#!/usr/bin/env python3
"""Authentication flow tracker — cross-protocol auth event correlation.

Tracks Kerberos (AS-REQ/TGS-REQ), NTLM (across SMB/HTTP/LDAP/MSSQL),
SSH connections, and LDAP binds.  Builds per-user authentication timelines
and identifies service accounts.
"""

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

from modules.base import BaseModule

logger = logging.getLogger(__name__)

# Protocol ports
KERBEROS_PORT = 88
SMB_PORT = 445
LDAP_PORT = 389
LDAPS_PORT = 636
RDP_PORT = 3389
SSH_PORT = 22
MSSQL_PORT = 1433

TCP_PROTO = 6
UDP_PROTO = 17

# Kerberos message types (in AS-REQ body)
KRB_AS_REQ = 10
KRB_AS_REP = 11
KRB_TGS_REQ = 12
KRB_TGS_REP = 13

# NTLM signature
NTLMSSP_SIGNATURE = b'NTLMSSP\x00'
NTLM_NEGOTIATE = 1
NTLM_CHALLENGE = 2
NTLM_AUTHENTICATE = 3

# Service account detection threshold
SERVICE_ACCOUNT_THRESHOLD = 5

_DB_INIT = """
CREATE TABLE IF NOT EXISTS auth_events (
    id             INTEGER PRIMARY KEY AUTOINCREMENT,
    timestamp      REAL NOT NULL,
    source_ip      TEXT NOT NULL,
    dest_ip        TEXT NOT NULL,
    protocol       TEXT NOT NULL,
    auth_type      TEXT NOT NULL,
    username       TEXT DEFAULT '',
    domain         TEXT DEFAULT '',
    target_service TEXT DEFAULT '',
    success        INTEGER DEFAULT -1
);
CREATE INDEX IF NOT EXISTS idx_auth_user ON auth_events(username);
CREATE INDEX IF NOT EXISTS idx_auth_ts ON auth_events(timestamp);
CREATE INDEX IF NOT EXISTS idx_auth_src ON auth_events(source_ip);
"""

FLUSH_INTERVAL = 30


class AuthFlowTracker(BaseModule):
    """Track authentication events across Kerberos, NTLM, SSH, LDAP."""

    name = "auth_flow_tracker"
    module_type = "passive"
    priority = 100
    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()
        # Track per-user service targets for service account detection
        self._user_targets: dict[str, set] = defaultdict(set)
        self._stats = {
            "packets_processed": 0,
            "kerberos_events": 0,
            "ntlm_events": 0,
            "ssh_events": 0,
            "ldap_bind_events": 0,
            "service_accounts_detected": 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("AuthFlowTracker requires _capture_bus in config")
            return

        # BPF for auth protocol ports
        bpf = (
            "port 88 or port 445 or port 389 or port 636 "
            "or port 3389 or port 22 or port 1433"
        )
        self._sub_queue = self._capture_bus.subscribe(
            name=self.name, bpf_filter=bpf, queue_depth=5000
        )

        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-auth-read"
        )
        self._read_thread.start()

        self._flush_thread = threading.Thread(
            target=self._flush_loop, daemon=True, name="sensor-auth-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("AuthFlowTracker started")

    def stop(self) -> None:
        if not self._running:
            return
        self._running = False
        if self._capture_bus:
            self._capture_bus.unsubscribe(self.name)
        for t in (self._read_thread, self._flush_thread):
            if t and t.is_alive():
                t.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("AuthFlowTracker 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:
        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_packet(pkt, ts)
            except Exception:
                logger.debug("Error processing auth packet", exc_info=True)

    def _process_packet(self, pkt: bytes, ts: float) -> None:
        """Route packet to appropriate protocol parser."""
        if len(pkt) < 34:
            return

        ethertype = struct.unpack("!H", pkt[12:14])[0]
        eth_offset = 14
        if ethertype == 0x8100:
            if len(pkt) < 38:
                return
            ethertype = struct.unpack("!H", pkt[16:18])[0]
            eth_offset = 18

        if ethertype != 0x0800:
            return

        ip_header = pkt[eth_offset:]
        if len(ip_header) < 20:
            return

        ihl = (ip_header[0] & 0x0F) * 4
        ip_proto = ip_header[9]
        src_ip = socket.inet_ntoa(ip_header[12:16])
        dst_ip = socket.inet_ntoa(ip_header[16:20])

        if ip_proto == TCP_PROTO and len(ip_header) >= ihl + 4:
            src_port, dst_port = struct.unpack("!HH", ip_header[ihl:ihl + 4])
            # TCP data offset
            if len(ip_header) < ihl + 12:
                return
            tcp_data_off = ((ip_header[ihl + 12] >> 4) & 0xF) * 4
            payload = ip_header[ihl + tcp_data_off:]
            if not payload:
                return

            port = dst_port
            if dst_port == KERBEROS_PORT or src_port == KERBEROS_PORT:
                self._parse_kerberos(payload, src_ip, dst_ip, ts)
            elif dst_port == SMB_PORT or src_port == SMB_PORT:
                self._parse_ntlm_in_payload(payload, src_ip, dst_ip, "smb", ts)
            elif dst_port in (LDAP_PORT, LDAPS_PORT) or src_port in (LDAP_PORT, LDAPS_PORT):
                self._parse_ldap_bind(payload, src_ip, dst_ip, ts)
                self._parse_ntlm_in_payload(payload, src_ip, dst_ip, "ldap", ts)
            elif dst_port == SSH_PORT or src_port == SSH_PORT:
                self._parse_ssh(payload, src_ip, dst_ip, ts)
            elif dst_port == RDP_PORT or src_port == RDP_PORT:
                self._parse_ntlm_in_payload(payload, src_ip, dst_ip, "rdp", ts)
            elif dst_port == MSSQL_PORT or src_port == MSSQL_PORT:
                self._parse_ntlm_in_payload(payload, src_ip, dst_ip, "mssql", ts)

        elif ip_proto == UDP_PROTO and len(ip_header) >= ihl + 8:
            src_port, dst_port = struct.unpack("!HH", ip_header[ihl:ihl + 4])
            payload = ip_header[ihl + 8:]
            if dst_port == KERBEROS_PORT or src_port == KERBEROS_PORT:
                self._parse_kerberos(payload, src_ip, dst_ip, ts)

    # ------------------------------------------------------------------
    # Kerberos parser
    # ------------------------------------------------------------------

    def _parse_kerberos(self, payload: bytes, src_ip: str, dst_ip: str,
                        ts: float) -> None:
        """Parse Kerberos AS-REQ and TGS-REQ to extract principal, realm, SPN."""
        if len(payload) < 10:
            return

        # Kerberos uses ASN.1 DER encoding
        # Application tag: AS-REQ=[APPLICATION 10], TGS-REQ=[APPLICATION 12]
        tag = payload[0]
        if tag & 0x1F == 0x1F:
            # Long-form tag, skip for now
            return

        app_class = (tag >> 6) & 0x03
        if app_class != 1:  # APPLICATION class
            return

        msg_type = tag & 0x1F
        auth_type = ""
        if msg_type == KRB_AS_REQ:
            auth_type = "kerberos_as_req"
        elif msg_type == KRB_TGS_REQ:
            auth_type = "kerberos_tgs_req"
        else:
            return

        # Extract principal and realm from ASN.1 — simplified extraction
        # Look for common realm and principal patterns in the raw bytes
        realm = self._extract_krb_string(payload, b'\x1b')  # GeneralString
        principal = self._extract_krb_string(payload, b'\x1b', skip=1)

        self._stats["kerberos_events"] += 1
        self._record_auth(
            ts, src_ip, dst_ip, "kerberos", auth_type,
            username=principal, domain=realm,
            target_service=principal if msg_type == KRB_TGS_REQ else "",
        )

    def _extract_krb_string(self, data: bytes, tag: bytes, skip: int = 0) -> str:
        """Extract a GeneralString value from ASN.1 data (simplified)."""
        search_tag = tag[0]
        pos = 0
        found = 0
        while pos < len(data) - 2:
            if data[pos] == search_tag:
                length = data[pos + 1]
                if length > 0 and length < 128 and pos + 2 + length <= len(data):
                    if found >= skip:
                        try:
                            return data[pos + 2:pos + 2 + length].decode("ascii", errors="replace")
                        except Exception:
                            return ""
                    found += 1
            pos += 1
        return ""

    # ------------------------------------------------------------------
    # NTLM parser
    # ------------------------------------------------------------------

    def _parse_ntlm_in_payload(self, payload: bytes, src_ip: str, dst_ip: str,
                                protocol: str, ts: float) -> None:
        """Search for NTLMSSP messages embedded in any protocol payload."""
        idx = payload.find(NTLMSSP_SIGNATURE)
        if idx == -1 or idx + 12 > len(payload):
            return

        ntlm_data = payload[idx:]
        msg_type = struct.unpack("<I", ntlm_data[8:12])[0]

        if msg_type == NTLM_AUTHENTICATE and len(ntlm_data) >= 88:
            # Type 3: Authenticate message
            username, domain, workstation = self._parse_ntlm_auth(ntlm_data)
            if username:
                self._stats["ntlm_events"] += 1
                self._record_auth(
                    ts, src_ip, dst_ip, protocol, "ntlm_auth",
                    username=username, domain=domain,
                )
        elif msg_type == NTLM_CHALLENGE and len(ntlm_data) >= 56:
            # Type 2: Challenge — extract target name
            target = self._parse_ntlm_challenge_target(ntlm_data)
            if target:
                self._stats["ntlm_events"] += 1
                self._record_auth(
                    ts, dst_ip, src_ip, protocol, "ntlm_challenge",
                    domain=target,
                )

    def _parse_ntlm_auth(self, data: bytes) -> tuple:
        """Parse NTLM Type 3 Authenticate message for username, domain, workstation."""
        if len(data) < 72:
            return ("", "", "")

        try:
            # LmChallengeResponse: offset 12
            # NtChallengeResponse: offset 20
            # DomainName: length(2), maxlen(2), offset(4) at byte 28
            domain_len = struct.unpack("<H", data[28:30])[0]
            domain_off = struct.unpack("<I", data[32:36])[0]

            # UserName: length(2), maxlen(2), offset(4) at byte 36
            user_len = struct.unpack("<H", data[36:38])[0]
            user_off = struct.unpack("<I", data[40:44])[0]

            # Workstation: length(2), maxlen(2), offset(4) at byte 44
            ws_len = struct.unpack("<H", data[44:46])[0]
            ws_off = struct.unpack("<I", data[48:52])[0]

            domain = data[domain_off:domain_off + domain_len].decode("utf-16-le", errors="replace") if domain_len else ""
            username = data[user_off:user_off + user_len].decode("utf-16-le", errors="replace") if user_len else ""
            workstation = data[ws_off:ws_off + ws_len].decode("utf-16-le", errors="replace") if ws_len else ""

            return (username, domain, workstation)
        except Exception:
            return ("", "", "")

    def _parse_ntlm_challenge_target(self, data: bytes) -> str:
        """Parse NTLM Type 2 Challenge for target name."""
        if len(data) < 24:
            return ""
        try:
            target_len = struct.unpack("<H", data[12:14])[0]
            target_off = struct.unpack("<I", data[16:20])[0]
            if target_len and target_off + target_len <= len(data):
                return data[target_off:target_off + target_len].decode("utf-16-le", errors="replace")
        except Exception:
            pass
        return ""

    # ------------------------------------------------------------------
    # SSH parser
    # ------------------------------------------------------------------

    def _parse_ssh(self, payload: bytes, src_ip: str, dst_ip: str,
                   ts: float) -> None:
        """Parse SSH version exchange and detect auth activity."""
        # SSH version string starts with "SSH-"
        if payload[:4] == b'SSH-':
            try:
                version_line = payload.split(b'\r\n')[0].decode("ascii", errors="replace")
            except Exception:
                version_line = ""

            self._stats["ssh_events"] += 1
            self._record_auth(
                ts, src_ip, dst_ip, "ssh", "version_exchange",
                target_service=version_line,
            )

    # ------------------------------------------------------------------
    # LDAP bind parser
    # ------------------------------------------------------------------

    def _parse_ldap_bind(self, payload: bytes, src_ip: str, dst_ip: str,
                         ts: float) -> None:
        """Parse LDAP simple bind request to extract username."""
        if len(payload) < 10:
            return

        # LDAP messages are BER-encoded
        # BindRequest: APPLICATION[0] -> sequence: version, name, auth
        # Look for a simple bind: tag 0x60 (APPLICATION CONSTRUCTED 0)
        if payload[0] != 0x30:  # Not a SEQUENCE
            return

        # Search for BindRequest tag (0x60) within the message
        idx = payload.find(b'\x60')
        if idx == -1 or idx + 10 > len(payload):
            return

        # Skip the BindRequest tag and length
        bind_data = payload[idx + 1:]
        if not bind_data:
            return

        # Try to read the length
        bind_len, len_size = self._read_ber_length(bind_data)
        if bind_len <= 0:
            return

        bind_body = bind_data[len_size:]

        # Version (INTEGER tag 0x02)
        if len(bind_body) < 3 or bind_body[0] != 0x02:
            return
        ver_len = bind_body[1]
        name_start = 2 + ver_len

        # Name (OCTET STRING tag 0x04)
        if len(bind_body) <= name_start + 2:
            return
        if bind_body[name_start] != 0x04:
            return
        name_len = bind_body[name_start + 1]
        if name_len > 0 and name_start + 2 + name_len <= len(bind_body):
            username = bind_body[name_start + 2:name_start + 2 + name_len].decode(
                "utf-8", errors="replace"
            )
            if username and username not in ("", "anonymous"):
                self._stats["ldap_bind_events"] += 1
                self._record_auth(
                    ts, src_ip, dst_ip, "ldap", "simple_bind",
                    username=username,
                )

    @staticmethod
    def _read_ber_length(data: bytes) -> tuple:
        """Read BER length encoding. Returns (length, bytes_consumed)."""
        if not data:
            return (0, 0)
        first = data[0]
        if first < 0x80:
            return (first, 1)
        num_bytes = first & 0x7F
        if num_bytes == 0 or len(data) < 1 + num_bytes:
            return (0, 0)
        length = 0
        for i in range(num_bytes):
            length = (length << 8) | data[1 + i]
        return (length, 1 + num_bytes)

    # ------------------------------------------------------------------
    # Recording
    # ------------------------------------------------------------------

    def _record_auth(self, ts: float, src_ip: str, dst_ip: str,
                     protocol: str, auth_type: str, username: str = "",
                     domain: str = "", target_service: str = "",
                     success: int = -1) -> None:
        """Buffer an auth event for batch insert."""
        with self._buffer_lock:
            self._write_buffer.append((
                ts, src_ip, dst_ip, protocol, auth_type,
                username, domain, target_service, success,
            ))

        # Track service account patterns
        if username:
            key = f"{domain}\\{username}" if domain else username
            self._user_targets[key].add(dst_ip)
            if len(self._user_targets[key]) >= SERVICE_ACCOUNT_THRESHOLD:
                # Only count once per user crossing threshold
                current_count = len([
                    u for u, targets in self._user_targets.items()
                    if len(targets) >= SERVICE_ACCOUNT_THRESHOLD
                ])
                if current_count > self._stats["service_accounts_detected"]:
                    self._stats["service_accounts_detected"] = current_count

    # ------------------------------------------------------------------
    # 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, "auth_flow_tracker.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 _flush_buffer(self) -> None:
        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 auth_events
                       (timestamp, source_ip, dest_ip, protocol, auth_type,
                        username, domain, target_service, success)
                       VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)""",
                    batch,
                )
        except Exception:
            logger.exception("Failed to flush auth events")

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