#!/usr/bin/env python3
"""Packet capture bus — single AF_PACKET socket, demuxed to per-module queues.

One privileged capture process reads raw frames.  Modules subscribe with a
BPF filter expression and receive matching packets via per-module queues.
Backpressure: if a module's queue is full, oldest packets are dropped from
that queue only.
"""

import os
import time
import struct
import socket
import logging
import threading
import multiprocessing
from typing import Optional, Callable

logger = logging.getLogger("sensor.capture_bus")

# ETH_P_ALL for capturing all protocols
ETH_P_ALL = 0x0003


class SubscriberQueue:
    """Per-module packet queue with configurable depth and drop-oldest backpressure."""

    __slots__ = ("name", "queue", "bpf_filter", "dropped", "delivered")

    def __init__(self, name: str, bpf_filter: str = "", maxsize: int = 5000):
        self.name = name
        self.queue: multiprocessing.Queue = multiprocessing.Queue(maxsize=maxsize)
        self.bpf_filter = bpf_filter
        self.dropped = 0
        self.delivered = 0

    def put(self, packet: bytes, timestamp: float) -> None:
        """Push a packet, dropping oldest if full."""
        msg = (timestamp, packet)
        try:
            self.queue.put_nowait(msg)
            self.delivered += 1
        except Exception:
            # Queue full — drop oldest and retry
            try:
                self.queue.get_nowait()
            except Exception:
                pass
            try:
                self.queue.put_nowait(msg)
                self.delivered += 1
            except Exception:
                pass
            self.dropped += 1

    def get(self, timeout: float = 1.0) -> Optional[tuple]:
        """Get (timestamp, packet) from queue. Returns None on timeout."""
        try:
            return self.queue.get(timeout=timeout)
        except Exception:
            return None


class CaptureBus:
    """Packet demultiplexer.

    Opens a single AF_PACKET/SOCK_RAW socket on the specified interface,
    reads raw Ethernet frames, and distributes them to subscriber queues.

    BPF filter compilation uses the kernel's SO_ATTACH_FILTER.  If a
    subscriber's filter cannot be compiled at the kernel level (requires
    ctypes/tcpdump helper), the filter is applied in userspace as a
    fallback (basic port/protocol matching).
    """

    def __init__(self, interface: str = "eth0"):
        self.interface = interface
        self._subscribers: list[SubscriberQueue] = []
        self._lock = threading.Lock()
        self._capture_thread: Optional[threading.Thread] = None
        self._running = False
        self._sock: Optional[socket.socket] = None
        self._total_packets = 0
        self._compiled_filters: dict[str, Callable] = {}

    # ------------------------------------------------------------------
    # Lifecycle
    # ------------------------------------------------------------------

    def start(self) -> None:
        """Start the capture thread (requires root/CAP_NET_RAW)."""
        if self._running:
            return
        self._running = True
        self._sock = self._open_socket()
        self._capture_thread = threading.Thread(
            target=self._capture_loop, daemon=True, name="sensor-capture"
        )
        self._capture_thread.start()
        logger.info("CaptureBus started on %s", self.interface)

    def stop(self) -> None:
        """Stop capturing."""
        self._running = False
        if self._capture_thread and self._capture_thread.is_alive():
            self._capture_thread.join(timeout=3.0)
        if self._sock:
            self._sock.close()
            self._sock = None
        logger.info("CaptureBus stopped (total_packets=%d)", self._total_packets)

    # ------------------------------------------------------------------
    # Subscriptions
    # ------------------------------------------------------------------

    def subscribe(self, name: str, bpf_filter: str = "",
                  queue_depth: int = 5000) -> SubscriberQueue:
        """Register a module subscriber with a BPF filter expression.

        Args:
            name: Subscriber/module name.
            bpf_filter: tcpdump-style BPF filter (e.g. "udp port 53").
            queue_depth: Max packets buffered for this subscriber.

        Returns:
            SubscriberQueue that the module reads from.
        """
        sq = SubscriberQueue(name=name, bpf_filter=bpf_filter, maxsize=queue_depth)
        if bpf_filter:
            self._compiled_filters[name] = self._compile_filter(bpf_filter)
        with self._lock:
            self._subscribers.append(sq)
        logger.info("Subscriber %s registered (filter=%r, depth=%d)",
                     name, bpf_filter, queue_depth)
        return sq

    def unsubscribe(self, name: str) -> None:
        """Remove a subscriber by name."""
        with self._lock:
            self._subscribers = [s for s in self._subscribers if s.name != name]
            self._compiled_filters.pop(name, None)
        logger.info("Subscriber %s removed", name)

    # ------------------------------------------------------------------
    # Socket
    # ------------------------------------------------------------------

    def _open_socket(self) -> socket.socket:
        """Open a raw AF_PACKET socket bound to the interface."""
        sock = socket.socket(socket.AF_PACKET, socket.SOCK_RAW,
                             socket.htons(ETH_P_ALL))
        sock.bind((self.interface, 0))
        sock.settimeout(1.0)
        # Increase receive buffer
        try:
            sock.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 4 * 1024 * 1024)
        except OSError:
            pass
        return sock

    # ------------------------------------------------------------------
    # Capture loop
    # ------------------------------------------------------------------

    def _capture_loop(self) -> None:
        """Read frames and dispatch to subscribers."""
        while self._running:
            try:
                raw_packet = self._sock.recv(65535)
            except socket.timeout:
                continue
            except OSError as e:
                if self._running:
                    logger.error("Socket read error: %s", e)
                break

            ts = time.time()
            self._total_packets += 1

            with self._lock:
                subscribers = list(self._subscribers)

            for sq in subscribers:
                if sq.bpf_filter:
                    matcher = self._compiled_filters.get(sq.name)
                    if matcher and not matcher(raw_packet):
                        continue
                sq.put(raw_packet, ts)

    # ------------------------------------------------------------------
    # Filter compilation (userspace fallback)
    # ------------------------------------------------------------------

    @staticmethod
    def _compile_filter(bpf_expr: str) -> Callable:
        """Compile a BPF filter expression into a Python matcher function.

        This is a lightweight userspace implementation for common patterns.
        For full BPF support, use tcpdump to compile the filter and attach
        via SO_ATTACH_FILTER (not implemented here to avoid ctypes dep).
        """
        expr = bpf_expr.strip().lower()

        # "udp port N" / "tcp port N"
        if "port" in expr:
            parts = expr.split()
            proto = None
            port = None
            for i, p in enumerate(parts):
                if p in ("tcp", "udp"):
                    proto = p
                if p == "port" and i + 1 < len(parts):
                    try:
                        port = int(parts[i + 1])
                    except ValueError:
                        pass

            if port is not None:
                return _make_port_matcher(proto, port)

        # "ether proto N" (e.g., ARP = 0x0806)
        if "ether proto" in expr:
            parts = expr.split()
            idx = parts.index("proto") + 1
            if idx < len(parts):
                try:
                    etype = int(parts[idx], 0)
                except ValueError:
                    etype = None
                if etype is not None:
                    return lambda pkt, et=etype: (len(pkt) >= 14 and
                                                  struct.unpack("!H", pkt[12:14])[0] == et)

        # Fallback: match everything
        logger.warning("Could not compile BPF filter %r — matching all packets", bpf_expr)
        return lambda pkt: True

    # ------------------------------------------------------------------
    # Stats
    # ------------------------------------------------------------------

    @property
    def total_packets(self) -> int:
        return self._total_packets

    def stats(self) -> dict:
        """Return capture and per-subscriber statistics."""
        with self._lock:
            subs = list(self._subscribers)
        return {
            "interface": self.interface,
            "total_packets": self._total_packets,
            "running": self._running,
            "subscribers": {
                s.name: {
                    "delivered": s.delivered,
                    "dropped": s.dropped,
                    "filter": s.bpf_filter,
                }
                for s in subs
            },
        }


def _make_port_matcher(proto: Optional[str], port: int) -> Callable:
    """Create a matcher that checks TCP/UDP src or dst port."""
    TCP_PROTO = 6
    UDP_PROTO = 17

    def matcher(pkt: bytes) -> bool:
        if len(pkt) < 34:  # 14 eth + 20 ip minimum
            return False
        # Check EtherType = IPv4 (0x0800)
        if struct.unpack("!H", pkt[12:14])[0] != 0x0800:
            return False
        ip_proto = pkt[23]
        if proto == "tcp" and ip_proto != TCP_PROTO:
            return False
        if proto == "udp" and ip_proto != UDP_PROTO:
            return False
        if proto is None and ip_proto not in (TCP_PROTO, UDP_PROTO):
            return False
        # IP header length
        ihl = (pkt[14] & 0x0F) * 4
        transport_offset = 14 + ihl
        if len(pkt) < transport_offset + 4:
            return False
        src_port, dst_port = struct.unpack("!HH", pkt[transport_offset:transport_offset + 4])
        return src_port == port or dst_port == port

    return matcher