main_forensics.py

import os
import time
import json
import logging
import requests
import networkx as nx
from datetime import datetime
from pyvis.network import Network
from dotenv import load_dotenv

load_dotenv()

# --- CONFIGURATION & SINK LIST ---
USDT_CONTRACT = "TR7NHqjeKQxGTCi8q8ZY4pL8otSzgjLj6t"
API_KEY = os.getenv("TRONSCAN_API_KEY")

# Major Exchange Hot Wallets (Sinks)
SINK_ADDRESSES = {
    "TwmTmsVy5hCdSwyth3vWJ1a8uR9g56S8qG": "Binance-Hot-1",
    "TAUN6FwrnwwmaEqYcckffC7wYmbaS6c767": "Binance-Hot-2",
    "TNPeeaaFB7K9cmo4uQpcU32zGK8G1NYqeL": "Huobi-Hot",
    "TM1zzNDZD2DPASbKcgdVoTYhfmYgtfwx9R": "OKX-Hot",
    "TLsV52sRDL79HXGGm9yzwKgccDfJrwpaPf": "Kraken-Hot",
    "TKzxdSv2FZKQrEqkKVgp5DcwEXBEKMg2Ax": "SunSwap-Router"
}

class TronForensics:
    def __init__(self, api_key, cache_file="blockchain_cache.json"):
        self.api_key = api_key
        self.cache_file = cache_file
        self.headers = {"TRON-PRO-API-KEY": api_key}
        self.last_call_time = 0
        self.graph = nx.DiGraph()
        self.cache = self._load_cache()
        self.start_address = None  # Track the starting point

    def _load_cache(self):
        if os.path.exists(self.cache_file):
            try:
                with open(self.cache_file, 'r') as f:
                    content = f.read().strip()
                    if not content:  # Empty file
                        return {}
                    return json.loads(content)
            except json.JSONDecodeError as e:
                logging.warning(f"Cache file corrupted, starting fresh: {e}")
                return {}
        return {}

    def _save_cache(self):
        with open(self.cache_file, 'w') as f:
            json.dump(self.cache, f)

    def rate_limited_get(self, url, params):
        """Ensures we stay under 5 calls/sec (0.2s delay)."""
        elapsed = time.time() - self.last_call_time
        if elapsed < 0.25: # Slightly more than 0.2 to be safe
            time.sleep(0.25 - elapsed)
        
        response = requests.get(url, params=params, headers=self.headers)
        self.last_call_time = time.time()
        
        if response.status_code == 429:
            logging.warning("Rate limit hit! Cooling down...")
            time.sleep(5)
            return self.rate_limited_get(url, params)
        return response.json()

    def get_usdt_transfers(self, address, direction="out"):
        """Fetches transfers with dust filtering and caching."""
        if address in self.cache:
            return self.cache[address]

        url = "https://apilist.tronscan.org/api/token_trc20/transfers"
        params = {
            "limit": 50, # High limit to save calls
            "start": 0,
            "contract_address": USDT_CONTRACT,
            "relatedAddress": address
        }
        
        data = self.rate_limited_get(url, params)
        transfers = data.get("token_transfers", [])
        
        # Clean & Filter Data
        results = []
        for tx in transfers:
            amount = float(tx.get("quant", 0)) / 1_000_000
            
            # EDGE CASE 1: Filter Dust/Spam (Less than 5 USDT is usually noise)
            if amount < 5.0: continue
            
            # EDGE CASE 2: Ensure it's the real USDT contract
            if tx.get("contract_address") != USDT_CONTRACT: continue

            results.append({
                "from": tx["from_address"],
                "to": tx["to_address"],
                "amount": amount,
                "ts": tx["block_ts"]
            })
        
        self.cache[address] = results
        self._save_cache()
        return results

    def trace(self, start_address, depth=3):
        self.start_address = start_address  # Remember the starting point
        queue = [(start_address, 0)]
        visited = set()

        while queue:
            addr, current_depth = queue.pop(0)
            if current_depth >= depth or addr in SINK_ADDRESSES or addr in visited:
                continue
            
            visited.add(addr)
            transfers = self.get_usdt_transfers(addr)
            
            for tx in transfers:
                # Add to graph
                self.graph.add_edge(tx["from"], tx["to"], weight=tx["amount"])
                
                # Recursive step
                if tx["to"] not in visited:
                    queue.append((tx["to"], current_depth + 1))

    def export_interactive(self, filename="ponzi_map.html"):
        net = Network(
            height="900px", 
            width="100%", 
            bgcolor="#0a0a0a", 
            font_color="white", 
            directed=True,
            heading="🕵️ USDT Money Flow Analysis"
        )
        
        # Better physics for clearer visualization
        net.barnes_hut(
            gravity=-80000,
            central_gravity=0.3,
            spring_length=200,
            spring_strength=0.001,
            damping=0.09
        )
        
        for node in self.graph.nodes():
            # Default: Intermediate wallet
            color = "#3498db"  # Blue
            size = 25
            shape = "dot"
            label = f"Wallet\n{node[:6]}...{node[-6:]}"
            title = f"<b>Intermediate Wallet</b><br>{node}<br>Click to view on Tronscan"
            border_color = "#2980b9"
            
            # Starting address (Suspect)
            if node == self.start_address:
                color = "#e74c3c"  # Red
                size = 40
                shape = "star"
                label = f"🎯 SUSPECT\n{node[:6]}...{node[-6:]}"
                title = f"<b>⚠️ INVESTIGATION TARGET</b><br>{node}<br>This is where the trace started"
                border_color = "#c0392b"
            
            # Exchange/Sink address
            elif node in SINK_ADDRESSES:
                color = "#2ecc71"  # Green
                size = 35
                shape = "box"
                label = f"🏦 {SINK_ADDRESSES[node]}\n{node[:6]}..."
                title = f"<b>💰 EXCHANGE WALLET</b><br>{SINK_ADDRESSES[node]}<br>{node}<br>Money likely cashed out here"
                border_color = "#27ae60"
            
            net.add_node(
                node, 
                label=label, 
                title=title, 
                color={"background": color, "border": border_color},
                size=size,
                shape=shape,
                borderWidth=3,
                font={"size": 14, "color": "white", "bold": True}
            )

        # Add edges with better styling
        for u, v, data in self.graph.edges(data=True):
            amount = data['weight']
            
            # Make edge thickness proportional to amount
            width = min(1 + (amount / 1000), 10)  # Cap at 10px
            
            # Color code by amount
            if amount > 10000:
                edge_color = "#e74c3c"  # Red for large amounts
            elif amount > 1000:
                edge_color = "#f39c12"  # Orange for medium
            else:
                edge_color = "#95a5a6"  # Gray for small
            
            label = f"${amount:,.2f}"
            title = f"<b>Transfer: ${amount:,.2f} USDT</b><br>From: {u}<br>To: {v}"
            
            net.add_edge(
                u, 
                v, 
                value=width,
                title=title,
                label=label,
                color=edge_color,
                arrows={"to": {"enabled": True, "scaleFactor": 1.2}},
                font={"size": 12, "color": "white", "strokeWidth": 0}
            )
        
        # Add custom legend and instructions
        legend_html = """
        <div style="position: fixed; top: 80px; left: 20px; background: rgba(0,0,0,0.85); 
                    padding: 20px; border-radius: 10px; border: 2px solid #3498db; 
                    font-family: Arial; color: white; z-index: 1000; max-width: 300px;">
            <h3 style="margin-top: 0; color: #3498db;">📊 Legend</h3>
            
            <div style="margin: 10px 0;">
                <span style="color: #e74c3c; font-size: 20px;">⭐</span>
                <strong style="color: #e74c3c;">RED STAR</strong> = Investigation Target (Suspect)
            </div>
            
            <div style="margin: 10px 0;">
                <span style="color: #3498db; font-size: 20px;">●</span>
                <strong style="color: #3498db;">BLUE CIRCLES</strong> = Intermediate Wallets
            </div>
            
            <div style="margin: 10px 0;">
                <span style="color: #2ecc71; font-size: 20px;">■</span>
                <strong style="color: #2ecc71;">GREEN BOXES</strong> = Exchange Wallets (Cash-out points)
            </div>
            
            <hr style="border-color: #34495e;">
            
            <h4 style="color: #f39c12; margin-bottom: 5px;">💸 Transfer Amounts</h4>
            <div style="font-size: 12px; line-height: 1.6;">
                <span style="color: #e74c3c;">●</span> Red arrow = Over $10,000<br>
                <span style="color: #f39c12;">●</span> Orange arrow = $1,000 - $10,000<br>
                <span style="color: #95a5a6;">●</span> Gray arrow = Under $1,000
            </div>
            
            <hr style="border-color: #34495e;">
            
            <h4 style="color: #3498db; margin-bottom: 5px;">🎮 Controls</h4>
            <div style="font-size: 12px; line-height: 1.6;">
                • <strong>Drag</strong> nodes to rearrange<br>
                • <strong>Scroll</strong> to zoom in/out<br>
                • <strong>Click</strong> nodes for details<br>
                • <strong>Hover</strong> over arrows for amounts
            </div>
        </div>
        """
        
        # Write HTML without opening browser
        net.write_html(filename, open_browser=False, notebook=False)
        
        # Add legend to the HTML file
        with open(filename, 'r', encoding='utf-8') as f:
            html_content = f.read()
        
        # Insert legend before closing body tag
        html_content = html_content.replace('</body>', f'{legend_html}</body>')
        
        with open(filename, 'w', encoding='utf-8') as f:
            f.write(html_content)
        
        print(f"\n✅ Investigation complete! Open '{filename}' in your browser to view the network.")
        print(f"📍 Starting point (Suspect): {self.start_address}")
        print(f"🔗 Total wallets traced: {len(self.graph.nodes())}")
        print(f"💸 Total transfers tracked: {len(self.graph.edges())}")

if __name__ == "__main__":
    investigator = TronForensics(api_key=API_KEY)
    # Replace with actual suspect address
    target = "TYCR4XDMvFZ6LXCkSaMWYnULZeeod8GCAM" 
    investigator.trace(target, depth=2)
    investigator.export_interactive()