#!/usr/bin/env python3
"""VLAN discovery — 802.1Q tag detection, DTP/CDP/LLDP/STP parsing.

Passively identifies VLAN infrastructure by parsing layer-2 control protocols.
Detects 802.1Q tagged frames, DTP trunk negotiation, 802.1X/EAPOL NAC,
CDP/LLDP neighbor announcements, and STP BPDUs.
"""

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

# Ethertypes and protocol IDs
ETHERTYPE_8021Q = 0x8100
ETHERTYPE_8021AD = 0x88A8  # QinQ
ETHERTYPE_EAPOL = 0x888E
ETHERTYPE_LLDP = 0x88CC

# CDP/DTP use SNAP encapsulation with LLC header (DSAP=DSAP=0xAA, SSAP=0xAA, ctrl=0x03)
# CDP OUI: 0x00000C, protocol: 0x2000
# DTP OUI: 0x00000C, protocol: 0x2004
# STP uses LLC (DSAP=0x42, SSAP=0x42)

# CDP TLV types
CDP_TLV_DEVICE_ID = 0x0001
CDP_TLV_ADDRESS = 0x0002
CDP_TLV_PORT_ID = 0x0003
CDP_TLV_CAPABILITIES = 0x0004
CDP_TLV_SOFTWARE_VERSION = 0x0005
CDP_TLV_PLATFORM = 0x0006
CDP_TLV_NATIVE_VLAN = 0x000A
CDP_TLV_MANAGEMENT_ADDR = 0x0016

# LLDP TLV types
LLDP_TLV_END = 0
LLDP_TLV_CHASSIS_ID = 1
LLDP_TLV_PORT_ID = 2
LLDP_TLV_TTL = 3
LLDP_TLV_PORT_DESC = 4
LLDP_TLV_SYSTEM_NAME = 5
LLDP_TLV_SYSTEM_DESC = 6
LLDP_TLV_MGMT_ADDR = 8

_DB_INIT = """
CREATE TABLE IF NOT EXISTS vlans (
    id          INTEGER PRIMARY KEY AUTOINCREMENT,
    vlan_id     INTEGER NOT NULL,
    name        TEXT DEFAULT '',
    source      TEXT NOT NULL,
    switch_name TEXT DEFAULT '',
    switch_port TEXT DEFAULT '',
    first_seen  REAL NOT NULL,
    UNIQUE(vlan_id, source, switch_name, switch_port)
);
CREATE INDEX IF NOT EXISTS idx_vlans_vlan_id ON vlans(vlan_id);
"""

FLUSH_INTERVAL = 30


class VLANDiscovery(BaseModule):
    """Passively discover VLANs via 802.1Q tags and L2 protocol parsing."""

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

    def __init__(self, bus, state, config, engine=None):
        super().__init__(bus, state, config, engine)
        self._capture_bus = None
        self._sub_queue = None
        self._read_thread: Optional[threading.Thread] = None
        self._flush_thread: Optional[threading.Thread] = None
        self._db_path = self._resolve_db_path()
        self._db_conn: Optional[sqlite3.Connection] = None
        self._write_buffer: list[tuple] = []
        self._buffer_lock = threading.Lock()
        self._stats = {
            "dot1q_frames": 0,
            "cdp_frames": 0,
            "lldp_frames": 0,
            "dtp_frames": 0,
            "stp_bpdus": 0,
            "eapol_frames": 0,
            "vlans_found": 0,
            "packets_processed": 0,
        }

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

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

        self._init_db()
        self._capture_bus = self.config.get("_capture_bus")
        if self._capture_bus is None:
            logger.error("VLANDiscovery requires _capture_bus in config")
            return

        # Broad filter: we need 802.1Q (any ethertype), CDP/DTP (SNAP), LLDP, STP, EAPOL
        # Compile a catch-all because these span multiple ethertypes and LLC frames
        self._sub_queue = self._capture_bus.subscribe(
            name=self.name, bpf_filter="", queue_depth=3000
        )

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

        self._read_thread = threading.Thread(
            target=self._read_loop, daemon=True, name="sensor-vlan-read"
        )
        self._read_thread.start()

        self._flush_thread = threading.Thread(
            target=self._flush_loop, daemon=True, name="sensor-vlan-flush"
        )
        self._flush_thread.start()

        self.state.set_module_status(self.name, "running", pid=self._pid)
        self.bus.emit("MODULE_STARTED", {"module": self.name}, source_module=self.name)
        logger.info("VLANDiscovery started")

    def stop(self) -> None:
        if not self._running:
            return
        self._running = False
        if self._capture_bus and self._sub_queue:
            self._capture_bus.unsubscribe(self.name)
        if self._read_thread and self._read_thread.is_alive():
            self._read_thread.join(timeout=5.0)
        if self._flush_thread and self._flush_thread.is_alive():
            self._flush_thread.join(timeout=5.0)
        self._flush_buffer()
        if self._db_conn:
            self._db_conn.close()
            self._db_conn = None
        self.state.set_module_status(self.name, "stopped")
        self.bus.emit("MODULE_STOPPED", {"module": self.name}, source_module=self.name)
        logger.info("VLANDiscovery stopped — stats: %s", self._stats)

    def status(self) -> dict:
        return {
            "running": self._running,
            "pid": self._pid,
            "uptime": time.time() - self._start_time if self._start_time else 0,
            **self._stats,
        }

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

    # ------------------------------------------------------------------
    # Packet processing
    # ------------------------------------------------------------------

    def _read_loop(self) -> None:
        """Read packets from capture bus and dispatch to parsers."""
        while self._running:
            result = self._sub_queue.get(timeout=1.0)
            if result is None:
                continue
            ts, pkt = result
            self._stats["packets_processed"] += 1
            try:
                self._process_frame(pkt, ts)
            except Exception:
                logger.debug("Error processing frame", exc_info=True)

    def _process_frame(self, pkt: bytes, ts: float) -> None:
        """Identify and dispatch frame to appropriate parser."""
        if len(pkt) < 14:
            return

        dst_mac = pkt[0:6]
        ethertype = struct.unpack("!H", pkt[12:14])[0]

        # 802.1Q tagged frame
        if ethertype == ETHERTYPE_8021Q or ethertype == ETHERTYPE_8021AD:
            self._parse_dot1q(pkt, ts)
            return

        # LLDP
        if ethertype == ETHERTYPE_LLDP:
            self._parse_lldp(pkt, ts)
            return

        # EAPOL (802.1X)
        if ethertype == ETHERTYPE_EAPOL:
            self._stats["eapol_frames"] += 1
            self._record_vlan(0, "eapol_detected", "", "", ts)
            return

        # Check for LLC/SNAP frames (ethertype field is actually length if < 0x0600)
        if ethertype <= 0x05DC and len(pkt) >= 22:
            self._parse_llc_snap(pkt, ts)

    def _parse_dot1q(self, pkt: bytes, ts: float) -> None:
        """Parse 802.1Q tagged frame to extract VLAN ID."""
        if len(pkt) < 18:
            return
        tci = struct.unpack("!H", pkt[14:16])[0]
        vlan_id = tci & 0x0FFF
        if vlan_id > 0:
            self._stats["dot1q_frames"] += 1
            self._record_vlan(vlan_id, "802.1q", "", "", ts)

    def _parse_llc_snap(self, pkt: bytes, ts: float) -> None:
        """Parse LLC/SNAP encapsulated frames for CDP, DTP, STP."""
        # Minimum: 14 (eth) + 3 (LLC) = 17 bytes for STP check
        if len(pkt) < 17:
            return

        dsap = pkt[14]
        ssap = pkt[15]
        ctrl = pkt[16]

        # STP BPDU: DSAP=0x42, SSAP=0x42
        if dsap == 0x42 and ssap == 0x42:
            self._parse_stp(pkt, ts)
            return

        # SNAP: DSAP=0xAA, SSAP=0xAA, Ctrl=0x03
        if dsap == 0xAA and ssap == 0xAA and ctrl == 0x03 and len(pkt) >= 22:
            oui = pkt[17:20]
            proto = struct.unpack("!H", pkt[20:22])[0]

            # CDP: OUI 0x00000C, protocol 0x2000
            if oui == b'\x00\x00\x0c' and proto == 0x2000:
                self._parse_cdp(pkt[22:], ts)
                return

            # DTP: OUI 0x00000C, protocol 0x2004
            if oui == b'\x00\x00\x0c' and proto == 0x2004:
                self._parse_dtp(pkt[22:], ts)
                return

    def _parse_cdp(self, payload: bytes, ts: float) -> None:
        """Parse CDP TLVs to extract switch name, port, VLAN, platform."""
        if len(payload) < 4:
            return
        self._stats["cdp_frames"] += 1

        # CDP header: version(1), TTL(1), checksum(2)
        offset = 4
        device_id = ""
        port_id = ""
        platform = ""
        native_vlan = 0

        while offset + 4 <= len(payload):
            tlv_type = struct.unpack("!H", payload[offset:offset + 2])[0]
            tlv_len = struct.unpack("!H", payload[offset + 2:offset + 4])[0]
            if tlv_len < 4 or offset + tlv_len > len(payload):
                break
            tlv_data = payload[offset + 4:offset + tlv_len]

            if tlv_type == CDP_TLV_DEVICE_ID:
                device_id = tlv_data.decode("ascii", errors="replace").strip('\x00')
            elif tlv_type == CDP_TLV_PORT_ID:
                port_id = tlv_data.decode("ascii", errors="replace").strip('\x00')
            elif tlv_type == CDP_TLV_PLATFORM:
                platform = tlv_data.decode("ascii", errors="replace").strip('\x00')
            elif tlv_type == CDP_TLV_NATIVE_VLAN and len(tlv_data) >= 2:
                native_vlan = struct.unpack("!H", tlv_data[:2])[0]

            offset += tlv_len

        if native_vlan > 0:
            self._record_vlan(native_vlan, "cdp", device_id, port_id, ts)

        if device_id:
            self.bus.emit("VLAN_DETECTED", {
                "source": "cdp",
                "switch_name": device_id,
                "switch_port": port_id,
                "platform": platform,
                "native_vlan": native_vlan,
            }, source_module=self.name)

    def _parse_dtp(self, payload: bytes, ts: float) -> None:
        """Parse DTP frames — Cisco VLAN trunk negotiation."""
        self._stats["dtp_frames"] += 1
        # DTP presence alone is noteworthy (trunk negotiation active)
        self.bus.emit("VLAN_DETECTED", {
            "source": "dtp",
            "detail": "DTP trunk negotiation detected — VLAN hopping may be possible",
        }, source_module=self.name)

    def _parse_lldp(self, pkt: bytes, ts: float) -> None:
        """Parse LLDP TLVs for system name, port description, management address."""
        if len(pkt) < 16:
            return
        self._stats["lldp_frames"] += 1

        # LLDP payload starts after Ethernet header (14 bytes)
        offset = 14
        system_name = ""
        port_desc = ""
        mgmt_addr = ""

        while offset + 2 <= len(pkt):
            tlv_header = struct.unpack("!H", pkt[offset:offset + 2])[0]
            tlv_type = (tlv_header >> 9) & 0x7F
            tlv_len = tlv_header & 0x01FF
            offset += 2

            if tlv_type == LLDP_TLV_END or offset + tlv_len > len(pkt):
                break

            tlv_data = pkt[offset:offset + tlv_len]

            if tlv_type == LLDP_TLV_SYSTEM_NAME:
                system_name = tlv_data.decode("utf-8", errors="replace")
            elif tlv_type == LLDP_TLV_PORT_DESC:
                port_desc = tlv_data.decode("utf-8", errors="replace")
            elif tlv_type == LLDP_TLV_MGMT_ADDR and tlv_len >= 6:
                addr_len = tlv_data[0]
                addr_subtype = tlv_data[1]
                if addr_subtype == 1 and addr_len >= 5:  # IPv4
                    addr_bytes = tlv_data[2:6]
                    mgmt_addr = ".".join(str(b) for b in addr_bytes)

            offset += tlv_len

        if system_name:
            self.bus.emit("VLAN_DETECTED", {
                "source": "lldp",
                "system_name": system_name,
                "port_desc": port_desc,
                "mgmt_addr": mgmt_addr,
            }, source_module=self.name)

    def _parse_stp(self, pkt: bytes, ts: float) -> None:
        """Parse STP BPDU to extract root bridge ID and priority."""
        # LLC header at 14, STP BPDU starts at 17
        if len(pkt) < 38:
            return
        self._stats["stp_bpdus"] += 1

        bpdu = pkt[17:]
        if len(bpdu) < 21:
            return

        # Protocol ID (2), version (1), type (1)
        proto_id = struct.unpack("!H", bpdu[0:2])[0]
        if proto_id != 0x0000:
            return

        # Flags (1 byte at offset 4)
        # Root bridge ID: 8 bytes at offset 5 (priority 2 bytes + MAC 6 bytes)
        if len(bpdu) >= 13:
            root_priority = struct.unpack("!H", bpdu[5:7])[0]
            root_mac = ":".join(f"{b:02x}" for b in bpdu[7:13])
            # Bridge ID: 8 bytes at offset 13
            bridge_priority = struct.unpack("!H", bpdu[13:15])[0] if len(bpdu) >= 15 else 0

            self.bus.emit("VLAN_DETECTED", {
                "source": "stp",
                "root_bridge_mac": root_mac,
                "root_priority": root_priority,
                "bridge_priority": bridge_priority,
            }, source_module=self.name)

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

    def _resolve_db_path(self) -> str:
        base = self.config.get("db_dir", os.path.join(
            os.path.expanduser("~"), ".implant"
        ))
        return os.path.join(base, "vlan_discovery.db")

    def _init_db(self) -> None:
        Path(self._db_path).parent.mkdir(parents=True, exist_ok=True)
        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(_DB_INIT)
        self._db_conn.commit()

    def _record_vlan(self, vlan_id: int, source: str, switch_name: str,
                     switch_port: str, ts: float) -> None:
        """Buffer a VLAN record for batch insert."""
        with self._buffer_lock:
            self._write_buffer.append((vlan_id, "", source, switch_name, switch_port, ts))

    def _flush_buffer(self) -> None:
        """Write buffered records to SQLite."""
        with self._buffer_lock:
            batch = list(self._write_buffer)
            self._write_buffer.clear()

        if not batch or not self._db_conn:
            return

        try:
            with self._db_conn:
                self._db_conn.executemany(
                    """INSERT INTO vlans (vlan_id, name, source, switch_name, switch_port, first_seen)
                       VALUES (?, ?, ?, ?, ?, ?)
                       ON CONFLICT(vlan_id, source, switch_name, switch_port) DO NOTHING""",
                    batch,
                )
            new_count = len(batch)
            self._stats["vlans_found"] += new_count
        except Exception:
            logger.exception("Failed to flush VLAN records")

    def _flush_loop(self) -> None:
        """Periodically flush write buffer to disk."""
        while self._running:
            time.sleep(FLUSH_INTERVAL)
            try:
                self._flush_buffer()
            except Exception:
                logger.exception("Flush loop error")
