"""
Entity Graph Visualization - Interactive Fund Flow for RugMaps + Scans
=======================================================================
Generates interactive entity relationship graphs showing wallet connections,
fund flows, and cluster relationships. SVG output for embedding in scan results
and RugMaps pages.
Premium feature: Visual forensics that sells in screenshots.
"""
import logging
from dataclasses import dataclass, field
from typing import Any
logger = logging.getLogger("entity.graph")
@dataclass
class GraphNode:
id: str
label: str
node_type: str = "wallet" # "wallet", "token", "contract", "entity", "cex"
risk_score: float = 0.0
chain: str = ""
size: int = 20
color: str = "#8B5CF6"
metadata: dict[str, Any] = field(default_factory=dict)
@dataclass
class GraphEdge:
source: str
target: str
edge_type: str = "transfer" # "transfer", "deploy", "fund", "cluster", "counterparty"
value_usd: float = 0.0
tx_count: int = 1
color: str = "#4B5563"
width: float = 1.0
def build_entity_graph(
center_address: str,
chain: str = "ethereum",
transactions: list[dict] | None = None,
labels: dict[str, Any] | None = None,
cluster_members: list[str] | None = None,
depth: int = 2,
) -> dict[str, Any]:
"""Build an entity relationship graph from wallet data.
Returns graph data suitable for visualization (SVG or Canvas).
"""
nodes: dict[str, GraphNode] = {}
edges: list[GraphEdge] = []
# ── Center node ──
center_id = f"{chain}:{center_address.lower()}"
nodes[center_id] = GraphNode(
id=center_id,
label=_shorten_addr(center_address),
node_type="wallet",
chain=chain,
size=30,
color="#22D3EE", # Cyan for center
metadata={"address": center_address, "chain": chain},
)
# ── Process transactions ──
txs = transactions or []
counterparties = set()
for tx in txs[:200]: # Cap at 200 transactions
if not isinstance(tx, dict):
continue
counterparty = tx.get("counterparty") or tx.get("to") or tx.get("from")
if not counterparty or counterparty.lower() == center_address.lower():
continue
cp_id = f"{tx.get('chain', chain)}:{counterparty.lower()}"
counterparties.add(cp_id)
value = float(tx.get("value_usd", 0) or 0)
tx_type = tx.get("type", "transfer")
is_outgoing = tx.get("direction") == "out" or tx.get("from", "").lower() == center_address.lower()
# Add counterparty node if new
if cp_id not in nodes:
# Determine node type and color
node_type, color = _classify_node(counterparty, labels or {}, tx)
nodes[cp_id] = GraphNode(
id=cp_id,
label=_shorten_addr(counterparty),
node_type=node_type,
chain=tx.get("chain", chain),
size=18 if node_type == "wallet" else 22,
color=color,
metadata={"address": counterparty},
)
# Add edge
edge_color = "#EF4444" if is_outgoing else "#22C55E" # Red=out, Green=in
edge_width = min(max(value / 10000, 1), 6) if value > 0 else 1
edges.append(
GraphEdge(
source=center_id if is_outgoing else cp_id,
target=cp_id if is_outgoing else center_id,
edge_type=tx_type,
value_usd=value,
tx_count=1,
color=edge_color,
width=edge_width,
)
)
# ── Add cluster members (if available) ──
if cluster_members:
cluster_color = "#F59E0B" # Amber for cluster
for member in cluster_members[:20]:
if member.lower() == center_address.lower():
continue
member_id = f"{chain}:{member.lower()}"
if member_id not in nodes:
nodes[member_id] = GraphNode(
id=member_id,
label=_shorten_addr(member),
node_type="cluster_member",
color=cluster_color,
size=16,
)
edges.append(
GraphEdge(
source=center_id,
target=member_id,
edge_type="cluster",
color=cluster_color,
width=1.5,
)
)
# ── Add label-based nodes (CEX, scams, etc.) ──
if labels:
label_data = labels.get(center_address, labels)
if isinstance(label_data, dict):
label_type = label_data.get("label_type") or label_data.get("type", "")
entity = label_data.get("entity") or label_data.get("name", "")
if label_type in ("cex", "exchange"):
nodes[center_id].node_type = "cex"
nodes[center_id].color = "#F59E0B" # Amber
nodes[center_id].label = entity or "Exchange"
elif label_type in ("scam", "phish", "hack"):
nodes[center_id].node_type = "scam"
nodes[center_id].color = "#EF4444" # Red
nodes[center_id].risk_score = 90
nodes[center_id].label = entity or "Scam"
elif label_type in ("mixer", "tumbler"):
nodes[center_id].node_type = "mixer"
nodes[center_id].color = "#8B5CF6" # Purple
nodes[center_id].risk_score = 70
# ── Risk-based coloring ──
for node in nodes.values():
if node.risk_score >= 80:
node.color = "#EF4444" # Red
elif node.risk_score >= 60:
node.color = "#F97316" # Orange
elif node.risk_score >= 40:
node.color = "#EAB308" # Yellow
# ── Convert to graph format ──
return {
"center": {"id": center_id, "address": center_address, "chain": chain},
"nodes": [
{
"id": n.id,
"label": n.label,
"type": n.node_type,
"chain": n.chain,
"risk_score": n.risk_score,
"size": n.size,
"color": n.color,
"metadata": n.metadata,
}
for n in nodes.values()
],
"edges": [
{
"source": e.source,
"target": e.target,
"type": e.edge_type,
"value_usd": e.value_usd,
"tx_count": e.tx_count,
"color": e.color,
"width": e.width,
}
for e in edges
],
"stats": {
"total_nodes": len(nodes),
"total_edges": len(edges),
"total_value_usd": sum(e.value_usd for e in edges),
"entity_types": list({n.node_type for n in nodes.values()}),
"risk_distribution": {
"high": sum(1 for n in nodes.values() if n.risk_score >= 60),
"medium": sum(1 for n in nodes.values() if 30 <= n.risk_score < 60),
"low": sum(1 for n in nodes.values() if n.risk_score < 30),
},
},
}
def generate_svg(graph_data: dict[str, Any], width: int = 800, height: int = 600) -> str:
"""Generate an SVG visualization of the entity graph.
Uses a simple force-directed layout algorithm.
"""
import math
nodes = graph_data.get("nodes", [])
edges = graph_data.get("edges", [])
center = graph_data.get("center", {})
center_id = center.get("id", "")
if not nodes:
return ''
# Simple radial layout: center in middle, others in concentric circles
cx, cy = width // 2, height // 2
# Separate center from others
other_nodes = [n for n in nodes if n["id"] != center_id]
center_node = next((n for n in nodes if n["id"] == center_id), None)
# Arrange other nodes in circles
positions = {}
if center_node:
positions[center_id] = (cx, cy)
rings = [
(150, 0.0), # radius, rotation offset
(250, 0.3),
(350, 0.15),
]
node_idx = 0
for ring_radius, rotation in rings:
nodes_in_ring = other_nodes[node_idx : node_idx + max(8, len(other_nodes) // len(rings))]
n = len(nodes_in_ring)
for i, node in enumerate(nodes_in_ring):
angle = (2 * math.pi * i / max(n, 1)) + rotation
x = cx + ring_radius * math.cos(angle)
y = cy + ring_radius * math.sin(angle)
positions[node["id"]] = (x, y)
node_idx += n
if node_idx >= len(other_nodes):
break
# Build SVG
svg_parts = [
f'")
return "\n".join(svg_parts)
def _classify_node(address: str, labels: dict[str, Any], tx: dict) -> tuple[str, str]:
"""Classify a counterparty node by type and color."""
addr_lower = address.lower()
# Check labels
label_info = labels.get(address, labels.get(addr_lower, {}))
if isinstance(label_info, dict):
label_type = label_info.get("label_type", label_info.get("type", ""))
if label_type in ("cex", "exchange"):
return ("cex", "#F59E0B")
if label_type in ("scam", "phish", "hack", "exploit"):
return ("scam", "#EF4444")
if label_type in ("mixer", "tumbler", "tornado"):
return ("mixer", "#8B5CF6")
if label_type in ("defi", "protocol"):
return ("protocol", "#06B6D4")
# Check transaction patterns
tx_type = tx.get("type", "")
if tx_type == "contract_creation" or tx_type == "deploy":
return ("contract", "#10B981")
if tx_type == "swap":
return ("dex", "#06B6D4")
return ("wallet", "#8B5CF6")
def _shorten_addr(addr: str) -> str:
"""Shorten address for display."""
if not addr:
return "?"
if len(addr) <= 12:
return addr
return f"{addr[:6]}...{addr[-4:]}"