bigbrother/core/bus.py

#!/usr/bin/env python3
"""Multiprocess-safe event bus with pub/sub pattern.

Uses multiprocessing.Queue for cross-process event delivery.
A dedicated dispatcher thread reads from the queue and fans out
to registered subscriber callbacks.
"""

import time
import threading
import logging
import multiprocessing
from dataclasses import dataclass, field, asdict
from typing import Callable, Optional

logger = logging.getLogger("bb.bus")

# Canonical event types
EVENT_TYPES = frozenset([
    "CREDENTIAL_FOUND", "HOST_DISCOVERED", "PCAP_ROTATED",
    "MODULE_STARTED", "MODULE_STOPPED", "MODULE_ERROR",
    "RESOURCE_WARNING", "RESOURCE_CRITICAL", "KILL_SWITCH",
    "CONNECTIVITY_CHANGED", "VLAN_DETECTED", "HANDSHAKE_CAPTURED",
    "TICKET_HARVESTED", "CLOUD_TOKEN_FOUND", "BEACON_DETECTED",
    "CHANGE_DETECTED", "THERMAL_WARNING", "SCOPE_VIOLATION",
    "TOOL_CRASHED", "TOOL_RESTARTED",
])


@dataclass
class Event:
    """Single event on the bus."""
    event_type: str
    payload: dict
    timestamp: float = field(default_factory=time.time)
    source_module: str = ""

    def to_dict(self) -> dict:
        return asdict(self)


class EventBus:
    """Multiprocess-safe publish/subscribe event bus.

    Subscribers register a callback with an optional event_type filter.
    Publishing pushes an Event into a shared mp.Queue.  A background
    dispatcher thread reads from the queue and invokes matching callbacks.
    """

    def __init__(self, maxsize: int = 10000):
        self._queue: multiprocessing.Queue = multiprocessing.Queue(maxsize=maxsize)
        self._subscribers: dict[str, list[Callable]] = {}  # event_type -> [callbacks]
        self._wildcard_subscribers: list[Callable] = []     # receive all events
        self._lock = threading.Lock()
        self._dispatcher_thread: Optional[threading.Thread] = None
        self._running = False
        self._event_count = 0

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

    def start(self) -> None:
        """Start the dispatcher thread."""
        if self._running:
            return
        self._running = True
        self._dispatcher_thread = threading.Thread(
            target=self._dispatch_loop, daemon=True, name="bb-bus-dispatch"
        )
        self._dispatcher_thread.start()
        logger.info("EventBus dispatcher started")

    def stop(self) -> None:
        """Stop the dispatcher thread. Drains remaining events first."""
        if not self._running:
            return
        self._running = False
        # Push sentinel to unblock queue.get()
        try:
            self._queue.put_nowait(None)
        except Exception:
            pass
        if self._dispatcher_thread and self._dispatcher_thread.is_alive():
            self._dispatcher_thread.join(timeout=5.0)
        logger.info("EventBus dispatcher stopped (dispatched %d events)", self._event_count)

    # ------------------------------------------------------------------
    # Pub / Sub
    # ------------------------------------------------------------------

    def subscribe(self, callback: Callable, event_type: Optional[str] = None) -> None:
        """Register a callback for events.

        Args:
            callback: Callable that takes a single Event argument.
            event_type: If set, only events of this type are delivered.
                        If None, the callback receives ALL events.
        """
        with self._lock:
            if event_type is None:
                self._wildcard_subscribers.append(callback)
            else:
                self._subscribers.setdefault(event_type, []).append(callback)

    def unsubscribe(self, callback: Callable, event_type: Optional[str] = None) -> None:
        """Remove a previously registered callback."""
        with self._lock:
            if event_type is None:
                try:
                    self._wildcard_subscribers.remove(callback)
                except ValueError:
                    pass
            else:
                subs = self._subscribers.get(event_type, [])
                try:
                    subs.remove(callback)
                except ValueError:
                    pass

    def publish(self, event: Event) -> None:
        """Publish an event onto the bus (multiprocess-safe).

        Non-blocking: if the queue is full the event is dropped and a
        warning is logged.
        """
        try:
            self._queue.put_nowait(event)
        except Exception:
            logger.warning("EventBus queue full — dropped %s from %s",
                           event.event_type, event.source_module)

    def emit(self, event_type: str, payload: dict, source_module: str = "") -> None:
        """Convenience wrapper: build an Event and publish it."""
        self.publish(Event(
            event_type=event_type,
            payload=payload,
            source_module=source_module,
        ))

    # ------------------------------------------------------------------
    # Internal dispatcher
    # ------------------------------------------------------------------

    def _dispatch_loop(self) -> None:
        """Background loop: pull events from the mp.Queue, fan out to callbacks."""
        while self._running:
            try:
                event = self._queue.get(timeout=0.5)
            except Exception:
                continue

            if event is None:  # sentinel
                continue

            self._event_count += 1

            with self._lock:
                targets = list(self._wildcard_subscribers)
                targets.extend(self._subscribers.get(event.event_type, []))

            for cb in targets:
                try:
                    cb(event)
                except Exception:
                    logger.exception("Subscriber callback error for %s", event.event_type)

    # ------------------------------------------------------------------
    # Introspection
    # ------------------------------------------------------------------

    @property
    def event_count(self) -> int:
        return self._event_count

    @property
    def queue_size(self) -> int:
        try:
            return self._queue.qsize()
        except NotImplementedError:
            return -1

    def get_queue(self) -> multiprocessing.Queue:
        """Return the underlying mp.Queue for sharing with child processes."""
        return self._queue