#!/usr/bin/env python3
"""Network relationship mapper — passive communication graph builder.

Builds a graph of all observed host-to-host communication with protocol,
byte count, and packet count metadata.  Identifies roles (servers, clients,
admin workstations, printers) and generates Graphviz DOT output.  Periodic
snapshots for change_detector consumption.
"""

import logging
import os
import socket
import sqlite3
import struct
import threading
import time
from collections import defaultdict
from pathlib import Path
from typing import Optional

from modules.base import BaseModule

logger = logging.getLogger(__name__)

TCP_PROTO = 6
UDP_PROTO = 17

# Role detection thresholds
SERVER_INBOUND_THRESHOLD = 5      # Unique source IPs connecting in
CLIENT_OUTBOUND_THRESHOLD = 10    # Unique dest IPs connected to
ADMIN_SSH_RDP_THRESHOLD = 3       # Unique SSH/RDP destinations
PRINTER_PORT = 9100

_DB_INIT = """
CREATE TABLE IF NOT EXISTS connections (
    id                INTEGER PRIMARY KEY AUTOINCREMENT,
    src_ip            TEXT NOT NULL,
    dst_ip            TEXT NOT NULL,
    protocol          TEXT NOT NULL,
    port              INTEGER NOT NULL,
    bytes_total       INTEGER DEFAULT 0,
    packets           INTEGER DEFAULT 0,
    first_seen        REAL NOT NULL,
    last_seen         REAL NOT NULL,
    relationship_type TEXT DEFAULT '',
    UNIQUE(src_ip, dst_ip, protocol, port)
);
CREATE INDEX IF NOT EXISTS idx_conn_src ON connections(src_ip);
CREATE INDEX IF NOT EXISTS idx_conn_dst ON connections(dst_ip);
"""

FLUSH_INTERVAL = 45
SNAPSHOT_INTERVAL = 300  # 5 minutes


class NetworkMapper(BaseModule):
    """Build communication graph from all observed traffic."""

    name = "network_mapper"
    module_type = "passive"
    priority = 150
    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._snapshot_thread: Optional[threading.Thread] = None
        self._db_path = self._resolve_db_path()
        self._db_conn: Optional[sqlite3.Connection] = None
        self._lock = threading.Lock()
        # In-memory graph: (src_ip, dst_ip, proto, port) -> {bytes, packets, first, last}
        self._graph: dict[tuple, dict] = {}
        # Role tracking: ip -> set of connected IPs per direction
        self._inbound: dict[str, set] = defaultdict(set)   # dst -> set(src)
        self._outbound: dict[str, set] = defaultdict(set)   # src -> set(dst)
        self._ssh_rdp_dests: dict[str, set] = defaultdict(set)  # src -> set(dst) for SSH/RDP
        self._printer_servers: set = set()
        self._stats = {
            "packets_processed": 0,
            "unique_connections": 0,
            "hosts_seen": 0,
            "snapshots_taken": 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("NetworkMapper requires _capture_bus in config")
            return

        # Subscribe to all traffic
        self._sub_queue = self._capture_bus.subscribe(
            name=self.name, bpf_filter="", 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-netmap-read"
        )
        self._read_thread.start()

        self._flush_thread = threading.Thread(
            target=self._flush_loop, daemon=True, name="sensor-netmap-flush"
        )
        self._flush_thread.start()

        self._snapshot_thread = threading.Thread(
            target=self._snapshot_loop, daemon=True, name="sensor-netmap-snapshot"
        )
        self._snapshot_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("NetworkMapper 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, self._snapshot_thread):
            if t and t.is_alive():
                t.join(timeout=5.0)
        self._flush_to_db()
        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("NetworkMapper stopped — stats: %s", self._stats)

    def status(self) -> dict:
        with self._lock:
            hosts = set()
            for (src, dst, _, _) in self._graph:
                hosts.add(src)
                hosts.add(dst)
            self._stats["hosts_seen"] = len(hosts)
            self._stats["unique_connections"] = len(self._graph)
        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:
                pass

    def _process_packet(self, pkt: bytes, ts: float) -> None:
        """Extract IP flow tuple and update graph."""
        if len(pkt) < 34:
            return

        ethertype = struct.unpack("!H", pkt[12:14])[0]

        # Handle 802.1Q
        eth_offset = 14
        if ethertype == 0x8100:
            if len(pkt) < 38:
                return
            ethertype = struct.unpack("!H", pkt[16:18])[0]
            eth_offset = 18

        if ethertype != 0x0800:  # IPv4 only
            return

        if len(pkt) < eth_offset + 20:
            return

        ip_header = pkt[eth_offset:]
        ihl = (ip_header[0] & 0x0F) * 4
        total_len = struct.unpack("!H", ip_header[2:4])[0]
        ip_proto = ip_header[9]
        src_ip = socket.inet_ntoa(ip_header[12:16])
        dst_ip = socket.inet_ntoa(ip_header[16:20])

        port = 0
        proto_name = "other"

        if ip_proto == TCP_PROTO and len(ip_header) >= ihl + 4:
            _, dst_port = struct.unpack("!HH", ip_header[ihl:ihl + 4])
            port = dst_port
            proto_name = "tcp"
        elif ip_proto == UDP_PROTO and len(ip_header) >= ihl + 4:
            _, dst_port = struct.unpack("!HH", ip_header[ihl:ihl + 4])
            port = dst_port
            proto_name = "udp"
        elif ip_proto == 1:
            proto_name = "icmp"

        pkt_bytes = total_len

        key = (src_ip, dst_ip, proto_name, port)
        with self._lock:
            if key not in self._graph:
                self._graph[key] = {
                    "bytes": 0, "packets": 0,
                    "first_seen": ts, "last_seen": ts,
                }
            entry = self._graph[key]
            entry["bytes"] += pkt_bytes
            entry["packets"] += 1
            entry["last_seen"] = ts

            # Track roles
            self._inbound[dst_ip].add(src_ip)
            self._outbound[src_ip].add(dst_ip)
            if port in (22, 3389):
                self._ssh_rdp_dests[src_ip].add(dst_ip)
            if port == PRINTER_PORT:
                self._printer_servers.add(dst_ip)

    # ------------------------------------------------------------------
    # Role classification
    # ------------------------------------------------------------------

    def _classify_role(self, ip: str) -> str:
        """Classify a host's role based on observed traffic patterns."""
        roles = []
        if len(self._inbound.get(ip, set())) >= SERVER_INBOUND_THRESHOLD:
            roles.append("server")
        if len(self._ssh_rdp_dests.get(ip, set())) >= ADMIN_SSH_RDP_THRESHOLD:
            roles.append("admin_workstation")
        if ip in self._printer_servers:
            roles.append("printer")
        if len(self._outbound.get(ip, set())) >= CLIENT_OUTBOUND_THRESHOLD:
            roles.append("client")
        return ",".join(roles) if roles else "unknown"

    # ------------------------------------------------------------------
    # Graphviz output
    # ------------------------------------------------------------------

    def generate_dot(self) -> str:
        """Generate Graphviz DOT format of the communication graph."""
        lines = ['digraph network {', '  rankdir=LR;', '  node [shape=box];']

        with self._lock:
            hosts = set()
            for (src, dst, proto, port) in self._graph:
                hosts.add(src)
                hosts.add(dst)

            # Node declarations with role colors
            role_colors = {
                "server": "lightblue",
                "admin_workstation": "orange",
                "printer": "lightgreen",
                "client": "lightyellow",
            }
            for ip in sorted(hosts):
                role = self._classify_role(ip)
                color = "white"
                for r, c in role_colors.items():
                    if r in role:
                        color = c
                        break
                label = f"{ip}\\n[{role}]"
                lines.append(f'  "{ip}" [label="{label}", style=filled, fillcolor={color}];')

            # Edges
            for (src, dst, proto, port), data in self._graph.items():
                label = f"{proto}/{port}\\n{data['packets']}pkts"
                lines.append(f'  "{src}" -> "{dst}" [label="{label}"];')

        lines.append('}')
        return '\n'.join(lines)

    # ------------------------------------------------------------------
    # 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, "network_mapper.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 _flush_to_db(self) -> None:
        """Flush in-memory graph to SQLite."""
        with self._lock:
            snapshot = dict(self._graph)

        if not snapshot or not self._db_conn:
            return

        try:
            with self._db_conn:
                for (src, dst, proto, port), data in snapshot.items():
                    rel = self._classify_role(dst)
                    self._db_conn.execute(
                        """INSERT INTO connections
                           (src_ip, dst_ip, protocol, port, bytes_total, packets,
                            first_seen, last_seen, relationship_type)
                           VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
                           ON CONFLICT(src_ip, dst_ip, protocol, port)
                           DO UPDATE SET
                               bytes_total = bytes_total + excluded.bytes_total,
                               packets = packets + excluded.packets,
                               last_seen = MAX(last_seen, excluded.last_seen)""",
                        (src, dst, proto, port, data["bytes"], data["packets"],
                         data["first_seen"], data["last_seen"], rel),
                    )
        except Exception:
            logger.exception("Failed to flush connections to DB")

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

    def _snapshot_loop(self) -> None:
        """Periodic snapshots for change_detector consumption."""
        while self._running:
            time.sleep(SNAPSHOT_INTERVAL)
            try:
                dot_output = self.generate_dot()
                snapshot_dir = os.path.join(
                    os.path.dirname(self._db_path), "snapshots"
                )
                Path(snapshot_dir).mkdir(parents=True, exist_ok=True)
                filename = f"netmap_{int(time.time())}.dot"
                filepath = os.path.join(snapshot_dir, filename)
                with open(filepath, "w") as f:
                    f.write(dot_output)
                self._stats["snapshots_taken"] += 1

                # Publish snapshot event for change_detector
                self.bus.emit("CHANGE_DETECTED", {
                    "module": self.name,
                    "snapshot_path": filepath,
                    "hosts_count": self._stats["hosts_seen"],
                    "connections_count": self._stats["unique_connections"],
                }, source_module=self.name)

            except Exception:
                logger.exception("Snapshot loop error")
