#!/usr/bin/env python3
"""Passive OS fingerprinting via p0f-style TCP analysis and protocol headers.

Analyzes:
  - TCP SYN/SYN-ACK: TTL, window size, DF flag, MSS, SACK, TCP timestamps
  - HTTP User-Agent headers
  - DHCP vendor class identifiers
  - SMB dialect negotiation
  - SSH version strings

Multiple signal sources produce confidence scoring per host. Results are
merged with the hosts table from host_discovery.
"""

import logging
import os
import re
import socket
import sqlite3
import struct
import threading
import time
from pathlib import Path
from typing import Optional

from modules.base import BaseModule

logger = logging.getLogger(__name__)

# p0f-style TCP signature database (built-in fallback)
# Format: (ttl_range, window_size, df, mss_range, sack, timestamps) -> (os_family, os_version)
TCP_SIG_DB = [
    # Linux signatures
    {"ttl": (64, 64), "df": True, "window": (5720, 65535), "mss": (1360, 1460),
     "sack": True, "ts": True, "os_family": "Linux", "os_version": "2.6+"},
    # Windows signatures
    {"ttl": (128, 128), "df": True, "window": (8192, 65535), "mss": (1360, 1460),
     "sack": True, "ts": False, "os_family": "Windows", "os_version": "7/10/Server"},
    {"ttl": (128, 128), "df": True, "window": (8192, 8192), "mss": (1360, 1460),
     "sack": True, "ts": False, "os_family": "Windows", "os_version": "XP/2003"},
    # macOS / iOS
    {"ttl": (64, 64), "df": True, "window": (65535, 65535), "mss": (1360, 1460),
     "sack": True, "ts": True, "os_family": "macOS", "os_version": "10.x+"},
    # FreeBSD
    {"ttl": (64, 64), "df": True, "window": (65535, 65535), "mss": (1360, 1460),
     "sack": True, "ts": True, "os_family": "FreeBSD", "os_version": ""},
    # Cisco IOS
    {"ttl": (255, 255), "df": False, "window": (4128, 4128), "mss": (536, 536),
     "sack": False, "ts": False, "os_family": "Cisco", "os_version": "IOS"},
    # Solaris
    {"ttl": (255, 255), "df": False, "window": (49232, 49232), "mss": (1360, 1460),
     "sack": False, "ts": True, "os_family": "Solaris", "os_version": "10+"},
]

# HTTP User-Agent patterns
UA_PATTERNS = [
    (re.compile(r"Windows NT 10\.0"), "Windows", "10/11"),
    (re.compile(r"Windows NT 6\.3"), "Windows", "8.1"),
    (re.compile(r"Windows NT 6\.2"), "Windows", "8"),
    (re.compile(r"Windows NT 6\.1"), "Windows", "7"),
    (re.compile(r"Windows NT 5\.1"), "Windows", "XP"),
    (re.compile(r"Mac OS X (\d+[._]\d+)"), "macOS", ""),
    (re.compile(r"Linux"), "Linux", ""),
    (re.compile(r"Ubuntu"), "Linux", "Ubuntu"),
    (re.compile(r"Android (\d+)"), "Android", ""),
    (re.compile(r"iPhone OS (\d+)"), "iOS", ""),
    (re.compile(r"iPad.*OS (\d+)"), "iPadOS", ""),
    (re.compile(r"CrOS"), "ChromeOS", ""),
]


class OSFingerprint(BaseModule):
    """Passive OS fingerprinting via TCP/protocol analysis."""

    name = "os_fingerprint"
    module_type = "passive"
    priority = 150
    requires_root = True
    requires_capture_bus = True

    BATCH_SIZE = 200
    FLUSH_INTERVAL = 120

    def __init__(self, bus, state, config, engine=None):
        super().__init__(bus, state, config, engine)
        self._capture_bus = None
        self._sub_queue = None
        self._reader_thread: Optional[threading.Thread] = None
        self._flusher_thread: Optional[threading.Thread] = None
        self._buffer = []
        self._buffer_lock = threading.Lock()
        self._db_path = ""
        self._db_conn: Optional[sqlite3.Connection] = None
        self._total_fingerprints = 0
        # Per-IP confidence tracker: ip -> {os_family: {method: confidence}}
        self._ip_os_scores = {}
        self._scores_lock = threading.Lock()
        # External signature database
        self._os_sigs = []

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

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

        self._capture_bus = self.config.get("capture_bus")
        if not self._capture_bus:
            logger.error("OSFingerprint requires capture_bus in config")
            return

        base_dir = self.config.get("data_dir", os.path.expanduser("~/.implant"))
        self._db_path = os.path.join(base_dir, "os_fingerprints.db")
        Path(os.path.dirname(self._db_path)).mkdir(parents=True, exist_ok=True)
        self._init_db()

        # Load external OS signatures
        sigs_path = self.config.get("os_sigs_db", "")
        if sigs_path and os.path.isfile(sigs_path):
            self._load_os_sigs(sigs_path)

        # Subscribe to SYN packets + HTTP/SMB/SSH for protocol fingerprinting
        # Using a broad filter; the module parses selectively
        bpf = "tcp[tcpflags] & (tcp-syn) != 0"
        self._sub_queue = self._capture_bus.subscribe(
            name=self.name, bpf_filter=bpf, queue_depth=8000
        )

        # Second subscription for application-layer fingerprints
        # We use a broader capture for HTTP/SMB/SSH — but since capture_bus
        # only supports one subscription per module name, we parse app-layer
        # from the same stream by also checking non-SYN packets that match
        # We subscribe with an empty filter to get everything and filter in code
        # Actually, let's keep SYN filter and add a second subscriber
        self._capture_bus.unsubscribe(self.name)
        self._sub_queue = self._capture_bus.subscribe(
            name=self.name,
            bpf_filter="tcp",  # All TCP — we filter SYN and app-layer in code
            queue_depth=10000,
        )

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

        self._reader_thread = threading.Thread(
            target=self._read_packets, daemon=True, name="sensor-osfp-reader"
        )
        self._reader_thread.start()

        self._flusher_thread = threading.Thread(
            target=self._flush_loop, daemon=True, name="sensor-osfp-flusher"
        )
        self._flusher_thread.start()

        self.state.set_module_status(self.name, "running", pid=os.getpid())
        logger.info("OSFingerprint started — passive TCP/protocol analysis")

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

        if self._capture_bus:
            self._capture_bus.unsubscribe(self.name)

        self._flush_buffer()
        if self._db_conn:
            self._db_conn.close()
            self._db_conn = None

        self.state.set_module_status(self.name, "stopped")
        logger.info("OSFingerprint stopped — %d fingerprints collected", self._total_fingerprints)

    def status(self) -> dict:
        with self._scores_lock:
            unique_hosts = len(self._ip_os_scores)
        return {
            "running": self._running,
            "pid": self._pid,
            "uptime": time.time() - self._start_time if self._start_time else 0,
            "total_fingerprints": self._total_fingerprints,
            "unique_hosts": unique_hosts,
            "buffer_size": len(self._buffer),
        }

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

    # ------------------------------------------------------------------
    # Database
    # ------------------------------------------------------------------

    def _init_db(self) -> None:
        self._db_conn = sqlite3.connect(self._db_path, check_same_thread=False)
        self._db_conn.execute("PRAGMA journal_mode=WAL")
        self._db_conn.execute("PRAGMA synchronous=NORMAL")
        self._db_conn.executescript("""
            CREATE TABLE IF NOT EXISTS os_fingerprints (
                id          INTEGER PRIMARY KEY AUTOINCREMENT,
                ip          TEXT NOT NULL,
                method      TEXT NOT NULL,
                signature   TEXT,
                os_family   TEXT,
                os_version  TEXT,
                confidence  REAL,
                timestamp   REAL NOT NULL
            );
            CREATE INDEX IF NOT EXISTS idx_osfp_ip ON os_fingerprints(ip);
            CREATE INDEX IF NOT EXISTS idx_osfp_family ON os_fingerprints(os_family);
        """)
        self._db_conn.commit()

    def _load_os_sigs(self, path: str) -> None:
        """Load OS signature database from SQLite."""
        try:
            conn = sqlite3.connect(path)
            rows = conn.execute(
                "SELECT ttl, window, df, mss, sack, timestamps, os_family, os_version "
                "FROM signatures"
            ).fetchall()
            for row in rows:
                self._os_sigs.append({
                    "ttl": (row[0], row[0]),
                    "window": (row[1], row[1]),
                    "df": bool(row[2]),
                    "mss": (row[3], row[3]),
                    "sack": bool(row[4]),
                    "ts": bool(row[5]),
                    "os_family": row[6],
                    "os_version": row[7] or "",
                })
            conn.close()
            logger.info("Loaded %d OS signatures from %s", len(self._os_sigs), path)
        except Exception:
            logger.warning("Failed to load OS signatures from %s", path)

    # ------------------------------------------------------------------
    # Packet reader
    # ------------------------------------------------------------------

    def _read_packets(self) -> None:
        while self._running:
            result = self._sub_queue.get(timeout=1.0)
            if result is None:
                continue
            ts, raw_packet = result
            try:
                self._parse_packet(ts, raw_packet)
            except Exception:
                pass

    def _parse_packet(self, ts: float, raw: bytes) -> None:
        if len(raw) < 14:
            return

        eth_type = struct.unpack("!H", raw[12:14])[0]
        ip_offset = 14
        if eth_type == 0x8100:
            if len(raw) < 18:
                return
            eth_type = struct.unpack("!H", raw[16:18])[0]
            ip_offset = 18

        if eth_type != 0x0800:
            return

        if len(raw) < ip_offset + 20:
            return

        ip_hdr = raw[ip_offset:]
        ihl = (ip_hdr[0] & 0x0F) * 4
        ip_proto = ip_hdr[9]
        if ip_proto != 6:
            return

        src_ip = socket.inet_ntoa(ip_hdr[12:16])
        dst_ip = socket.inet_ntoa(ip_hdr[16:20])
        ip_ttl = ip_hdr[8]
        ip_flags = struct.unpack("!H", ip_hdr[6:8])[0]
        df_flag = bool(ip_flags & 0x4000)

        tcp_offset = ip_offset + ihl
        if len(raw) < tcp_offset + 20:
            return

        tcp_hdr = raw[tcp_offset:]
        src_port, dst_port = struct.unpack("!HH", tcp_hdr[:4])
        tcp_flags = tcp_hdr[13]
        window = struct.unpack("!H", tcp_hdr[14:16])[0]
        tcp_hdr_len = ((tcp_hdr[12] >> 4) & 0xF) * 4

        syn = bool(tcp_flags & 0x02)
        ack = bool(tcp_flags & 0x10)

        # SYN or SYN-ACK — do TCP fingerprinting
        if syn:
            self._fingerprint_tcp_syn(
                ts, src_ip, ip_ttl, df_flag, window,
                tcp_hdr[:tcp_hdr_len], syn_ack=ack,
            )

        # Application-layer fingerprinting for non-SYN packets with payload
        if not syn:
            payload = raw[tcp_offset + tcp_hdr_len:]
            if payload:
                if src_port == 80 or dst_port == 80 or src_port == 8080 or dst_port == 8080:
                    self._fingerprint_http(ts, src_ip, dst_ip, src_port, payload)
                elif src_port == 22 or dst_port == 22:
                    self._fingerprint_ssh(ts, src_ip, src_port, payload)
                elif src_port == 445 or dst_port == 445:
                    self._fingerprint_smb(ts, src_ip, src_port, payload)

    # ------------------------------------------------------------------
    # TCP SYN fingerprinting (p0f-style)
    # ------------------------------------------------------------------

    def _fingerprint_tcp_syn(self, ts: float, ip: str, ttl: int,
                             df: bool, window: int, tcp_header: bytes,
                             syn_ack: bool = False) -> None:
        """Analyze TCP SYN/SYN-ACK options for OS fingerprinting."""
        # Parse TCP options
        mss = 0
        has_sack = False
        has_timestamps = False
        wscale = 0

        if len(tcp_header) > 20:
            opt_offset = 20
            while opt_offset < len(tcp_header):
                opt_kind = tcp_header[opt_offset]
                if opt_kind == 0:  # End of options
                    break
                if opt_kind == 1:  # NOP
                    opt_offset += 1
                    continue
                if opt_offset + 1 >= len(tcp_header):
                    break
                opt_len = tcp_header[opt_offset + 1]
                if opt_len < 2 or opt_offset + opt_len > len(tcp_header):
                    break

                if opt_kind == 2 and opt_len == 4:  # MSS
                    mss = struct.unpack("!H", tcp_header[opt_offset + 2:opt_offset + 4])[0]
                elif opt_kind == 3 and opt_len == 3:  # Window Scale
                    wscale = tcp_header[opt_offset + 2]
                elif opt_kind == 4:  # SACK Permitted
                    has_sack = True
                elif opt_kind == 8:  # Timestamps
                    has_timestamps = True

                opt_offset += opt_len

        # Normalize TTL to nearest power-of-2 boundary
        initial_ttl = self._normalize_ttl(ttl)

        # Build signature string
        sig = f"ttl:{initial_ttl}:win:{window}:mss:{mss}:df:{int(df)}:sack:{int(has_sack)}:ts:{int(has_timestamps)}:wscale:{wscale}"

        # Match against signature database
        os_family, os_version, confidence = self._match_tcp_signature(
            initial_ttl, window, df, mss, has_sack, has_timestamps,
        )

        if os_family:
            method = "tcp_syn_ack" if syn_ack else "tcp_syn"
            self._record_fingerprint(ts, ip, method, sig, os_family, os_version, confidence)

    def _match_tcp_signature(self, ttl: int, window: int, df: bool,
                             mss: int, sack: bool, timestamps: bool) -> tuple:
        """Match TCP parameters against signature database."""
        best_match = ("", "", 0.0)
        best_score = 0

        sig_sources = self._os_sigs if self._os_sigs else TCP_SIG_DB

        for sig in sig_sources:
            score = 0
            total = 6

            ttl_lo, ttl_hi = sig["ttl"]
            if ttl_lo <= ttl <= ttl_hi:
                score += 2  # TTL is weighted higher

            win_lo, win_hi = sig["window"]
            if win_lo <= window <= win_hi:
                score += 1

            if sig["df"] == df:
                score += 1

            mss_lo, mss_hi = sig["mss"]
            if mss_lo <= mss <= mss_hi:
                score += 1

            if sig["sack"] == sack:
                score += 0.5

            if sig["ts"] == timestamps:
                score += 0.5

            if score > best_score:
                best_score = score
                confidence = score / total
                best_match = (sig["os_family"], sig["os_version"], confidence)

        return best_match

    @staticmethod
    def _normalize_ttl(ttl: int) -> int:
        """Estimate initial TTL from observed TTL."""
        if ttl <= 32:
            return 32
        elif ttl <= 64:
            return 64
        elif ttl <= 128:
            return 128
        else:
            return 255

    # ------------------------------------------------------------------
    # Application-layer fingerprinting
    # ------------------------------------------------------------------

    def _fingerprint_http(self, ts: float, src_ip: str, dst_ip: str,
                          src_port: int, payload: bytes) -> None:
        """Extract OS info from HTTP User-Agent headers."""
        try:
            text = payload[:4096].decode("utf-8", errors="ignore")
        except Exception:
            return

        ua_match = re.search(r"User-Agent:\s*(.+?)(?:\r\n|\n)", text, re.IGNORECASE)
        if not ua_match:
            return

        ua = ua_match.group(1).strip()
        # The User-Agent is from the client (request sender)
        # If src_port is ephemeral (>1024), this is a client
        if src_port > 1024:
            fp_ip = src_ip
        else:
            fp_ip = dst_ip

        for pattern, os_family, os_version in UA_PATTERNS:
            m = pattern.search(ua)
            if m:
                version = os_version
                if not version and m.lastindex:
                    version = m.group(1).replace("_", ".")
                self._record_fingerprint(
                    ts, fp_ip, "http_ua", ua[:200],
                    os_family, version, 0.7,
                )
                break

    def _fingerprint_ssh(self, ts: float, ip: str, src_port: int,
                         payload: bytes) -> None:
        """Extract OS info from SSH version string."""
        try:
            text = payload[:256].decode("ascii", errors="ignore")
        except Exception:
            return

        if not text.startswith("SSH-"):
            return

        # SSH version string: SSH-2.0-OpenSSH_8.9p1 Ubuntu-3ubuntu0.1
        version_str = text.strip()
        # The SSH server sends its banner — src_port should be 22
        fp_ip = ip if src_port == 22 else ip

        os_family = ""
        os_version = ""

        if "Ubuntu" in version_str:
            os_family = "Linux"
            os_version = "Ubuntu"
        elif "Debian" in version_str:
            os_family = "Linux"
            os_version = "Debian"
        elif "FreeBSD" in version_str:
            os_family = "FreeBSD"
        elif "OpenSSH" in version_str:
            # Generic OpenSSH — likely Linux or BSD
            os_family = "Linux/BSD"

        if os_family:
            self._record_fingerprint(
                ts, fp_ip, "ssh_banner", version_str[:200],
                os_family, os_version, 0.8,
            )

    def _fingerprint_smb(self, ts: float, ip: str, src_port: int,
                         payload: bytes) -> None:
        """Extract OS info from SMB negotiate response."""
        # SMB2 header: 0xFE 'S' 'M' 'B'
        if len(payload) < 68:
            return

        # Look for SMB2 header
        smb2_offset = payload.find(b"\xfeSMB")
        if smb2_offset < 0:
            # Try SMB1
            smb1_offset = payload.find(b"\xffSMB")
            if smb1_offset >= 0 and src_port == 445:
                self._record_fingerprint(
                    ts, ip, "smb_dialect", "SMB1",
                    "Windows", "XP/2003 or Samba", 0.5,
                )
            return

        # SMB2 negotiate response from server (src_port 445)
        if src_port != 445:
            return

        smb2_data = payload[smb2_offset:]
        if len(smb2_data) < 68:
            return

        # SMB2 header is 64 bytes, then negotiate response
        # Dialect at offset 4-5 of negotiate response (after 64-byte header)
        neg_response = smb2_data[64:]
        if len(neg_response) < 6:
            return

        # struct_size(2) + security_mode(2) + dialect_revision(2)
        dialect = struct.unpack("<H", neg_response[4:6])[0]

        dialect_map = {
            0x0202: ("Windows", "Vista/2008 or Samba"),
            0x0210: ("Windows", "7/2008R2"),
            0x0300: ("Windows", "8/2012"),
            0x0302: ("Windows", "8.1/2012R2"),
            0x0311: ("Windows", "10/2016+"),
        }

        if dialect in dialect_map:
            os_family, os_version = dialect_map[dialect]
            self._record_fingerprint(
                ts, ip, "smb_dialect", f"SMB{dialect:#06x}",
                os_family, os_version, 0.75,
            )

    # ------------------------------------------------------------------
    # Record management
    # ------------------------------------------------------------------

    def _record_fingerprint(self, ts: float, ip: str, method: str,
                            signature: str, os_family: str,
                            os_version: str, confidence: float) -> None:
        """Record an OS fingerprint observation."""
        self._total_fingerprints += 1

        # Update per-IP confidence scoring
        with self._scores_lock:
            if ip not in self._ip_os_scores:
                self._ip_os_scores[ip] = {}
            scores = self._ip_os_scores[ip]
            if os_family not in scores:
                scores[os_family] = {}
            scores[os_family][method] = confidence

        record = (ip, method, signature[:500], os_family, os_version, confidence, ts)
        with self._buffer_lock:
            self._buffer.append(record)
            if len(self._buffer) >= self.BATCH_SIZE:
                self._flush_buffer()

    def get_best_guess(self, ip: str) -> dict:
        """Get the highest-confidence OS guess for an IP."""
        with self._scores_lock:
            scores = self._ip_os_scores.get(ip, {})

        if not scores:
            return {"os_family": "", "confidence": 0.0}

        # Aggregate confidence per OS family
        best_family = ""
        best_conf = 0.0
        for os_family, methods in scores.items():
            # Average confidence across methods, boosted by method count
            avg_conf = sum(methods.values()) / len(methods)
            method_bonus = min(len(methods) * 0.1, 0.3)  # Up to 30% bonus
            total_conf = min(avg_conf + method_bonus, 1.0)
            if total_conf > best_conf:
                best_conf = total_conf
                best_family = os_family

        return {"os_family": best_family, "confidence": round(best_conf, 2)}

    # ------------------------------------------------------------------
    # Buffer flush
    # ------------------------------------------------------------------

    def _flush_loop(self) -> None:
        while self._running:
            time.sleep(self.FLUSH_INTERVAL)
            try:
                self._flush_buffer()
            except Exception:
                logger.exception("OS fingerprint flush error")

    def _flush_buffer(self) -> None:
        with self._buffer_lock:
            batch = list(self._buffer)
            self._buffer.clear()

        if not batch or not self._db_conn:
            return

        try:
            self._db_conn.executemany(
                "INSERT INTO os_fingerprints "
                "(ip, method, signature, os_family, os_version, confidence, timestamp) "
                "VALUES (?, ?, ?, ?, ?, ?, ?)",
                batch,
            )
            self._db_conn.commit()
        except Exception:
            logger.exception("Failed to flush %d OS fingerprints", len(batch))
