#!/usr/bin/env python3
"""SMB file access monitor — passive SMB2/3 session and file tracking.

Parses SMB2/3 Tree Connect requests (share names), Create requests (file paths),
and Read/Write operations.  Tracks per-user share access and detects
GPP/SYSVOL access patterns indicating potential Group Policy Preference
password exposure.  SMB3 encrypted sessions are opaque.
"""

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__)

# SMB2 command IDs
SMB2_NEGOTIATE = 0x0000
SMB2_SESSION_SETUP = 0x0001
SMB2_TREE_CONNECT = 0x0003
SMB2_CREATE = 0x0005
SMB2_READ = 0x0008
SMB2_WRITE = 0x0009

# SMB2 header magic
SMB2_MAGIC = b'\xfeSMB'

# GPP-related paths (case-insensitive matching)
GPP_PATTERNS = (
    "groups.xml", "services.xml", "scheduledtasks.xml",
    "datasources.xml", "printers.xml", "drives.xml",
    "policies", "sysvol",
)

_DB_INIT = """
CREATE TABLE IF NOT EXISTS smb_access (
    id           INTEGER PRIMARY KEY AUTOINCREMENT,
    timestamp    REAL NOT NULL,
    source_ip    TEXT NOT NULL,
    username     TEXT DEFAULT '',
    share_name   TEXT DEFAULT '',
    file_path    TEXT DEFAULT '',
    operation    TEXT NOT NULL,
    smb_version  TEXT DEFAULT ''
);
CREATE INDEX IF NOT EXISTS idx_smb_user ON smb_access(username);
CREATE INDEX IF NOT EXISTS idx_smb_share ON smb_access(share_name);
CREATE INDEX IF NOT EXISTS idx_smb_ts ON smb_access(timestamp);
"""

FLUSH_INTERVAL = 30


class SMBMonitor(BaseModule):
    """Monitor SMB2/3 file share access passively."""

    name = "smb_monitor"
    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 session ID -> username mapping
        self._sessions: dict[tuple, str] = {}  # (src_ip, session_id) -> username
        # Track tree ID -> share name mapping
        self._trees: dict[tuple, str] = {}     # (src_ip, tree_id) -> share_name
        self._stats = {
            "packets_processed": 0,
            "tree_connects": 0,
            "file_creates": 0,
            "reads": 0,
            "writes": 0,
            "gpp_access_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("SMBMonitor requires _capture_bus in config")
            return

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

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

    def _process_packet(self, pkt: bytes, ts: float) -> None:
        """Extract TCP payload and parse SMB2."""
        if len(pkt) < 54:  # eth(14) + ip(20) + tcp(20) minimum
            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

        # SMB uses NetBIOS Session Service: 1-byte type + 3-byte length
        # Type 0x00 = session message
        if payload[0] != 0x00:
            return
        nb_len = struct.unpack("!I", b'\x00' + payload[1:4])[0]
        smb_data = payload[4:]

        if len(smb_data) < 64:
            return

        # Determine direction: request goes TO port 445
        client_ip = src_ip if dst_port == 445 else dst_ip

        self._parse_smb2(smb_data, client_ip, ts)

    def _parse_smb2(self, data: bytes, client_ip: str, ts: float) -> None:
        """Parse an SMB2 message header and dispatch to command handler."""
        if data[:4] != SMB2_MAGIC:
            return

        if len(data) < 64:
            return

        # SMB2 header (64 bytes)
        header_len = struct.unpack("<H", data[4:6])[0]
        command = struct.unpack("<H", data[12:14])[0]
        flags = struct.unpack("<I", data[16:20])[0]
        is_response = bool(flags & 0x00000001)

        session_id = struct.unpack("<Q", data[40:48])[0]
        tree_id = struct.unpack("<I", data[36:40])[0]

        smb_body = data[64:]

        if command == SMB2_TREE_CONNECT and is_response and len(smb_body) >= 8:
            # Tree Connect Response doesn't contain the share name,
            # but we see the path in the request
            pass
        elif command == SMB2_TREE_CONNECT and not is_response:
            self._handle_tree_connect_request(smb_body, client_ip, session_id, tree_id, ts)
        elif command == SMB2_CREATE and not is_response:
            self._handle_create_request(smb_body, client_ip, session_id, tree_id, ts)
        elif command == SMB2_READ and not is_response:
            self._stats["reads"] += 1
            share = self._trees.get((client_ip, tree_id), "")
            self._record_access(ts, client_ip, "", share, "", "read", "smb2")
        elif command == SMB2_WRITE and not is_response:
            self._stats["writes"] += 1
            share = self._trees.get((client_ip, tree_id), "")
            self._record_access(ts, client_ip, "", share, "", "write", "smb2")

    def _handle_tree_connect_request(self, body: bytes, client_ip: str,
                                      session_id: int, tree_id: int,
                                      ts: float) -> None:
        """Parse SMB2 Tree Connect request to extract share name."""
        # Tree Connect Request body (SMB 3.1.1):
        # StructureSize(2) + Reserved/Flags(2) + PathOffset(2) + PathLength(2)
        if len(body) < 8:
            return

        path_offset = struct.unpack("<H", body[4:6])[0]
        path_length = struct.unpack("<H", body[6:8])[0]

        # PathOffset is relative to the beginning of the SMB2 header
        # Since body starts at offset 64, adjust
        # But we receive body starting at offset 0, so the path is at:
        # (path_offset - 64) in body terms, OR it may follow immediately
        # In practice, path follows the 8-byte struct at body[8:]
        if path_length > 0 and len(body) >= 8 + path_length:
            try:
                share_path = body[8:8 + path_length].decode("utf-16-le", errors="replace")
            except Exception:
                share_path = ""

            # Extract just the share name from \\server\share
            share_name = share_path.rsplit("\\", 1)[-1] if "\\" in share_path else share_path

            self._trees[(client_ip, tree_id)] = share_name
            self._stats["tree_connects"] += 1

            username = self._sessions.get((client_ip, session_id), "")
            self._record_access(ts, client_ip, username, share_name, "", "tree_connect", "smb2")

            # Check for SYSVOL/NETLOGON access
            if share_name.lower() in ("sysvol", "netlogon"):
                logger.info("SYSVOL/NETLOGON access from %s — potential GPP exposure", client_ip)

    def _handle_create_request(self, body: bytes, client_ip: str,
                                session_id: int, tree_id: int,
                                ts: float) -> None:
        """Parse SMB2 Create request to extract file path."""
        # Create Request: StructureSize(2) + SecurityFlags(1) + RequestedOplockLevel(1)
        # + ImpersonationLevel(4) + SmbCreateFlags(8) + Reserved(8)
        # + DesiredAccess(4) + FileAttributes(4) + ShareAccess(4)
        # + CreateDisposition(4) + CreateOptions(4) + NameOffset(2) + NameLength(2)
        if len(body) < 56:
            return

        name_offset = struct.unpack("<H", body[44:46])[0]
        name_length = struct.unpack("<H", body[46:48])[0]

        if name_length == 0:
            return

        # Name follows the fixed-size structure
        # name_offset is relative to SMB2 header start, body starts at +64
        name_start = 56  # After the fixed-size fields in Create body
        if len(body) >= name_start + name_length:
            try:
                file_path = body[name_start:name_start + name_length].decode(
                    "utf-16-le", errors="replace"
                )
            except Exception:
                file_path = ""

            if file_path:
                self._stats["file_creates"] += 1
                share = self._trees.get((client_ip, tree_id), "")
                username = self._sessions.get((client_ip, session_id), "")
                self._record_access(ts, client_ip, username, share, file_path, "create", "smb2")

                # GPP detection
                file_lower = file_path.lower()
                for pattern in GPP_PATTERNS:
                    if pattern in file_lower:
                        self._stats["gpp_access_detected"] += 1
                        emit_credential_found(self.bus, self.name, {
                            "source": "smb_gpp_access",
                            "client_ip": client_ip,
                            "share": share,
                            "file_path": file_path,
                            "detail": "GPP/SYSVOL file access — may contain cleartext passwords",
                        })
                        break

    # ------------------------------------------------------------------
    # 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, "smb_monitor.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_access(self, ts: float, source_ip: str, username: str,
                       share_name: str, file_path: str, operation: str,
                       smb_version: str) -> None:
        with self._buffer_lock:
            self._write_buffer.append((
                ts, source_ip, username, share_name, file_path, operation, smb_version
            ))

    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 smb_access
                       (timestamp, source_ip, username, share_name, file_path,
                        operation, smb_version)
                       VALUES (?, ?, ?, ?, ?, ?, ?)""",
                    batch,
                )
        except Exception:
            logger.exception("Failed to flush SMB access records")

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