#!/usr/bin/env python3
"""Operator audit trail — append-only log with HMAC integrity chain.

Logs all operator actions: SSH sessions, CLI commands, module activations,
config changes.  Each row includes an HMAC of the previous row, forming
a tamper-evident chain.  Subscribes to MODULE_STARTED, MODULE_STOPPED,
and KILL_SWITCH events.
"""

import hashlib
import hmac
import json
import logging
import os
import sqlite3
import subprocess
import threading
import time
from typing import Optional

from modules.base import BaseModule

logger = logging.getLogger(__name__)

_SCHEMA = """
CREATE TABLE IF NOT EXISTS audit_log (
    id          INTEGER PRIMARY KEY AUTOINCREMENT,
    timestamp   REAL    NOT NULL,
    action      TEXT    NOT NULL,
    operator    TEXT    DEFAULT '',
    details     TEXT    DEFAULT '',
    hmac        TEXT    NOT NULL
);

CREATE INDEX IF NOT EXISTS idx_audit_ts     ON audit_log(timestamp);
CREATE INDEX IF NOT EXISTS idx_audit_action ON audit_log(action);
"""

# Action types
ACTION_MODULE_START   = "module_start"
ACTION_MODULE_STOP    = "module_stop"
ACTION_KILL_SWITCH    = "kill_switch"
ACTION_CONFIG_CHANGE  = "config_change"
ACTION_SSH_SESSION    = "ssh_session"
ACTION_CLI_COMMAND    = "cli_command"
ACTION_CRED_EXPORT    = "cred_export"
ACTION_DATA_EXFIL     = "data_exfil"
ACTION_BASELINE_BUILD = "baseline_build"
ACTION_MANUAL         = "manual"


class OperatorAudit(BaseModule):
    """Append-only operator audit trail with HMAC integrity chain."""

    name = "operator_audit"
    module_type = "intel"
    priority = 50  # Critical — always runs
    requires_root = False

    def __init__(self, bus, state, config, engine=None):
        super().__init__(bus, state, config, engine)
        self._db_path = ""
        self._conn: Optional[sqlite3.Connection] = None
        self._lock = threading.Lock()
        self._hmac_key = b""
        self._last_hmac = ""
        self._entry_count = 0

    # ------------------------------------------------------------------
    # BaseModule interface
    # ------------------------------------------------------------------

    def start(self) -> None:
        if self._running:
            return

        base_dir = self.config.get("data_dir", os.path.expanduser("~/.implant"))
        self._db_path = os.path.join(base_dir, "operator_audit.db")
        os.makedirs(os.path.dirname(self._db_path), exist_ok=True)

        self._conn = sqlite3.connect(self._db_path, check_same_thread=False)
        self._conn.execute("PRAGMA journal_mode=WAL")
        self._conn.execute("PRAGMA synchronous=NORMAL")
        self._conn.row_factory = sqlite3.Row
        self._conn.executescript(_SCHEMA)

        # HMAC key: from config or derive from machine ID
        key_config = self.config.get("audit_hmac_key", "")
        if key_config:
            self._hmac_key = key_config.encode() if isinstance(key_config, str) else key_config
        else:
            self._hmac_key = self._derive_machine_key()

        # Load last HMAC from chain
        self._last_hmac = self._get_last_hmac()

        # Subscribe to events
        self.bus.subscribe(self._on_module_started, "MODULE_STARTED")
        self.bus.subscribe(self._on_module_stopped, "MODULE_STOPPED")
        self.bus.subscribe(self._on_kill_switch, "KILL_SWITCH")

        self._running = True
        self._pid = os.getpid()
        self._start_time = time.time()
        self.state.set_module_status(self.name, "running", pid=self._pid)

        # Log our own start
        self.log_action(
            ACTION_MODULE_START,
            details="OperatorAudit started",
        )

        logger.info(
            "OperatorAudit started — db=%s, chain_length=%d",
            self._db_path, self._entry_count,
        )

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

        # Log our own stop before shutting down
        self.log_action(
            ACTION_MODULE_STOP,
            details="OperatorAudit stopping",
        )

        self._running = False

        self.bus.unsubscribe(self._on_module_started, "MODULE_STARTED")
        self.bus.unsubscribe(self._on_module_stopped, "MODULE_STOPPED")
        self.bus.unsubscribe(self._on_kill_switch, "KILL_SWITCH")

        if self._conn:
            self._conn.close()
            self._conn = None

        self.state.set_module_status(self.name, "stopped")
        logger.info(
            "OperatorAudit stopped — %d entries in chain", self._entry_count
        )

    def status(self) -> dict:
        return {
            "running": self._running,
            "pid": self._pid,
            "uptime": time.time() - self._start_time if self._start_time else 0,
            "entries": self._entry_count,
            "chain_valid": self.verify_chain(),
        }

    def configure(self, config: dict) -> None:
        old_config = json.dumps(self.config, sort_keys=True, default=str)
        self.config.update(config)
        new_config = json.dumps(self.config, sort_keys=True, default=str)

        if old_config != new_config:
            self.log_action(
                ACTION_CONFIG_CHANGE,
                details=f"Config updated",
            )

    # ------------------------------------------------------------------
    # Event handlers
    # ------------------------------------------------------------------

    def _on_module_started(self, event) -> None:
        p = event.payload
        module_name = p.get("module", event.source_module or "unknown")
        if module_name == self.name:
            return  # Skip our own start event (already logged)
        self.log_action(
            ACTION_MODULE_START,
            details=f"Module started: {module_name} (pid={p.get('pid', '?')})",
        )

    def _on_module_stopped(self, event) -> None:
        p = event.payload
        module_name = p.get("module", event.source_module or "unknown")
        if module_name == self.name:
            return
        self.log_action(
            ACTION_MODULE_STOP,
            details=f"Module stopped: {module_name}",
        )

    def _on_kill_switch(self, event) -> None:
        p = event.payload
        self.log_action(
            ACTION_KILL_SWITCH,
            details=f"KILL SWITCH activated: reason={p.get('reason', 'unknown')}",
        )

    # ------------------------------------------------------------------
    # Core logging
    # ------------------------------------------------------------------

    def log_action(self, action: str, operator: str = "",
                   details: str = "") -> None:
        """Log an operator action with HMAC chain integrity.

        Args:
            action: Action type (use ACTION_* constants).
            operator: Operator identifier (SSH key fingerprint, username, etc).
            details: Free-text description of the action.
        """
        if not self._conn:
            return

        if not operator:
            operator = self._detect_operator()

        ts = time.time()

        # Compute HMAC: H(key, previous_hmac | timestamp | action | operator | details)
        msg = f"{self._last_hmac}|{ts}|{action}|{operator}|{details}"
        entry_hmac = hmac.new(
            self._hmac_key, msg.encode(), hashlib.sha256
        ).hexdigest()

        with self._lock:
            try:
                self._conn.execute(
                    """INSERT INTO audit_log
                       (timestamp, action, operator, details, hmac)
                       VALUES (?, ?, ?, ?, ?)""",
                    (ts, action, operator, details, entry_hmac),
                )
                self._conn.commit()
                self._last_hmac = entry_hmac
                self._entry_count += 1
            except Exception:
                logger.exception("Failed to log audit action: %s", action)

    def log_ssh_session(self, remote_ip: str, key_fingerprint: str = "",
                        event: str = "connect") -> None:
        """Log an SSH session event."""
        self.log_action(
            ACTION_SSH_SESSION,
            operator=key_fingerprint or remote_ip,
            details=f"SSH {event} from {remote_ip} "
                    f"(key={key_fingerprint or 'password'})",
        )

    def log_cli_command(self, command: str, operator: str = "") -> None:
        """Log a CLI command execution."""
        self.log_action(
            ACTION_CLI_COMMAND,
            operator=operator,
            details=f"Command: {command}",
        )

    # ------------------------------------------------------------------
    # HMAC chain
    # ------------------------------------------------------------------

    def _derive_machine_key(self) -> bytes:
        """Derive HMAC key from machine-specific identifiers."""
        components = []

        # CPU serial (RPi)
        try:
            with open("/proc/cpuinfo", "r") as f:
                for line in f:
                    if line.startswith("Serial"):
                        components.append(line.strip())
                        break
        except (IOError, PermissionError):
            pass

        # Machine ID
        try:
            with open("/etc/machine-id", "r") as f:
                components.append(f.read().strip())
        except (IOError, PermissionError):
            pass

        if not components:
            components.append("bigbrother-default-key")

        combined = "|".join(components)
        return hashlib.sha256(combined.encode()).digest()

    def _get_last_hmac(self) -> str:
        """Get the HMAC of the last entry in the chain."""
        with self._lock:
            row = self._conn.execute(
                "SELECT hmac FROM audit_log ORDER BY id DESC LIMIT 1"
            ).fetchone()

            count_row = self._conn.execute(
                "SELECT COUNT(*) FROM audit_log"
            ).fetchone()
            self._entry_count = count_row[0] if count_row else 0

        if row:
            return row["hmac"]
        return ""  # Genesis — empty for first entry

    def _detect_operator(self) -> str:
        """Attempt to detect the current operator identity."""
        # Check SSH connection
        ssh_client = os.environ.get("SSH_CLIENT", "")
        if ssh_client:
            parts = ssh_client.split()
            remote_ip = parts[0] if parts else "unknown"

            # Try to get SSH key fingerprint
            try:
                result = subprocess.run(
                    ["who", "-m"],
                    capture_output=True, timeout=2,
                )
                who = result.stdout.decode().strip()
                if who:
                    return f"ssh:{remote_ip}:{who.split()[0]}"
            except Exception:
                pass

            return f"ssh:{remote_ip}"

        # Local operator
        user = os.environ.get("USER", os.environ.get("LOGNAME", "local"))
        return f"local:{user}"

    # ------------------------------------------------------------------
    # Chain verification
    # ------------------------------------------------------------------

    def verify_chain(self) -> bool:
        """Verify the entire HMAC chain integrity.

        Returns:
            True if the chain is valid, False if tampered.
        """
        if not self._conn:
            return False

        with self._lock:
            rows = self._conn.execute(
                "SELECT timestamp, action, operator, details, hmac "
                "FROM audit_log ORDER BY id ASC"
            ).fetchall()

        if not rows:
            return True  # Empty chain is valid

        prev_hmac = ""
        for row in rows:
            ts = row["timestamp"]
            action = row["action"]
            operator = row["operator"]
            details = row["details"]
            stored_hmac = row["hmac"]

            msg = f"{prev_hmac}|{ts}|{action}|{operator}|{details}"
            computed = hmac.new(
                self._hmac_key, msg.encode(), hashlib.sha256
            ).hexdigest()

            if computed != stored_hmac:
                logger.error(
                    "HMAC chain broken at ts=%.3f action=%s — "
                    "expected %s, got %s",
                    ts, action, computed[:16], stored_hmac[:16],
                )
                return False

            prev_hmac = stored_hmac

        return True

    # ------------------------------------------------------------------
    # Export
    # ------------------------------------------------------------------

    def export_log(self, since: float = 0, limit: int = 10000) -> list:
        """Export audit log entries for engagement reports."""
        with self._lock:
            rows = self._conn.execute(
                """SELECT * FROM audit_log
                   WHERE timestamp > ?
                   ORDER BY id ASC LIMIT ?""",
                (since, limit),
            ).fetchall()

        return [dict(r) for r in rows]

    def export_json(self, since: float = 0) -> str:
        """Export audit log as JSON string."""
        entries = self.export_log(since)
        return json.dumps(entries, indent=2, default=str)

    def get_summary(self) -> dict:
        """Return audit log summary."""
        with self._lock:
            rows = self._conn.execute(
                """SELECT action, COUNT(*) FROM audit_log
                   GROUP BY action ORDER BY COUNT(*) DESC"""
            ).fetchall()

            operators = self._conn.execute(
                """SELECT DISTINCT operator FROM audit_log
                   WHERE operator != ''"""
            ).fetchall()

        return {
            "total_entries": self._entry_count,
            "chain_valid": self.verify_chain(),
            "by_action": {a: c for a, c in rows},
            "operators": [r[0] for r in operators],
        }