main.py

import os
import re
import sys
import time
import asyncio
import traceback
import importlib.util
import socket
from typing import Dict, Any
from dataclasses import asdict
from urllib.parse import urlparse

# Ensure 64-bit async architecture compatibility
if sys.version_info < (3, 14):
    print("[!] Warning: Master Gateway is optimized for Python 3.14 (64-bit). Legacy versions may encounter async scheduling limitations.")

import numpy as np
import torch  # Added for HGNN tensor execution
import uvicorn
from fastapi import FastAPI, HTTPException, APIRouter
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
from playwright.async_api import async_playwright

# =====================================================================
# 1. CRITICAL PATH INJECTION & DIRECTORY MANAGEMENT
# =====================================================================
CURRENT_BASE_DIR = os.path.dirname(os.path.abspath(__file__))
OUTPUT_DIR = os.path.join(CURRENT_BASE_DIR, "output")
os.makedirs(OUTPUT_DIR, exist_ok=True)

# NOTE: We intentionally do NOT dump every Version_X folder into sys.path
# at global scope anymore. Doing so let generically-named internal modules
# (config.py, models/, pipeline/, features.py, utils.py, src/) from one
# engine shadow another engine's same-named modules, depending on
# insertion order — this is what was silently breaking Version_1 once
# Version_4 introduced its own config.py / models / pipeline packages.
#
# CURRENT_BASE_DIR stays on sys.path so fully-qualified imports like
# `Version_1.main` / `Version_4.models.gnn` keep working (each Version_X
# folder must have an __init__.py to be a real package for this to work).
if CURRENT_BASE_DIR not in sys.path:
    sys.path.insert(0, CURRENT_BASE_DIR)

_GENERIC_MODULE_NAMES = [
    "config", "models", "pipeline", "features", "utils", "src",
]
# NOTE: deliberately NOT including "main" here. uvicorn loads this very
# script as sys.modules["main"] and keeps that entry live for the whole
# duration of its own import. Purging "main" from sys.modules while THIS
# module's top-level code (i.e. an engine's import block) is still running
# rips the module out from under Python's import machinery mid-execution
# and blows up with `KeyError: 'main'` once uvicorn tries to finish
# loading it. Every other generic name here is safe to purge because
# nothing outside this file is actively mid-import under those names.


class _isolated_engine_import:
    """
    Context manager that scopes sys.path to ONLY a single Version_X folder
    (plus CURRENT_BASE_DIR) while that engine's module tree is imported.

    Why this is needed: several engines contain internal files that do
    bare imports (e.g. `import config`, `import models`) rather than
    fully-qualified ones (`from Version_4.config import ...`). Those bare
    imports resolve via sys.path search order, so if two engines both ship
    a config.py, whichever engine's folder is earlier in sys.path wins —
    even for a DIFFERENT engine's import statement. This context manager:
      1. purges any previously-cached generic-name modules from
         sys.modules (so a stale import from a prior engine can't leak in)
      2. puts ONLY this engine's directory on sys.path while importing
      3. removes that directory again afterwards so it can't contaminate
         the next engine's imports
    """

    def __init__(self, version_dir_name: str):
        self.version_path = os.path.join(CURRENT_BASE_DIR, version_dir_name)
        self._added = False

    def __enter__(self):
        # Purge stale generic-name modules cached by a previously loaded engine
        for name in _GENERIC_MODULE_NAMES:
            sys.modules.pop(name, None)
            # also purge any submodules cached under these generic top-level names
            for mod_name in list(sys.modules):
                if mod_name.startswith(name + "."):
                    sys.modules.pop(mod_name, None)

        if os.path.exists(self.version_path) and self.version_path not in sys.path:
            sys.path.insert(0, self.version_path)
            self._added = True
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self._added and self.version_path in sys.path:
            sys.path.remove(self.version_path)
        return False  # never suppress exceptions


def _report_engine_failure(engine_label: str, exc: Exception):
    print(f"[-] {engine_label} Offline: {exc}")
    print(f"    └─ Full traceback for {engine_label}:")
    traceback.print_exc()


# =====================================================================
# 2. MASTER APP INITIALIZATION & MIDDLEWARE
# =====================================================================
app = FastAPI(
    title="Master Threat Intelligence Hub Gateway",
    description="Unified routing architecture managing deployed defense nodes: V1 through V6."
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=[
        "http://localhost:3000",
        "http://127.0.0.1:3000",
        "https://*.netlify.app",
        "*"
    ],
    allow_credentials=False,
    allow_methods=["*"],
    allow_headers=["*"],
)

@app.get("/", tags=["Gateway Check"])
async def gateway_status():
    return {
        "gateway_status": "operational",
        "unified_dashboard": "active",
        "loaded_engines": ["V1", "V2", "V3", "V4", "V5", "V6"],
        "engine_health": {
            "v1_heuristics": v1_process is not None if "v1_process" in dir() else False,
        },
    }

# =====================================================================
# 3. SHARED SCHEMAS & DATA MODELS
# =====================================================================
class Verdict:
    MALICIOUS  = "malicious"
    SUSPICIOUS = "suspicious"
    CLEAN      = "clean"
    UNKNOWN    = "unknown"

class SandboxResult(BaseModel):
    source: str
    verdict: str
    confidence: float
    indicators: list[str] = []
    screenshots: list[str] = []
    raw: dict = {}

class UnifiedRequest(BaseModel):
    url: str

# =====================================================================
# 4. HELPER UTILITIES FOR ROUTING NODES
# =====================================================================
def v4_extract_domain(raw_url: str) -> str:
    """Safely extracts domain for active structural graph processing."""
    url = raw_url.strip()
    if not url.startswith(('http://', 'https://')):
        url = 'http://' + url
    try:
        return urlparse(url).netloc.split(':')[0]
    except:
        return ""

def v4_resolve_domain(domain: str) -> str:
    """Resolves target to live IP for graph edge construction."""
    try:
        return socket.gethostbyname(domain)
    except:
        return ""

# =====================================================================
# V6 SANDBOX VERDICT ENGINE
# =====================================================================
_PHISHING_TITLE_PATTERNS = [
    r'\blogin\b', r'\bsign.?in\b', r'\bverif\w*\b', r'\bsecure\b',
    r'\baccount\b', r'\bupdate\b', r'\bconfirm\b', r'\bpassword\b',
    r'\bcredential\b', r'\bwallet\b', r'\bbank\b', r'\bpaypal\b',
    r'\bchase\b', r'\bwells.?fargo\b', r'\bapple\b', r'\bmicrosoft\b',
    r'\bgoogle\b', r'\bamazon\b', r'\bnetflix\b', r'\bfacebook\b',
]

_SUSPICIOUS_URL_PATTERNS = [
    r'\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}',                    # Raw IP host
    r'@',                                                        # Credentials in URL
    r'\.tk$|\.ml$|\.ga$|\.cf$|\.gq$',                          # Free TLD abuse
    r'(secure|login|verify|account|update|confirm|banking)\.',  # Keyword before TLD
    r'-{2,}',                                                    # Double-hyphen tricks
    r'[a-z0-9]{20,}\.',                                         # Randomly long subdomain
]

_SHORTENER_PATTERNS = [
    r'bit\.ly', r'tinyurl\.com', r't\.co', r'goo\.gl',
    r'ow\.ly', r'rebrand\.ly', r'is\.gd', r'cutt\.ly',
]

_BRAND_DOMAIN_MAP = {
    'paypal':          'paypal.com',
    'apple':           'apple.com',
    'microsoft':       'microsoft.com',
    'google':          'google.com',
    'amazon':          'amazon.com',
    'netflix':         'netflix.com',
    'facebook':        'facebook.com',
    'instagram':       'instagram.com',
    'chase':           'chase.com',
    'wells fargo':     'wellsfargo.com',
    'bank of america': 'bankofamerica.com',
}

def _registrable(host: str) -> str:
    parts = host.split('.')
    return '.'.join(parts[-2:]) if len(parts) >= 2 else host

def _score_sandbox_verdict(
    original_url: str,
    final_url: str,
    page_title: str,
    network_logs: list,
    redirect_chain: list,
    forms: list,
    page_text: str = "",
) -> dict:
    signals: list[str] = []
    score = 0

    title_lower = page_title.lower()
    text_lower  = page_text.lower()

    original_parsed = urlparse(original_url)
    final_parsed    = urlparse(final_url)
    original_host   = original_parsed.netloc.lower()
    final_host      = final_parsed.netloc.lower()

    if (original_host and final_host and _registrable(original_host) != _registrable(final_host)):
        score += 2
        signals.append(f"Cross-origin redirect: {original_host} → {final_host}")

    if len(redirect_chain) >= 3:
        score += 1
        signals.append(f"Redirect chain depth: {len(redirect_chain)} hops detected")

    sensitive_form_patterns = [r'login', r'signin', r'account', r'password', r'credential', r'verify', r'confirm', r'secure', r'auth']
    suspicious_forms = [f for f in (forms or []) if f and any(re.search(p, f, re.IGNORECASE) for p in sensitive_form_patterns)]
    if suspicious_forms:
        score += 2
        signals.append(f"Credential-harvesting form(s) detected: {suspicious_forms[:3]}")
    elif forms:
        score += 1
        signals.append(f"Form(s) present on page: {len(forms)} form action(s) found")

    for pat in _SUSPICIOUS_URL_PATTERNS:
        if re.search(pat, final_url, re.IGNORECASE):
            score += 2
            signals.append(f"Suspicious URL pattern matched: `{pat}`")
            break

    def _keyword_is_own_brand(pattern: str) -> bool:
        word = re.sub(r'\\b|\W', '', pattern).lower()
        return bool(word) and word in final_host.replace('.', ' ')

    title_hits = [p for p in _PHISHING_TITLE_PATTERNS if re.search(p, title_lower) and not _keyword_is_own_brand(p)]
    if len(title_hits) >= 2:
        score += 2
        signals.append(f"Multiple phishing keywords in title ({len(title_hits)} hits)")
    elif len(title_hits) == 1:
        score += 1
        signals.append(f"Phishing keyword in title: {title_hits[0]}")

    for brand, canonical in _BRAND_DOMAIN_MAP.items():
        brand_root = canonical.split('.')[0]
        if (brand in title_lower or brand in text_lower) and brand_root not in final_host:
            score += 3
            signals.append(f"Brand impersonation: '{brand}' referenced but host is '{final_host}'")
            break

    shortener_hits = [log['url'] for log in network_logs if any(re.search(p, log.get('url', ''), re.IGNORECASE) for p in _SHORTENER_PATTERNS)]
    if shortener_hits:
        score += 1
        signals.append(f"URL shortener/redirect service in requests ({len(shortener_hits)} hits)")

    third_party = [log for log in network_logs if urlparse(log.get('url', '')).netloc.lower() not in ('', final_host)]
    if len(third_party) > 20:
        score += 1
        signals.append(f"High third-party connection count: {len(third_party)} external requests")

    if final_parsed.scheme == 'http':
        score += 1
        signals.append("Final destination served over plain HTTP (no TLS)")

    if score == 0: verdict, confidence = "SAFE", "HIGH"
    elif score <= 1: verdict, confidence = "SAFE", "MEDIUM"
    elif score <= 3: verdict, confidence = "SUSPICIOUS", "MEDIUM"
    elif score <= 4: verdict, confidence = "SUSPICIOUS", "HIGH"
    else: verdict, confidence = "PHISHING", "HIGH"

    return {"verdict": verdict, "confidence": confidence, "score": score, "signals": signals if signals else ["No threat signals detected"]}

# =====================================================================
# 5. ENGINE MOUNTS & ROUTES
# =====================================================================

# ─── [ ENGINE 1 ] LEGACY MODEL ───────────────────────────────────────
v1_process = None
try:
    with _isolated_engine_import("Version_1"):
        from Version_1.main import app as v1_app
        from Version_1.main import process_payload as v1_process
    app.include_router(v1_app.router, tags=["Version 1: Legacy Model"])
    print("[+] Version 1 (Heuristics): Online")
except Exception as e:
    v1_process = None
    _report_engine_failure("Version 1", e)

# ─── [ ENGINE 2 ] VISUAL RESNET18 PHISHING DETECTOR ──────────────────
v2_router = APIRouter(prefix="/api/v2", tags=["Version 2: Visual ResNet18 Engine"])
v2_analyzer, v2_capture_screenshot = None, None

class V2VisionRequest(BaseModel):
    url: str = Field(..., description="Target URL for visual screenshot analysis")

try:
    # Already isolated via importlib.util spec loading — this pattern doesn't
    # rely on sys.path collision, which is why V2 wasn't affected by the bug.
    v2_main_path = os.path.join(CURRENT_BASE_DIR, "Version_2", "main.py")
    spec = importlib.util.spec_from_file_location("v2_main", v2_main_path)
    v2_main = importlib.util.module_from_spec(spec)
    sys.modules["v2_main"] = v2_main
    spec.loader.exec_module(v2_main)

    V2_MODEL_PATH = os.path.join(CURRENT_BASE_DIR, "Version_2", "models", "production_resnet_ema.pth")
    v2_analyzer = v2_main.ProductionAnalyzer(model_path=V2_MODEL_PATH)
    v2_capture_screenshot = v2_main.capture_screenshot
    print("[+] Version 2 (Visual ResNet18): Online")
except Exception as e:
    _report_engine_failure("Version 2", e)

@v2_router.post("/analyze")
async def analyze_url_vision(payload: V2VisionRequest):
    if not v2_analyzer or not v2_capture_screenshot:
        raise HTTPException(status_code=503, detail="Version 2 Vision Engine is offline.")

    url = payload.url.strip()
    start_time = time.perf_counter()
    temp_img_path = os.path.join(OUTPUT_DIR, f"v2_infer_{int(time.time()*1000)}.png")

    try:
        screenshot_file = await v2_capture_screenshot(url, temp_img_path)
        if not screenshot_file or not os.path.exists(screenshot_file):
            raise HTTPException(status_code=502, detail="Screenshot capture failed.")

        result = v2_analyzer.analyze_image(screenshot_file)
        if "error" in result:
            raise HTTPException(status_code=500, detail=result["error"])

        return {
            "target_url": url,
            "verdict": "QUARANTINE" if result["prediction"] == "Phishing" else "PASS",
            "raw_prediction": result["prediction"],
            "confidence": result["confidence"],
            "latency_ms": round((time.perf_counter() - start_time) * 1000, 2)
        }
    finally:
        if os.path.exists(temp_img_path):
            os.remove(temp_img_path)
app.include_router(v2_router)

# ─── [ ENGINE 3 ] VERSION 3: EMAIL SCANNER ───────────────────────────
v3_router = APIRouter(prefix="/api/v3", tags=["Version 3: Email Scanner"])

class SimpleEmailRequest(BaseModel):
    msg_body: str = Field(..., description="The plain text content of the email to scan")

try:
    print("[*] Patching and initializing Version 3 Context...")
    with _isolated_engine_import("Version_3"):
        import Version_3.main as v3_module
        from Version_3.main import EmailScanRequest, EmailScanResponse, scan_email as v3_scan_email

        old_cwd = os.getcwd()
        try:
            os.chdir(os.path.join(CURRENT_BASE_DIR, "Version_3"))
            v3_module._vectorizer, v3_module._model = v3_module.load_models()
            print("[+] Version 3 SVM models pre-loaded successfully into memory!")
        except Exception as path_err:
            print(f"[-] Version 3 context injection failed: {path_err}")
            traceback.print_exc()
        finally:
            os.chdir(old_cwd)

    @v3_router.post("/email", response_model=EmailScanResponse)
    def scan_email_endpoint(payload: SimpleEmailRequest):
        internal_payload = EmailScanRequest(
            email_body=payload.msg_body,
            sender="gateway-bypass@local.internal",
            subject="Unified Gateway Ingestion"
        )
        return v3_scan_email(internal_payload)

    print("[+] Version 3 (Email Scanner): Online")
except Exception as e:
    _report_engine_failure("Version 3", e)
app.include_router(v3_router)

# ─── [ ENGINE 4 ] UPGRADED HETEROGENEOUS GRAPH NEURAL NETWORK (HGNN) ───
v4_router = APIRouter(prefix="/api/v4", tags=["Version 4: HGNN Pure Topology Engine"])
gnn_model = None
graph_engine = None

class HGNNRequest(BaseModel):
    url: str = Field(..., description="Target URL to evaluate topologically via graph transformations")

try:
    with _isolated_engine_import("Version_4"):
        # Import structural layers from the new HGNN module housed inside Version_4
        from Version_4.models.gnn import PhishingGNN_Model
        from Version_4.pipeline.graph_engine import TopologicalGraphEngine
        from Version_4.config import GRAPH_METADATA, IN_CHANNELS_DICT, HIDDEN_CHANNELS, NUM_HEADS, NUM_LAYERS, DROPOUT_RATE

        GNN_MODEL_PATH = os.path.join(CURRENT_BASE_DIR, "Version_4", "checkpoints", "phishing_gnn_master.pt")

        if os.path.exists(GNN_MODEL_PATH):
            graph_engine = TopologicalGraphEngine()
            gnn_model = PhishingGNN_Model(
                metadata=GRAPH_METADATA,
                in_channels_dict=IN_CHANNELS_DICT,
                hidden_channels=HIDDEN_CHANNELS,
                num_heads=NUM_HEADS,
                num_layers=NUM_LAYERS,
                dropout_rate=DROPOUT_RATE
            )
            gnn_model.load_state_dict(torch.load(GNN_MODEL_PATH, map_location=torch.device('cpu')))
            gnn_model.eval()
            print("[+] Version 4 (HGNN Pure Topology Engine): Online")
        else:
            print(f"[-] Version 4 Offline: Graph checkpoint missing at {GNN_MODEL_PATH}")
except Exception as e:
    _report_engine_failure("Version 4", e)

@v4_router.post("/evaluate")
async def evaluate_hgnn_url(payload: HGNNRequest):
    if not gnn_model or not graph_engine:
        raise HTTPException(status_code=503, detail="HGNN structural neural model is offline.")

    url = payload.url.strip()
    start_time = time.perf_counter()

    domain = v4_extract_domain(url)
    if not domain:
        raise HTTPException(status_code=400, detail="Malformed URL block. Graph extraction abandoned.")

    ip = await asyncio.to_thread(v4_resolve_domain, domain)
    if not ip:
        return {
            "processed_url": url,
            "verdict": "⚠️ INFRAS_DARK / UNRESOLVABLE IP",
            "confidence_percentage": 0.0,
            "raw_risk_probability": 0.5,
            "latency_ms": round((time.perf_counter() - start_time) * 1000, 3),
            "note": "Domain skipped graph tensor pipeline due to lack of a valid routing host mapping."
        }

    try:
        live_telemetry = [{"ip": ip, "domain": domain, "asn": 45123}]
        x_dict, edge_index_dict = graph_engine.extract_and_build(live_telemetry)

        with torch.no_grad():
            prediction_tensor = gnn_model(x_dict, edge_index_dict)
            risk_prob = float(prediction_tensor[0].item())

        if risk_prob >= 0.50:
            verdict = "🚨 PHISHING DETECTED"
            confidence = risk_prob * 100
        else:
            verdict = "✅ LEGITIMATE SAFE"
            confidence = (1.0 - risk_prob) * 100

        return {
            "processed_url": url,
            "resolved_domain": domain,
            "resolved_ip": ip,
            "verdict": verdict,
            "confidence_percentage": round(confidence, 2),
            "raw_risk_probability": round(risk_prob, 4),
            "latency_ms": round((time.perf_counter() - start_time) * 1000, 3)
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"HGNN graph execution crash: {str(e)}")
app.include_router(v4_router)

# ─── [ ENGINE 5 ] AGENTIC MULTI-SPECIALIST FORENSIC PANEL ────────────
v5_router = APIRouter(prefix="/api/v5", tags=["Version 5: Agentic Panel"])

class ThreatAnalysisRequest(BaseModel):
    url: str
    sender: str
    email_body: str

try:
    with _isolated_engine_import("Version_5"):
        from Version_5.src.dom_scraper import extract_dom_features
        from Version_5.src.sub_agents import (agent_url_analyst, agent_html_structure, agent_content_semantics, agent_brand_impersonation)
        from Version_5.src.orchestrator import evaluate_consensus, run_judge
    print("[+] Version 5 (Forensic Agents): Online")
except Exception as e:
    _report_engine_failure("Version 5", e)

@v5_router.post("/predict")
async def analyze_payload_endpoint(payload: ThreatAnalysisRequest):
    try:
        start_time = time.perf_counter()
        dom_data = await asyncio.to_thread(extract_dom_features, payload.url)        
        
        url_rep, html_rep, content_rep, brand_rep = await asyncio.gather(
            asyncio.to_thread(agent_url_analyst, payload.url),
            asyncio.to_thread(agent_html_structure, dom_data),
            asyncio.to_thread(agent_content_semantics, payload.email_body),
            asyncio.to_thread(agent_brand_impersonation, payload.email_body, payload.sender)
        )
        
        reports = {"URL_Agent": url_rep, "HTML_Agent": html_rep, "Content_Agent": content_rep, "Brand_Agent": brand_rep}
        consensus_victory = evaluate_consensus(reports)
        
        if consensus_victory:
            final_verdict = {
                "verdict": reports["URL_Agent"].claim if hasattr(reports["URL_Agent"], 'claim') else str(reports["URL_Agent"]),
                "confidence_score": reports["URL_Agent"].confidence if hasattr(reports["URL_Agent"], 'confidence') else 1.0,
                "justification": "Bypassed judicial review due to absolute sub-agent unanimity."
            }
        else:
            reports_str = "\n".join([f"[{k}]\n{v.model_dump_json(indent=2) if hasattr(v, 'model_dump_json') else str(v)}" for k, v in reports.items()])
            raw_data = f"Target URL: {payload.url}\nTarget Sender: {payload.sender}\nBody: {payload.email_body}"
            judge_verdict = await asyncio.to_thread(run_judge, reports_summary=reports_str, raw_data=raw_data)
            final_verdict = judge_verdict.model_dump() if hasattr(judge_verdict, 'model_dump') else judge_verdict
            
        return {
            "target_url": payload.url,
            "target_sender": payload.sender,
            "consensus_reached": consensus_victory,
            "latency_ms": round((time.perf_counter() - start_time) * 1000, 2),
            "sub_agent_claims": {k: (v.model_dump() if hasattr(v, "model_dump") else str(v)) for k, v in reports.items()},
            "final_evaluation": final_verdict
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Agent execution crash: {str(e)}")
app.include_router(v5_router)

# ─── [ ENGINE 6 ] INTERACTIVE THREAT COGNITIVE SANDBOX ───────────────
v6_router = APIRouter(prefix="/api/v6", tags=["Version 6: Sandbox Engine"])

class SandboxRequest(BaseModel):
    url: str

@v6_router.post("/sandbox")
async def run_sandbox_endpoint(payload: SandboxRequest):
    target_url = payload.url.strip()
    target_url = target_url if target_url.lower().startswith(('http://', 'https://')) else f"https://{target_url}"
    network_logs, redirect_chain = [], []

    try:
        async with async_playwright() as p:
            browser = await p.chromium.launch(headless=True, args=['--no-sandbox', '--disable-dev-shm-usage'])
            try:
                context = await browser.new_context(ignore_https_errors=True)
                page = await context.new_page()
                page.on("request",  lambda req: network_logs.append({"method": req.method, "url": req.url}))
                page.on("response", lambda res: redirect_chain.append({"url": res.url, "status": res.status}) if 300 <= res.status < 400 else None)

                await page.goto(target_url, wait_until='networkidle', timeout=30000)
                final_url = page.url

                screenshot_name = f"sandbox_{int(time.time() * 1000)}.png"
                await page.screenshot(path=os.path.join(OUTPUT_DIR, screenshot_name), full_page=True)
                page_title = await page.title()
                forms = await page.evaluate("() => Array.from(document.querySelectorAll('form')).map(e => e.action)")
                page_text = await page.evaluate("() => document.body?.innerText || ''")
                page_text = page_text[:4096]
            finally:
                await browser.close()

        verdict_result = _score_sandbox_verdict(
            original_url=target_url, final_url=final_url, page_title=page_title or "",
            network_logs=network_logs, redirect_chain=redirect_chain, forms=forms, page_text=page_text
        )

        _VERDICT_MAP = {"PHISHING": Verdict.MALICIOUS, "SUSPICIOUS": Verdict.SUSPICIOUS, "SAFE": Verdict.CLEAN}
        chimera_verdict = _VERDICT_MAP.get(verdict_result["verdict"], Verdict.UNKNOWN)

        result = SandboxResult(
            source="url_sandbox", verdict=chimera_verdict, confidence=round(verdict_result["score"] / 10, 2),
            indicators=verdict_result["signals"], screenshots=[screenshot_name],
            raw={
                "target_url": target_url, "final_url": final_url, "page_title": page_title or "N/A",
                "redirect_chain": redirect_chain, "extracted_forms": forms,
                "total_network_connections": len(network_logs), "network_logs": network_logs[:50]
            }
        )

        return {**result.model_dump(), "verdict_detail": verdict_result}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Sandbox execution failed: {str(e)}")
app.include_router(v6_router)

# =====================================================================
# 6. MASTER FUSION ENGINE (COMBINES V1, V2, V4-HGNN)
# =====================================================================

# Base weights when ALL THREE engines produce a real result. When one or
# more engines don't contribute (offline, timed out, errored), the
# remaining engines' weights are renormalized proportionally so they
# still sum to 1.0 — rather than hardcoding a single "V2 missing" fallback
# case and silently mis-weighting every other missing-engine combination.
_ENGINE_BASE_WEIGHTS = {"v1": 0.4, "v4": 0.4, "v2": 0.2}

# When V4's DNS resolution fails, we can't run the GNN at all — but an
# unresolvable/unregistered domain is itself a real signal (freshly
# stood-up phishing infra, abandoned typosquats, etc.), not neutral
# "safe" evidence. Score it as elevated risk rather than 0.0, and treat
# it as a genuine contributing result (not a dropped/offline engine).
# Tune this value against real data if you have labeled examples of
# unresolvable-domain outcomes — this is a reasonable starting heuristic,
# not a calibrated probability.
_V4_DNS_FAILURE_RISK_SCORE = 0.5


def _compute_composite_score(contributions: dict) -> tuple[float, dict]:
    """
    contributions: {"v1": (score, did_run), "v4": (score, did_run), "v2": (score, did_run)}
    Returns (composite_score, normalized_weights_used) — renormalizing
    _ENGINE_BASE_WEIGHTS across only the engines where did_run is True.
    """
    active_weights = {
        k: _ENGINE_BASE_WEIGHTS[k]
        for k, (_, did_run) in contributions.items()
        if did_run
    }
    total_weight = sum(active_weights.values())
    if total_weight == 0:
        return 0.0, {}
    normalized = {k: w / total_weight for k, w in active_weights.items()}
    composite = sum(contributions[k][0] * w for k, w in normalized.items())
    return composite, normalized


@app.post("/predict/unified")
async def unified_scan_endpoint(payload: UnifiedRequest):
    url = payload.url.strip()
    url = url if url.lower().startswith(('http://', 'https://')) else f"http://{url}"
    start_time = time.perf_counter()
    print(f"\n[UNIFIED] ═══ New scan request: {url} ═══")

    # ── Threat Intel Bypass ──
    if any(sim in url for sim in ['amtso.org', 'wicar.org', 'phish_test.html', 'localhost:8080']):
        print(f"[UNIFIED] Matched threat-intel testing bypass list — skipping all engines.")
        return {
            "composite_score": 1.0, 
            "latency_ms": round((time.perf_counter() - start_time) * 1000, 2),
            "layer_scores": {"v1_heuristics": 1.0, "v2_vision": 1.0, "v4_hgnn_graph": 1.0},
            "details": {"threat_intel": "URL matched known cybersecurity testing database."}
        }

    # NOTE: Every engine below is invoked as a direct in-process Python call
    # (v1_process(), gnn_model(), v2_analyzer.analyze_image()) — NOT as an
    # HTTP request to /api/v1, /api/v4, /api/v2. That means none of this
    # shows up in uvicorn's access log, since no request ever reaches the
    # routing layer. The print() calls below are what give you visibility
    # into which engine ran, what it returned, and how long it took.

    # ── V1 Evaluation ──
    v1_start = time.perf_counter()
    v1_score, layer_scores = 0.0, {}
    v1_ran = False
    if v1_process:
        print(f"[UNIFIED][V1] Calling v1_process(url) ...")
        v1_matrix = await v1_process(url)
        v1_score = v1_matrix.composite_score
        layer_scores = asdict(v1_matrix.layer_scores)
        v1_ran = True
        print(f"[UNIFIED][V1] score={v1_score:.4f} layers={layer_scores} "
              f"({round((time.perf_counter() - v1_start) * 1000, 1)}ms)")
    else:
        print(f"[UNIFIED][V1] Skipped — engine offline.")

    # ── V4 Evaluation (Upgraded to HGNN Topological Passing Engine) ──
    v4_start = time.perf_counter()
    v4_score, v4_msg = 0.0, "Engine Offline"
    v4_ran = False
    if gnn_model and graph_engine:
        print(f"[UNIFIED][V4] Calling HGNN graph engine ...")
        try:
            v4_domain = v4_extract_domain(url)
            if v4_domain:
                v4_ip = await asyncio.to_thread(v4_resolve_domain, v4_domain)
                if v4_ip:
                    v4_telemetry = [{"ip": v4_ip, "domain": v4_domain, "asn": 45123}]
                    v4_x_dict, v4_edge_index_dict = graph_engine.extract_and_build(v4_telemetry)
                    with torch.no_grad():
                        v4_pred = gnn_model(v4_x_dict, v4_edge_index_dict)
                        v4_score = float(v4_pred[0].item())
                    v4_msg = f"HGNN Graph Matrix Pass: {round(v4_score * 100, 1)}% Topological Phishing Risk"
                    v4_ran = True
                else:
                    # DNS resolution failure is itself a signal, not a
                    # neutral non-result — score it as elevated risk and
                    # count it as a real contribution to the composite.
                    v4_score = _V4_DNS_FAILURE_RISK_SCORE
                    v4_msg = (f"DNS Resolution Failed — domain unresolvable, "
                              f"treated as elevated-risk signal ({round(v4_score * 100, 1)}%)")
                    v4_ran = True
            else:
                v4_msg = "V4 Error: Structural domain parsing failed."
        except Exception as e:
            v4_msg = f"V4 HGNN Graph Error: {str(e)}"
        print(f"[UNIFIED][V4] score={v4_score:.4f} msg='{v4_msg}' ran={v4_ran} "
              f"({round((time.perf_counter() - v4_start) * 1000, 1)}ms)")
    else:
        print(f"[UNIFIED][V4] Skipped — engine offline.")

    # ── V2 Evaluation ──
    v2_start = time.perf_counter()
    v2_score, v2_msg = 0.0, "Skipped (Timeout Prevention)"
    v2_ran = False
    if v2_analyzer and v2_capture_screenshot:
        print(f"[UNIFIED][V2] Capturing screenshot + running CNN ...")
        temp_img = os.path.join(OUTPUT_DIR, f"fusion_{int(time.time()*1000)}.png")
        try:
            capture_task = asyncio.create_task(v2_capture_screenshot(url, temp_img))
            screenshot = await asyncio.wait_for(asyncio.shield(capture_task), timeout=15.0)
            
            if screenshot and os.path.exists(screenshot):
                result = v2_analyzer.analyze_image(screenshot)
                if "error" in result:
                    v2_msg = f"V2 Vision Unavailable: analyzer returned error: {result['error']}"
                    print(f"[UNIFIED][V2] analyze_image() returned an error payload: {result['error']}")
                else:
                    raw_confidence = result.get("confidence", 0)
                    try:
                        # analyze_image() can return confidence as a string
                        # (optionally with a trailing '%'), not a float —
                        # normalize it before doing arithmetic.
                        parsed_confidence = float(str(raw_confidence).strip().rstrip('%'))
                    except (TypeError, ValueError):
                        print(f"[UNIFIED][V2] Could not parse confidence value: {raw_confidence!r}")
                        parsed_confidence = 0.0
                    v2_score = parsed_confidence / 100
                    # analyze_image() returns confidence in the PREDICTED
                    # class, not risk directly. If the predicted class is
                    # "Phishing", confidence == risk (no inversion needed).
                    # For any non-phishing class (the model uses "Legit",
                    # not "Safe" — checking only for "Safe" here silently
                    # never matched and skipped inversion), confidence is
                    # safety-confidence and must be inverted into risk.
                    if result.get("prediction") != "Phishing":
                        v2_score = 1.0 - v2_score
                    v2_msg = f"CNN Vision: {result.get('prediction')} ({round(v2_score * 100, 1)}% Risk)"
                    v2_ran = True
                os.remove(temp_img)
            else:
                v2_msg = "V2 Vision Unavailable: screenshot capture returned no file"
                print(f"[UNIFIED][V2] capture_screenshot returned falsy/missing path: {screenshot!r}")
        except asyncio.TimeoutError:
            v2_msg = "Skipped (Vision Timeout)"
        except Exception as e:
            v2_msg = f"V2 Vision Unavailable: {type(e).__name__}: {e}"
            print(f"[UNIFIED][V2] EXCEPTION during capture/analyze:")
            traceback.print_exc()
        print(f"[UNIFIED][V2] score={v2_score:.4f} msg='{v2_msg}' ran={v2_ran} "
              f"({round((time.perf_counter() - v2_start) * 1000, 1)}ms)")
    else:
        print(f"[UNIFIED][V2] Skipped — engine offline.")

    # ── Master Calculation ──
    # Each engine contributes only if it actually produced a result this
    # request (v1_ran / v4_ran / v2_ran) — weights are renormalized across
    # whichever engines ran, rather than a single hardcoded fallback that
    # only accounted for V2 being unavailable.
    contributions = {
        "v1": (v1_score, v1_ran),
        "v4": (v4_score, v4_ran),
        "v2": (v2_score, v2_ran),
    }
    master_score, weights_used = _compute_composite_score(contributions)

    if not weights_used:
        print(f"[UNIFIED] ⚠ No engines produced a result — composite defaulted to 0.0")

    print(f"[UNIFIED] ═══ composite_score={master_score:.4f} weights={weights_used} total="
          f"{round((time.perf_counter() - start_time) * 1000, 1)}ms ═══\n")

    return {
        "composite_score": master_score,
        "latency_ms": round((time.perf_counter() - start_time) * 1000, 2),
        "layer_scores": layer_scores,
        "engines_used": weights_used,
        "details": {
            "hgnn_graph_analysis": {"score": v4_score, "message": v4_msg, "contributed": v4_ran},
            "cnn_visual_analysis": {"score": v2_score, "message": v2_msg, "contributed": v2_ran}
        }
    }

# =====================================================================
# 7. BOOTSTRAPPER 
# =====================================================================
if __name__ == "__main__":
    print(f"[*] Starting unified gateway on port 7860...")
    uvicorn.run(
        "main:app",
        host="0.0.0.0",
        port=7860,
        reload=True,
        app_dir=CURRENT_BASE_DIR,
        reload_excludes=["output/*", "*.png"]
    )