#!/usr/bin/env python3
"""Passive host discovery via broadcast/multicast protocol observation.

Parses:
  - ARP requests/replies (IP-MAC mapping)
  - DHCP request/ACK (hostname, vendor class, option 55 fingerprinting)
  - mDNS/Bonjour (hostnames, services)
  - NetBIOS name queries (NBNS port 137)
  - SSDP/UPnP announcements (port 1900)
  - LLMNR queries (port 5355)

Publishes HOST_DISCOVERED events. Correlates multiple signals per host
for IP + MAC + hostname + vendor (OUI) + OS guess.
"""

import logging
import os
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__)


class HostDiscovery(BaseModule):
    """Build host inventory from passive network observation."""

    name = "host_discovery"
    module_type = "passive"
    priority = 100
    requires_root = True
    requires_capture_bus = True

    BATCH_SIZE = 200
    FLUSH_INTERVAL = 60

    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
        # In-memory host table: ip -> host_info dict
        self._hosts = {}
        self._hosts_lock = threading.Lock()
        self._pending_updates = []
        self._update_lock = threading.Lock()
        self._db_path = ""
        self._db_conn: Optional[sqlite3.Connection] = None
        self._total_hosts = 0
        # OUI lookup cache: first 3 bytes hex -> vendor
        self._oui_cache = {}
        # DHCP fingerprint cache: option55 -> os_guess
        self._dhcp_fp_cache = {}

    # ------------------------------------------------------------------
    # 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("HostDiscovery 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, "hosts.db")
        Path(os.path.dirname(self._db_path)).mkdir(parents=True, exist_ok=True)
        self._init_db()

        # Load OUI database
        oui_path = self.config.get("oui_db", "")
        if oui_path and os.path.isfile(oui_path):
            self._load_oui_db(oui_path)

        # Load DHCP fingerprint database
        dhcp_fp_path = self.config.get("dhcp_fingerprints_db", "")
        if dhcp_fp_path and os.path.isfile(dhcp_fp_path):
            self._load_dhcp_fingerprints(dhcp_fp_path)

        # Load BPF filter
        bpf_path = self.config.get("bpf_filter_path", "")
        bpf_filter = ""
        if bpf_path and os.path.isfile(bpf_path):
            with open(bpf_path) as f:
                lines = [l.strip() for l in f if l.strip() and not l.startswith("#")]
                bpf_filter = " ".join(lines)
        if not bpf_filter:
            bpf_filter = "arp or port 67 or port 68 or port 5353 or port 1900 or port 5355 or port 137 or port 138"

        self._sub_queue = self._capture_bus.subscribe(
            name=self.name, bpf_filter=bpf_filter, 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-hostdisc-reader"
        )
        self._reader_thread.start()

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

        self.state.set_module_status(self.name, "running", pid=os.getpid())
        logger.info("HostDiscovery started")

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

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

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

        self.state.set_module_status(self.name, "stopped")
        logger.info("HostDiscovery stopped — %d hosts discovered", self._total_hosts)

    def status(self) -> dict:
        with self._hosts_lock:
            host_count = len(self._hosts)
        return {
            "running": self._running,
            "pid": self._pid,
            "uptime": time.time() - self._start_time if self._start_time else 0,
            "total_hosts": host_count,
            "oui_entries": len(self._oui_cache),
            "dhcp_fingerprints": len(self._dhcp_fp_cache),
        }

    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 hosts (
                id               INTEGER PRIMARY KEY AUTOINCREMENT,
                ip               TEXT NOT NULL,
                mac              TEXT,
                hostname         TEXT,
                vendor           TEXT,
                os_guess         TEXT,
                dhcp_fingerprint TEXT,
                first_seen       REAL NOT NULL,
                last_seen        REAL NOT NULL,
                source           TEXT
            );
            CREATE UNIQUE INDEX IF NOT EXISTS idx_host_ip ON hosts(ip);
            CREATE INDEX IF NOT EXISTS idx_host_mac ON hosts(mac);
            CREATE INDEX IF NOT EXISTS idx_host_hostname ON hosts(hostname);
        """)
        self._db_conn.commit()

    def _load_oui_db(self, path: str) -> None:
        """Load OUI vendor database. Expected format: 'AA:BB:CC<tab>Vendor Name' per line."""
        try:
            conn = sqlite3.connect(path)
            rows = conn.execute("SELECT oui, vendor FROM oui").fetchall()
            for oui, vendor in rows:
                self._oui_cache[oui.upper().replace(":", "").replace("-", "")] = vendor
            conn.close()
            logger.info("Loaded %d OUI entries", len(self._oui_cache))
        except Exception:
            # Try plain text format
            try:
                with open(path) as f:
                    for line in f:
                        line = line.strip()
                        if not line or line.startswith("#"):
                            continue
                        parts = line.split("\t", 1)
                        if len(parts) == 2:
                            oui = parts[0].upper().replace(":", "").replace("-", "")
                            self._oui_cache[oui] = parts[1]
                logger.info("Loaded %d OUI entries from text", len(self._oui_cache))
            except Exception:
                logger.warning("Failed to load OUI database from %s", path)

    def _load_dhcp_fingerprints(self, path: str) -> None:
        """Load DHCP option 55 fingerprint database."""
        try:
            conn = sqlite3.connect(path)
            rows = conn.execute("SELECT fingerprint, os_name FROM fingerprints").fetchall()
            for fp, os_name in rows:
                self._dhcp_fp_cache[fp] = os_name
            conn.close()
            logger.info("Loaded %d DHCP fingerprints", len(self._dhcp_fp_cache))
        except Exception:
            try:
                with open(path) as f:
                    for line in f:
                        line = line.strip()
                        if not line or line.startswith("#"):
                            continue
                        parts = line.split("\t", 1)
                        if len(parts) == 2:
                            self._dhcp_fp_cache[parts[0]] = parts[1]
                logger.info("Loaded %d DHCP fingerprints from text", len(self._dhcp_fp_cache))
            except Exception:
                logger.warning("Failed to load DHCP fingerprints from %s", path)

    # ------------------------------------------------------------------
    # OUI lookup
    # ------------------------------------------------------------------

    def _lookup_oui(self, mac: str) -> str:
        """Look up vendor from MAC address OUI (first 3 octets)."""
        oui = mac.upper().replace(":", "").replace("-", "")[:6]
        return self._oui_cache.get(oui, "")

    # ------------------------------------------------------------------
    # Host update
    # ------------------------------------------------------------------

    def _update_host(self, ts: float, ip: str, mac: str = "", hostname: str = "",
                     vendor: str = "", os_guess: str = "",
                     dhcp_fingerprint: str = "", source: str = "") -> None:
        """Update or create host entry. Merges new data with existing."""
        if not ip or ip == "0.0.0.0" or ip.startswith("255."):
            return

        is_new = False
        with self._hosts_lock:
            if ip not in self._hosts:
                self._hosts[ip] = {
                    "ip": ip, "mac": "", "hostname": "", "vendor": "",
                    "os_guess": "", "dhcp_fingerprint": "",
                    "first_seen": ts, "last_seen": ts, "source": source,
                }
                is_new = True

            host = self._hosts[ip]
            host["last_seen"] = ts

            if mac and not host["mac"]:
                host["mac"] = mac
                if not vendor:
                    vendor = self._lookup_oui(mac)
            if hostname and not host["hostname"]:
                host["hostname"] = hostname
            if vendor and not host["vendor"]:
                host["vendor"] = vendor
            if os_guess and not host["os_guess"]:
                host["os_guess"] = os_guess
            if dhcp_fingerprint and not host["dhcp_fingerprint"]:
                host["dhcp_fingerprint"] = dhcp_fingerprint
            if source:
                existing = host.get("source", "")
                if source not in existing:
                    host["source"] = f"{existing},{source}" if existing else source

        if is_new:
            self._total_hosts += 1
            self.bus.emit("HOST_DISCOVERED", {
                "ip": ip, "mac": mac, "hostname": hostname,
                "vendor": vendor, "source": source,
            }, source_module=self.name)

    # ------------------------------------------------------------------
    # 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:
        """Route packet to appropriate protocol parser."""
        if len(raw) < 14:
            return

        eth_type = struct.unpack("!H", raw[12:14])[0]
        src_mac = ":".join(f"{b:02x}" for b in raw[6:12])

        if eth_type == 0x0806:  # ARP
            self._parse_arp(ts, raw, src_mac)
            return

        ip_offset = 14
        if eth_type == 0x8100:  # VLAN
            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]
        src_ip = socket.inet_ntoa(ip_hdr[12:16])
        dst_ip = socket.inet_ntoa(ip_hdr[16:20])

        if ip_proto == 17:  # UDP
            udp_offset = ip_offset + ihl
            if len(raw) < udp_offset + 8:
                return
            src_port, dst_port = struct.unpack("!HH", raw[udp_offset:udp_offset + 4])
            payload = raw[udp_offset + 8:]

            if dst_port == 67 or dst_port == 68:
                self._parse_dhcp(ts, payload, src_mac)
            elif dst_port == 5353 or src_port == 5353:
                self._parse_mdns(ts, src_ip, src_mac, payload)
            elif dst_port == 137 or src_port == 137:
                self._parse_nbns(ts, src_ip, src_mac, payload)
            elif dst_port == 1900:
                self._parse_ssdp(ts, src_ip, src_mac, payload)
            elif dst_port == 5355 or src_port == 5355:
                self._parse_llmnr(ts, src_ip, src_mac, payload)

            # Register any UDP source
            self._update_host(ts, src_ip, mac=src_mac, source="traffic")

    # ------------------------------------------------------------------
    # Protocol parsers
    # ------------------------------------------------------------------

    def _parse_arp(self, ts: float, raw: bytes, eth_src_mac: str) -> None:
        """Parse ARP request/reply for IP-MAC mapping."""
        if len(raw) < 42:  # 14 eth + 28 ARP
            return

        arp_data = raw[14:]
        # ARP: htype(2) ptype(2) hlen(1) plen(1) oper(2) sha(6) spa(4) tha(6) tpa(4)
        htype, ptype, hlen, plen, oper = struct.unpack("!HHBBH", arp_data[:8])

        if htype != 1 or ptype != 0x0800:
            return

        sender_mac = ":".join(f"{b:02x}" for b in arp_data[8:14])
        sender_ip = socket.inet_ntoa(arp_data[14:18])
        target_ip = socket.inet_ntoa(arp_data[24:28])

        if sender_ip != "0.0.0.0":
            self._update_host(ts, sender_ip, mac=sender_mac, source="arp")

    def _parse_dhcp(self, ts: float, payload: bytes, src_mac: str) -> None:
        """Parse DHCP request/ACK for hostname, vendor class, option 55."""
        if len(payload) < 240:
            return

        # DHCP: op(1) htype(1) hlen(1) hops(1) xid(4) secs(2) flags(2)
        # ciaddr(4) yiaddr(4) siaddr(4) giaddr(4) chaddr(16) ...
        op = payload[0]
        yiaddr = socket.inet_ntoa(payload[16:20])
        chaddr = ":".join(f"{b:02x}" for b in payload[28:34])

        # Parse DHCP options (start at offset 240, after magic cookie)
        if payload[236:240] != b"\x63\x82\x53\x63":
            return

        hostname = ""
        vendor_class = ""
        option_55 = ""
        assigned_ip = yiaddr if yiaddr != "0.0.0.0" else ""
        msg_type = 0

        offset = 240
        while offset < len(payload):
            opt = payload[offset]
            if opt == 255:  # End
                break
            if opt == 0:  # Pad
                offset += 1
                continue

            if offset + 1 >= len(payload):
                break
            opt_len = payload[offset + 1]
            offset += 2

            if offset + opt_len > len(payload):
                break

            opt_data = payload[offset:offset + opt_len]

            if opt == 53 and opt_len == 1:  # Message Type
                msg_type = opt_data[0]
            elif opt == 12:  # Hostname
                hostname = opt_data.decode("ascii", errors="replace").rstrip("\x00")
            elif opt == 60:  # Vendor Class
                vendor_class = opt_data.decode("ascii", errors="replace").rstrip("\x00")
            elif opt == 55:  # Parameter Request List
                option_55 = ",".join(str(b) for b in opt_data)
            elif opt == 50 and opt_len == 4:  # Requested IP
                if not assigned_ip:
                    assigned_ip = socket.inet_ntoa(opt_data)

            offset += opt_len

        # DHCP fingerprint lookup
        os_guess = ""
        if option_55:
            os_guess = self._dhcp_fp_cache.get(option_55, "")

        if assigned_ip:
            self._update_host(
                ts, assigned_ip, mac=chaddr, hostname=hostname,
                vendor=vendor_class, os_guess=os_guess,
                dhcp_fingerprint=option_55, source="dhcp",
            )
        elif chaddr:
            # No IP yet, but we can log the MAC
            pass

    def _parse_mdns(self, ts: float, src_ip: str, src_mac: str,
                    payload: bytes) -> None:
        """Parse mDNS for hostname discovery."""
        if len(payload) < 12:
            return

        # mDNS uses standard DNS format on port 5353
        _, flags, qdcount, ancount, _, _ = struct.unpack("!HHHHHH", payload[:12])
        is_response = bool(flags & 0x8000)

        offset = 12

        # Parse questions
        for _ in range(qdcount):
            name, offset = self._read_dns_name(payload, offset)
            if offset + 4 > len(payload):
                return
            offset += 4  # qtype + qclass

        # Parse answers (for responses)
        if is_response:
            for _ in range(ancount):
                if offset >= len(payload):
                    break
                name, offset = self._read_dns_name(payload, offset)
                if offset + 10 > len(payload):
                    break
                rtype, _, ttl, rdlength = struct.unpack("!HHIH", payload[offset:offset + 10])
                offset += 10
                if offset + rdlength > len(payload):
                    break

                # Extract hostname from .local names
                if name and name.endswith(".local"):
                    hostname = name.rsplit(".local", 1)[0]
                    if hostname:
                        self._update_host(ts, src_ip, mac=src_mac,
                                          hostname=hostname, source="mdns")

                offset += rdlength
        else:
            # For queries, the querier's presence is noted
            self._update_host(ts, src_ip, mac=src_mac, source="mdns")

    def _parse_nbns(self, ts: float, src_ip: str, src_mac: str,
                    payload: bytes) -> None:
        """Parse NetBIOS Name Service queries/responses."""
        if len(payload) < 12:
            return

        _, flags, qdcount, ancount, _, _ = struct.unpack("!HHHHHH", payload[:12])
        is_response = bool(flags & 0x8000)

        offset = 12

        # Parse question names
        for _ in range(qdcount):
            if offset >= len(payload):
                break
            nbname, offset = self._read_nbns_name(payload, offset)
            if offset + 4 > len(payload):
                break
            offset += 4  # qtype + qclass

            if nbname:
                self._update_host(ts, src_ip, mac=src_mac,
                                  hostname=nbname, source="nbns")

        # Parse answer names
        if is_response:
            for _ in range(ancount):
                if offset >= len(payload):
                    break
                nbname, offset = self._read_nbns_name(payload, offset)
                if offset + 10 > len(payload):
                    break
                _, _, _, rdlength = struct.unpack("!HHIH", payload[offset:offset + 10])
                offset += 10

                if nbname:
                    self._update_host(ts, src_ip, mac=src_mac,
                                      hostname=nbname, source="nbns")

                if offset + rdlength > len(payload):
                    break
                offset += rdlength

    @staticmethod
    def _read_nbns_name(data: bytes, offset: int) -> tuple:
        """Decode a NetBIOS encoded name. Returns (name, new_offset)."""
        if offset >= len(data):
            return "", offset

        length = data[offset]
        offset += 1

        if length != 32:
            # Skip non-standard length
            return "", offset + length + 1  # +1 for trailing null length byte

        if offset + 32 > len(data):
            return "", offset

        encoded = data[offset:offset + 32]
        offset += 32

        # Skip trailing length byte
        if offset < len(data):
            offset += 1

        # Decode: each pair of bytes encodes one character
        name_chars = []
        for i in range(0, 32, 2):
            ch = ((encoded[i] - ord('A')) << 4) | (encoded[i + 1] - ord('A'))
            if 32 <= ch < 127:
                name_chars.append(chr(ch))
        name = "".join(name_chars).rstrip()

        return name, offset

    def _parse_ssdp(self, ts: float, src_ip: str, src_mac: str,
                    payload: bytes) -> None:
        """Parse SSDP/UPnP announcements for device info."""
        try:
            text = payload.decode("utf-8", errors="ignore")
        except Exception:
            return

        # SSDP uses HTTP-like headers
        server = ""
        usn = ""
        for line in text.split("\r\n"):
            lower = line.lower()
            if lower.startswith("server:"):
                server = line.split(":", 1)[1].strip()
            elif lower.startswith("usn:"):
                usn = line.split(":", 1)[1].strip()

        vendor = server if server else ""
        hostname = ""
        if usn:
            # USN often contains device UUID
            pass

        self._update_host(ts, src_ip, mac=src_mac, vendor=vendor, source="ssdp")

    def _parse_llmnr(self, ts: float, src_ip: str, src_mac: str,
                     payload: bytes) -> None:
        """Parse LLMNR queries for hostname discovery."""
        if len(payload) < 12:
            return

        # LLMNR uses DNS message format
        _, flags, qdcount, ancount, _, _ = struct.unpack("!HHHHHH", payload[:12])
        is_response = bool(flags & 0x8000)

        offset = 12
        for _ in range(qdcount):
            name, offset = self._read_dns_name(payload, offset)
            if offset + 4 > len(payload):
                break
            offset += 4

            if name:
                if is_response:
                    self._update_host(ts, src_ip, mac=src_mac,
                                      hostname=name, source="llmnr")
                else:
                    # Querier looking for this name
                    self._update_host(ts, src_ip, mac=src_mac, source="llmnr")

    @staticmethod
    def _read_dns_name(data: bytes, offset: int) -> tuple:
        """Read a DNS-format name. Returns (name, new_offset)."""
        parts = []
        jumped = False
        saved_offset = offset
        jumps = 0

        while offset < len(data):
            length = data[offset]
            if length == 0:
                offset += 1
                break
            if (length & 0xC0) == 0xC0:
                if offset + 1 >= len(data):
                    break
                ptr = struct.unpack("!H", data[offset:offset + 2])[0] & 0x3FFF
                if not jumped:
                    saved_offset = offset + 2
                    jumped = True
                offset = ptr
                jumps += 1
                if jumps > 10:
                    break
                continue
            offset += 1
            if offset + length > len(data):
                break
            parts.append(data[offset:offset + length].decode("utf-8", errors="replace"))
            offset += length

        name = ".".join(parts) if parts else ""
        return (name, saved_offset if jumped else offset)

    # ------------------------------------------------------------------
    # Flush hosts to SQLite
    # ------------------------------------------------------------------

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

    def _flush_hosts(self) -> None:
        """Write all in-memory hosts to SQLite."""
        with self._hosts_lock:
            hosts_snapshot = list(self._hosts.values())

        if not hosts_snapshot or not self._db_conn:
            return

        try:
            for h in hosts_snapshot:
                self._db_conn.execute(
                    """INSERT INTO hosts (ip, mac, hostname, vendor, os_guess,
                       dhcp_fingerprint, first_seen, last_seen, source)
                       VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
                       ON CONFLICT(ip) DO UPDATE SET
                           mac = COALESCE(NULLIF(excluded.mac, ''), hosts.mac),
                           hostname = COALESCE(NULLIF(excluded.hostname, ''), hosts.hostname),
                           vendor = COALESCE(NULLIF(excluded.vendor, ''), hosts.vendor),
                           os_guess = COALESCE(NULLIF(excluded.os_guess, ''), hosts.os_guess),
                           dhcp_fingerprint = COALESCE(NULLIF(excluded.dhcp_fingerprint, ''), hosts.dhcp_fingerprint),
                           last_seen = excluded.last_seen,
                           source = excluded.source
                    """,
                    (h["ip"], h["mac"], h["hostname"], h["vendor"],
                     h["os_guess"], h["dhcp_fingerprint"],
                     h["first_seen"], h["last_seen"], h["source"]),
                )
            self._db_conn.commit()
        except Exception:
            logger.exception("Failed to flush %d hosts", len(hosts_snapshot))
