#!/usr/bin/env python3
"""Per-user activity timeline — aggregates login events, service access,
file operations, and web activity into chronological user profiles.

Ingests data from dns_logger, auth_flow_tracker, smb_monitor, and
credential_sniffer via the event bus and periodic state queries.
Identifies work hours, admin activity windows, and service account behavior.
"""

import json
import logging
import os
import sqlite3
import threading
import time
from collections import defaultdict
from typing import Optional

from modules.base import BaseModule

logger = logging.getLogger("sensor.intel.user_timeline")

_SCHEMA = """
CREATE TABLE IF NOT EXISTS user_events (
    id        INTEGER PRIMARY KEY AUTOINCREMENT,
    timestamp REAL    NOT NULL,
    username  TEXT    NOT NULL,
    domain    TEXT    DEFAULT '',
    event_type TEXT   NOT NULL,
    source_ip TEXT,
    target_ip TEXT,
    service   TEXT,
    detail    TEXT    DEFAULT '',
    source_module TEXT DEFAULT ''
);

CREATE INDEX IF NOT EXISTS idx_ue_user  ON user_events(username);
CREATE INDEX IF NOT EXISTS idx_ue_ts    ON user_events(timestamp);
CREATE INDEX IF NOT EXISTS idx_ue_type  ON user_events(event_type);

CREATE TABLE IF NOT EXISTS user_profiles (
    username       TEXT PRIMARY KEY,
    domain         TEXT DEFAULT '',
    first_seen     REAL,
    last_seen      REAL,
    event_count    INTEGER DEFAULT 0,
    is_admin       INTEGER DEFAULT 0,
    is_service_acct INTEGER DEFAULT 0,
    typical_hours  TEXT DEFAULT '',
    workstations   TEXT DEFAULT '',
    services_used  TEXT DEFAULT '',
    notes          TEXT DEFAULT ''
);
"""

# Event types
EVENT_LOGIN = "login"
EVENT_LOGOUT = "logout"
EVENT_AUTH_FAIL = "auth_fail"
EVENT_SERVICE_ACCESS = "service_access"
EVENT_FILE_ACCESS = "file_access"
EVENT_WEB_BROWSE = "web_browse"
EVENT_ADMIN_ACTION = "admin_action"
EVENT_CRED_CAPTURED = "cred_captured"

# Admin indicators: services or patterns that suggest admin activity
_ADMIN_SERVICES = {
    "rdp", "ssh", "winrm", "wmi", "psremoting", "dcom",
    "ldap", "kerberos", "smb_admin", "rpc",
}

_ADMIN_USERNAMES = {
    "administrator", "admin", "root", "domain admin",
}

# Service account patterns
_SVC_PATTERNS = (
    "svc_", "svc-", "service_", "sa_", "task_", "app_",
    "sql_", "iis_", "backup_", "scan_", "nessus", "splunk",
    "crowdstrike", "sccm", "wsus",
)


class UserTimeline(BaseModule):
    """Build per-user activity timelines from network observation data."""

    name = "user_timeline"
    module_type = "intel"
    priority = 200
    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._ingest_thread: Optional[threading.Thread] = None
        self._event_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, "user_timeline.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)

        # Subscribe to credential and host events for user correlation
        self.bus.subscribe(self._on_credential_found, "CREDENTIAL_FOUND")
        self.bus.subscribe(self._on_host_discovered, "HOST_DISCOVERED")

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

        # Periodic ingestion from other module state
        self._ingest_thread = threading.Thread(
            target=self._periodic_ingest, daemon=True,
            name="sensor-user-timeline-ingest",
        )
        self._ingest_thread.start()

        self.state.set_module_status(self.name, "running", pid=self._pid)
        logger.info("UserTimeline started — db=%s", self._db_path)

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

        self.bus.unsubscribe(self._on_credential_found, "CREDENTIAL_FOUND")
        self.bus.unsubscribe(self._on_host_discovered, "HOST_DISCOVERED")

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

        self.state.set_module_status(self.name, "stopped")
        logger.info("UserTimeline stopped — %d events recorded", self._event_count)

    def status(self) -> dict:
        user_count = 0
        if self._conn:
            try:
                with self._lock:
                    row = self._conn.execute(
                        "SELECT COUNT(DISTINCT username) FROM user_events"
                    ).fetchone()
                    user_count = row[0] if row else 0
            except Exception:
                pass

        return {
            "running": self._running,
            "pid": self._pid,
            "uptime": time.time() - self._start_time if self._start_time else 0,
            "total_events": self._event_count,
            "tracked_users": user_count,
        }

    def configure(self, config: dict) -> None:
        self.config.update(config)

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

    def _on_credential_found(self, event) -> None:
        """Record credential capture as a user event."""
        p = event.payload
        username = p.get("username", "")
        if not username:
            return

        self._record_event(
            username=username,
            domain=p.get("domain", ""),
            event_type=EVENT_CRED_CAPTURED,
            source_ip=p.get("source_ip", ""),
            target_ip=p.get("target_ip", ""),
            service=p.get("service", ""),
            detail=f"Type: {p.get('cred_type', 'unknown')}",
            source_module=event.source_module,
        )

    def _on_host_discovered(self, event) -> None:
        """Check for user information in host discovery data."""
        p = event.payload
        # Some host discovery may include logged-in user information
        username = p.get("logged_in_user", "")
        if username:
            self._record_event(
                username=username,
                domain=p.get("domain", ""),
                event_type=EVENT_LOGIN,
                source_ip=p.get("ip", ""),
                target_ip="",
                service="workstation",
                detail=f"Active on {p.get('hostname', p.get('ip', ''))}",
                source_module=event.source_module,
            )

    # ------------------------------------------------------------------
    # Event recording
    # ------------------------------------------------------------------

    def _record_event(self, username: str, domain: str, event_type: str,
                      source_ip: str = "", target_ip: str = "",
                      service: str = "", detail: str = "",
                      source_module: str = "",
                      timestamp: float = None) -> None:
        """Record a user activity event."""
        if not username:
            return

        ts = timestamp or time.time()
        self._event_count += 1

        with self._lock:
            try:
                self._conn.execute(
                    """INSERT INTO user_events
                       (timestamp, username, domain, event_type, source_ip,
                        target_ip, service, detail, source_module)
                       VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)""",
                    (ts, username, domain, event_type, source_ip,
                     target_ip, service, detail, source_module),
                )

                # Update user profile
                is_admin = 1 if (
                    service.lower() in _ADMIN_SERVICES
                    or username.lower() in _ADMIN_USERNAMES
                    or event_type == EVENT_ADMIN_ACTION
                ) else 0

                is_svc = 1 if any(
                    username.lower().startswith(p) for p in _SVC_PATTERNS
                ) else 0

                self._conn.execute(
                    """INSERT INTO user_profiles
                       (username, domain, first_seen, last_seen, event_count,
                        is_admin, is_service_acct)
                       VALUES (?, ?, ?, ?, 1, ?, ?)
                       ON CONFLICT(username) DO UPDATE SET
                           last_seen = MAX(excluded.last_seen, user_profiles.last_seen),
                           event_count = user_profiles.event_count + 1,
                           is_admin = MAX(excluded.is_admin, user_profiles.is_admin),
                           is_service_acct = MAX(excluded.is_service_acct, user_profiles.is_service_acct)
                    """,
                    (username, domain, ts, ts, is_admin, is_svc),
                )

                self._conn.commit()
            except Exception:
                logger.exception("Failed to record user event for %s", username)

    def ingest_auth_event(self, username: str, domain: str, source_ip: str,
                          target_ip: str, service: str, success: bool,
                          timestamp: float = None) -> None:
        """Ingest an authentication event from external modules."""
        event_type = EVENT_LOGIN if success else EVENT_AUTH_FAIL
        detail = "success" if success else "failure"
        self._record_event(
            username=username, domain=domain, event_type=event_type,
            source_ip=source_ip, target_ip=target_ip, service=service,
            detail=detail, source_module="auth_flow_tracker",
            timestamp=timestamp,
        )

    def ingest_file_access(self, username: str, source_ip: str,
                           file_path: str, action: str = "read",
                           timestamp: float = None) -> None:
        """Ingest a file access event (from smb_monitor)."""
        self._record_event(
            username=username, domain="", event_type=EVENT_FILE_ACCESS,
            source_ip=source_ip, target_ip="", service="smb",
            detail=f"{action}: {file_path}",
            source_module="smb_monitor", timestamp=timestamp,
        )

    def ingest_web_activity(self, source_ip: str, domain: str,
                            url: str = "", timestamp: float = None) -> None:
        """Ingest web browsing activity (correlated to user by IP)."""
        # Look up username by source_ip from recent events
        username = self._resolve_user_by_ip(source_ip)
        if not username:
            username = f"host:{source_ip}"

        self._record_event(
            username=username, domain="", event_type=EVENT_WEB_BROWSE,
            source_ip=source_ip, target_ip="", service="web",
            detail=url or domain,
            source_module="dns_logger", timestamp=timestamp,
        )

    def _resolve_user_by_ip(self, ip: str) -> str:
        """Try to resolve a username from a source IP using recent login events."""
        with self._lock:
            try:
                row = self._conn.execute(
                    """SELECT username FROM user_events
                       WHERE source_ip = ? AND event_type = ?
                       ORDER BY timestamp DESC LIMIT 1""",
                    (ip, EVENT_LOGIN),
                ).fetchone()
                return row["username"] if row else ""
            except Exception:
                return ""

    # ------------------------------------------------------------------
    # Periodic state ingestion
    # ------------------------------------------------------------------

    def _periodic_ingest(self) -> None:
        """Periodically pull data from other module state stores."""
        interval = self.config.get("timeline_ingest_interval", 120)
        while self._running:
            time.sleep(interval)
            try:
                self._ingest_from_auth_tracker()
                self._ingest_from_smb_monitor()
                self._ingest_from_dns_logger()
                self._update_user_profiles()
            except Exception:
                logger.exception("Periodic user timeline ingest failed")

    def _ingest_from_auth_tracker(self) -> None:
        """Pull authentication events from auth_flow_tracker state."""
        last_ts = self.state.get(self.name, "auth_tracker_last_ts")
        last_ts = float(last_ts) if last_ts else 0

        auth_json = self.state.get("auth_flow_tracker", "recent_events")
        if not auth_json:
            return

        try:
            events = json.loads(auth_json)
            for evt in events:
                ts = evt.get("timestamp", 0)
                if ts <= last_ts:
                    continue
                self.ingest_auth_event(
                    username=evt.get("username", ""),
                    domain=evt.get("domain", ""),
                    source_ip=evt.get("source_ip", ""),
                    target_ip=evt.get("target_ip", ""),
                    service=evt.get("service", ""),
                    success=evt.get("success", True),
                    timestamp=ts,
                )
                last_ts = max(last_ts, ts)

            self.state.set(self.name, "auth_tracker_last_ts", str(last_ts))
        except (json.JSONDecodeError, TypeError):
            pass

    def _ingest_from_smb_monitor(self) -> None:
        """Pull SMB file access events from smb_monitor state."""
        last_ts = self.state.get(self.name, "smb_monitor_last_ts")
        last_ts = float(last_ts) if last_ts else 0

        smb_json = self.state.get("smb_monitor", "recent_file_ops")
        if not smb_json:
            return

        try:
            events = json.loads(smb_json)
            for evt in events:
                ts = evt.get("timestamp", 0)
                if ts <= last_ts:
                    continue
                self.ingest_file_access(
                    username=evt.get("username", ""),
                    source_ip=evt.get("source_ip", ""),
                    file_path=evt.get("path", ""),
                    action=evt.get("action", "read"),
                    timestamp=ts,
                )
                last_ts = max(last_ts, ts)

            self.state.set(self.name, "smb_monitor_last_ts", str(last_ts))
        except (json.JSONDecodeError, TypeError):
            pass

    def _ingest_from_dns_logger(self) -> None:
        """Pull DNS query data for web activity correlation."""
        last_ts = self.state.get(self.name, "dns_logger_last_ts")
        last_ts = float(last_ts) if last_ts else 0

        dns_json = self.state.get("dns_logger", "recent_queries")
        if not dns_json:
            return

        try:
            queries = json.loads(dns_json)
            for q in queries:
                ts = q.get("timestamp", 0)
                if ts <= last_ts:
                    continue
                self.ingest_web_activity(
                    source_ip=q.get("client_ip", ""),
                    domain=q.get("query_name", ""),
                    timestamp=ts,
                )
                last_ts = max(last_ts, ts)

            self.state.set(self.name, "dns_logger_last_ts", str(last_ts))
        except (json.JSONDecodeError, TypeError):
            pass

    def _update_user_profiles(self) -> None:
        """Refresh computed fields on user profiles (work hours, workstations, etc)."""
        with self._lock:
            try:
                users = self._conn.execute(
                    "SELECT DISTINCT username FROM user_events"
                ).fetchall()

                for (username,) in users:
                    # Compute typical hours
                    rows = self._conn.execute(
                        """SELECT timestamp FROM user_events
                           WHERE username = ?""",
                        (username,),
                    ).fetchall()

                    hours = defaultdict(int)
                    for (ts,) in rows:
                        hour = time.localtime(ts).tm_hour
                        hours[hour] += 1

                    # Top active hours (>10% of activity)
                    total = sum(hours.values())
                    if total > 0:
                        typical = sorted(
                            [h for h, c in hours.items() if c / total > 0.1]
                        )
                    else:
                        typical = []

                    # Workstations (source IPs used for login)
                    ws_rows = self._conn.execute(
                        """SELECT DISTINCT source_ip FROM user_events
                           WHERE username = ? AND event_type = ? AND source_ip != ''""",
                        (username, EVENT_LOGIN),
                    ).fetchall()
                    workstations = [r[0] for r in ws_rows]

                    # Services used
                    svc_rows = self._conn.execute(
                        """SELECT DISTINCT service FROM user_events
                           WHERE username = ? AND service != ''""",
                        (username,),
                    ).fetchall()
                    services = [r[0] for r in svc_rows]

                    self._conn.execute(
                        """UPDATE user_profiles
                           SET typical_hours = ?, workstations = ?, services_used = ?
                           WHERE username = ?""",
                        (json.dumps(typical), json.dumps(workstations),
                         json.dumps(services), username),
                    )

                self._conn.commit()
            except Exception:
                logger.exception("Failed to update user profiles")

    # ------------------------------------------------------------------
    # Query interface
    # ------------------------------------------------------------------

    def get_timeline(self, username: str, limit: int = 200,
                     since: float = 0) -> list:
        """Get chronological activity timeline for a user."""
        with self._lock:
            rows = self._conn.execute(
                """SELECT * FROM user_events
                   WHERE username = ? AND timestamp > ?
                   ORDER BY timestamp DESC LIMIT ?""",
                (username, since, limit),
            ).fetchall()

        return [dict(r) for r in rows]

    def get_active_users(self, hours: float = 1.0) -> list:
        """Get users active within the given time window."""
        since = time.time() - (hours * 3600)
        with self._lock:
            rows = self._conn.execute(
                """SELECT username, domain, COUNT(*) as event_count,
                          MAX(timestamp) as last_active
                   FROM user_events
                   WHERE timestamp > ?
                   GROUP BY username
                   ORDER BY last_active DESC""",
                (since,),
            ).fetchall()

        return [dict(r) for r in rows]

    def get_admin_users(self) -> list:
        """Get all users identified as having admin privileges."""
        with self._lock:
            rows = self._conn.execute(
                """SELECT * FROM user_profiles
                   WHERE is_admin = 1
                   ORDER BY last_seen DESC"""
            ).fetchall()

        return [dict(r) for r in rows]

    def get_service_accounts(self) -> list:
        """Get all detected service accounts."""
        with self._lock:
            rows = self._conn.execute(
                """SELECT * FROM user_profiles
                   WHERE is_service_acct = 1
                   ORDER BY event_count DESC"""
            ).fetchall()

        return [dict(r) for r in rows]

    def get_user_profile(self, username: str) -> Optional[dict]:
        """Get computed profile for a user."""
        with self._lock:
            row = self._conn.execute(
                "SELECT * FROM user_profiles WHERE username = ?",
                (username,),
            ).fetchone()

        if not row:
            return None

        profile = dict(row)
        # Parse JSON fields
        for field in ("typical_hours", "workstations", "services_used"):
            try:
                profile[field] = json.loads(profile.get(field, "[]"))
            except (json.JSONDecodeError, TypeError):
                profile[field] = []

        return profile

    def get_user_work_hours(self, username: str) -> dict:
        """Analyze and return a user's typical work hours."""
        with self._lock:
            rows = self._conn.execute(
                """SELECT timestamp FROM user_events
                   WHERE username = ?""",
                (username,),
            ).fetchall()

        if not rows:
            return {"username": username, "hours": {}, "typical_start": None, "typical_end": None}

        hours = defaultdict(int)
        for (ts,) in rows:
            hour = time.localtime(ts).tm_hour
            hours[hour] += 1

        total = sum(hours.values())
        active_hours = sorted([h for h, c in hours.items() if c / total > 0.05])

        return {
            "username": username,
            "hours": dict(hours),
            "active_hours": active_hours,
            "typical_start": active_hours[0] if active_hours else None,
            "typical_end": active_hours[-1] if active_hours else None,
            "total_events": total,
        }