"""
SAFE-MCP Penetration Testing Engine

Simulates attack patterns from the SAFE-MCP threat taxonomy against Navil's
anomaly detectors and policy engine, verifying that each detector correctly
identifies the corresponding attack scenario.

Usage:
    from navil.pentest import PentestEngine
    engine = PentestEngine(anomaly_detector, policy_engine)
    report = engine.run_all()
"""

from __future__ import annotations

import logging
from dataclasses import asdict, dataclass
from datetime import datetime, timedelta, timezone
from typing import Any

from navil.anomaly_detector import BehavioralAnomalyDetector

logger = logging.getLogger(__name__)


@dataclass
class PentestResult:
    """Result of a single pentest scenario."""

    scenario: str
    description: str
    attack_steps: list[str]
    detected: bool
    alerts_fired: list[dict[str, Any]]
    severity: str
    verdict: str  # PASS, FAIL, PARTIAL

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


# ── Scenario Registry ────────────────────────────────────────

SCENARIOS: dict[str, dict[str, str]] = {
    "reconnaissance": {
        "description": "Excessive tools/list probing to enumerate available tools",
        "target_detector": "RECONNAISSANCE",
        "expected_severity": "MEDIUM",
    },
    "persistence": {
        "description": "Bot-like periodic reconnection at regular intervals",
        "target_detector": "PERSISTENCE",
        "expected_severity": "HIGH",
    },
    "defense_evasion": {
        "description": "Large encoded payload to evade content inspection",
        "target_detector": "DEFENSE_EVASION",
        "expected_severity": "HIGH",
    },
    "lateral_movement": {
        "description": "Cross-server tool chaining across multiple MCP servers",
        "target_detector": "LATERAL_MOVEMENT",
        "expected_severity": "HIGH",
    },
    "c2_beaconing": {
        "description": "Command & control beaconing with regular timing and small payloads",
        "target_detector": "COMMAND_AND_CONTROL",
        "expected_severity": "CRITICAL",
    },
    "supply_chain": {
        "description": "Calling an unregistered tool not in the server's tools/list",
        "target_detector": "SUPPLY_CHAIN",
        "expected_severity": "CRITICAL",
    },
    "rug_pull": {
        "description": "Sudden shift to new tools not in baseline behavior",
        "target_detector": "RUG_PULL",
        "expected_severity": "HIGH",
    },
    "data_exfiltration": {
        "description": "Data volume spike exceeding 5x the established baseline",
        "target_detector": "DATA_EXFILTRATION",
        "expected_severity": "CRITICAL",
    },
    "privilege_escalation": {
        "description": "First-time access to sensitive admin tools",
        "target_detector": "PRIVILEGE_ESCALATION",
        "expected_severity": "CRITICAL",
    },
    "rate_spike": {
        "description": "Sudden 4x increase in invocation rate",
        "target_detector": "RATE_SPIKE",
        "expected_severity": "MEDIUM",
    },
    "policy_bypass": {
        "description": "Attempt to call a tool explicitly denied by policy",
        "target_detector": "POLICY",
        "expected_severity": "POLICY",
    },
}


class PentestEngine:
    """Simulates SAFE-MCP attack scenarios against Navil detectors."""

    AGENT = "pentest-agent"

    def __init__(
        self,
        anomaly_detector: BehavioralAnomalyDetector,
        policy_engine: Any,
    ) -> None:
        self._detector_template = anomaly_detector
        self._policy_engine = policy_engine

    # ── Public API ───────────────────────────────────────────

    def run_all(self) -> dict[str, Any]:
        """Run all attack scenarios and return a summary report."""
        results: list[dict[str, Any]] = []
        passed = 0
        failed = 0
        partial = 0

        for name in SCENARIOS:
            result = self.run_scenario(name)
            results.append(result.to_dict())
            if result.verdict == "PASS":
                passed += 1
            elif result.verdict == "FAIL":
                failed += 1
            else:
                partial += 1

        total = len(results)
        detection_rate = round((passed / total) * 100, 1) if total > 0 else 0.0

        return {
            "status": "completed",
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "total_scenarios": total,
            "passed": passed,
            "failed": failed,
            "partial": partial,
            "detection_rate": detection_rate,
            "results": results,
        }

    def run_scenario(self, name: str) -> PentestResult:
        """Run a single named scenario."""
        if name not in SCENARIOS:
            raise ValueError(f"Unknown scenario: {name}. Available: {', '.join(SCENARIOS.keys())}")

        scenario_info = SCENARIOS[name]
        method = getattr(self, f"_attack_{name}")
        return method(scenario_info)

    # ── Fresh Detector ───────────────────────────────────────

    def _fresh_detector(self) -> BehavioralAnomalyDetector:
        """Create a fresh detector so scenarios don't interfere."""
        return BehavioralAnomalyDetector()

    # ── Attack Scenarios ─────────────────────────────────────

    def _attack_reconnaissance(self, info: dict[str, str]) -> PentestResult:
        """Simulate excessive tools/list probing (>5 in 10 min)."""
        det = self._fresh_detector()
        steps = []

        for i in range(7):
            det.record_invocation(
                agent_name=self.AGENT,
                tool_name="__tools_list__",
                action="tools/list",
                duration_ms=10,
                is_list_tools=True,
            )
            steps.append(f"Called tools/list ({i + 1}/7)")

        alerts = [a for a in det.alerts if a.anomaly_type == "RECONNAISSANCE"]
        detected = len(alerts) >= 1

        return PentestResult(
            scenario="reconnaissance",
            description=info["description"],
            attack_steps=steps,
            detected=detected,
            alerts_fired=[self._alert_to_dict(a) for a in alerts],
            severity=info["expected_severity"],
            verdict="PASS" if detected else "FAIL",
        )

    def _attack_persistence(self, info: dict[str, str]) -> PentestResult:
        """Simulate bot-like periodic calls at exact 30s intervals."""
        det = self._fresh_detector()
        steps = []

        base = datetime.now(timezone.utc) - timedelta(minutes=10)
        for i in range(8):
            ts = base + timedelta(seconds=i * 30)
            det.record_invocation(
                agent_name=self.AGENT,
                tool_name="heartbeat",
                action="ping",
                duration_ms=5,
                data_accessed_bytes=0,
                success=True,
                timestamp=ts.isoformat(),
            )
            steps.append(f"Heartbeat at t+{i * 30}s (interval=30s)")

        det._detect_persistence(self.AGENT)
        alerts = [a for a in det.alerts if a.anomaly_type == "PERSISTENCE"]
        detected = len(alerts) >= 1

        return PentestResult(
            scenario="persistence",
            description=info["description"],
            attack_steps=steps,
            detected=detected,
            alerts_fired=[self._alert_to_dict(a) for a in alerts],
            severity=info["expected_severity"],
            verdict="PASS" if detected else "FAIL",
        )

    def _attack_defense_evasion(self, info: dict[str, str]) -> PentestResult:
        """Simulate large encoded payload (>5000 bytes)."""
        det = self._fresh_detector()
        steps = []

        det.record_invocation(
            agent_name=self.AGENT,
            tool_name="execute",
            action="run",
            duration_ms=100,
            arguments_size_bytes=6000,
        )
        steps.append("Sent tool call with 6000-byte argument payload")

        alerts = [a for a in det.alerts if a.anomaly_type == "DEFENSE_EVASION"]
        detected = len(alerts) >= 1

        return PentestResult(
            scenario="defense_evasion",
            description=info["description"],
            attack_steps=steps,
            detected=detected,
            alerts_fired=[self._alert_to_dict(a) for a in alerts],
            severity=info["expected_severity"],
            verdict="PASS" if detected else "FAIL",
        )

    def _attack_lateral_movement(self, info: dict[str, str]) -> PentestResult:
        """Simulate talking to >3 MCP servers in 5 min."""
        det = self._fresh_detector()
        steps = []

        servers = [
            "http://mcp-alpha:3000",
            "http://mcp-beta:3000",
            "http://mcp-gamma:3000",
            "http://mcp-delta:3000",
        ]
        for server in servers:
            det.record_invocation(
                agent_name=self.AGENT,
                tool_name="query",
                action="tools/call",
                duration_ms=50,
                target_server=server,
            )
            steps.append(f"Called tools/call on {server}")

        alerts = [a for a in det.alerts if a.anomaly_type == "LATERAL_MOVEMENT"]
        detected = len(alerts) >= 1

        return PentestResult(
            scenario="lateral_movement",
            description=info["description"],
            attack_steps=steps,
            detected=detected,
            alerts_fired=[self._alert_to_dict(a) for a in alerts],
            severity=info["expected_severity"],
            verdict="PASS" if detected else "FAIL",
        )

    def _attack_c2_beaconing(self, info: dict[str, str]) -> PentestResult:
        """Simulate C2 beaconing: regular intervals + small consistent responses."""
        det = self._fresh_detector()
        steps = []

        base = datetime.now(timezone.utc) - timedelta(minutes=5)
        for i in range(8):
            ts = base + timedelta(seconds=i * 10)
            det.record_invocation(
                agent_name=self.AGENT,
                tool_name="status",
                action="check",
                duration_ms=20,
                data_accessed_bytes=0,
                success=True,
                response_size_bytes=256,
                timestamp=ts.isoformat(),
            )
            steps.append(f"Beacon at t+{i * 10}s (response=256B)")

        det._detect_command_and_control(self.AGENT)
        alerts = [a for a in det.alerts if a.anomaly_type == "COMMAND_AND_CONTROL"]
        detected = len(alerts) >= 1

        return PentestResult(
            scenario="c2_beaconing",
            description=info["description"],
            attack_steps=steps,
            detected=detected,
            alerts_fired=[self._alert_to_dict(a) for a in alerts],
            severity=info["expected_severity"],
            verdict="PASS" if detected else "FAIL",
        )

    def _attack_supply_chain(self, info: dict[str, str]) -> PentestResult:
        """Simulate calling an unregistered tool on a known server."""
        det = self._fresh_detector()
        steps = []

        server = "http://mcp-server:3000"
        det.register_server_tools(server, ["read", "write", "list"])
        steps.append(f"Registered tools [read, write, list] for {server}")

        det.record_invocation(
            agent_name=self.AGENT,
            tool_name="inject_backdoor",
            action="tools/call",
            duration_ms=100,
            target_server=server,
        )
        steps.append("Called unregistered tool 'inject_backdoor' on registered server")

        alerts = [a for a in det.alerts if a.anomaly_type == "SUPPLY_CHAIN"]
        detected = len(alerts) >= 1

        return PentestResult(
            scenario="supply_chain",
            description=info["description"],
            attack_steps=steps,
            detected=detected,
            alerts_fired=[self._alert_to_dict(a) for a in alerts],
            severity=info["expected_severity"],
            verdict="PASS" if detected else "FAIL",
        )

    def _attack_rug_pull(self, info: dict[str, str]) -> PentestResult:
        """Simulate sudden shift to >3 new tools in 1 hour."""
        det = self._fresh_detector()
        steps = []

        # Build baseline with standard tools
        baseline_time = datetime.now(timezone.utc) - timedelta(hours=2)
        for i in range(20):
            ts = baseline_time + timedelta(minutes=i * 5)
            det.record_invocation(
                agent_name=self.AGENT,
                tool_name="read_file",
                action="read",
                duration_ms=50,
                data_accessed_bytes=100,
                success=True,
                timestamp=ts.isoformat(),
            )
        det._build_baseline(self.AGENT)
        steps.append("Built baseline with 20 'read_file' invocations over 2 hours")

        # Now suddenly use 4 new tools
        new_tools = ["admin_panel", "credential_dump", "network_scan", "exfil_data"]
        for tool in new_tools:
            det.record_invocation(
                agent_name=self.AGENT,
                tool_name=tool,
                action="execute",
                duration_ms=100,
                data_accessed_bytes=500,
            )
            steps.append(f"Accessed new tool: {tool}")

        alerts = [a for a in det.alerts if a.anomaly_type == "RUG_PULL"]
        detected = len(alerts) >= 1

        return PentestResult(
            scenario="rug_pull",
            description=info["description"],
            attack_steps=steps,
            detected=detected,
            alerts_fired=[self._alert_to_dict(a) for a in alerts],
            severity=info["expected_severity"],
            verdict="PASS" if detected else "FAIL",
        )

    def _attack_data_exfiltration(self, info: dict[str, str]) -> PentestResult:
        """Simulate data volume spike >5x baseline."""
        det = self._fresh_detector()
        steps = []

        # Build baseline with modest data access
        baseline_time = datetime.now(timezone.utc) - timedelta(hours=2)
        for i in range(10):
            ts = baseline_time + timedelta(minutes=i * 10)
            det.record_invocation(
                agent_name=self.AGENT,
                tool_name="read_data",
                action="read",
                duration_ms=50,
                data_accessed_bytes=100,
                success=True,
                timestamp=ts.isoformat(),
            )
        det._build_baseline(self.AGENT)
        steps.append("Built baseline: 10 invocations, 100B each (baseline ~1000B)")

        # Exfiltrate 6x baseline volume
        det.record_invocation(
            agent_name=self.AGENT,
            tool_name="read_data",
            action="read",
            duration_ms=50,
            data_accessed_bytes=6000,
        )
        steps.append("Accessed 6000 bytes in single call (6x baseline volume)")

        alerts = [a for a in det.alerts if a.anomaly_type == "DATA_EXFILTRATION"]
        detected = len(alerts) >= 1

        return PentestResult(
            scenario="data_exfiltration",
            description=info["description"],
            attack_steps=steps,
            detected=detected,
            alerts_fired=[self._alert_to_dict(a) for a in alerts],
            severity=info["expected_severity"],
            verdict="PASS" if detected else "FAIL",
        )

    def _attack_privilege_escalation(self, info: dict[str, str]) -> PentestResult:
        """Simulate first-time access to sensitive admin tools."""
        det = self._fresh_detector()
        steps = []

        # Build baseline with safe tools
        baseline_time = datetime.now(timezone.utc) - timedelta(hours=2)
        for i in range(10):
            ts = baseline_time + timedelta(minutes=i * 10)
            det.record_invocation(
                agent_name=self.AGENT,
                tool_name="read_logs",
                action="read",
                duration_ms=50,
                data_accessed_bytes=100,
                success=True,
                timestamp=ts.isoformat(),
            )
        det._build_baseline(self.AGENT)
        steps.append("Built baseline with safe 'read_logs' tool only")

        # Access sensitive admin tool
        det.record_invocation(
            agent_name=self.AGENT,
            tool_name="admin_panel",
            action="admin",
            duration_ms=50,
        )
        steps.append("Accessed sensitive tool: admin_panel (first-time)")

        alerts = [a for a in det.alerts if a.anomaly_type == "PRIVILEGE_ESCALATION"]
        detected = len(alerts) >= 1

        return PentestResult(
            scenario="privilege_escalation",
            description=info["description"],
            attack_steps=steps,
            detected=detected,
            alerts_fired=[self._alert_to_dict(a) for a in alerts],
            severity=info["expected_severity"],
            verdict="PASS" if detected else "FAIL",
        )

    def _attack_rate_spike(self, info: dict[str, str]) -> PentestResult:
        """Simulate 4x baseline call rate in 30 min."""
        det = self._fresh_detector()
        steps = []

        # Build baseline with moderate rate (~10 calls / 2 hours)
        baseline_time = datetime.now(timezone.utc) - timedelta(hours=2)
        for i in range(10):
            ts = baseline_time + timedelta(minutes=i * 12)
            det.record_invocation(
                agent_name=self.AGENT,
                tool_name="query",
                action="read",
                duration_ms=50,
                data_accessed_bytes=50,
                success=True,
                timestamp=ts.isoformat(),
            )
        det._build_baseline(self.AGENT)
        steps.append("Built baseline: 10 calls over 2 hours (~5/hr)")

        # Spike: 40 calls in 30 min (4x rate)
        for _ in range(40):
            det.record_invocation(
                agent_name=self.AGENT,
                tool_name="query",
                action="read",
                duration_ms=10,
                data_accessed_bytes=10,
            )
        steps.append("Fired 40 calls in rapid succession (4x baseline rate)")

        alerts = [a for a in det.alerts if a.anomaly_type == "RATE_SPIKE"]
        detected = len(alerts) >= 1

        return PentestResult(
            scenario="rate_spike",
            description=info["description"],
            attack_steps=steps,
            detected=detected,
            alerts_fired=[self._alert_to_dict(a) for a in alerts],
            severity=info["expected_severity"],
            verdict="PASS" if detected else "FAIL",
        )

    def _attack_policy_bypass(self, info: dict[str, str]) -> PentestResult:
        """Simulate calling a tool explicitly denied by policy."""
        steps = []

        # Configure policy to deny a tool — set directly on the engine
        from navil.policy_engine import PolicyEngine

        pe = PolicyEngine()
        pe.policy = {
            "version": "1.0",
            "agents": {
                self.AGENT: {
                    "tools_allowed": ["read_file", "write_file"],
                    "tools_denied": ["destroy_database"],
                    "rate_limit_per_hour": 1000,
                    "data_clearance": "PUBLIC",
                }
            },
            "tools": {},
        }
        steps.append("Loaded policy denying 'destroy_database' tool")

        # Attempt to call denied tool
        allowed, reason = pe.check_tool_call(
            agent_name=self.AGENT,
            tool_name="destroy_database",
            action="delete",
        )
        steps.append(f"Called check_tool_call('destroy_database') -> allowed={allowed}")

        detected = not allowed
        alert_dict = {
            "anomaly_type": "POLICY",
            "severity": "POLICY",
            "agent_name": self.AGENT,
            "description": reason,
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "evidence": [f"Policy denied: {reason}"],
            "recommended_action": "Review agent permissions",
        }

        return PentestResult(
            scenario="policy_bypass",
            description=info["description"],
            attack_steps=steps,
            detected=detected,
            alerts_fired=[alert_dict] if detected else [],
            severity="POLICY",
            verdict="PASS" if detected else "FAIL",
        )

    # ── Helpers ──────────────────────────────────────────────

    @staticmethod
    def _alert_to_dict(alert: Any) -> dict[str, Any]:
        """Convert an AnomalyAlert to a plain dict."""
        if hasattr(alert, "__dataclass_fields__"):
            return {
                "anomaly_type": alert.anomaly_type,
                "severity": alert.severity,
                "agent_name": alert.agent_name,
                "description": alert.description,
                "timestamp": alert.timestamp,
                "evidence": alert.evidence,
                "recommended_action": alert.recommended_action,
                "confidence": getattr(alert, "confidence", 0.0),
            }
        return dict(alert) if isinstance(alert, dict) else {"raw": str(alert)}
