#!/usr/bin/env python3
"""IDS self-testing module: assess detection risk of planned actions against
Snort/Suricata community rules before activation."""

import logging
import os
import re
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional

from modules.base import BaseModule

logger = logging.getLogger(__name__)

# Detection risk levels
RISK_LOW = "LOW"
RISK_MEDIUM = "MEDIUM"
RISK_HIGH = "HIGH"
RISK_CRITICAL = "CRITICAL"

# Action-to-risk baseline mapping (before rule checking)
ACTION_RISK_BASELINE = {
    # Passive (LOW)
    "passive_sniff": RISK_LOW,
    "dns_logging": RISK_LOW,
    "host_discovery_passive": RISK_LOW,
    "credential_sniff": RISK_LOW,
    "pcap_capture": RISK_LOW,
    # Medium (detectable with tuned IDS)
    "arp_spoof": RISK_MEDIUM,
    "dhcp_spoof": RISK_MEDIUM,
    "dns_spoof": RISK_MEDIUM,
    "nbns_spoof": RISK_MEDIUM,
    "llmnr_spoof": RISK_MEDIUM,
    "mdns_spoof": RISK_MEDIUM,
    "evil_twin": RISK_MEDIUM,
    "wpad_spoof": RISK_MEDIUM,
    "vlan_hop": RISK_MEDIUM,
    # High (commonly detected)
    "responder": RISK_HIGH,
    "ntlm_relay": RISK_HIGH,
    "kerberoast": RISK_HIGH,
    "smb_relay": RISK_HIGH,
    "http_proxy_inject": RISK_HIGH,
    "port_scan": RISK_HIGH,
    # Critical (almost certainly detected)
    "mitmproxy_ssl_intercept": RISK_CRITICAL,
    "ssl_strip": RISK_CRITICAL,
    "dns_hijack_external": RISK_CRITICAL,
    "exploit_delivery": RISK_CRITICAL,
}

# Bettercap caplet action to risk mapping
CAPLET_ACTION_RISK = {
    "net.probe": RISK_MEDIUM,
    "net.sniff": RISK_LOW,
    "arp.spoof": RISK_MEDIUM,
    "dns.spoof": RISK_MEDIUM,
    "dhcp6.spoof": RISK_MEDIUM,
    "http.proxy": RISK_HIGH,
    "https.proxy": RISK_CRITICAL,
    "tcp.proxy": RISK_HIGH,
    "wifi.deauth": RISK_HIGH,
    "wifi.ap": RISK_MEDIUM,
    "any.proxy": RISK_HIGH,
    "hid.inject": RISK_CRITICAL,
}


@dataclass
class RiskAssessment:
    """Result of an IDS risk assessment."""
    action: str
    risk_level: str
    matching_sids: list[str] = field(default_factory=list)
    matching_rules: list[str] = field(default_factory=list)
    recommendation: str = ""
    timestamp: float = field(default_factory=time.time)

    def to_dict(self) -> dict:
        return {
            "action": self.action,
            "risk_level": self.risk_level,
            "matching_sids": self.matching_sids,
            "matching_rules": self.matching_rules,
            "recommendation": self.recommendation,
            "timestamp": self.timestamp,
        }


class IDSTester(BaseModule):
    """On-demand IDS risk assessment: check planned actions against known
    Snort/Suricata signatures before module activation."""

    name = "ids_tester"
    module_type = "stealth"
    priority = -100
    requires_root = False

    def __init__(self, bus, state, config, engine=None):
        super().__init__(bus, state, config, engine)
        self._rules_dir = config.get("rules_dir", "data/ids_rules")
        self._rules: list[dict] = []
        self._last_assessment: Optional[RiskAssessment] = None
        self._assessments: list[RiskAssessment] = []

    # ------------------------------------------------------------------
    # BaseModule interface
    # ------------------------------------------------------------------

    def start(self) -> None:
        self._running = True
        self._start_time = time.time()
        self._pid = os.getpid()

        self._load_rules()

        self.state.set_module_status(self.name, "running", pid=self._pid)
        self.bus.emit("MODULE_STARTED", {"module": self.name}, source_module=self.name)
        logger.info("IDSTester ready — %d rules loaded", len(self._rules))

    def stop(self) -> None:
        self._running = False
        self.state.set_module_status(self.name, "stopped")
        self.bus.emit("MODULE_STOPPED", {"module": self.name}, source_module=self.name)
        logger.info("IDSTester stopped")

    def status(self) -> dict:
        return {
            "running": self._running,
            "pid": self._pid,
            "uptime": time.time() - self._start_time if self._start_time else 0,
            "rules_loaded": len(self._rules),
            "last_assessment_time": self._last_assessment.timestamp if self._last_assessment else None,
            "last_assessment_action": self._last_assessment.action if self._last_assessment else None,
            "last_assessment_risk": self._last_assessment.risk_level if self._last_assessment else None,
            "total_assessments": len(self._assessments),
        }

    def configure(self, config: dict) -> None:
        if "rules_dir" in config:
            self._rules_dir = config["rules_dir"]
            self._load_rules()

    # ------------------------------------------------------------------
    # Public API: risk assessment
    # ------------------------------------------------------------------

    def assess_risk(self, action: str) -> dict:
        """Assess detection risk for a planned action.

        Args:
            action: Action identifier (e.g., 'arp_spoof', 'responder', 'mitmproxy_ssl_intercept')

        Returns:
            dict with risk_level, matching_sids, recommendation
        """
        # Start with baseline risk
        baseline_risk = ACTION_RISK_BASELINE.get(action, RISK_MEDIUM)

        # Check against loaded IDS rules
        matching_sids = []
        matching_rules = []
        keywords = self._action_to_keywords(action)

        for rule in self._rules:
            rule_text = rule.get("raw", "").lower()
            rule_msg = rule.get("msg", "").lower()
            for keyword in keywords:
                if keyword in rule_text or keyword in rule_msg:
                    sid = rule.get("sid", "unknown")
                    matching_sids.append(str(sid))
                    matching_rules.append(rule.get("msg", "unknown"))
                    break

        # Escalate risk if IDS rules match
        final_risk = baseline_risk
        if matching_sids:
            risk_order = [RISK_LOW, RISK_MEDIUM, RISK_HIGH, RISK_CRITICAL]
            current_idx = risk_order.index(baseline_risk)
            escalation = min(len(matching_sids), 2)  # Max +2 levels
            final_risk = risk_order[min(current_idx + escalation, len(risk_order) - 1)]

        recommendation = self._generate_recommendation(action, final_risk, len(matching_sids))

        assessment = RiskAssessment(
            action=action,
            risk_level=final_risk,
            matching_sids=matching_sids[:20],  # Cap at 20
            matching_rules=matching_rules[:20],
            recommendation=recommendation,
        )

        self._last_assessment = assessment
        self._assessments.append(assessment)
        # Keep only last 100 assessments
        if len(self._assessments) > 100:
            self._assessments = self._assessments[-100:]

        logger.info(
            "Risk assessment: %s -> %s (%d matching SIDs)",
            action, final_risk, len(matching_sids),
        )

        return assessment.to_dict()

    def assess_caplet(self, caplet_content: str) -> list[dict]:
        """Assess detection risk for bettercap caplet actions.

        Args:
            caplet_content: Raw caplet file content

        Returns:
            List of risk assessment dicts, one per detected action
        """
        results = []
        for line in caplet_content.splitlines():
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            # Check each known caplet action
            for action_prefix, risk in CAPLET_ACTION_RISK.items():
                if line.startswith(action_prefix) or f" {action_prefix}" in line:
                    assessment = self.assess_risk(action_prefix.replace(".", "_"))
                    assessment["caplet_line"] = line
                    results.append(assessment)
                    break
        return results

    def preflight_check(self, module_name: str, actions: list[str]) -> dict:
        """Pre-flight validation: assess all actions a module will perform.

        Args:
            module_name: Name of the module to be activated
            actions: List of action identifiers the module will perform

        Returns:
            dict with overall_risk, action_risks, go_nogo recommendation
        """
        action_risks = []
        risk_order = [RISK_LOW, RISK_MEDIUM, RISK_HIGH, RISK_CRITICAL]
        max_risk_idx = 0

        for action in actions:
            result = self.assess_risk(action)
            action_risks.append(result)
            try:
                idx = risk_order.index(result["risk_level"])
                max_risk_idx = max(max_risk_idx, idx)
            except ValueError:
                pass

        overall_risk = risk_order[max_risk_idx]

        # go/nogo based on overall risk
        if overall_risk == RISK_CRITICAL:
            go_nogo = "NO-GO"
            summary = f"Module '{module_name}' has CRITICAL detection risk — activation NOT recommended"
        elif overall_risk == RISK_HIGH:
            go_nogo = "CAUTION"
            summary = f"Module '{module_name}' has HIGH detection risk — activate only if IDS evasion is confirmed"
        elif overall_risk == RISK_MEDIUM:
            go_nogo = "GO"
            summary = f"Module '{module_name}' has MEDIUM detection risk — proceed with monitoring"
        else:
            go_nogo = "GO"
            summary = f"Module '{module_name}' has LOW detection risk — safe to activate"

        return {
            "module": module_name,
            "overall_risk": overall_risk,
            "go_nogo": go_nogo,
            "summary": summary,
            "action_risks": action_risks,
        }

    # ------------------------------------------------------------------
    # Internal: rule loading
    # ------------------------------------------------------------------

    def _load_rules(self) -> None:
        """Load Snort/Suricata rules from the rules directory."""
        self._rules = []
        rules_path = Path(self._rules_dir)
        if not rules_path.is_dir():
            logger.info("IDS rules directory not found: %s", self._rules_dir)
            return

        rule_files = list(rules_path.glob("*.rules"))
        if not rule_files:
            logger.info("No .rules files in %s", self._rules_dir)
            return

        for rule_file in rule_files:
            try:
                with open(rule_file, "r", errors="replace") as f:
                    for line in f:
                        line = line.strip()
                        if not line or line.startswith("#"):
                            continue
                        parsed = self._parse_rule(line)
                        if parsed:
                            self._rules.append(parsed)
            except Exception:
                logger.exception("Failed to parse rules file: %s", rule_file)

        logger.info("Loaded %d IDS rules from %d files", len(self._rules), len(rule_files))

    @staticmethod
    def _parse_rule(line: str) -> Optional[dict]:
        """Parse a single Snort/Suricata rule line into a dict."""
        # Extract SID
        sid_match = re.search(r"sid:\s*(\d+)", line)
        sid = sid_match.group(1) if sid_match else None

        # Extract message
        msg_match = re.search(r'msg:\s*"([^"]*)"', line)
        msg = msg_match.group(1) if msg_match else ""

        # Extract classtype
        class_match = re.search(r"classtype:\s*([^;]+)", line)
        classtype = class_match.group(1).strip() if class_match else ""

        if not sid:
            return None

        return {
            "sid": sid,
            "msg": msg,
            "classtype": classtype,
            "raw": line,
        }

    # ------------------------------------------------------------------
    # Internal: keyword mapping
    # ------------------------------------------------------------------

    @staticmethod
    def _action_to_keywords(action: str) -> list[str]:
        """Map action identifiers to IDS rule keywords for matching."""
        keyword_map = {
            "arp_spoof": ["arp", "spoof", "poison", "gratuitous arp", "arp-scan"],
            "dns_spoof": ["dns", "spoof", "dns poison", "dns hijack"],
            "dhcp_spoof": ["dhcp", "rogue dhcp", "dhcp spoof"],
            "responder": ["responder", "llmnr", "nbns", "nbt-ns", "mdns", "wpad", "netbios"],
            "ntlm_relay": ["ntlm", "relay", "smb relay", "ntlmrelay"],
            "smb_relay": ["smb", "relay", "smb relay"],
            "kerberoast": ["kerberos", "kerberoast", "tgs-rep", "spn"],
            "mitmproxy_ssl_intercept": ["mitm", "ssl intercept", "tls intercept", "ssl strip"],
            "ssl_strip": ["ssl strip", "sslstrip", "hsts bypass"],
            "port_scan": ["port scan", "nmap", "syn scan", "tcp scan"],
            "evil_twin": ["evil twin", "rogue ap", "fake ap", "deauth"],
            "nbns_spoof": ["nbns", "nbt-ns", "netbios"],
            "llmnr_spoof": ["llmnr", "multicast dns"],
            "mdns_spoof": ["mdns", "multicast dns", "avahi"],
            "wpad_spoof": ["wpad", "proxy auto", "proxy autoconfig"],
            "vlan_hop": ["vlan", "802.1q", "dtp", "double tag"],
            "http_proxy_inject": ["http inject", "http proxy", "http mitm"],
            "net_probe": ["arp scan", "net probe", "network scan", "host discovery"],
            "net_sniff": ["sniff", "promiscuous", "pcap"],
            "wifi_deauth": ["deauth", "disassoc", "wifi attack"],
            "wifi_ap": ["rogue ap", "fake ap", "evil twin"],
        }
        # Normalize action name
        action_lower = action.lower().replace(".", "_")
        keywords = keyword_map.get(action_lower, [action_lower.replace("_", " ")])
        # Always include the raw action name
        keywords.append(action_lower.replace("_", " "))
        return list(set(keywords))

    @staticmethod
    def _generate_recommendation(action: str, risk: str, sid_count: int) -> str:
        """Generate a human-readable recommendation."""
        if risk == RISK_CRITICAL:
            return (
                f"CRITICAL: '{action}' matches {sid_count} IDS signatures. "
                "Do NOT activate without confirmed IDS blind spots or rule suppression. "
                "Consider alternative approaches or wait for off-hours."
            )
        elif risk == RISK_HIGH:
            return (
                f"HIGH: '{action}' is commonly detected ({sid_count} matching SIDs). "
                "Verify target network has no Suricata/Snort/Zeek monitoring, or use "
                "traffic mimicry to reduce signature exposure."
            )
        elif risk == RISK_MEDIUM:
            return (
                f"MEDIUM: '{action}' may be detected by tuned IDS ({sid_count} matching SIDs). "
                "Proceed with caution — monitor for alerts and be ready to back off."
            )
        else:
            return (
                f"LOW: '{action}' has minimal detection footprint ({sid_count} matching SIDs). "
                "Safe to proceed with standard OPSEC."
            )
