#!/usr/bin/env python3
"""Passive TLS SNI extraction from ClientHello messages.

Subscribes to capture_bus with BPF "tcp port 443". Parses TLS ClientHello
handshake messages and extracts the Server Name Indication (SNI) extension.
Handles fragmented ClientHello across multiple TCP segments via a reassembly
buffer keyed on (src_ip, src_port, dst_ip, dst_port).
"""

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

from modules.base import BaseModule

logger = logging.getLogger(__name__)

# TLS record types
TLS_HANDSHAKE = 22
# TLS handshake types
TLS_CLIENT_HELLO = 1
# TLS extension type for SNI
EXT_SNI = 0x0000

# TLS version mapping
TLS_VERSIONS = {
    0x0300: "SSLv3",
    0x0301: "TLS 1.0",
    0x0302: "TLS 1.1",
    0x0303: "TLS 1.2",
    0x0304: "TLS 1.3",
}

# Max reassembly buffer entries (prevent memory leak from orphaned streams)
MAX_REASSEMBLY_ENTRIES = 10000
# Max age for reassembly buffer entries (seconds)
REASSEMBLY_TIMEOUT = 30


class TLSSNIExtractor(BaseModule):
    """Extract SNI hostnames from TLS ClientHello messages."""

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

    BATCH_SIZE = 500
    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
        self._buffer = []
        self._buffer_lock = threading.Lock()
        self._db_path = ""
        self._db_conn: Optional[sqlite3.Connection] = None
        self._total_sni = 0
        # TCP reassembly buffer: (src,sport,dst,dport) -> (data, timestamp)
        self._reassembly: OrderedDict = OrderedDict()
        self._reassembly_lock = threading.Lock()

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

        self._sub_queue = self._capture_bus.subscribe(
            name=self.name, bpf_filter="tcp port 443", queue_depth=8000
        )

        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-sni-reader"
        )
        self._reader_thread.start()

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

        self.state.set_module_status(self.name, "running", pid=os.getpid())
        logger.info("TLSSNIExtractor started — monitoring TLS on port 443")

    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("TLSSNIExtractor stopped — %d SNIs extracted", self._total_sni)

    def status(self) -> dict:
        return {
            "running": self._running,
            "pid": self._pid,
            "uptime": time.time() - self._start_time if self._start_time else 0,
            "total_sni": self._total_sni,
            "buffer_size": len(self._buffer),
            "reassembly_entries": len(self._reassembly),
        }

    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 tls_sni (
                id          INTEGER PRIMARY KEY AUTOINCREMENT,
                timestamp   REAL NOT NULL,
                source_ip   TEXT NOT NULL,
                dest_ip     TEXT NOT NULL,
                dest_port   INTEGER NOT NULL,
                sni         TEXT NOT NULL,
                tls_version TEXT
            );
            CREATE INDEX IF NOT EXISTS idx_sni_source ON tls_sni(source_ip);
            CREATE INDEX IF NOT EXISTS idx_sni_hostname ON tls_sni(sni);
            CREATE INDEX IF NOT EXISTS idx_sni_ts ON tls_sni(timestamp);
        """)
        self._db_conn.commit()

    # ------------------------------------------------------------------
    # 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:
        """Parse Ethernet -> IPv4 -> TCP -> TLS ClientHello."""
        if len(raw) < 14:
            return

        eth_type = struct.unpack("!H", raw[12:14])[0]
        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]
        if ip_proto != 6:  # TCP only
            return

        src_ip = socket.inet_ntoa(ip_hdr[12:16])
        dst_ip = socket.inet_ntoa(ip_hdr[16:20])

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

        src_port, dst_port = struct.unpack("!HH", raw[tcp_offset:tcp_offset + 4])
        tcp_hdr_len = ((raw[tcp_offset + 12] >> 4) & 0xF) * 4
        payload_offset = tcp_offset + tcp_hdr_len
        payload = raw[payload_offset:]

        if not payload:
            return

        # Try to extract SNI from this payload (possibly with reassembly)
        flow_key = (src_ip, src_port, dst_ip, dst_port)
        self._process_tls_payload(ts, flow_key, payload, src_ip, dst_ip, dst_port)

    def _process_tls_payload(self, ts: float, flow_key: tuple,
                             payload: bytes, src_ip: str, dst_ip: str,
                             dst_port: int) -> None:
        """Parse TLS records from TCP payload, with reassembly for fragments."""
        # Check if we have buffered data for this flow
        with self._reassembly_lock:
            if flow_key in self._reassembly:
                prev_data, _ = self._reassembly.pop(flow_key)
                payload = prev_data + payload

        # Try to parse TLS record
        offset = 0
        while offset < len(payload):
            if len(payload) - offset < 5:
                # Not enough for TLS record header — buffer for reassembly
                self._buffer_fragment(flow_key, payload[offset:], ts)
                return

            content_type = payload[offset]
            tls_version = struct.unpack("!H", payload[offset + 1:offset + 3])[0]
            record_length = struct.unpack("!H", payload[offset + 3:offset + 5])[0]

            if content_type != TLS_HANDSHAKE:
                return

            if record_length > 16384:  # Max TLS record size
                return

            if len(payload) - offset - 5 < record_length:
                # Fragmented — buffer remaining data
                self._buffer_fragment(flow_key, payload[offset:], ts)
                return

            record_data = payload[offset + 5:offset + 5 + record_length]
            sni, version_str = self._parse_client_hello(record_data, tls_version)
            if sni:
                self._total_sni += 1
                record = (ts, src_ip, dst_ip, dst_port, sni, version_str)
                with self._buffer_lock:
                    self._buffer.append(record)
                    if len(self._buffer) >= self.BATCH_SIZE:
                        self._flush_buffer()
                return  # Only care about the first ClientHello per connection

            offset += 5 + record_length

    def _buffer_fragment(self, flow_key: tuple, data: bytes, ts: float) -> None:
        """Buffer a TCP fragment for reassembly."""
        with self._reassembly_lock:
            self._reassembly[flow_key] = (data, ts)
            # Evict oldest entries if buffer too large
            while len(self._reassembly) > MAX_REASSEMBLY_ENTRIES:
                self._reassembly.popitem(last=False)
            # Evict expired entries periodically
            now = time.time()
            expired = [k for k, (_, t) in self._reassembly.items()
                       if now - t > REASSEMBLY_TIMEOUT]
            for k in expired:
                self._reassembly.pop(k, None)

    @staticmethod
    def _parse_client_hello(data: bytes, record_tls_version: int) -> tuple:
        """Parse a TLS Handshake record to extract SNI from ClientHello.

        Returns (sni_hostname, tls_version_string) or (None, None).
        """
        if len(data) < 4:
            return None, None

        handshake_type = data[0]
        if handshake_type != TLS_CLIENT_HELLO:
            return None, None

        # Handshake length (3 bytes)
        hs_length = struct.unpack("!I", b"\x00" + data[1:4])[0]
        if len(data) < 4 + hs_length:
            return None, None

        offset = 4

        # ClientHello: version(2) + random(32) + session_id_len(1) + ...
        if offset + 34 > len(data):
            return None, None

        client_version = struct.unpack("!H", data[offset:offset + 2])[0]
        offset += 2 + 32  # skip version + random

        # Session ID
        if offset >= len(data):
            return None, None
        session_id_len = data[offset]
        offset += 1 + session_id_len

        # Cipher suites
        if offset + 2 > len(data):
            return None, None
        cipher_suites_len = struct.unpack("!H", data[offset:offset + 2])[0]
        offset += 2 + cipher_suites_len

        # Compression methods
        if offset >= len(data):
            return None, None
        comp_methods_len = data[offset]
        offset += 1 + comp_methods_len

        # Extensions
        if offset + 2 > len(data):
            return None, None
        extensions_len = struct.unpack("!H", data[offset:offset + 2])[0]
        offset += 2

        extensions_end = offset + extensions_len
        if extensions_end > len(data):
            extensions_end = len(data)

        # Determine TLS version — check for supported_versions extension (0x002b)
        # which overrides the record-layer version for TLS 1.3
        detected_version = client_version
        sni_hostname = None

        ext_offset = offset
        while ext_offset + 4 <= extensions_end:
            ext_type = struct.unpack("!H", data[ext_offset:ext_offset + 2])[0]
            ext_len = struct.unpack("!H", data[ext_offset + 2:ext_offset + 4])[0]
            ext_data_start = ext_offset + 4
            ext_data_end = ext_data_start + ext_len

            if ext_data_end > extensions_end:
                break

            if ext_type == EXT_SNI:
                # SNI extension: list_length(2), type(1)=0 (hostname), name_length(2), name
                sni_data = data[ext_data_start:ext_data_end]
                sni_hostname = _parse_sni_extension(sni_data)

            elif ext_type == 0x002B:
                # supported_versions extension — pick highest version
                sv_data = data[ext_data_start:ext_data_end]
                if len(sv_data) >= 1:
                    sv_list_len = sv_data[0]
                    sv_offset = 1
                    max_ver = 0
                    while sv_offset + 2 <= 1 + sv_list_len and sv_offset + 2 <= len(sv_data):
                        ver = struct.unpack("!H", sv_data[sv_offset:sv_offset + 2])[0]
                        if ver > max_ver:
                            max_ver = ver
                        sv_offset += 2
                    if max_ver > 0:
                        detected_version = max_ver

            ext_offset = ext_data_end

        version_str = TLS_VERSIONS.get(detected_version, f"0x{detected_version:04x}")
        return sni_hostname, version_str

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

    def _flush_loop(self) -> None:
        while self._running:
            time.sleep(self.FLUSH_INTERVAL)
            try:
                self._flush_buffer()
            except Exception:
                logger.exception("SNI 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 tls_sni (timestamp, source_ip, dest_ip, dest_port, sni, tls_version) "
                "VALUES (?, ?, ?, ?, ?, ?)",
                batch,
            )
            self._db_conn.commit()
        except Exception:
            logger.exception("Failed to flush %d SNI records", len(batch))


def _parse_sni_extension(data: bytes) -> Optional[str]:
    """Parse the SNI extension data to extract the hostname."""
    if len(data) < 5:
        return None

    sni_list_len = struct.unpack("!H", data[0:2])[0]
    offset = 2

    while offset + 3 <= len(data) and offset < 2 + sni_list_len:
        name_type = data[offset]
        name_len = struct.unpack("!H", data[offset + 1:offset + 3])[0]
        offset += 3

        if offset + name_len > len(data):
            break

        if name_type == 0:  # host_name
            try:
                return data[offset:offset + name_len].decode("ascii")
            except UnicodeDecodeError:
                return data[offset:offset + name_len].decode("utf-8", errors="replace")

        offset += name_len

    return None
