tracker.py

import heapq
from concurrent.futures import ThreadPoolExecutor

import networkx as nx
from api_clients import fetch_btc_address_data

class BlockchainTracker:
    """
    Handles graph-based representation of blockchain transaction flows and traces paths.
    Uses NetworkX DiGraph internally to map connections and analyze paths.
    """
    def __init__(self, trail_data, source, target):
        self.trail_data = trail_data
        self.source = source
        self.target = target
        self.G = nx.DiGraph()
        self.malicious_path = []
        self.total_hops = 0
        self.path_amount = 0.0
        self.build_graph()
        self.find_path()

    @classmethod
    def from_dynamic_btc_trace(cls, seed_address, limit=5, max_hops=2, branch_limit=2):
        """
        Dynamically query Blockchain.info API and build a multi-hop, branching BlockchainTracker.

        Fetches are batched and run concurrently one BFS hop level at a time (instead of
        one address at a time), and each address is only ever fetched once per trace via
        `address_cache`. This preserves the exact same traversal order and endpoint-selection
        logic as before -- it just removes redundant round trips and lets independent
        network calls within a level overlap instead of queueing behind each other.
        """
        trail_data = {}
        address_cache = {}

        def fetch_cached(address):
            if address not in address_cache:
                address_cache[address] = fetch_btc_address_data(address, limit=limit)
            return address_cache[address]

        # 1. Fetch seed transactions to determine flow direction
        seed_data = fetch_cached(seed_address)
        if not seed_data or "txs" not in seed_data or not seed_data["txs"]:
            return cls({}, seed_address, seed_address)

        txs = seed_data["txs"]

        send_count = 0
        recv_count = 0
        for tx in txs:
            is_sender = False
            for inp in tx.get("inputs", []):
                prev_out = inp.get("prev_out")
                if prev_out and prev_out.get("addr") == seed_address:
                    is_sender = True
                    break
            if is_sender:
                send_count += 1
            else:
                recv_count += 1

        trace_outgoing = send_count >= recv_count

        visited = [seed_address]
        visited_set = {seed_address}

        endpoint = seed_address
        endpoint_hop = -1
        max_flow_val = 0.0

        def add_neighbor(neighbors, address, value, transaction):
            info = neighbors.setdefault(
                address,
                {"amount": 0.0, "transaction_ids": [], "transaction_records": []},
            )
            info["amount"] += value
            transaction_id = transaction.get("hash") or transaction.get("txid")
            if transaction_id and transaction_id not in info["transaction_ids"]:
                info["transaction_ids"].append(transaction_id)
                info["transaction_records"].append(
                    {
                        "id": transaction_id,
                        "block_height": transaction.get("block_height"),
                        "timestamp": transaction.get("time"),
                    }
                )

        # BFS by level: every address in `current_level` is at the same hop depth,
        # so their fetches are independent of one another and safe to run in parallel.
        current_level = [seed_address]
        hop = 0
        while current_level and hop < max_hops:
            to_fetch = [addr for addr in current_level if addr not in address_cache]
            if to_fetch:
                with ThreadPoolExecutor(max_workers=min(8, len(to_fetch))) as pool:
                    fetched = pool.map(
                        lambda address: fetch_btc_address_data(address, limit=limit),
                        to_fetch,
                    )
                    for addr, data in zip(to_fetch, fetched):
                        address_cache[addr] = data

            next_level = []
            for current_addr in current_level:
                addr_data = address_cache.get(current_addr)
                if not addr_data or "txs" not in addr_data or not addr_data["txs"]:
                    continue

                node_txs = addr_data["txs"]
                neighbors = {}  # target_addr -> amount and transaction provenance

                for tx in node_txs:
                    if trace_outgoing:
                        is_sender = False
                        for inp in tx.get("inputs", []):
                            prev_out = inp.get("prev_out")
                            if prev_out and prev_out.get("addr") == current_addr:
                                is_sender = True
                                break
                        if is_sender:
                            for out in tx.get("out", []):
                                recipient = out.get("addr")
                                val = out.get("value", 0) / 1e8
                                if recipient and recipient != current_addr:
                                    add_neighbor(neighbors, recipient, val, tx)
                    else:
                        is_receiver = False
                        for out in tx.get("out", []):
                            if out.get("addr") == current_addr:
                                is_receiver = True
                                break
                        if is_receiver:
                            for inp in tx.get("inputs", []):
                                prev_out = inp.get("prev_out")
                                if prev_out:
                                    sender = prev_out.get("addr")
                                    val = prev_out.get("value", 0) / 1e8
                                    if sender and sender != current_addr:
                                        add_neighbor(neighbors, sender, val, tx)

                if not neighbors:
                    continue

                # Sort neighbors by value transferred (highest flow first)
                sorted_neighbors = sorted(
                    neighbors.items(), key=lambda x: x[1]["amount"], reverse=True
                )
                top_neighbors = sorted_neighbors[:branch_limit]

                # Record in trail_data
                if trace_outgoing:
                    trail_data[current_addr] = [
                        (addr, info["amount"], info) for addr, info in neighbors.items()
                    ]
                    for addr, info in top_neighbors:
                        if addr not in visited_set:
                            visited_set.add(addr)
                            visited.append(addr)
                            next_level.append(addr)
                            # Record target as the deepest leaf in the highest flow path
                            if (
                                hop + 1 > endpoint_hop
                                or (
                                    hop + 1 == endpoint_hop
                                    and info["amount"] > max_flow_val
                                )
                            ):
                                endpoint_hop = hop + 1
                                max_flow_val = info["amount"]
                                endpoint = addr
                else:
                    for sender, info in neighbors.items():
                        if sender not in trail_data:
                            trail_data[sender] = []
                        trail_data[sender].append(
                            (current_addr, info["amount"], info)
                        )
                    for addr, info in top_neighbors:
                        if addr not in visited_set:
                            visited_set.add(addr)
                            visited.append(addr)
                            next_level.append(addr)
                            if (
                                hop + 1 > endpoint_hop
                                or (
                                    hop + 1 == endpoint_hop
                                    and info["amount"] > max_flow_val
                                )
                            ):
                                endpoint_hop = hop + 1
                                max_flow_val = info["amount"]
                                endpoint = addr

            current_level = next_level
            hop += 1

        # Define source & target for path calculations
        if trace_outgoing:
            source = seed_address
            target = endpoint if endpoint != seed_address else (visited[1] if len(visited) > 1 else seed_address)
        else:
            source = endpoint if endpoint != seed_address else (visited[1] if len(visited) > 1 else seed_address)
            target = seed_address

        return cls(trail_data, source, target)

    def build_graph(self):
        """
        Populate the NetworkX DiGraph with edges and amounts from the trail data.
        """
        for src, dests in self.trail_data.items():
            for edge in dests:
                dest, amount = edge[:2]
                metadata = edge[2] if len(edge) > 2 and isinstance(edge[2], dict) else {}
                transaction_ids = list(metadata.get("transaction_ids", []))
                transaction_records = list(metadata.get("transaction_records", []))
                if self.G.has_edge(src, dest):
                    self.G[src][dest]["amount"] += amount
                    existing_ids = self.G[src][dest].setdefault("transaction_ids", [])
                    self.G[src][dest]["transaction_ids"] = list(
                        dict.fromkeys(existing_ids + transaction_ids)
                    )
                    existing_records = self.G[src][dest].setdefault(
                        "transaction_records", []
                    )
                    known_ids = {record.get("id") for record in existing_records}
                    self.G[src][dest]["transaction_records"].extend(
                        record
                        for record in transaction_records
                        if record.get("id") not in known_ids
                    )
                    self.G[src][dest]["transaction_count"] = len(
                        self.G[src][dest]["transaction_ids"]
                    )
                else:
                    self.G.add_edge(
                        src,
                        dest,
                        amount=amount,
                        transaction_ids=transaction_ids,
                        transaction_records=transaction_records,
                        transaction_count=len(transaction_ids),
                    )

    def find_path(self):
        """
        Compute the strongest fund-flow path from source to target.

        The minimum edge amount is used as the path's bottleneck because a
        trace should prefer the route that preserves the most value across
        every hop, rather than the route with the fewest nodes. A widest-path
        search avoids enumerating every simple path in a branching graph.
        """
        try:
            if self.source == self.target:
                self.malicious_path = [self.source]
                self.total_hops = 0
                self.path_amount = 0.0
                return

            best_capacity = {self.source: float("inf")}
            best_hops = {self.source: 0}
            predecessors = {}
            queue = [(float("-inf"), 0, self.source)]

            while queue:
                negative_capacity, hops, current = heapq.heappop(queue)
                capacity = -negative_capacity
                if capacity != best_capacity.get(current) or hops != best_hops.get(current):
                    continue
                if current == self.target:
                    break

                for neighbor, data in self.G[current].items():
                    edge_amount = data.get("amount", 0.0)
                    candidate_capacity = min(capacity, edge_amount)
                    candidate_hops = hops + 1
                    current_capacity = best_capacity.get(neighbor, 0.0)
                    current_hops = best_hops.get(neighbor, float("inf"))
                    if (
                        candidate_capacity > current_capacity
                        or (
                            candidate_capacity == current_capacity
                            and candidate_hops < current_hops
                        )
                    ):
                        best_capacity[neighbor] = candidate_capacity
                        best_hops[neighbor] = candidate_hops
                        predecessors[neighbor] = current
                        heapq.heappush(
                            queue, (-candidate_capacity, candidate_hops, neighbor)
                        )

            if self.target in best_capacity:
                path = [self.target]
                while path[-1] != self.source:
                    path.append(predecessors[path[-1]])
                self.malicious_path = list(reversed(path))
                self.path_amount = best_capacity[self.target]
            else:
                self.malicious_path = []
                self.path_amount = 0.0
            self.total_hops = len(self.malicious_path) - 1
        except (nx.NetworkXNoPath, nx.NodeNotFound):
            self.malicious_path = []
            self.total_hops = 0
            self.path_amount = 0.0
        except Exception as e:
            print(f"Error computing graph path: {e}")
            self.malicious_path = []
            self.total_hops = 0
            self.path_amount = 0.0

    def get_metrics(self):
        """
        Get tracking metrics.
        """
        path_transaction_ids = []
        for source, target in zip(self.malicious_path, self.malicious_path[1:]):
            for transaction_id in self.G[source][target].get("transaction_ids", []):
                if transaction_id not in path_transaction_ids:
                    path_transaction_ids.append(transaction_id)

        edges = [
            {
                "source": source,
                "target": target,
                "amount": data.get("amount", 0.0),
                "transaction_ids": list(data.get("transaction_ids", [])),
                "transaction_records": list(data.get("transaction_records", [])),
                "transaction_count": data.get("transaction_count", 0),
            }
            for source, target, data in self.G.edges(data=True)
        ]
        return {
            "hops": self.total_hops,
            "path": self.malicious_path,
            "path_amount": self.path_amount,
            "path_transaction_ids": path_transaction_ids,
            "edges": edges,
            "start": self.source,
            "end": self.target,
            "last_traced_address": (
                self.malicious_path[-2]
                if len(self.malicious_path) > 1 else self.source
            )
        }