#!/usr/bin/env python3
"""Security posture mapper — detect security infrastructure by traffic patterns.

Identifies EDR, SIEM, vulnerability scanners, honeypots, NAC, and network
monitoring tools by observing DNS queries, HTTP traffic, and port patterns.
Gates active module risk assessment — if strong security tooling is detected,
active operations should increase caution.
"""

import json
import logging
import os
import re
import sqlite3
import threading
import time
from typing import Optional

from modules.base import BaseModule

logger = logging.getLogger(__name__)

_SCHEMA = """
CREATE TABLE IF NOT EXISTS security_tools (
    id               INTEGER PRIMARY KEY AUTOINCREMENT,
    tool_type        TEXT    NOT NULL,
    tool_name        TEXT    NOT NULL,
    evidence         TEXT    NOT NULL,
    confidence       REAL    DEFAULT 0.5,
    target_ip        TEXT    DEFAULT '',
    target_hostname  TEXT    DEFAULT '',
    first_seen       REAL    NOT NULL,
    last_seen        REAL    NOT NULL,
    impacts_active_ops INTEGER DEFAULT 1,
    details          TEXT    DEFAULT '',
    UNIQUE(tool_type, tool_name, target_ip)
);

CREATE INDEX IF NOT EXISTS idx_st_type ON security_tools(tool_type);
CREATE INDEX IF NOT EXISTS idx_st_name ON security_tools(tool_name);
"""

# -----------------------------------------------------------------------
# Detection signatures
# -----------------------------------------------------------------------

# DNS patterns: (regex, tool_type, tool_name, confidence, impacts_active)
_DNS_SIGNATURES = [
    # EDR / Endpoint Protection
    (re.compile(r"ts01-[a-z0-9]+\.cloudsink\.net$", re.I),
     "edr", "CrowdStrike Falcon", 0.95, True),
    (re.compile(r"(lfodown|lfoup)\.crowdstrike\.com$", re.I),
     "edr", "CrowdStrike Falcon", 0.95, True),
    (re.compile(r"\.crowdstrike\.com$", re.I),
     "edr", "CrowdStrike Falcon", 0.9, True),
    (re.compile(r"[a-z0-9]+\.sentinelone\.net$", re.I),
     "edr", "SentinelOne", 0.95, True),
    (re.compile(r"management.*\.sentinelone\.net$", re.I),
     "edr", "SentinelOne", 0.95, True),
    (re.compile(r"\.cybereason\.com$", re.I),
     "edr", "Cybereason", 0.9, True),
    (re.compile(r"\.carbonblack\.(com|io)$", re.I),
     "edr", "Carbon Black", 0.9, True),
    (re.compile(r"\.cb\.carbonblack\.io$", re.I),
     "edr", "Carbon Black Cloud", 0.95, True),
    (re.compile(r"(wdcp|wdcpalt)\.microsoft\.com$", re.I),
     "edr", "Microsoft Defender ATP", 0.85, True),
    (re.compile(r"\.securitycenter\.windows\.com$", re.I),
     "edr", "Microsoft Defender ATP", 0.9, True),
    (re.compile(r"\.endpoint\.security\.microsoft\.com$", re.I),
     "edr", "Microsoft Defender for Endpoint", 0.95, True),
    (re.compile(r"\.sophos\.(com|net)$", re.I),
     "edr", "Sophos", 0.8, True),
    (re.compile(r"\.cylance\.(com|io)$", re.I),
     "edr", "Cylance/BlackBerry Protect", 0.85, True),
    (re.compile(r"\.trellix\.com$", re.I),
     "edr", "Trellix (McAfee/FireEye)", 0.85, True),
    (re.compile(r"\.fireeye\.com$", re.I),
     "edr", "FireEye/Mandiant", 0.85, True),

    # SIEM
    (re.compile(r"\.splunkcloud\.com$", re.I),
     "siem", "Splunk Cloud", 0.9, False),
    (re.compile(r"input.*\.splunkcloud\.com$", re.I),
     "siem", "Splunk Cloud HEC", 0.95, False),
    (re.compile(r"\.sumologic\.com$", re.I),
     "siem", "Sumo Logic", 0.85, False),
    (re.compile(r"\.logz\.io$", re.I),
     "siem", "Logz.io", 0.85, False),
    (re.compile(r"\.elastic-cloud\.com$", re.I),
     "siem", "Elastic Cloud", 0.85, False),
    (re.compile(r"\.datadoghq\.com$", re.I),
     "siem", "Datadog", 0.8, False),

    # Vulnerability Scanners
    (re.compile(r"(plugins|sensor)\.nessus\.org$", re.I),
     "vuln_scanner", "Nessus", 0.9, False),
    (re.compile(r"\.tenable\.(com|io)$", re.I),
     "vuln_scanner", "Tenable", 0.85, False),
    (re.compile(r"\.qualys\.com$", re.I),
     "vuln_scanner", "Qualys", 0.9, False),
    (re.compile(r"\.rapid7\.com$", re.I),
     "vuln_scanner", "Rapid7 InsightVM", 0.85, False),

    # Honeypots / Deception
    (re.compile(r"\.canary\.tools$", re.I),
     "honeypot", "Thinkst Canary", 0.95, True),
    (re.compile(r"canarytokens\..*\.com$", re.I),
     "honeypot", "Canary Tokens", 0.9, True),
    (re.compile(r"\.illusive-networks\.com$", re.I),
     "honeypot", "Illusive Networks", 0.85, True),
    (re.compile(r"\.attivo\.com$", re.I),
     "honeypot", "Attivo Networks", 0.85, True),

    # NAC
    (re.compile(r"\.forescout\.com$", re.I),
     "nac", "Forescout", 0.85, True),
    (re.compile(r"ise[\w-]*\.(local|internal|corp|lan)$", re.I),
     "nac", "Cisco ISE", 0.8, True),

    # Threat Intelligence
    (re.compile(r"\.virustotal\.com$", re.I),
     "threat_intel", "VirusTotal", 0.7, False),
    (re.compile(r"\.shodan\.io$", re.I),
     "threat_intel", "Shodan", 0.7, False),
]

# Port-based signatures: port -> (tool_type, tool_name, confidence, impacts_active)
_PORT_SIGNATURES = {
    8088:  ("siem", "Splunk HEC", 0.7, False),
    8089:  ("siem", "Splunk Management", 0.7, False),
    9997:  ("siem", "Splunk Forwarder", 0.75, False),
    1514:  ("siem", "Wazuh Agent", 0.7, False),
    1515:  ("siem", "Wazuh Enrollment", 0.75, False),
    55000: ("siem", "Wazuh API", 0.8, False),
    8834:  ("vuln_scanner", "Nessus Scanner", 0.85, False),
    3790:  ("vuln_scanner", "Rapid7 Metasploit Pro/Nexpose", 0.75, False),
    9390:  ("vuln_scanner", "OpenVAS/Greenbone", 0.8, False),
    9392:  ("vuln_scanner", "Greenbone GSA", 0.8, False),
    1812:  ("nac", "RADIUS Authentication", 0.6, True),
    1813:  ("nac", "RADIUS Accounting", 0.6, True),
    3799:  ("nac", "RADIUS CoA (Change of Authorization)", 0.7, True),
}

# HTTP User-Agent patterns
_UA_SIGNATURES = [
    (re.compile(r"Nessus", re.I), "vuln_scanner", "Nessus", 0.9),
    (re.compile(r"Qualys", re.I), "vuln_scanner", "Qualys", 0.9),
    (re.compile(r"InsightVM|Rapid7", re.I), "vuln_scanner", "Rapid7", 0.85),
    (re.compile(r"Nikto", re.I), "vuln_scanner", "Nikto", 0.8),
    (re.compile(r"OpenVAS", re.I), "vuln_scanner", "OpenVAS", 0.85),
    (re.compile(r"CrowdStrike", re.I), "edr", "CrowdStrike", 0.7),
    (re.compile(r"Splunk", re.I), "siem", "Splunk", 0.7),
]

# T-Pot honeypot: multiple decoy services on unusual port ranges
_TPOT_PORT_CLUSTERS = {
    # T-Pot maps real honeypots to high ports; seeing many of these is a sign
    64295, 64297, 64300,  # T-Pot management
}


class SecurityPosture(BaseModule):
    """Map the target network's security tooling and posture."""

    name = "security_posture"
    module_type = "intel"
    priority = 100
    requires_root = False

    def __init__(self, bus, state, config, engine=None):
        super().__init__(bus, state, config, engine)
        self._db_path = ""
        self._conn: Optional[sqlite3.Connection] = None
        self._lock = threading.Lock()
        self._scan_thread: Optional[threading.Thread] = None
        self._detection_count = 0

        # Cache of risk assessment
        self._risk_level = "unknown"  # low, medium, high, critical
        self._active_ops_safe = True

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

    def start(self) -> None:
        if self._running:
            return

        base_dir = self.config.get("data_dir", os.path.expanduser("~/.implant"))
        self._db_path = os.path.join(base_dir, "security_posture.db")
        os.makedirs(os.path.dirname(self._db_path), exist_ok=True)

        self._conn = sqlite3.connect(self._db_path, check_same_thread=False)
        self._conn.execute("PRAGMA journal_mode=WAL")
        self._conn.execute("PRAGMA synchronous=NORMAL")
        self._conn.row_factory = sqlite3.Row
        self._conn.executescript(_SCHEMA)

        # Subscribe to host discovery for port-based detection
        self.bus.subscribe(self._on_host_discovered, "HOST_DISCOVERED")

        self._running = True
        self._pid = os.getpid()
        self._start_time = time.time()

        self._scan_thread = threading.Thread(
            target=self._periodic_scan, daemon=True,
            name="sensor-security-posture",
        )
        self._scan_thread.start()

        # Initial risk assessment from any previously stored data
        self._reassess_risk()

        self.state.set_module_status(self.name, "running", pid=self._pid)
        logger.info("SecurityPosture started — db=%s", self._db_path)

    def stop(self) -> None:
        if not self._running:
            return
        self._running = False

        self.bus.unsubscribe(self._on_host_discovered, "HOST_DISCOVERED")

        if self._conn:
            self._conn.close()
            self._conn = None

        self.state.set_module_status(self.name, "stopped")
        logger.info(
            "SecurityPosture stopped — %d tools detected, risk=%s",
            self._detection_count, self._risk_level,
        )

    def status(self) -> dict:
        tool_count = 0
        if self._conn:
            try:
                with self._lock:
                    row = self._conn.execute(
                        "SELECT COUNT(*) FROM security_tools"
                    ).fetchone()
                    tool_count = row[0] if row else 0
            except Exception:
                pass

        return {
            "running": self._running,
            "pid": self._pid,
            "uptime": time.time() - self._start_time if self._start_time else 0,
            "detected_tools": tool_count,
            "risk_level": self._risk_level,
            "active_ops_safe": self._active_ops_safe,
        }

    def configure(self, config: dict) -> None:
        self.config.update(config)

    # ------------------------------------------------------------------
    # Detection functions
    # ------------------------------------------------------------------

    def analyze_dns_query(self, query_name: str, client_ip: str = "",
                          resolved_ip: str = "") -> Optional[dict]:
        """Analyze a DNS query for security tool signatures."""
        for pattern, tool_type, tool_name, confidence, impacts in _DNS_SIGNATURES:
            if pattern.search(query_name):
                det = self._record_detection(
                    tool_type=tool_type,
                    tool_name=tool_name,
                    evidence=f"DNS: {query_name}",
                    confidence=confidence,
                    target_ip=client_ip,
                    target_hostname=query_name,
                    impacts_active=impacts,
                    details=json.dumps({
                        "query": query_name,
                        "client_ip": client_ip,
                        "resolved_ip": resolved_ip,
                    }),
                )
                return det
        return None

    def analyze_port(self, target_ip: str, port: int,
                     hostname: str = "") -> Optional[dict]:
        """Check if a port suggests security tooling."""
        if port in _PORT_SIGNATURES:
            tool_type, tool_name, confidence, impacts = _PORT_SIGNATURES[port]
            return self._record_detection(
                tool_type=tool_type,
                tool_name=tool_name,
                evidence=f"Port {port} open on {target_ip}",
                confidence=confidence,
                target_ip=target_ip,
                target_hostname=hostname,
                impacts_active=impacts,
            )

        # T-Pot cluster detection
        if port in _TPOT_PORT_CLUSTERS:
            return self._record_detection(
                tool_type="honeypot",
                tool_name="T-Pot",
                evidence=f"T-Pot management port {port} on {target_ip}",
                confidence=0.8,
                target_ip=target_ip,
                target_hostname=hostname,
                impacts_active=True,
            )

        return None

    def analyze_user_agent(self, user_agent: str, source_ip: str = "") -> Optional[dict]:
        """Check HTTP User-Agent for security scanner signatures."""
        for pattern, tool_type, tool_name, confidence in _UA_SIGNATURES:
            if pattern.search(user_agent):
                return self._record_detection(
                    tool_type=tool_type,
                    tool_name=tool_name,
                    evidence=f"User-Agent: {user_agent[:100]}",
                    confidence=confidence,
                    target_ip=source_ip,
                    impacts_active=False,
                    details=json.dumps({"user_agent": user_agent}),
                )
        return None

    def analyze_radius_traffic(self, src_ip: str, dst_ip: str,
                                dst_port: int) -> Optional[dict]:
        """Detect 802.1X/RADIUS NAC infrastructure."""
        if dst_port in (1812, 1813, 3799):
            return self._record_detection(
                tool_type="nac",
                tool_name="RADIUS/802.1X",
                evidence=f"RADIUS traffic {src_ip} -> {dst_ip}:{dst_port}",
                confidence=0.8,
                target_ip=dst_ip,
                impacts_active=True,
                details=json.dumps({
                    "src": src_ip, "dst": dst_ip, "port": dst_port
                }),
            )
        return None

    def _on_host_discovered(self, event) -> None:
        """Check discovered host ports for security tools."""
        p = event.payload
        ip = p.get("ip", "")
        hostname = p.get("hostname", "")
        ports = p.get("open_ports", [])

        for port in ports:
            self.analyze_port(ip, port, hostname)

    # ------------------------------------------------------------------
    # Database operations
    # ------------------------------------------------------------------

    def _record_detection(self, tool_type: str, tool_name: str,
                          evidence: str, confidence: float,
                          target_ip: str = "", target_hostname: str = "",
                          impacts_active: bool = True,
                          details: str = "") -> dict:
        """Record a security tool detection."""
        now = time.time()
        detection = {
            "tool_type": tool_type,
            "tool_name": tool_name,
            "evidence": evidence,
            "confidence": confidence,
            "target_ip": target_ip,
            "impacts_active_ops": impacts_active,
        }

        with self._lock:
            try:
                self._conn.execute(
                    """INSERT INTO security_tools
                       (tool_type, tool_name, evidence, confidence,
                        target_ip, target_hostname, first_seen, last_seen,
                        impacts_active_ops, details)
                       VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                       ON CONFLICT(tool_type, tool_name, target_ip) DO UPDATE SET
                           last_seen = excluded.last_seen,
                           confidence = MAX(security_tools.confidence, excluded.confidence),
                           evidence = security_tools.evidence || '; ' || excluded.evidence
                    """,
                    (tool_type, tool_name, evidence, confidence,
                     target_ip, target_hostname, now, now,
                     1 if impacts_active else 0, details),
                )
                self._conn.commit()
                self._detection_count += 1
            except Exception:
                logger.exception("Failed to record security tool detection")
                return detection

        logger.warning(
            "Security tool detected: %s (%s) at %s — confidence=%.2f, "
            "impacts_active=%s",
            tool_name, tool_type, target_ip or "unknown", confidence,
            impacts_active,
        )

        # Reassess risk after each new detection
        self._reassess_risk()

        return detection

    # ------------------------------------------------------------------
    # Risk assessment
    # ------------------------------------------------------------------

    def _reassess_risk(self) -> None:
        """Recalculate overall security risk level based on detected tools."""
        if not self._conn:
            return

        with self._lock:
            try:
                tools = self._conn.execute(
                    "SELECT tool_type, tool_name, confidence, impacts_active_ops "
                    "FROM security_tools"
                ).fetchall()
            except Exception:
                return

        if not tools:
            self._risk_level = "low"
            self._active_ops_safe = True
            return

        has_edr = any(t["tool_type"] == "edr" and t["confidence"] >= 0.7
                      for t in tools)
        has_siem = any(t["tool_type"] == "siem" and t["confidence"] >= 0.7
                       for t in tools)
        has_nac = any(t["tool_type"] == "nac" and t["confidence"] >= 0.7
                      for t in tools)
        has_honeypot = any(t["tool_type"] == "honeypot" and t["confidence"] >= 0.7
                           for t in tools)
        has_vuln_scanner = any(t["tool_type"] == "vuln_scanner"
                                and t["confidence"] >= 0.7 for t in tools)

        # Score-based risk assessment
        score = 0
        if has_edr:
            score += 40
        if has_siem:
            score += 20
        if has_nac:
            score += 25
        if has_honeypot:
            score += 30
        if has_vuln_scanner:
            score += 10

        if score >= 60:
            self._risk_level = "critical"
            self._active_ops_safe = False
        elif score >= 40:
            self._risk_level = "high"
            self._active_ops_safe = False
        elif score >= 20:
            self._risk_level = "medium"
            self._active_ops_safe = True  # with caution
        else:
            self._risk_level = "low"
            self._active_ops_safe = True

        # Persist for other modules
        self.state.set(self.name, "risk_level", self._risk_level)
        self.state.set(self.name, "active_ops_safe", json.dumps(self._active_ops_safe))
        self.state.set(self.name, "has_edr", json.dumps(has_edr))
        self.state.set(self.name, "has_siem", json.dumps(has_siem))
        self.state.set(self.name, "has_nac", json.dumps(has_nac))
        self.state.set(self.name, "has_honeypot", json.dumps(has_honeypot))

    # ------------------------------------------------------------------
    # Periodic scanning
    # ------------------------------------------------------------------

    def _periodic_scan(self) -> None:
        """Pull DNS data from other modules and analyze."""
        interval = self.config.get("posture_scan_interval", 120)
        while self._running:
            time.sleep(interval)
            try:
                self._scan_dns_data()
                self._reassess_risk()
            except Exception:
                logger.exception("Security posture scan failed")

    def _scan_dns_data(self) -> None:
        """Pull recent DNS queries from dns_logger and analyze for security tools."""
        last_ts = self.state.get(self.name, "dns_scan_last_ts")
        last_ts = float(last_ts) if last_ts else 0

        dns_json = self.state.get("dns_logger", "recent_queries")
        if not dns_json:
            return

        try:
            queries = json.loads(dns_json)
            max_ts = last_ts
            for q in queries:
                ts = q.get("timestamp", 0)
                if ts <= last_ts:
                    continue
                self.analyze_dns_query(
                    query_name=q.get("query_name", ""),
                    client_ip=q.get("client_ip", ""),
                    resolved_ip=q.get("resolved_ip", ""),
                )
                max_ts = max(max_ts, ts)

            self.state.set(self.name, "dns_scan_last_ts", str(max_ts))
        except (json.JSONDecodeError, TypeError):
            pass

    # ------------------------------------------------------------------
    # Query interface
    # ------------------------------------------------------------------

    def get_risk_level(self) -> str:
        """Return current overall risk level."""
        return self._risk_level

    def is_active_ops_safe(self) -> bool:
        """Return whether active operations are considered safe."""
        return self._active_ops_safe

    def get_detected_tools(self, tool_type: str = None) -> list:
        """Get all detected security tools."""
        with self._lock:
            if tool_type:
                rows = self._conn.execute(
                    """SELECT * FROM security_tools
                       WHERE tool_type = ?
                       ORDER BY confidence DESC""",
                    (tool_type,),
                ).fetchall()
            else:
                rows = self._conn.execute(
                    """SELECT * FROM security_tools
                       ORDER BY confidence DESC"""
                ).fetchall()

        return [dict(r) for r in rows]

    def get_posture_summary(self) -> dict:
        """Return a full security posture summary."""
        with self._lock:
            rows = self._conn.execute(
                """SELECT tool_type, COUNT(*) as count,
                          AVG(confidence) as avg_confidence,
                          MAX(impacts_active_ops) as any_impacts_active
                   FROM security_tools
                   GROUP BY tool_type
                   ORDER BY count DESC"""
            ).fetchall()

        return {
            "risk_level": self._risk_level,
            "active_ops_safe": self._active_ops_safe,
            "total_tools": sum(r[1] for r in rows),
            "by_type": {
                r[0]: {
                    "count": r[1],
                    "avg_confidence": round(r[2], 2),
                    "impacts_active": bool(r[3]),
                }
                for r in rows
            },
        }
