#!/usr/bin/env python3
"""Passive DNS query logger — builds per-host browsing history from wire.

Subscribes to capture_bus with BPF "port 53". Parses DNS query/response
packets using struct (no scapy dependency). Buffers records and batch-flushes
to SQLite every 60 seconds or 1000 records.

Also flags DoH connections to known resolvers on port 443 as DNS blind spots.
"""

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.dns_logger")

# Well-known DoH resolver IPs
DOH_RESOLVERS = frozenset([
    "1.1.1.1", "1.0.0.1",             # Cloudflare
    "8.8.8.8", "8.8.4.4",             # Google
    "9.9.9.9", "149.112.112.112",     # Quad9
    "208.67.222.222", "208.67.220.220",  # OpenDNS
])

# DNS query type map
QTYPES = {
    1: "A", 2: "NS", 5: "CNAME", 6: "SOA", 12: "PTR",
    15: "MX", 16: "TXT", 28: "AAAA", 33: "SRV", 35: "NAPTR",
    43: "DS", 46: "RRSIG", 47: "NSEC", 48: "DNSKEY",
    52: "TLSA", 65: "HTTPS", 257: "CAA", 255: "ANY",
}


class DNSLogger(BaseModule):
    """Log all DNS queries and responses per source IP."""

    name = "dns_logger"
    module_type = "passive"
    priority = 100
    requires_root = True

    BATCH_SIZE = 1000
    FLUSH_INTERVAL = 60  # seconds

    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_queries = 0
        self._doh_detections = 0

    # ------------------------------------------------------------------
    # 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("DNSLogger requires capture_bus in config")
            return

        # Database setup
        base_dir = self.config.get("data_dir", os.path.expanduser("~/.bigbrother"))
        self._db_path = os.path.join(base_dir, "dns_queries.db")
        Path(os.path.dirname(self._db_path)).mkdir(parents=True, exist_ok=True)
        self._init_db()

        # Subscribe to capture bus for DNS traffic
        self._sub_queue = self._capture_bus.subscribe(
            name=self.name, bpf_filter="port 53", queue_depth=10000
        )

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

        # Packet reader thread
        self._reader_thread = threading.Thread(
            target=self._read_packets, daemon=True, name="bb-dns-reader"
        )
        self._reader_thread.start()

        # Periodic flusher thread
        self._flusher_thread = threading.Thread(
            target=self._flush_loop, daemon=True, name="bb-dns-flusher"
        )
        self._flusher_thread.start()

        self.state.set_module_status(self.name, "running", pid=os.getpid())
        logger.info("DNSLogger started — listening for DNS on port 53")

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

        # Unsubscribe from capture bus
        if self._capture_bus:
            self._capture_bus.unsubscribe(self.name)

        # Final flush
        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("DNSLogger stopped — %d total queries logged", self._total_queries)

    def status(self) -> dict:
        return {
            "running": self._running,
            "pid": self._pid,
            "uptime": time.time() - self._start_time if self._start_time else 0,
            "total_queries": self._total_queries,
            "doh_detections": self._doh_detections,
            "buffer_size": len(self._buffer),
        }

    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 dns_queries (
                id          INTEGER PRIMARY KEY AUTOINCREMENT,
                timestamp   REAL NOT NULL,
                source_ip   TEXT NOT NULL,
                domain      TEXT NOT NULL,
                query_type  TEXT,
                response_ips TEXT,
                ttl         INTEGER,
                is_doh      INTEGER DEFAULT 0
            );
            CREATE INDEX IF NOT EXISTS idx_dns_source ON dns_queries(source_ip);
            CREATE INDEX IF NOT EXISTS idx_dns_domain ON dns_queries(domain);
            CREATE INDEX IF NOT EXISTS idx_dns_ts ON dns_queries(timestamp);
        """)
        self._db_conn.commit()

    # ------------------------------------------------------------------
    # Packet reader
    # ------------------------------------------------------------------

    def _read_packets(self) -> None:
        """Pull packets from capture bus queue and parse DNS."""
        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  # Malformed packets are silently dropped

    def _parse_packet(self, ts: float, raw: bytes) -> None:
        """Parse an Ethernet frame containing a DNS packet."""
        if len(raw) < 14:
            return

        # Ethernet header
        eth_type = struct.unpack("!H", raw[12:14])[0]
        if eth_type == 0x8100:  # 802.1Q VLAN
            if len(raw) < 18:
                return
            eth_type = struct.unpack("!H", raw[16:18])[0]
            ip_offset = 18
        else:
            ip_offset = 14

        if eth_type != 0x0800:  # IPv4 only
            return

        if len(raw) < ip_offset + 20:
            return

        # IPv4 header
        ip_header = raw[ip_offset:]
        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])

        # UDP (17) or TCP (6) transport
        transport_offset = ip_offset + ihl

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

            # Check for DoH: known resolvers on port 443
            if dst_port == 443 and dst_ip in DOH_RESOLVERS:
                self._doh_detections += 1
                record = (ts, src_ip, f"[DoH:{dst_ip}]", "DoH", "", 0, 1)
                with self._buffer_lock:
                    self._buffer.append(record)
                    if len(self._buffer) >= self.BATCH_SIZE:
                        self._flush_buffer()
                return

            if src_port != 53 and dst_port != 53:
                return

        elif ip_proto == 6:  # TCP DNS (rare, used for large responses)
            if len(raw) < transport_offset + 20:
                return
            src_port, dst_port = struct.unpack("!HH", raw[transport_offset:transport_offset + 4])
            if src_port != 53 and dst_port != 53:
                return
            tcp_header_len = ((raw[transport_offset + 12] >> 4) & 0xF) * 4
            dns_offset = transport_offset + tcp_header_len
            # TCP DNS has 2-byte length prefix
            if len(raw) < dns_offset + 2:
                return
            dns_offset += 2
        else:
            return

        # Parse DNS payload
        dns_data = raw[dns_offset:]
        if len(dns_data) < 12:
            return

        self._parse_dns(ts, src_ip, dst_ip, src_port, dst_port, dns_data)

    def _parse_dns(self, ts: float, src_ip: str, dst_ip: str,
                   src_port: int, dst_port: int, data: bytes) -> None:
        """Parse DNS header + question/answer sections."""
        # DNS header: ID(2) FLAGS(2) QDCOUNT(2) ANCOUNT(2) NSCOUNT(2) ARCOUNT(2)
        txid, flags, qdcount, ancount, nscount, arcount = struct.unpack(
            "!HHHHHH", data[:12]
        )
        is_response = bool(flags & 0x8000)
        offset = 12

        # Parse questions
        domains = []
        query_types = []
        for _ in range(qdcount):
            domain, offset = self._read_name(data, offset)
            if offset + 4 > len(data):
                return
            qtype, qclass = struct.unpack("!HH", data[offset:offset + 4])
            offset += 4
            if domain:
                domains.append(domain)
                query_types.append(QTYPES.get(qtype, str(qtype)))

        # Parse answers (responses only)
        response_ips = []
        min_ttl = 0
        if is_response:
            for _ in range(ancount):
                if offset >= len(data):
                    break
                _name, offset = self._read_name(data, offset)
                if offset + 10 > len(data):
                    break
                rtype, _rclass, ttl, rdlength = struct.unpack(
                    "!HHIH", data[offset:offset + 10]
                )
                offset += 10
                if offset + rdlength > len(data):
                    break

                if rtype == 1 and rdlength == 4:  # A record
                    ip = socket.inet_ntoa(data[offset:offset + 4])
                    response_ips.append(ip)
                elif rtype == 28 and rdlength == 16:  # AAAA record
                    try:
                        ip = socket.inet_ntop(socket.AF_INET6, data[offset:offset + 16])
                        response_ips.append(ip)
                    except Exception:
                        pass

                if ttl > 0:
                    min_ttl = ttl if min_ttl == 0 else min(min_ttl, ttl)

                offset += rdlength

        # Build records for each queried domain
        # For queries: source_ip is the querier
        # For responses: dst_ip is the querier (source is the DNS server)
        querier_ip = src_ip if not is_response else dst_ip

        for i, domain in enumerate(domains):
            if not domain or domain == ".":
                continue
            qtype = query_types[i] if i < len(query_types) else ""
            resp_str = ",".join(response_ips) if response_ips else ""

            record = (ts, querier_ip, domain, qtype, resp_str, min_ttl, 0)
            self._total_queries += 1

            with self._buffer_lock:
                self._buffer.append(record)
                if len(self._buffer) >= self.BATCH_SIZE:
                    self._flush_buffer()

    @staticmethod
    def _read_name(data: bytes, offset: int, max_jumps: int = 10) -> tuple:
        """Read a DNS name with pointer compression. 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:
                # Pointer
                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 > max_jumps:
                    break
                continue

            offset += 1
            if offset + length > len(data):
                break
            try:
                parts.append(data[offset:offset + length].decode("utf-8", errors="replace"))
            except Exception:
                parts.append("?")
            offset += length

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

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

    def _flush_loop(self) -> None:
        """Periodically flush DNS record buffer to SQLite."""
        while self._running:
            time.sleep(self.FLUSH_INTERVAL)
            try:
                self._flush_buffer()
            except Exception:
                logger.exception("DNS flush error")

    def _flush_buffer(self) -> None:
        """Write buffered DNS records to SQLite."""
        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 dns_queries (timestamp, source_ip, domain, query_type, response_ips, ttl, is_doh) "
                "VALUES (?, ?, ?, ?, ?, ?, ?)",
                batch,
            )
            self._db_conn.commit()
        except Exception:
            logger.exception("Failed to flush %d DNS records", len(batch))