#!/usr/bin/env python3
"""Database protocol interceptor — passive database query and auth monitoring.

Parses login packets and query metadata for:
- MSSQL TDS (1433): TDS7 Login packet, SQL batch text
- MySQL (3306): Handshake, Login Request, COM_QUERY
- PostgreSQL (5432): Startup message, Simple Query
- Redis (6379): AUTH command, key operations
- MongoDB (27017): OP_MSG saslStart, find operations

Logs queries but NOT full result sets to avoid excessive data.
"""

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("bb.passive.db_interceptor")

# TDS packet types
TDS_SQL_BATCH = 0x01
TDS_LOGIN7 = 0x10
TDS_PRELOGIN = 0x12

# MySQL commands
MYSQL_COM_QUERY = 0x03
MYSQL_COM_INIT_DB = 0x02

# PostgreSQL message types
PG_STARTUP = 0x00  # No type byte — identified by structure
PG_SIMPLE_QUERY = ord('Q')
PG_PASSWORD = ord('p')

# Redis inline delimiters
REDIS_CRLF = b'\r\n'

# MongoDB opcodes
MONGO_OP_MSG = 2013

_DB_INIT = """
CREATE TABLE IF NOT EXISTS db_queries (
    id         INTEGER PRIMARY KEY AUTOINCREMENT,
    timestamp  REAL NOT NULL,
    source_ip  TEXT NOT NULL,
    dest_ip    TEXT NOT NULL,
    db_type    TEXT NOT NULL,
    username   TEXT DEFAULT '',
    database   TEXT DEFAULT '',
    query_text TEXT DEFAULT '',
    operation  TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_dbq_type ON db_queries(db_type);
CREATE INDEX IF NOT EXISTS idx_dbq_user ON db_queries(username);
CREATE INDEX IF NOT EXISTS idx_dbq_ts ON db_queries(timestamp);
"""

FLUSH_INTERVAL = 45
MAX_QUERY_LEN = 2000  # Truncate queries beyond this


class DBInterceptor(BaseModule):
    """Passively intercept database authentication and queries."""

    name = "db_interceptor"
    module_type = "passive"
    priority = 150
    requires_root = 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 = {
            "packets_processed": 0,
            "mssql_events": 0,
            "mysql_events": 0,
            "postgres_events": 0,
            "redis_events": 0,
            "mongodb_events": 0,
            "logins_captured": 0,
            "queries_captured": 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("DBInterceptor requires _capture_bus in config")
            return

        bpf = "port 1433 or port 3306 or port 5432 or port 6379 or port 27017"
        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="bb-dbi-read"
        )
        self._read_thread.start()

        self._flush_thread = threading.Thread(
            target=self._flush_loop, daemon=True, name="bb-dbi-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("DBInterceptor 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("DBInterceptor 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 DB packet", exc_info=True)

    def _process_packet(self, pkt: bytes, ts: float) -> None:
        """Route packet to the correct database parser based on port."""
        if len(pkt) < 54:
            return

        ethertype = struct.unpack("!H", pkt[12:14])[0]
        eth_offset = 14
        if ethertype == 0x8100:
            if len(pkt) < 58:
                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]
        if ip_proto != 6:  # TCP only
            return

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

        if len(ip_header) < ihl + 20:
            return
        src_port, dst_port = struct.unpack("!HH", ip_header[ihl:ihl + 4])
        tcp_data_off = ((ip_header[ihl + 12] >> 4) & 0xF) * 4
        payload = ip_header[ihl + tcp_data_off:]

        if len(payload) < 4:
            return

        # Route by destination port (client -> server)
        if dst_port == 1433 or src_port == 1433:
            self._parse_mssql(payload, src_ip, dst_ip, dst_port == 1433, ts)
        elif dst_port == 3306 or src_port == 3306:
            self._parse_mysql(payload, src_ip, dst_ip, dst_port == 3306, ts)
        elif dst_port == 5432 or src_port == 5432:
            self._parse_postgres(payload, src_ip, dst_ip, dst_port == 5432, ts)
        elif dst_port == 6379 or src_port == 6379:
            self._parse_redis(payload, src_ip, dst_ip, dst_port == 6379, ts)
        elif dst_port == 27017 or src_port == 27017:
            self._parse_mongodb(payload, src_ip, dst_ip, dst_port == 27017, ts)

    # ------------------------------------------------------------------
    # MSSQL TDS parser
    # ------------------------------------------------------------------

    def _parse_mssql(self, payload: bytes, src_ip: str, dst_ip: str,
                     is_to_server: bool, ts: float) -> None:
        """Parse MSSQL TDS packets."""
        if len(payload) < 8:
            return

        # TDS header: type(1), status(1), length(2), SPID(2), packet(1), window(1)
        tds_type = payload[0]
        tds_len = struct.unpack("!H", payload[2:4])[0]

        if tds_type == TDS_LOGIN7 and is_to_server and len(payload) >= 36:
            self._parse_tds_login7(payload[8:], src_ip, dst_ip, ts)
        elif tds_type == TDS_SQL_BATCH and is_to_server:
            # SQL batch: header group + SQL text (all Unicode LE after 8-byte TDS header)
            self._parse_tds_sql_batch(payload[8:], src_ip, dst_ip, ts)

    def _parse_tds_login7(self, data: bytes, src_ip: str, dst_ip: str,
                          ts: float) -> None:
        """Parse TDS7 Login packet for username, hostname, app name."""
        if len(data) < 36:
            return

        self._stats["mssql_events"] += 1
        self._stats["logins_captured"] += 1

        try:
            # TDS7 Login: fixed header of 36 bytes then variable data
            # HostName offset(2) + length(2) at offset 8
            # UserName offset(2) + length(2) at offset 16
            # AppName offset(2) + length(2) at offset 24
            # ServerName offset(2) + length(2) at offset 28
            # Database offset(2) + length(2) at offset 40 (but we have 36-byte minimum)

            hostname_off = struct.unpack("<H", data[8:10])[0]
            hostname_len = struct.unpack("<H", data[10:12])[0]

            username_off = struct.unpack("<H", data[16:18])[0]
            username_len = struct.unpack("<H", data[18:20])[0]

            appname_off = struct.unpack("<H", data[24:26])[0]
            appname_len = struct.unpack("<H", data[26:28])[0]

            # Offsets are in chars (UTF-16LE = 2 bytes per char)
            username = self._read_utf16le(data, username_off, username_len)
            hostname = self._read_utf16le(data, hostname_off, hostname_len)
            appname = self._read_utf16le(data, appname_off, appname_len)

            database = ""
            if len(data) >= 44:
                db_off = struct.unpack("<H", data[40:42])[0]
                db_len = struct.unpack("<H", data[42:44])[0]
                database = self._read_utf16le(data, db_off, db_len)

            self._record_query(
                ts, src_ip, dst_ip, "mssql", username, database,
                f"LOGIN hostname={hostname} app={appname}", "login"
            )

            self.bus.emit("CREDENTIAL_FOUND", {
                "source": "mssql_tds_login",
                "username": username,
                "hostname": hostname,
                "source_ip": src_ip,
                "dest_ip": dst_ip,
                "database": database,
            }, source_module=self.name)

        except Exception:
            logger.debug("Failed to parse TDS Login7", exc_info=True)

    def _parse_tds_sql_batch(self, data: bytes, src_ip: str, dst_ip: str,
                              ts: float) -> None:
        """Parse TDS SQL Batch for query text."""
        if len(data) < 4:
            return

        # TDS SQL Batch: ALL_HEADERS (variable) + SQL text
        # ALL_HEADERS starts with TotalLength(4)
        total_headers_len = struct.unpack("<I", data[0:4])[0]
        if total_headers_len > len(data):
            total_headers_len = 0

        sql_bytes = data[total_headers_len:]
        if not sql_bytes:
            return

        try:
            query = sql_bytes.decode("utf-16-le", errors="replace")
        except Exception:
            query = sql_bytes.decode("ascii", errors="replace")

        query = query.strip('\x00').strip()[:MAX_QUERY_LEN]
        if query:
            self._stats["mssql_events"] += 1
            self._stats["queries_captured"] += 1
            self._record_query(ts, src_ip, dst_ip, "mssql", "", "", query, "query")

    @staticmethod
    def _read_utf16le(data: bytes, char_offset: int, char_length: int) -> str:
        """Read UTF-16LE string from TDS data using char offset/length."""
        byte_off = char_offset * 2
        byte_len = char_length * 2
        if byte_off + byte_len > len(data):
            return ""
        try:
            return data[byte_off:byte_off + byte_len].decode("utf-16-le", errors="replace")
        except Exception:
            return ""

    # ------------------------------------------------------------------
    # MySQL parser
    # ------------------------------------------------------------------

    def _parse_mysql(self, payload: bytes, src_ip: str, dst_ip: str,
                     is_to_server: bool, ts: float) -> None:
        """Parse MySQL protocol packets."""
        if len(payload) < 5:
            return

        # MySQL packet: length(3 LE) + sequence_id(1) + payload
        pkt_len = struct.unpack("<I", payload[0:3] + b'\x00')[0]
        seq_id = payload[3]
        mysql_data = payload[4:]

        if not mysql_data:
            return

        if not is_to_server:
            # Server -> client: check for Handshake (protocol version 10)
            if mysql_data[0] == 10 and seq_id == 0:
                self._parse_mysql_handshake(mysql_data, src_ip, dst_ip, ts)
            return

        # Client -> server
        cmd = mysql_data[0]
        if cmd == MYSQL_COM_QUERY and len(mysql_data) > 1:
            query = mysql_data[1:min(len(mysql_data), MAX_QUERY_LEN + 1)].decode(
                "utf-8", errors="replace"
            )
            self._stats["mysql_events"] += 1
            self._stats["queries_captured"] += 1
            self._record_query(ts, src_ip, dst_ip, "mysql", "", "", query, "query")

        elif seq_id == 1 and len(mysql_data) >= 32:
            # Login Request (HandshakeResponse)
            self._parse_mysql_login(mysql_data, src_ip, dst_ip, ts)

    def _parse_mysql_handshake(self, data: bytes, src_ip: str, dst_ip: str,
                                ts: float) -> None:
        """Parse MySQL server Handshake for version info."""
        # Protocol version (1) + server version (null-terminated)
        if len(data) < 5:
            return
        null_idx = data.find(b'\x00', 1)
        if null_idx < 0:
            return
        server_version = data[1:null_idx].decode("ascii", errors="replace")
        self._stats["mysql_events"] += 1
        # Record as informational
        self._record_query(
            ts, dst_ip, src_ip, "mysql", "", "",
            f"SERVER_HANDSHAKE version={server_version}", "handshake"
        )

    def _parse_mysql_login(self, data: bytes, src_ip: str, dst_ip: str,
                            ts: float) -> None:
        """Parse MySQL Login Request (HandshakeResponse41)."""
        if len(data) < 32:
            return

        try:
            # Client capabilities (4 bytes), max packet (4), charset (1),
            # reserved (23 bytes), username (null-terminated)
            username_start = 32
            null_idx = data.find(b'\x00', username_start)
            if null_idx < 0:
                return
            username = data[username_start:null_idx].decode("utf-8", errors="replace")

            # After username + null + auth data, there may be a database name
            database = ""
            # Skip auth response length + data
            pos = null_idx + 1
            if pos < len(data):
                auth_len = data[pos]
                pos += 1 + auth_len
                if pos < len(data):
                    db_null = data.find(b'\x00', pos)
                    if db_null > pos:
                        database = data[pos:db_null].decode("utf-8", errors="replace")

            self._stats["mysql_events"] += 1
            self._stats["logins_captured"] += 1
            self._record_query(ts, src_ip, dst_ip, "mysql", username, database, "", "login")

            self.bus.emit("CREDENTIAL_FOUND", {
                "source": "mysql_login",
                "username": username,
                "database": database,
                "source_ip": src_ip,
                "dest_ip": dst_ip,
            }, source_module=self.name)

        except Exception:
            logger.debug("Failed to parse MySQL login", exc_info=True)

    # ------------------------------------------------------------------
    # PostgreSQL parser
    # ------------------------------------------------------------------

    def _parse_postgres(self, payload: bytes, src_ip: str, dst_ip: str,
                        is_to_server: bool, ts: float) -> None:
        """Parse PostgreSQL protocol messages."""
        if not is_to_server or len(payload) < 8:
            return

        # Check for Startup Message: length(4) + protocol_version(4)
        # Protocol version 3.0 = 0x00030000
        if len(payload) >= 8:
            msg_len = struct.unpack("!I", payload[0:4])[0]
            proto_ver = struct.unpack("!I", payload[4:8])[0]

            if proto_ver == 0x00030000 and msg_len > 8:
                # Startup message: parse key=value pairs
                self._parse_pg_startup(payload[8:msg_len], src_ip, dst_ip, ts)
                return

        # Simple Query: type='Q' + length(4) + query_string
        if payload[0] == PG_SIMPLE_QUERY and len(payload) >= 6:
            query_len = struct.unpack("!I", payload[1:5])[0]
            query = payload[5:min(5 + query_len - 1, 5 + MAX_QUERY_LEN)].decode(
                "utf-8", errors="replace"
            ).rstrip('\x00')
            if query:
                self._stats["postgres_events"] += 1
                self._stats["queries_captured"] += 1
                self._record_query(ts, src_ip, dst_ip, "postgres", "", "", query, "query")

    def _parse_pg_startup(self, data: bytes, src_ip: str, dst_ip: str,
                           ts: float) -> None:
        """Parse PostgreSQL Startup Message key-value pairs."""
        username = ""
        database = ""

        pos = 0
        while pos < len(data) - 1:
            null_idx = data.find(b'\x00', pos)
            if null_idx < 0:
                break
            key = data[pos:null_idx].decode("utf-8", errors="replace")
            pos = null_idx + 1

            val_null = data.find(b'\x00', pos)
            if val_null < 0:
                break
            value = data[pos:val_null].decode("utf-8", errors="replace")
            pos = val_null + 1

            if key == "user":
                username = value
            elif key == "database":
                database = value

            if not key:  # Terminal null
                break

        if username:
            self._stats["postgres_events"] += 1
            self._stats["logins_captured"] += 1
            self._record_query(ts, src_ip, dst_ip, "postgres", username, database, "", "login")

            self.bus.emit("CREDENTIAL_FOUND", {
                "source": "postgres_startup",
                "username": username,
                "database": database,
                "source_ip": src_ip,
                "dest_ip": dst_ip,
            }, source_module=self.name)

    # ------------------------------------------------------------------
    # Redis parser
    # ------------------------------------------------------------------

    def _parse_redis(self, payload: bytes, src_ip: str, dst_ip: str,
                     is_to_server: bool, ts: float) -> None:
        """Parse Redis RESP protocol for AUTH commands and key operations."""
        if not is_to_server or len(payload) < 3:
            return

        # RESP: commands start with * (array) followed by count
        if payload[0:1] == b'*':
            args = self._parse_resp_array(payload)
            if not args:
                return

            cmd = args[0].upper() if args else ""

            if cmd == "AUTH":
                # AUTH [username] password
                password = args[-1] if len(args) >= 2 else ""
                username = args[1] if len(args) >= 3 else ""
                self._stats["redis_events"] += 1
                self._stats["logins_captured"] += 1
                self._record_query(ts, src_ip, dst_ip, "redis", username, "", "AUTH ***", "login")

                self.bus.emit("CREDENTIAL_FOUND", {
                    "source": "redis_auth",
                    "username": username,
                    "source_ip": src_ip,
                    "dest_ip": dst_ip,
                }, source_module=self.name)

            elif cmd in ("GET", "SET", "HGET", "HSET", "DEL", "KEYS",
                         "MGET", "MSET", "LPUSH", "RPUSH", "SADD", "ZADD"):
                query = " ".join(args[:3])  # Command + first 2 args max
                self._stats["redis_events"] += 1
                self._stats["queries_captured"] += 1
                self._record_query(
                    ts, src_ip, dst_ip, "redis", "", "",
                    query[:MAX_QUERY_LEN], "query"
                )

    def _parse_resp_array(self, data: bytes) -> list:
        """Parse a RESP array into a list of string arguments."""
        args = []
        lines = data.split(REDIS_CRLF)
        if not lines or not lines[0].startswith(b'*'):
            return args

        try:
            count = int(lines[0][1:])
        except ValueError:
            return args

        idx = 1
        while len(args) < count and idx < len(lines) - 1:
            if lines[idx].startswith(b'$'):
                try:
                    str_len = int(lines[idx][1:])
                except ValueError:
                    break
                idx += 1
                if idx < len(lines):
                    args.append(lines[idx].decode("utf-8", errors="replace")[:MAX_QUERY_LEN])
                idx += 1
            else:
                # Inline format
                args.append(lines[idx].decode("utf-8", errors="replace"))
                idx += 1

        return args

    # ------------------------------------------------------------------
    # MongoDB parser
    # ------------------------------------------------------------------

    def _parse_mongodb(self, payload: bytes, src_ip: str, dst_ip: str,
                       is_to_server: bool, ts: float) -> None:
        """Parse MongoDB wire protocol for auth and find operations."""
        if len(payload) < 16:
            return

        # MongoDB wire protocol header: length(4 LE) + requestID(4) + responseTo(4) + opCode(4)
        msg_len = struct.unpack("<I", payload[0:4])[0]
        opcode = struct.unpack("<I", payload[12:16])[0]

        if opcode != MONGO_OP_MSG or len(payload) < 21:
            return

        if not is_to_server:
            return

        # OP_MSG: flagBits(4) + sections
        # Section kind 0: body (BSON document)
        sections_start = 20
        if sections_start >= len(payload):
            return

        section_kind = payload[sections_start]
        if section_kind != 0:
            return

        bson_data = payload[sections_start + 1:]
        if len(bson_data) < 5:
            return

        # Minimal BSON parsing: look for command keys
        bson_str = bson_data.decode("utf-8", errors="replace")

        if "saslStart" in bson_str or "authenticate" in bson_str:
            self._stats["mongodb_events"] += 1
            self._stats["logins_captured"] += 1
            # Extract mechanism if visible
            mechanism = ""
            if "SCRAM-SHA" in bson_str:
                mechanism = "SCRAM-SHA"
            elif "PLAIN" in bson_str:
                mechanism = "PLAIN"
            self._record_query(
                ts, src_ip, dst_ip, "mongodb", "", "",
                f"AUTH mechanism={mechanism}", "login"
            )

        elif "find" in bson_str or "aggregate" in bson_str or "insert" in bson_str:
            # Extract collection name heuristically
            operation = "find" if "find" in bson_str else ("aggregate" if "aggregate" in bson_str else "insert")
            self._stats["mongodb_events"] += 1
            self._stats["queries_captured"] += 1
            self._record_query(
                ts, src_ip, dst_ip, "mongodb", "", "",
                f"{operation} (BSON)", "query"
            )

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

    def _resolve_db_path(self) -> str:
        base = self.config.get("db_dir", os.path.join(
            os.path.expanduser("~"), ".bigbrother"
        ))
        return os.path.join(base, "db_interceptor.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_query(self, ts: float, src_ip: str, dst_ip: str,
                      db_type: str, username: str, database: str,
                      query_text: str, operation: str) -> None:
        with self._buffer_lock:
            self._write_buffer.append((
                ts, src_ip, dst_ip, db_type, username, database,
                query_text[:MAX_QUERY_LEN], operation
            ))

    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 db_queries
                       (timestamp, source_ip, dest_ip, db_type, username,
                        database, query_text, operation)
                       VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
                    batch,
                )
        except Exception:
            logger.exception("Failed to flush DB queries")

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