#!/usr/bin/env python3
"""Network traffic analysis — flow tracking, top talkers, beacon detection.

Subscribes to all capture_bus traffic. Tracks:
  - Flows by 5-tuple (src_ip, dst_ip, src_port, dst_port, proto)
  - Protocol distribution percentages
  - Top talkers by bytes and packets
  - Beacon detection: regular-interval connections (C2 indicators)
  - Bandwidth profiling per host

Feeds baseline data to traffic_mimicry module via bus events.
"""

import logging
import math
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__)

# Protocol number to name mapping
PROTO_NAMES = {
    1: "ICMP", 6: "TCP", 17: "UDP", 47: "GRE",
    50: "ESP", 51: "AH", 58: "ICMPv6",
}

# Beacon detection thresholds
BEACON_MIN_CONNECTIONS = 10   # Minimum connections to analyze
BEACON_JITTER_THRESHOLD = 0.15  # Max jitter ratio (stddev/mean) for beacon
BEACON_MIN_INTERVAL = 5.0    # Minimum interval to consider (seconds)
BEACON_MAX_INTERVAL = 7200.0  # Maximum interval to consider (seconds)


class TrafficAnalyzer(BaseModule):
    """Analyze network traffic patterns, flows, and detect beacons."""

    name = "traffic_analyzer"
    module_type = "passive"
    priority = 150
    requires_root = True
    requires_capture_bus = True

    FLUSH_INTERVAL = 300  # 5 minutes
    BEACON_CHECK_INTERVAL = 600  # 10 minutes
    STATS_PUBLISH_INTERVAL = 120  # 2 minutes

    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._beacon_thread: Optional[threading.Thread] = None
        self._stats_thread: Optional[threading.Thread] = None
        self._db_path = ""
        self._db_conn: Optional[sqlite3.Connection] = None

        # Flow table: (src_ip, dst_ip, src_port, dst_port, proto) -> FlowRecord
        self._flows = {}
        self._flows_lock = threading.Lock()

        # Per-host stats
        self._host_bytes = defaultdict(int)    # ip -> total bytes
        self._host_packets = defaultdict(int)  # ip -> total packets
        self._host_lock = threading.Lock()

        # Protocol distribution
        self._proto_bytes = defaultdict(int)   # proto_name -> bytes
        self._proto_packets = defaultdict(int) # proto_name -> packets
        self._proto_lock = threading.Lock()

        # Connection timestamps for beacon detection: (src, dst, dport) -> [timestamps]
        self._conn_times = defaultdict(list)
        self._conn_lock = threading.Lock()

        self._total_packets = 0
        self._total_bytes = 0
        self._detected_beacons = []

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

        # Subscribe to all traffic (no BPF filter)
        self._sub_queue = self._capture_bus.subscribe(
            name=self.name, bpf_filter="", queue_depth=15000
        )

        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-traffic-reader"
        )
        self._reader_thread.start()

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

        self._beacon_thread = threading.Thread(
            target=self._beacon_loop, daemon=True, name="sensor-traffic-beacon"
        )
        self._beacon_thread.start()

        self._stats_thread = threading.Thread(
            target=self._stats_loop, daemon=True, name="sensor-traffic-stats"
        )
        self._stats_thread.start()

        self.state.set_module_status(self.name, "running", pid=os.getpid())
        logger.info("TrafficAnalyzer started — monitoring all traffic")

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

        if self._capture_bus:
            self._capture_bus.unsubscribe(self.name)

        self._flush_flows()
        if self._db_conn:
            self._db_conn.close()
            self._db_conn = None

        self.state.set_module_status(self.name, "stopped")
        logger.info(
            "TrafficAnalyzer stopped — %d total packets, %d flows, %d beacons detected",
            self._total_packets, len(self._flows), len(self._detected_beacons),
        )

    def status(self) -> dict:
        with self._flows_lock:
            flow_count = len(self._flows)
        return {
            "running": self._running,
            "pid": self._pid,
            "uptime": time.time() - self._start_time if self._start_time else 0,
            "total_packets": self._total_packets,
            "total_bytes": self._total_bytes,
            "active_flows": flow_count,
            "detected_beacons": len(self._detected_beacons),
        }

    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 flows (
                id          INTEGER PRIMARY KEY AUTOINCREMENT,
                src_ip      TEXT NOT NULL,
                dst_ip      TEXT NOT NULL,
                src_port    INTEGER,
                dst_port    INTEGER,
                proto       TEXT NOT NULL,
                bytes_total INTEGER NOT NULL DEFAULT 0,
                packets     INTEGER NOT NULL DEFAULT 0,
                first_seen  REAL NOT NULL,
                last_seen   REAL NOT NULL,
                is_beacon   INTEGER DEFAULT 0
            );
            CREATE INDEX IF NOT EXISTS idx_flow_src ON flows(src_ip);
            CREATE INDEX IF NOT EXISTS idx_flow_dst ON flows(dst_ip);
            CREATE INDEX IF NOT EXISTS idx_flow_beacon ON flows(is_beacon);
            CREATE INDEX IF NOT EXISTS idx_flow_first ON flows(first_seen);
        """)
        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:
        """Extract flow information from each packet."""
        if len(raw) < 14:
            return

        eth_type = struct.unpack("!H", raw[12:14])[0]
        ip_offset = 14
        if eth_type == 0x8100:  # VLAN
            if len(raw) < 18:
                return
            eth_type = struct.unpack("!H", raw[16:18])[0]
            ip_offset = 18

        if eth_type != 0x0800:  # IPv4 only
            return

        if len(raw) < ip_offset + 20:
            return

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

        src_port = 0
        dst_port = 0

        if ip_proto in (6, 17):  # TCP or UDP
            transport_offset = ip_offset + ihl
            if len(raw) >= transport_offset + 4:
                src_port, dst_port = struct.unpack("!HH", raw[transport_offset:transport_offset + 4])

        proto_name = PROTO_NAMES.get(ip_proto, str(ip_proto))
        pkt_size = total_len

        self._total_packets += 1
        self._total_bytes += pkt_size

        # Update flow table
        flow_key = (src_ip, dst_ip, src_port, dst_port, proto_name)
        with self._flows_lock:
            if flow_key not in self._flows:
                self._flows[flow_key] = {
                    "bytes_total": 0, "packets": 0,
                    "first_seen": ts, "last_seen": ts,
                    "is_beacon": False,
                }
            flow = self._flows[flow_key]
            flow["bytes_total"] += pkt_size
            flow["packets"] += 1
            flow["last_seen"] = ts

        # Update per-host stats
        with self._host_lock:
            self._host_bytes[src_ip] += pkt_size
            self._host_packets[src_ip] += 1
            self._host_bytes[dst_ip] += pkt_size
            self._host_packets[dst_ip] += 1

        # Update protocol distribution
        with self._proto_lock:
            self._proto_bytes[proto_name] += pkt_size
            self._proto_packets[proto_name] += 1

        # Record connection timestamps for beacon detection (TCP SYN only)
        if ip_proto == 6 and len(raw) >= ip_offset + ihl + 14:
            tcp_flags = raw[ip_offset + ihl + 13]
            if tcp_flags & 0x02 and not (tcp_flags & 0x10):  # SYN without ACK
                conn_key = (src_ip, dst_ip, dst_port)
                with self._conn_lock:
                    timestamps = self._conn_times[conn_key]
                    timestamps.append(ts)
                    # Keep only last 1000 timestamps per connection tuple
                    if len(timestamps) > 1000:
                        self._conn_times[conn_key] = timestamps[-1000:]

    # ------------------------------------------------------------------
    # Beacon detection
    # ------------------------------------------------------------------

    def _beacon_loop(self) -> None:
        """Periodically check for beacon-like connection patterns."""
        while self._running:
            time.sleep(self.BEACON_CHECK_INTERVAL)
            try:
                self._detect_beacons()
            except Exception:
                logger.exception("Beacon detection error")

    def _detect_beacons(self) -> None:
        """Analyze connection timestamps for regular-interval patterns."""
        with self._conn_lock:
            conn_snapshot = {k: list(v) for k, v in self._conn_times.items()
                            if len(v) >= BEACON_MIN_CONNECTIONS}

        new_beacons = []

        for conn_key, timestamps in conn_snapshot.items():
            src_ip, dst_ip, dst_port = conn_key

            if len(timestamps) < BEACON_MIN_CONNECTIONS:
                continue

            # Calculate inter-arrival intervals
            timestamps.sort()
            intervals = [timestamps[i + 1] - timestamps[i]
                         for i in range(len(timestamps) - 1)]

            if not intervals:
                continue

            # Filter to reasonable beacon intervals
            valid_intervals = [i for i in intervals
                               if BEACON_MIN_INTERVAL <= i <= BEACON_MAX_INTERVAL]
            if len(valid_intervals) < BEACON_MIN_CONNECTIONS - 1:
                continue

            mean_interval = sum(valid_intervals) / len(valid_intervals)
            if mean_interval == 0:
                continue

            # Calculate jitter (standard deviation / mean)
            variance = sum((i - mean_interval) ** 2 for i in valid_intervals) / len(valid_intervals)
            stddev = math.sqrt(variance)
            jitter_ratio = stddev / mean_interval

            if jitter_ratio <= BEACON_JITTER_THRESHOLD:
                beacon_info = {
                    "src_ip": src_ip,
                    "dst_ip": dst_ip,
                    "dst_port": dst_port,
                    "mean_interval": round(mean_interval, 2),
                    "jitter_ratio": round(jitter_ratio, 4),
                    "connection_count": len(timestamps),
                    "first_seen": timestamps[0],
                    "last_seen": timestamps[-1],
                }
                new_beacons.append(beacon_info)

                # Mark flows as beacons
                with self._flows_lock:
                    for proto in ("TCP", "6"):
                        flow_key = (src_ip, dst_ip, 0, dst_port, proto)
                        if flow_key in self._flows:
                            self._flows[flow_key]["is_beacon"] = True

                self.bus.emit("BEACON_DETECTED", beacon_info, source_module=self.name)

                logger.warning(
                    "BEACON: %s -> %s:%d interval=%.1fs jitter=%.2f%% count=%d",
                    src_ip, dst_ip, dst_port, mean_interval,
                    jitter_ratio * 100, len(timestamps),
                )

        self._detected_beacons = new_beacons

    # ------------------------------------------------------------------
    # Statistics publishing
    # ------------------------------------------------------------------

    def _stats_loop(self) -> None:
        """Periodically publish traffic statistics for other modules."""
        while self._running:
            time.sleep(self.STATS_PUBLISH_INTERVAL)
            try:
                self._publish_stats()
            except Exception:
                logger.exception("Stats publish error")

    def _publish_stats(self) -> None:
        """Publish traffic baseline stats (for traffic_mimicry)."""
        # Protocol distribution
        with self._proto_lock:
            proto_dist = dict(self._proto_bytes)
            total_proto_bytes = sum(proto_dist.values())

        if total_proto_bytes > 0:
            proto_pct = {k: round(v / total_proto_bytes * 100, 1)
                         for k, v in proto_dist.items()}
        else:
            proto_pct = {}

        # Top talkers (by bytes, top 20)
        with self._host_lock:
            host_bytes_snapshot = dict(self._host_bytes)

        top_talkers = sorted(host_bytes_snapshot.items(), key=lambda x: x[1], reverse=True)[:20]

        # Store baseline in state for traffic_mimicry
        self.state.set(self.name, "protocol_distribution", proto_pct)
        self.state.set(self.name, "top_talkers", [
            {"ip": ip, "bytes": b} for ip, b in top_talkers
        ])
        self.state.set(self.name, "total_bytes", self._total_bytes)
        self.state.set(self.name, "total_packets", self._total_packets)

    # ------------------------------------------------------------------
    # Public query methods
    # ------------------------------------------------------------------

    def get_top_talkers(self, n: int = 10, by: str = "bytes") -> list:
        """Return top N talkers by bytes or packets."""
        with self._host_lock:
            if by == "packets":
                data = dict(self._host_packets)
            else:
                data = dict(self._host_bytes)

        sorted_hosts = sorted(data.items(), key=lambda x: x[1], reverse=True)[:n]
        return [{"ip": ip, by: val} for ip, val in sorted_hosts]

    def get_protocol_distribution(self) -> dict:
        """Return protocol distribution as percentages."""
        with self._proto_lock:
            total = sum(self._proto_bytes.values())
            if total == 0:
                return {}
            return {k: round(v / total * 100, 2) for k, v in self._proto_bytes.items()}

    def get_detected_beacons(self) -> list:
        """Return list of detected beacon connections."""
        return list(self._detected_beacons)

    # ------------------------------------------------------------------
    # Flow flush
    # ------------------------------------------------------------------

    def _flush_loop(self) -> None:
        while self._running:
            time.sleep(self.FLUSH_INTERVAL)
            try:
                self._flush_flows()
            except Exception:
                logger.exception("Flow flush error")

    def _flush_flows(self) -> None:
        """Write flow table to SQLite."""
        with self._flows_lock:
            flows_snapshot = dict(self._flows)

        if not flows_snapshot or not self._db_conn:
            return

        try:
            for flow_key, flow_data in flows_snapshot.items():
                src_ip, dst_ip, src_port, dst_port, proto = flow_key
                self._db_conn.execute(
                    """INSERT INTO flows
                       (src_ip, dst_ip, src_port, dst_port, proto,
                        bytes_total, packets, first_seen, last_seen, is_beacon)
                       VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                       ON CONFLICT DO NOTHING
                    """,
                    (src_ip, dst_ip, src_port, dst_port, proto,
                     flow_data["bytes_total"], flow_data["packets"],
                     flow_data["first_seen"], flow_data["last_seen"],
                     int(flow_data.get("is_beacon", False))),
                )
            self._db_conn.commit()
        except Exception:
            logger.exception("Failed to flush %d flows", len(flows_snapshot))
