import gradio as gr
import plotly.graph_objects as go
import plotly.express as px
from plotly.subplots import make_subplots
import pandas as pd
from typing import Dict, List, Tuple, Any

# Color scheme for MITRE phases
PHASE_COLORS = {
    "initial_access": "#FF4444",      # Red
    "execution": "#FF8C00",           # Orange
    "credential_access": "#FFD700",   # Gold
    "lateral_movement": "#4169E1",    # Blue
    "privilege_escalation": "#9370DB", # Purple
    "persistence": "#20B2AA",         # Teal
    "command_control": "#FF1493",     # Deep Pink
    "exfiltration": "#1a1a1a",        # Black
    "impact": "#8B0000",              # Dark Red
}

PHASE_LABELS_EN = {
    "initial_access": "Initial Access",
    "execution": "Execution",
    "credential_access": "Credential Access",
    "lateral_movement": "Lateral Movement",
    "privilege_escalation": "Privilege Escalation",
    "persistence": "Persistence",
    "command_control": "Command & Control",
    "exfiltration": "Exfiltration",
    "impact": "Impact",
}

PHASE_LABELS_FR = {
    "initial_access": "Accès initial",
    "execution": "Exécution",
    "credential_access": "Accès aux identifiants",
    "lateral_movement": "Mouvement latéral",
    "privilege_escalation": "Escalade de privilèges",
    "persistence": "Persistance",
    "command_control": "Commande et contrôle",
    "exfiltration": "Exfiltration",
    "impact": "Impact",
}

# Comprehensive attack paths
ATTACK_PATHS = [
    {
        "id": "AP-001",
        "name_en": "Classic AD Domain Compromise",
        "name_fr": "Compromission de domaine AD classique",
        "category": "active_directory",
        "nodes": [
            {"id": "n1", "label_en": "Phishing Email", "label_fr": "Email de Phishing", "mitre": "T1566", "phase": "initial_access"},
            {"id": "n2", "label_en": "Macro Execution", "label_fr": "Exécution Macro", "mitre": "T1059", "phase": "execution"},
            {"id": "n3", "label_en": "Credential Dump (LSASS)", "label_fr": "Dump d'identifiants (LSASS)", "mitre": "T1003.001", "phase": "credential_access"},
            {"id": "n4", "label_en": "Lateral Movement (PsExec)", "label_fr": "Mouvement latéral (PsExec)", "mitre": "T1021.002", "phase": "lateral_movement"},
            {"id": "n5", "label_en": "Domain Admin Compromise", "label_fr": "Compromission Admin du domaine", "mitre": "T1078.002", "phase": "privilege_escalation"},
            {"id": "n6", "label_en": "DCSync for Persistence", "label_fr": "DCSync pour la persistance", "mitre": "T1003.006", "phase": "persistence"},
            {"id": "n7", "label_en": "Data Exfiltration", "label_fr": "Exfiltration de données", "mitre": "T1041", "phase": "exfiltration"},
        ],
        "edges": [("n1","n2"), ("n2","n3"), ("n3","n4"), ("n4","n5"), ("n5","n6"), ("n6","n7")],
        "description_en": "Attacker sends phishing email to user. Target opens macro-enabled document, triggering code execution. Malware dumps credentials from LSASS. Attacker uses stolen credentials for lateral movement to server. Domain admin account is compromised. DCSync is set up for persistent backdoor access. Sensitive data is exfiltrated.",
        "description_fr": "L'attaquant envoie un email de phishing à l'utilisateur. La cible ouvre un document avec macros, déclenchant l'exécution de code. Le malware vide les identifiants de LSASS. L'attaquant utilise les identifiants volés pour se déplacer latéralement. Le compte admin du domaine est compromis. DCSync est configuré pour un accès persistant. Les données sensibles sont exfiltrées.",
    },
    {
        "id": "AP-002",
        "name_en": "Web App to Internal Network",
        "name_fr": "Application Web au réseau interne",
        "category": "web_application",
        "nodes": [
            {"id": "n1", "label_en": "SQL Injection", "label_fr": "Injection SQL", "mitre": "T1190", "phase": "initial_access"},
            {"id": "n2", "label_en": "Web Shell Upload", "label_fr": "Upload de Web Shell", "mitre": "T1190", "phase": "execution"},
            {"id": "n3", "label_en": "Reverse Shell", "label_fr": "Reverse Shell", "mitre": "T1059", "phase": "execution"},
            {"id": "n4", "label_en": "Network Enumeration", "label_fr": "Énumération réseau", "mitre": "T1592", "phase": "command_control"},
            {"id": "n5", "label_en": "Database Access", "label_fr": "Accès à la base de données", "mitre": "T1005", "phase": "credential_access"},
            {"id": "n6", "label_en": "Internal Server Compromise", "label_fr": "Compromission du serveur interne", "mitre": "T1021.001", "phase": "lateral_movement"},
            {"id": "n7", "label_en": "Data Breach", "label_fr": "Fuite de données", "mitre": "T1020", "phase": "exfiltration"},
        ],
        "edges": [("n1","n2"), ("n2","n3"), ("n3","n4"), ("n4","n5"), ("n4","n6"), ("n5","n7"), ("n6","n7")],
        "description_en": "SQL injection vulnerability exploited to gain initial access. Attacker uploads web shell for code execution. Reverse shell established for interactive access. Network enumeration reveals internal infrastructure. Database accessed for credentials and data. Internal servers compromised for lateral movement. Sensitive data exfiltrated.",
        "description_fr": "Vulnérabilité d'injection SQL exploitée pour un accès initial. L'attaquant charge un web shell pour l'exécution de code. Shell inversé établi pour l'accès interactif. L'énumération du réseau révèle l'infrastructure interne. Base de données accédée pour les identifiants et les données. Serveurs internes compromis pour le mouvement latéral. Les données sensibles sont exfiltrées.",
    },
    {
        "id": "AP-003",
        "name_en": "Cloud Account Takeover",
        "name_fr": "Prise de contrôle du compte cloud",
        "category": "cloud",
        "nodes": [
            {"id": "n1", "label_en": "Credential Stuffing", "label_fr": "Credential Stuffing", "mitre": "T1110.004", "phase": "initial_access"},
            {"id": "n2", "label_en": "MFA Bypass/Disable", "label_fr": "Contournement/Désactivation MFA", "mitre": "T1556", "phase": "credential_access"},
            {"id": "n3", "label_en": "API Token Theft", "label_fr": "Vol de jeton API", "mitre": "T1528", "phase": "credential_access"},
            {"id": "n4", "label_en": "IAM Role Enumeration", "label_fr": "Énumération des rôles IAM", "mitre": "T1526", "phase": "persistence"},
            {"id": "n5", "label_en": "Privilege Escalation", "label_fr": "Escalade de privilèges", "mitre": "T1078.004", "phase": "privilege_escalation"},
            {"id": "n6", "label_en": "Data Access", "label_fr": "Accès aux données", "mitre": "T1526", "phase": "credential_access"},
            {"id": "n7", "label_en": "Cloud Data Exfiltration", "label_fr": "Exfiltration de données cloud", "mitre": "T1020.001", "phase": "exfiltration"},
        ],
        "edges": [("n1","n2"), ("n2","n3"), ("n3","n4"), ("n4","n5"), ("n5","n6"), ("n6","n7")],
        "description_en": "Attacker obtains credentials via credential stuffing from leaked database. MFA is bypassed or disabled by social engineering. API tokens are stolen for persistent access. IAM roles are enumerated to find privilege escalation paths. Privileges escalated to admin role. Cloud storage and databases accessed. Exfiltrate sensitive cloud data.",
        "description_fr": "L'attaquant obtient les identifiants via credential stuffing à partir de base de données divulguée. L'AMF est contournée ou désactivée par ingénierie sociale. Les jetons API sont volés pour un accès persistant. Les rôles IAM sont énumérés pour trouver des chemins d'escalade. Les privilèges sont escaladés vers le rôle administrateur. L'accès aux stockages et bases de données cloud. Exfiltration de données cloud sensibles.",
    },
    {
        "id": "AP-004",
        "name_en": "Ransomware Kill Chain",
        "name_fr": "Chaîne d'attaque Ransomware",
        "category": "network",
        "nodes": [
            {"id": "n1", "label_en": "Phishing with Dropper", "label_fr": "Phishing avec dropper", "mitre": "T1566", "phase": "initial_access"},
            {"id": "n2", "label_en": "Malware Execution", "label_fr": "Exécution du malware", "mitre": "T1204", "phase": "execution"},
            {"id": "n3", "label_en": "C2 Beaconing", "label_fr": "C2 Beaconing", "mitre": "T1071", "phase": "command_control"},
            {"id": "n4", "label_en": "Network Reconnaissance", "label_fr": "Reconnaissance réseau", "mitre": "T1087", "phase": "persistence"},
            {"id": "n5", "label_en": "Lateral Movement", "label_fr": "Mouvement latéral", "mitre": "T1021", "phase": "lateral_movement"},
            {"id": "n6", "label_en": "Backup Encryption", "label_fr": "Chiffrement des sauvegardes", "mitre": "T1486", "phase": "impact"},
            {"id": "n7", "label_en": "Full Encryption & Ransom", "label_fr": "Chiffrement complet & rançon", "mitre": "T1486", "phase": "impact"},
        ],
        "edges": [("n1","n2"), ("n2","n3"), ("n3","n4"), ("n4","n5"), ("n5","n6"), ("n6","n7")],
        "description_en": "Phishing email with malware dropper attachment. User executes malware. Malware establishes C2 communication with attacker infrastructure. Network reconnaissance to identify high-value targets. Lateral movement to file servers and backup systems. Backup systems encrypted to prevent recovery. Full encryption of data with ransom demand.",
        "description_fr": "Email de phishing avec pièce jointe dropper malveillant. L'utilisateur exécute le malware. Le malware établit la communication C2 avec l'infrastructure de l'attaquant. Reconnaissance réseau pour identifier les cibles de haute valeur. Mouvement latéral vers les serveurs de fichiers et les systèmes de sauvegarde. Chiffrement des systèmes de sauvegarde pour empêcher la récupération. Chiffrement complet des données avec demande de rançon.",
    },
    {
        "id": "AP-005",
        "name_en": "Supply Chain Attack",
        "name_fr": "Attaque de chaîne d'approvisionnement",
        "category": "network",
        "nodes": [
            {"id": "n1", "label_en": "Compromise Dependency", "label_fr": "Compromission de dépendance", "mitre": "T1195.001", "phase": "initial_access"},
            {"id": "n2", "label_en": "Code Execution on Build", "label_fr": "Exécution de code au build", "mitre": "T1059", "phase": "execution"},
            {"id": "n3", "label_en": "Malicious Payload Injection", "label_fr": "Injection de charge malveillante", "mitre": "T1195.002", "phase": "execution"},
            {"id": "n4", "label_en": "Distributed to Customers", "label_fr": "Distribué aux clients", "mitre": "T1195", "phase": "initial_access"},
            {"id": "n5", "label_en": "C2 Establishment", "label_fr": "Établissement C2", "mitre": "T1071", "phase": "command_control"},
            {"id": "n6", "label_en": "Persistence Mechanism", "label_fr": "Mécanisme de persistance", "mitre": "T1547", "phase": "persistence"},
            {"id": "n7", "label_en": "Data Exfiltration", "label_fr": "Exfiltration de données", "mitre": "T1041", "phase": "exfiltration"},
        ],
        "edges": [("n1","n2"), ("n2","n3"), ("n3","n4"), ("n4","n5"), ("n5","n6"), ("n6","n7")],
        "description_en": "Open-source dependency compromised through account takeover. Attacker injects malicious code into library. Payload included in build artifacts. Backdoor distributed to all customers. C2 communication established from infected systems. Persistence achieved through startup scripts. Data exfiltrated from multiple organizations.",
        "description_fr": "Dépendance open-source compromise via prise de contrôle de compte. L'attaquant injecte du code malveillant dans la bibliothèque. La charge utile incluse dans les artefacts de build. Porte dérobée distribuée à tous les clients. Communication C2 établie à partir des systèmes infectés. Persistance réalisée par des scripts de démarrage. Données exfiltrées de plusieurs organisations.",
    },
    {
        "id": "AP-006",
        "name_en": "Insider Threat - Data Theft",
        "name_fr": "Menace interne - Vol de données",
        "category": "insider_threat",
        "nodes": [
            {"id": "n1", "label_en": "Legitimate Access", "label_fr": "Accès légitime", "mitre": "T1078", "phase": "initial_access"},
            {"id": "n2", "label_en": "Escalate Access Permissions", "label_fr": "Escalader les autorisations d'accès", "mitre": "T1078.003", "phase": "privilege_escalation"},
            {"id": "n3", "label_en": "Data Discovery", "label_fr": "Découverte de données", "mitre": "T1005", "phase": "credential_access"},
            {"id": "n4", "label_en": "Copy to External Drive", "label_fr": "Copier sur lecteur externe", "mitre": "T1052", "phase": "exfiltration"},
            {"id": "n5", "label_en": "Upload to Personal Cloud", "label_fr": "Télécharger vers cloud personnel", "mitre": "T1020.001", "phase": "exfiltration"},
            {"id": "n6", "label_en": "Delete Logs & Cover Tracks", "label_fr": "Supprimer les journaux & couvrir les traces", "mitre": "T1070", "phase": "persistence"},
            {"id": "n7", "label_en": "Share with Competitors", "label_fr": "Partager avec concurrents", "mitre": "T1041", "phase": "impact"},
        ],
        "edges": [("n1","n2"), ("n2","n3"), ("n3","n4"), ("n3","n5"), ("n4","n6"), ("n5","n6"), ("n6","n7")],
        "description_en": "Disgruntled employee with legitimate system access. Escalates access by requesting additional permissions from manager. Discovers and identifies sensitive/proprietary data. Copies data to personal external storage device. Uploads data to personal cloud storage account. Deletes audit logs and system events to cover tracks. Shares trade secrets with competitors for financial gain.",
        "description_fr": "Employé mécontent ayant accès légitime aux systèmes. Escalade l'accès en demandant des autorisations supplémentaires au manager. Découvre et identifie les données sensibles/propriétaires. Copie les données sur un dispositif de stockage externe personnel. Télécharge les données vers un compte cloud personnel. Supprime les journaux d'audit et les événements système pour couvrir les traces. Partage les secrets commerciaux avec les concurrents pour un gain financier.",
    },
    {
        "id": "AP-007",
        "name_en": "NTLM Relay Attack",
        "name_fr": "Attaque par relais NTLM",
        "category": "active_directory",
        "nodes": [
            {"id": "n1", "label_en": "LLMNR/mDNS Poisoning", "label_fr": "Empoisonnement LLMNR/mDNS", "mitre": "T1557.002", "phase": "initial_access"},
            {"id": "n2", "label_en": "NTLM Auth Capture", "label_fr": "Capture d'authentification NTLM", "mitre": "T1557.001", "phase": "credential_access"},
            {"id": "n3", "label_en": "Relay to Target Server", "label_fr": "Relais vers serveur cible", "mitre": "T1557.002", "phase": "lateral_movement"},
            {"id": "n4", "label_en": "Impersonate User", "label_fr": "Usurper l'identité de l'utilisateur", "mitre": "T1134.003", "phase": "lateral_movement"},
            {"id": "n5", "label_en": "Shell Access", "label_fr": "Accès Shell", "mitre": "T1059", "phase": "execution"},
            {"id": "n6", "label_en": "Install Persistence", "label_fr": "Installer la persistance", "mitre": "T1547", "phase": "persistence"},
            {"id": "n7", "label_en": "Lateral Escalation", "label_fr": "Escalade latérale", "mitre": "T1021", "phase": "privilege_escalation"},
        ],
        "edges": [("n1","n2"), ("n2","n3"), ("n3","n4"), ("n4","n5"), ("n5","n6"), ("n6","n7")],
        "description_en": "Attacker on network poisons LLMNR/mDNS responses. Captures NTLM authentication exchanges from users. Relays captured NTLM hash to target server (file share, email, etc). Successfully authenticates as legitimate user. Gains shell access to compromised server. Installs persistent backdoor. Uses compromised account for privilege escalation.",
        "description_fr": "L'attaquant sur le réseau empoisonne les réponses LLMNR/mDNS. Capture les échanges d'authentification NTLM des utilisateurs. Relaye le hachage NTLM capturé vers le serveur cible (partage de fichiers, email, etc). S'authentifie avec succès en tant qu'utilisateur légitime. Obtient l'accès shell au serveur compromis. Installe une porte dérobée persistante. Utilise le compte compromis pour l'escalade de privilèges.",
    },
    {
        "id": "AP-008",
        "name_en": "Kerberos Attack Chain",
        "name_fr": "Chaîne d'attaque Kerberos",
        "category": "active_directory",
        "nodes": [
            {"id": "n1", "label_en": "SPN Enumeration", "label_fr": "Énumération SPN", "mitre": "T1087.002", "phase": "initial_access"},
            {"id": "n2", "label_en": "Service Ticket Request", "label_fr": "Demande de ticket de service", "mitre": "T1558.003", "phase": "credential_access"},
            {"id": "n3", "label_en": "Kerberoasting (Crack Hash)", "label_fr": "Kerberoasting (Crack Hash)", "mitre": "T1558.003", "phase": "credential_access"},
            {"id": "n4", "label_en": "Service Account Password", "label_fr": "Mot de passe du compte de service", "mitre": "T1110", "phase": "credential_access"},
            {"id": "n5", "label_en": "Pass-the-Hash Attack", "label_fr": "Attaque Pass-the-Hash", "mitre": "T1550.002", "phase": "lateral_movement"},
            {"id": "n6", "label_en": "DCSync Attack", "label_fr": "Attaque DCSync", "mitre": "T1003.006", "phase": "privilege_escalation"},
            {"id": "n7", "label_en": "Extract All Hashes", "label_fr": "Extraire tous les hachages", "mitre": "T1003", "phase": "exfiltration"},
        ],
        "edges": [("n1","n2"), ("n2","n3"), ("n3","n4"), ("n4","n5"), ("n5","n6"), ("n6","n7")],
        "description_en": "Enumerate Service Principal Names (SPNs) in Active Directory. Request service tickets for identified services. Extract hashes from tickets for offline cracking. Crack service account password through brute force. Use Pass-the-Hash technique to authenticate as service account. Execute DCSync to extract domain controller hashes. Obtain all user hashes for complete domain compromise.",
        "description_fr": "Énumérer les noms principaux de service (SPN) dans Active Directory. Demander des tickets de service pour les services identifiés. Extraire les hachages des tickets pour le craquage hors ligne. Craquer le mot de passe du compte de service par force brute. Utiliser la technique Pass-the-Hash pour s'authentifier en tant que compte de service. Exécuter DCSync pour extraire les hachages du contrôleur de domaine. Obtenir tous les hachages d'utilisateur pour une compromission complète du domaine.",
    },
    {
        "id": "AP-009",
        "name_en": "Web API Exploitation - Mass Data Access",
        "name_fr": "Exploitation d'API Web - Accès massif aux données",
        "category": "web_application",
        "nodes": [
            {"id": "n1", "label_en": "API Enumeration", "label_fr": "Énumération d'API", "mitre": "T1592.004", "phase": "initial_access"},
            {"id": "n2", "label_en": "Authentication Bypass", "label_fr": "Contournement d'authentification", "mitre": "T1556", "phase": "credential_access"},
            {"id": "n3", "label_en": "API Token Extraction", "label_fr": "Extraction de jeton API", "mitre": "T1528", "phase": "credential_access"},
            {"id": "n4", "label_en": "IDOR Vulnerability", "label_fr": "Vulnérabilité IDOR", "mitre": "T1190", "phase": "credential_access"},
            {"id": "n5", "label_en": "Parameter Tampering", "label_fr": "Manipulation de paramètres", "mitre": "T1007", "phase": "lateral_movement"},
            {"id": "n6", "label_en": "Bulk Data Export", "label_fr": "Export en masse de données", "mitre": "T1005", "phase": "exfiltration"},
            {"id": "n7", "label_en": "Sell on Dark Web", "label_fr": "Vendre sur le dark web", "mitre": "T1041", "phase": "impact"},
        ],
        "edges": [("n1","n2"), ("n2","n3"), ("n3","n4"), ("n4","n5"), ("n5","n6"), ("n6","n7")],
        "description_en": "Attacker discovers API endpoints through reverse engineering. Bypasses authentication using JWT manipulation or session fixation. Extracts API tokens from client-side code. Exploits IDOR vulnerability to access other users' data. Modifies parameters to access admin endpoints. Performs bulk export of sensitive customer data. Sells leaked data on dark web forums.",
        "description_fr": "L'attaquant découvre les points de terminaison de l'API par ingénierie inverse. Contourne l'authentification en utilisant la manipulation JWT ou la fixation de session. Extrait les jetons API du code côté client. Exploite la vulnérabilité IDOR pour accéder aux données d'autres utilisateurs. Modifie les paramètres pour accéder aux points de terminaison admin. Effectue un export en masse de données clients sensibles. Vend les données divulguées sur des forums du dark web.",
    },
    {
        "id": "AP-010",
        "name_en": "Zero-Day Exploitation & Persistence",
        "name_fr": "Exploitation de zéro-jour et persistance",
        "category": "network",
        "nodes": [
            {"id": "n1", "label_en": "Target Reconnaissance", "label_fr": "Reconnaissance de la cible", "mitre": "T1592", "phase": "initial_access"},
            {"id": "n2", "label_en": "Zero-Day Discovery", "label_fr": "Découverte de zéro-jour", "mitre": "T1190", "phase": "initial_access"},
            {"id": "n3", "label_en": "Exploit Execution", "label_fr": "Exécution de l'exploit", "mitre": "T1203", "phase": "execution"},
            {"id": "n4", "label_en": "Reverse Shell", "label_fr": "Reverse Shell", "mitre": "T1059", "phase": "execution"},
            {"id": "n5", "label_en": "C2 Infrastructure", "label_fr": "Infrastructure C2", "mitre": "T1071", "phase": "command_control"},
            {"id": "n6", "label_en": "Rootkit Installation", "label_fr": "Installation de Rootkit", "mitre": "T1014", "phase": "persistence"},
            {"id": "n7", "label_en": "Long-term Data Theft", "label_fr": "Vol de données à long terme", "mitre": "T1041", "phase": "exfiltration"},
        ],
        "edges": [("n1","n2"), ("n2","n3"), ("n3","n4"), ("n4","n5"), ("n5","n6"), ("n6","n7")],
        "description_en": "Attacker performs detailed reconnaissance on target organization. Discovers and acquires zero-day vulnerability in popular software. Crafts reliable exploit code. Targets high-value system through spear-phishing or watering hole attack. Reverse shell establishes interactive access. C2 infrastructure set up for command execution. Rootkit installed for kernel-level persistence. Harvests sensitive data over months without detection.",
        "description_fr": "L'attaquant effectue une reconnaissance détaillée de l'organisation cible. Découvre et acquiert une vulnérabilité zéro-jour dans un logiciel populaire. Crée un code d'exploitation fiable. Cible un système de haute valeur via une attaque de phishing ciblée ou de watering hole. Reverse shell établit un accès interactif. Infrastructure C2 configurée pour l'exécution de commandes. Rootkit installé pour la persistance au niveau du noyau. Récolte les données sensibles pendant des mois sans détection.",
    },
]

def create_network_graph(attack_path: Dict[str, Any], language: str = "en") -> go.Figure:
    """Create an interactive Plotly network graph for attack path visualization."""

    nodes = attack_path["nodes"]
    edges = attack_path["edges"]

    # Create positions using a hierarchical layout
    positions = {}
    phases_order = [
        "initial_access", "execution", "credential_access", "lateral_movement",
        "privilege_escalation", "persistence", "command_control", "exfiltration", "impact"
    ]

    phase_positions = {phase: [] for phase in phases_order}
    for node in nodes:
        phase = node["phase"]
        phase_positions[phase].append(node["id"])

    # Position nodes based on phase
    x_pos = {}
    y_pos = {}
    for phase_idx, phase in enumerate(phases_order):
        node_ids = phase_positions[phase]
        if node_ids:
            x = phase_idx * 2
            for node_idx, node_id in enumerate(node_ids):
                y = (node_idx - len(node_ids) / 2) * 1.5
                x_pos[node_id] = x
                y_pos[node_id] = y

    # Create edge traces
    edge_x = []
    edge_y = []
    for edge in edges:
        x0, y0 = x_pos.get(edge[0], 0), y_pos.get(edge[0], 0)
        x1, y1 = x_pos.get(edge[1], 0), y_pos.get(edge[1], 0)
        edge_x.append(x0)
        edge_x.append(x1)
        edge_x.append(None)
        edge_y.append(y0)
        edge_y.append(y1)
        edge_y.append(None)

    edge_trace = go.Scatter(
        x=edge_x, y=edge_y,
        mode='lines',
        line=dict(width=2.5, color='rgba(100, 100, 100, 0.5)'),
        hoverinfo='none',
        showlegend=False,
    )

    # Create node traces grouped by phase
    node_traces = {}
    phase_labels = PHASE_LABELS_FR if language == "fr" else PHASE_LABELS_EN

    for phase in set(node["phase"] for node in nodes):
        phase_nodes = [n for n in nodes if n["phase"] == phase]
        node_x = [x_pos[n["id"]] for n in phase_nodes]
        node_y = [y_pos[n["id"]] for n in phase_nodes]

        label_key = "label_fr" if language == "fr" else "label_en"
        hover_texts = [
            f"<b>{n[label_key]}</b><br>"
            f"MITRE: {n['mitre']}<br>"
            f"Phase: {phase_labels[n['phase']]}"
            for n in phase_nodes
        ]

        node_trace = go.Scatter(
            x=node_x, y=node_y,
            mode='markers+text',
            name=phase_labels[phase],
            marker=dict(
                size=30,
                color=PHASE_COLORS[phase],
                line=dict(width=2, color='white'),
            ),
            text=[n[label_key] for n in phase_nodes],
            textposition="middle center",
            textfont=dict(size=9, color='white', family='Arial Black'),
            hovertext=hover_texts,
            hoverinfo='text',
            showlegend=True,
        )
        node_traces[phase] = node_trace

    # Create figure
    fig = go.Figure(data=[edge_trace] + list(node_traces.values()))

    fig.update_layout(
        title=dict(
            text=attack_path[f"name_{language}"],
            font=dict(size=24, color='#1a1a1a', family='Arial Black'),
        ),
        showlegend=True,
        hovermode='closest',
        margin=dict(b=20, l=5, r=5, t=60),
        xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
        yaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
        plot_bgcolor='rgba(240, 240, 240, 1)',
        paper_bgcolor='white',
        font=dict(size=12, family='Arial'),
        height=600,
        width=1200,
    )

    return fig

def get_attack_stats(attack_path: Dict[str, Any], language: str = "en") -> str:
    """Generate statistics about the attack path."""
    nodes = attack_path["nodes"]
    edges = attack_path["edges"]

    unique_phases = set(node["phase"] for node in nodes)
    mitre_techniques = [node["mitre"] for node in nodes]

    phase_labels = PHASE_LABELS_FR if language == "fr" else PHASE_LABELS_EN

    title_stats = "ATTACK STATISTICS" if language == "en" else "STATISTIQUES D'ATTAQUE"
    lbl_stages = "Attack Stages" if language == "en" else "Étapes d'attaque"
    lbl_connections = "Connection Steps" if language == "en" else "Étapes de connexion"
    lbl_mitre = "MITRE Techniques Used" if language == "en" else "Techniques MITRE utilisées"
    lbl_phases = "Attack Phases" if language == "en" else "Phases d'attaque"
    lbl_involved = "Phases Involved" if language == "en" else "Phases impliquées"

    stats = f"""
    **{title_stats}**

    - **{lbl_stages}:** {len(nodes)}
    - **{lbl_connections}:** {len(edges)}
    - **{lbl_mitre}:** {len(set(mitre_techniques))}
    - **{lbl_phases}:** {len(unique_phases)}

    **{lbl_involved}:**
    """

    for phase in sorted(unique_phases):
        count = sum(1 for n in nodes if n["phase"] == phase)
        stats += f"\n- {phase_labels[phase]}: {count}"

    return stats

def create_phase_breakdown(attack_path: Dict[str, Any], language: str = "en") -> go.Figure:
    """Create a bar chart showing attack phase breakdown."""
    nodes = attack_path["nodes"]
    phase_labels = PHASE_LABELS_FR if language == "fr" else PHASE_LABELS_EN

    phase_counts = {}
    for node in nodes:
        phase = node["phase"]
        phase_counts[phase] = phase_counts.get(phase, 0) + 1

    phases_ordered = [p for p in [
        "initial_access", "execution", "credential_access", "lateral_movement",
        "privilege_escalation", "persistence", "command_control", "exfiltration", "impact"
    ] if p in phase_counts]

    counts = [phase_counts[p] for p in phases_ordered]
    colors = [PHASE_COLORS[p] for p in phases_ordered]
    labels = [phase_labels[p] for p in phases_ordered]

    fig = go.Figure(data=[
        go.Bar(
            x=labels,
            y=counts,
            marker=dict(color=colors, line=dict(color='white', width=2)),
            text=counts,
            textposition='auto',
            hovertemplate='<b>%{x}</b><br>Techniques: %{y}<extra></extra>',
        )
    ])

    fig.update_layout(
        title="Attack Phase Breakdown" if language == "en" else "Répartition des phases d'attaque",
        xaxis_title=f"{'MITRE Phase' if language == 'en' else 'Phase MITRE'}",
        yaxis_title=f"{'Techniques' if language == 'en' else 'Techniques'}",
        plot_bgcolor='rgba(240, 240, 240, 1)',
        paper_bgcolor='white',
        font=dict(size=11, family='Arial'),
        height=400,
        showlegend=False,
    )

    return fig

def get_node_details(attack_path: Dict[str, Any], language: str = "en") -> str:
    """Generate detailed information about each attack node."""
    nodes = attack_path["nodes"]
    phase_labels = PHASE_LABELS_FR if language == "fr" else PHASE_LABELS_EN
    label_key = "label_fr" if language == "fr" else "label_en"

    breakdown_title = "ATTACK STAGES BREAKDOWN" if language == "en" else "DÉTAILS DES ÉTAPES D'ATTAQUE"
    details = f"**{breakdown_title}**\n\n"

    for idx, node in enumerate(nodes, 1):
        phase_label = phase_labels[node["phase"]]
        details += f"{idx}. **{node[label_key]}** (MITRE: {node['mitre']})\n"
        details += f"   Phase: {phase_label}\n\n"

    return details

def interface():
    """Build the Gradio interface."""

    with gr.Blocks(theme=gr.themes.Soft(), title="Attack Path Visualizer") as demo:
        # Header
        with gr.Row():
            gr.Markdown("""
            # 🗺️ Attack Path Visualizer

            **Cybersecurity Attack Chain Visualization Tool**

            Explore 10 realistic attack scenarios with interactive flowcharts, MITRE framework mapping, and tactical insights.
            Perfect for threat modeling, red team training, and security awareness.
            """)

        with gr.Row():
            with gr.Column(scale=2):
                language = gr.Radio(
                    choices=[("English", "en"), ("Français", "fr")],
                    value="en",
                    label="Language",
                    scale=1,
                )
            with gr.Column(scale=3):
                attack_dropdown = gr.Dropdown(
                    choices=[(f"{ap['id']} - {ap['name_en']}", ap['id']) for ap in ATTACK_PATHS],
                    value=ATTACK_PATHS[0]["id"],
                    label="Select Attack Path",
                    scale=1,
                )

        # Main visualization
        attack_graph = gr.Plot(label="Attack Flowchart", scale=2)

        # Description and breakdown
        with gr.Row():
            with gr.Column(scale=1):
                attack_description = gr.Markdown("", label="Attack Description")
            with gr.Column(scale=1):
                attack_stats = gr.Markdown("", label="Statistics")

        # Phase breakdown chart
        phase_chart = gr.Plot(label="Phase Breakdown")

        # Node details
        node_details = gr.Markdown("", label="Stage Details")

        # Update function
        def update_visualization(selected_path_id, lang):
            attack_path = next((ap for ap in ATTACK_PATHS if ap["id"] == selected_path_id), ATTACK_PATHS[0])

            graph = create_network_graph(attack_path, lang)
            desc_key = f"description_{lang}"
            description = f"**{'Overview' if lang == 'en' else 'Aperçu'}**\n\n{attack_path[desc_key]}"
            stats = get_attack_stats(attack_path, lang)
            chart = create_phase_breakdown(attack_path, lang)
            details = get_node_details(attack_path, lang)

            return graph, description, stats, chart, details

        # Bind changes
        attack_dropdown.change(
            fn=update_visualization,
            inputs=[attack_dropdown, language],
            outputs=[attack_graph, attack_description, attack_stats, phase_chart, node_details],
        )

        language.change(
            fn=update_visualization,
            inputs=[attack_dropdown, language],
            outputs=[attack_graph, attack_description, attack_stats, phase_chart, node_details],
        )

        # Load initial visualization
        demo.load(
            fn=update_visualization,
            inputs=[attack_dropdown, language],
            outputs=[attack_graph, attack_description, attack_stats, phase_chart, node_details],
        )

        # Footer
        with gr.Row():
            gr.Markdown("""
            ---

            **Resources:**
            - [MITRE ATT&CK Framework](https://attack.mitre.org/)
            - [NIST Cybersecurity Framework](https://www.nist.gov/cyberframework)
            - [OWASP Security Top 10](https://owasp.org/www-project-top-ten/)

            **Educational Tool** - For authorized security testing and awareness training only.

            Built with Gradio | Made for LinkedIn | Open Source
            """)

    return demo

if __name__ == "__main__":
    demo = interface()
    demo.launch()