Add Phase 1 core infrastructure: event bus, state manager, engine, capture bus, tool manager, scheduler, resource monitor, kill switch
Core framework (8 files in core/) + BaseModule ABC (2 files in modules/): - bus.py: Multiprocess-safe pub/sub event bus using mp.Queue with background dispatcher thread, 20 canonical event types - state.py: SQLite persistent state with WAL mode, dedicated writer thread with 500ms coalesce, module status + generic key-value store - engine.py: Module lifecycle manager with multiprocess model, dependency resolution (topological sort), hardware tier detection (opi_zero3/pi_zero/generic), resource budgeting, scope enforcement, engagement phase gating - capture_bus.py: Single AF_PACKET socket packet demux with per-module subscriber queues, BPF filter matching, drop-oldest backpressure - tool_manager.py: Subprocess lifecycle for external tools (bettercap/tcpdump/Responder/mitmproxy/ntlmrelayx) with exponential backoff restart, /proc resource monitoring, process disguise via prctl, health check callbacks, graceful SIGTERM->SIGKILL shutdown - scheduler.py: Cron-like task scheduler with jitter, thread pool execution, enable/disable/run_now controls - resource_monitor.py: System + per-process RAM/CPU/disk/thermal monitoring, OOM kill lowest-priority module, thermal shedding at 70C - kill_switch.py: 7-phase wipe sequence (stop procs, corrective ARP, LUKS destroy, shred keys, zero-fill, clear RAM, reboot), boot flag for interrupted wipe resume, dead man's switch - modules/base.py: BaseModule ABC with start/stop/status/configure/health_check interface contract
This commit is contained in:
n0mad1k
@@ -0,0 +1,20 @@
|
||||
from core.bus import EventBus, Event
|
||||
from core.state import StateManager
|
||||
from core.engine import Engine
|
||||
from core.capture_bus import CaptureBus
|
||||
from core.tool_manager import ToolManager
|
||||
from core.scheduler import Scheduler
|
||||
from core.resource_monitor import ResourceMonitor
|
||||
from core.kill_switch import KillSwitch
|
||||
|
||||
__all__ = [
|
||||
"EventBus",
|
||||
"Event",
|
||||
"StateManager",
|
||||
"Engine",
|
||||
"CaptureBus",
|
||||
"ToolManager",
|
||||
"Scheduler",
|
||||
"ResourceMonitor",
|
||||
"KillSwitch",
|
||||
]
|
||||
+185
@@ -0,0 +1,185 @@
|
||||
#!/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
|
||||
@@ -0,0 +1,291 @@
|
||||
#!/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("bb.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="bb-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
|
||||
+433
@@ -0,0 +1,433 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Module lifecycle manager.
|
||||
|
||||
Each BigBrother module runs in its own multiprocessing.Process.
|
||||
The Engine handles:
|
||||
- Loading module classes from the modules/ tree
|
||||
- Dependency resolution (topological sort)
|
||||
- Start / stop / restart / start_all / stop_all
|
||||
- Resource budgeting per hardware tier
|
||||
- Interface conflict detection
|
||||
- Scope enforcement
|
||||
- Engagement phase gating (passive_only for N days, then active_allowed)
|
||||
"""
|
||||
|
||||
import os
|
||||
import time
|
||||
import signal
|
||||
import logging
|
||||
import importlib
|
||||
import multiprocessing
|
||||
from collections import defaultdict, deque
|
||||
from typing import Optional, Type
|
||||
|
||||
from core.bus import EventBus, Event
|
||||
from core.state import StateManager
|
||||
from modules.base import BaseModule
|
||||
|
||||
logger = logging.getLogger("bb.engine")
|
||||
|
||||
# Hardware tier definitions — module limits and resource ceilings
|
||||
HARDWARE_TIERS = {
|
||||
"opi_zero3": {
|
||||
"max_passive": 16,
|
||||
"max_active": 6,
|
||||
"max_ram_mb": 3072,
|
||||
"max_cpu_pct": 80,
|
||||
"allow_mitmproxy": True,
|
||||
"allow_bettercap": True,
|
||||
},
|
||||
"pi_zero": {
|
||||
"max_passive": 4,
|
||||
"max_active": 0,
|
||||
"max_ram_mb": 400,
|
||||
"max_cpu_pct": 70,
|
||||
"allow_mitmproxy": False,
|
||||
"allow_bettercap": False,
|
||||
},
|
||||
"generic": {
|
||||
"max_passive": 16,
|
||||
"max_active": 99,
|
||||
"max_ram_mb": 0, # 0 = no limit
|
||||
"max_cpu_pct": 90,
|
||||
"allow_mitmproxy": True,
|
||||
"allow_bettercap": True,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def detect_hardware_tier() -> str:
|
||||
"""Detect hardware tier from /proc/device-tree/model and /proc/cpuinfo."""
|
||||
model = ""
|
||||
try:
|
||||
with open("/proc/device-tree/model", "r") as f:
|
||||
model = f.read().strip().lower()
|
||||
except (FileNotFoundError, PermissionError):
|
||||
pass
|
||||
|
||||
if "orange pi zero3" in model or "orange pi zero 3" in model:
|
||||
return "opi_zero3"
|
||||
|
||||
# Check cpuinfo for known SoCs
|
||||
try:
|
||||
with open("/proc/cpuinfo", "r") as f:
|
||||
cpuinfo = f.read().lower()
|
||||
except (FileNotFoundError, PermissionError):
|
||||
cpuinfo = ""
|
||||
|
||||
if "allwinner" in cpuinfo and "h618" in cpuinfo:
|
||||
return "opi_zero3"
|
||||
|
||||
if "raspberry pi zero 2" in model or "bcm2710" in cpuinfo:
|
||||
return "pi_zero"
|
||||
|
||||
return "generic"
|
||||
|
||||
|
||||
def _topo_sort(modules: dict[str, "ModuleEntry"]) -> list[str]:
|
||||
"""Topological sort of modules by dependency. Returns ordered list of names."""
|
||||
in_degree: dict[str, int] = defaultdict(int)
|
||||
graph: dict[str, list[str]] = defaultdict(list)
|
||||
all_names = set(modules.keys())
|
||||
|
||||
for name, entry in modules.items():
|
||||
for dep in entry.module_class.dependencies:
|
||||
if dep in all_names:
|
||||
graph[dep].append(name)
|
||||
in_degree[name] += 1
|
||||
if name not in in_degree:
|
||||
in_degree[name] = 0
|
||||
|
||||
queue = deque(n for n, d in in_degree.items() if d == 0)
|
||||
order = []
|
||||
while queue:
|
||||
node = queue.popleft()
|
||||
order.append(node)
|
||||
for neighbor in graph[node]:
|
||||
in_degree[neighbor] -= 1
|
||||
if in_degree[neighbor] == 0:
|
||||
queue.append(neighbor)
|
||||
|
||||
if len(order) != len(all_names):
|
||||
missing = all_names - set(order)
|
||||
logger.error("Circular dependency detected involving: %s", missing)
|
||||
# Append remaining modules anyway (best-effort)
|
||||
order.extend(missing)
|
||||
|
||||
return order
|
||||
|
||||
|
||||
class ModuleEntry:
|
||||
"""Registry entry for a loaded module."""
|
||||
__slots__ = ("module_class", "config", "process", "started_at")
|
||||
|
||||
def __init__(self, module_class: Type[BaseModule], config: dict):
|
||||
self.module_class = module_class
|
||||
self.config = config
|
||||
self.process: Optional[multiprocessing.Process] = None
|
||||
self.started_at: Optional[float] = None
|
||||
|
||||
|
||||
def _module_runner(module_class: Type[BaseModule], bus_queue: multiprocessing.Queue,
|
||||
state_db_path: str, config: dict, engine_pipe) -> None:
|
||||
"""Entry point for each module process.
|
||||
|
||||
Instantiates the module, wires up a local EventBus that forwards
|
||||
events through the shared mp.Queue, and calls module.start().
|
||||
"""
|
||||
# Build a lightweight bus proxy that pushes to the shared queue
|
||||
class _BusProxy:
|
||||
def __init__(self, queue):
|
||||
self._q = queue
|
||||
|
||||
def publish(self, event):
|
||||
try:
|
||||
self._q.put_nowait(event)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def emit(self, event_type, payload, source_module=""):
|
||||
self.publish(Event(
|
||||
event_type=event_type,
|
||||
payload=payload,
|
||||
source_module=source_module,
|
||||
))
|
||||
|
||||
def subscribe(self, *a, **kw):
|
||||
pass # Subscriptions only work in the main process
|
||||
|
||||
bus_proxy = _BusProxy(bus_queue)
|
||||
state = StateManager(db_path=state_db_path)
|
||||
state.start()
|
||||
|
||||
module = module_class(bus=bus_proxy, state=state, config=config)
|
||||
|
||||
# Signal handlers for graceful shutdown
|
||||
def _handle_term(sig, frame):
|
||||
try:
|
||||
module.stop()
|
||||
except Exception:
|
||||
pass
|
||||
state.stop()
|
||||
os._exit(0)
|
||||
|
||||
signal.signal(signal.SIGTERM, _handle_term)
|
||||
signal.signal(signal.SIGINT, _handle_term)
|
||||
|
||||
try:
|
||||
module.start()
|
||||
except Exception:
|
||||
logger.exception("Module %s crashed in start()", module.name)
|
||||
bus_proxy.emit("MODULE_ERROR", {"module": module.name, "error": "start_crash"},
|
||||
source_module=module.name)
|
||||
finally:
|
||||
state.stop()
|
||||
|
||||
|
||||
class Engine:
|
||||
"""Module lifecycle manager.
|
||||
|
||||
Manages multiprocess module instances with dependency resolution,
|
||||
hardware tier enforcement, scope checking, and phase gating.
|
||||
"""
|
||||
|
||||
def __init__(self, bus: EventBus, state: StateManager, config: dict):
|
||||
self.bus = bus
|
||||
self.state = state
|
||||
self.config = config
|
||||
self._modules: dict[str, ModuleEntry] = {}
|
||||
self._tier = config.get("device", {}).get("platform", "auto")
|
||||
if self._tier == "auto":
|
||||
self._tier = detect_hardware_tier()
|
||||
self._tier_limits = HARDWARE_TIERS.get(self._tier, HARDWARE_TIERS["generic"])
|
||||
self._scope = config.get("scope", {})
|
||||
self._phase = config.get("engagement_phase", {}).get("mode", "passive_only")
|
||||
self._lock = multiprocessing.Lock()
|
||||
logger.info("Engine initialized (tier=%s, phase=%s)", self._tier, self._phase)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Module registration
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def register(self, module_class: Type[BaseModule], config: dict = None) -> None:
|
||||
"""Register a module class for management."""
|
||||
name = module_class.name
|
||||
if name in self._modules:
|
||||
logger.warning("Module %s already registered, replacing", name)
|
||||
self._modules[name] = ModuleEntry(module_class, config or {})
|
||||
logger.debug("Registered module: %s (type=%s, priority=%d)",
|
||||
name, module_class.module_type, module_class.priority)
|
||||
|
||||
def load_module(self, module_path: str, config: dict = None) -> None:
|
||||
"""Dynamically load a module from a dotted Python path.
|
||||
|
||||
e.g. "modules.passive.dns_logger.DnsLogger"
|
||||
"""
|
||||
parts = module_path.rsplit(".", 1)
|
||||
if len(parts) != 2:
|
||||
raise ValueError(f"Expected 'package.ClassName', got: {module_path}")
|
||||
mod = importlib.import_module(parts[0])
|
||||
cls = getattr(mod, parts[1])
|
||||
if not (isinstance(cls, type) and issubclass(cls, BaseModule)):
|
||||
raise TypeError(f"{module_path} is not a BaseModule subclass")
|
||||
self.register(cls, config)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Lifecycle
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def start(self, name: str) -> bool:
|
||||
"""Start a single module by name. Returns True on success."""
|
||||
with self._lock:
|
||||
entry = self._modules.get(name)
|
||||
if entry is None:
|
||||
logger.error("Module %s not registered", name)
|
||||
return False
|
||||
|
||||
if entry.process and entry.process.is_alive():
|
||||
logger.warning("Module %s already running", name)
|
||||
return True
|
||||
|
||||
# Phase gating
|
||||
if not self._phase_allows(entry.module_class):
|
||||
logger.warning("Module %s blocked by phase=%s", name, self._phase)
|
||||
return False
|
||||
|
||||
# Tier limits
|
||||
if not self._tier_allows(entry.module_class):
|
||||
logger.warning("Module %s blocked by tier=%s limits", name, self._tier)
|
||||
return False
|
||||
|
||||
# Dependency check
|
||||
for dep in entry.module_class.dependencies:
|
||||
dep_entry = self._modules.get(dep)
|
||||
if dep_entry is None or not (dep_entry.process and dep_entry.process.is_alive()):
|
||||
logger.error("Module %s requires %s which is not running", name, dep)
|
||||
return False
|
||||
|
||||
proc = multiprocessing.Process(
|
||||
target=_module_runner,
|
||||
args=(
|
||||
entry.module_class,
|
||||
self.bus.get_queue(),
|
||||
self.state._db_path,
|
||||
entry.config,
|
||||
None,
|
||||
),
|
||||
name=f"bb-{name}",
|
||||
daemon=True,
|
||||
)
|
||||
proc.start()
|
||||
entry.process = proc
|
||||
entry.started_at = time.time()
|
||||
|
||||
self.state.set_module_status(name, "running", pid=proc.pid)
|
||||
self.bus.emit("MODULE_STARTED", {"module": name, "pid": proc.pid},
|
||||
source_module="engine")
|
||||
logger.info("Started module %s (pid=%d)", name, proc.pid)
|
||||
return True
|
||||
|
||||
def stop(self, name: str, timeout: float = 5.0) -> bool:
|
||||
"""Stop a module. SIGTERM -> wait -> SIGKILL."""
|
||||
with self._lock:
|
||||
entry = self._modules.get(name)
|
||||
if entry is None or entry.process is None:
|
||||
return True
|
||||
|
||||
proc = entry.process
|
||||
if proc.is_alive():
|
||||
proc.terminate()
|
||||
proc.join(timeout=timeout)
|
||||
if proc.is_alive():
|
||||
logger.warning("Module %s did not stop, sending SIGKILL", name)
|
||||
proc.kill()
|
||||
proc.join(timeout=2.0)
|
||||
|
||||
entry.process = None
|
||||
entry.started_at = None
|
||||
self.state.set_module_status(name, "stopped")
|
||||
self.bus.emit("MODULE_STOPPED", {"module": name}, source_module="engine")
|
||||
logger.info("Stopped module %s", name)
|
||||
return True
|
||||
|
||||
def restart(self, name: str) -> bool:
|
||||
"""Stop then start a module."""
|
||||
self.stop(name)
|
||||
time.sleep(0.2)
|
||||
return self.start(name)
|
||||
|
||||
def start_all(self) -> dict[str, bool]:
|
||||
"""Start all registered modules in dependency order."""
|
||||
order = _topo_sort(self._modules)
|
||||
results = {}
|
||||
for name in order:
|
||||
results[name] = self.start(name)
|
||||
return results
|
||||
|
||||
def stop_all(self, timeout: float = 5.0) -> None:
|
||||
"""Stop all running modules in reverse dependency order."""
|
||||
order = _topo_sort(self._modules)
|
||||
for name in reversed(order):
|
||||
self.stop(name, timeout=timeout)
|
||||
|
||||
def get_status(self, name: str = None) -> dict:
|
||||
"""Get status of one module or all modules."""
|
||||
if name:
|
||||
entry = self._modules.get(name)
|
||||
if entry is None:
|
||||
return {"error": "not registered"}
|
||||
alive = entry.process.is_alive() if entry.process else False
|
||||
return {
|
||||
"name": name,
|
||||
"running": alive,
|
||||
"pid": entry.process.pid if entry.process else None,
|
||||
"uptime": time.time() - entry.started_at if entry.started_at and alive else 0,
|
||||
"type": entry.module_class.module_type,
|
||||
"priority": entry.module_class.priority,
|
||||
}
|
||||
|
||||
result = {}
|
||||
for n in self._modules:
|
||||
result[n] = self.get_status(n)
|
||||
return result
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Gating / enforcement
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _phase_allows(self, cls: Type[BaseModule]) -> bool:
|
||||
"""Check if current engagement phase allows this module type."""
|
||||
if self._phase == "passive_only":
|
||||
return cls.module_type in ("passive", "stealth", "intel", "connectivity")
|
||||
return True # active_allowed or unrestricted
|
||||
|
||||
def _tier_allows(self, cls: Type[BaseModule]) -> bool:
|
||||
"""Check if hardware tier limits allow this module."""
|
||||
limits = self._tier_limits
|
||||
running_by_type = self._count_running_by_type()
|
||||
|
||||
if cls.module_type == "passive":
|
||||
if running_by_type.get("passive", 0) >= limits["max_passive"]:
|
||||
return False
|
||||
elif cls.module_type == "active":
|
||||
if running_by_type.get("active", 0) >= limits["max_active"]:
|
||||
return False
|
||||
return True
|
||||
|
||||
def _count_running_by_type(self) -> dict[str, int]:
|
||||
"""Count running modules by type."""
|
||||
counts: dict[str, int] = defaultdict(int)
|
||||
for entry in self._modules.values():
|
||||
if entry.process and entry.process.is_alive():
|
||||
counts[entry.module_class.module_type] += 1
|
||||
return dict(counts)
|
||||
|
||||
def set_phase(self, phase: str) -> None:
|
||||
"""Switch engagement phase (passive_only | active_allowed)."""
|
||||
self._phase = phase
|
||||
logger.info("Engagement phase set to: %s", phase)
|
||||
|
||||
def check_scope(self, target: str) -> bool:
|
||||
"""Check if a target IP/hostname is within engagement scope."""
|
||||
if not self._scope.get("enabled", False):
|
||||
return True
|
||||
|
||||
import ipaddress
|
||||
|
||||
networks = self._scope.get("networks", [])
|
||||
hosts = self._scope.get("hosts", [])
|
||||
exclude = self._scope.get("exclude", [])
|
||||
|
||||
# Check exclusions first
|
||||
try:
|
||||
addr = ipaddress.ip_address(target)
|
||||
for exc in exclude:
|
||||
if addr in ipaddress.ip_network(exc, strict=False):
|
||||
return False
|
||||
for net in networks:
|
||||
if addr in ipaddress.ip_network(net, strict=False):
|
||||
return True
|
||||
except ValueError:
|
||||
pass # Not an IP, check hostname
|
||||
|
||||
if target in hosts:
|
||||
return True
|
||||
|
||||
domains = self._scope.get("domains", [])
|
||||
for dom in domains:
|
||||
if target.endswith("." + dom) or target == dom:
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
@property
|
||||
def tier(self) -> str:
|
||||
return self._tier
|
||||
|
||||
@property
|
||||
def tier_limits(self) -> dict:
|
||||
return dict(self._tier_limits)
|
||||
|
||||
@property
|
||||
def phase(self) -> str:
|
||||
return self._phase
|
||||
@@ -0,0 +1,473 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Kill switch — full evidence wipe and shutdown.
|
||||
|
||||
Sequence:
|
||||
1. Stop all modules + subprocesses (bettercap, tcpdump, Responder)
|
||||
2. Send corrective ARP if ARP spoof was active
|
||||
3. Destroy LUKS header (milliseconds)
|
||||
4. Shred encryption keys + credential database
|
||||
5. dd zero-fill data partition
|
||||
6. Clear RAM
|
||||
7. Reboot (boot flag auto-completes if interrupted)
|
||||
|
||||
Trigger methods: CLI, SSH, BLE, dead man's switch.
|
||||
Boot flag at /var/run/.bb_wipe ensures wipe completes if interrupted.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import signal
|
||||
import shutil
|
||||
import struct
|
||||
import socket
|
||||
import logging
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
from core.bus import EventBus
|
||||
|
||||
logger = logging.getLogger("bb.kill_switch")
|
||||
|
||||
# Boot flag — if this file exists on boot, resume wipe
|
||||
WIPE_FLAG = "/var/run/.bb_wipe"
|
||||
|
||||
# Paths to wipe (relative to install root or absolute)
|
||||
DEFAULT_SENSITIVE_PATHS = [
|
||||
"storage/creds/",
|
||||
"storage/config/",
|
||||
"storage/intel/",
|
||||
"storage/audit/",
|
||||
"storage/dns_logs/",
|
||||
"storage/sni_logs/",
|
||||
"storage/logs/",
|
||||
"storage/tool_logs/",
|
||||
"storage/pcaps/",
|
||||
"storage/baseline/",
|
||||
]
|
||||
|
||||
SENSITIVE_FILES = [
|
||||
"storage/config/encryption.key",
|
||||
"storage/creds/credentials.db",
|
||||
"storage/config/luks_header.bak",
|
||||
]
|
||||
|
||||
|
||||
class KillSwitch:
|
||||
"""Full wipe orchestrator.
|
||||
|
||||
Call execute() to perform the complete wipe sequence.
|
||||
Call check_boot_flag() on startup to resume an interrupted wipe.
|
||||
"""
|
||||
|
||||
def __init__(self, bus: EventBus, config: dict,
|
||||
engine=None, tool_manager=None):
|
||||
self.bus = bus
|
||||
self.config = config
|
||||
self.engine = engine
|
||||
self.tool_manager = tool_manager
|
||||
self._install_path = config.get("device", {}).get(
|
||||
"install_path", "/opt/.cache/bb"
|
||||
)
|
||||
self._luks_device = config.get("security", {}).get(
|
||||
"luks_device", "/dev/mapper/bb-data"
|
||||
)
|
||||
self._luks_partition = config.get("security", {}).get(
|
||||
"luks_partition", ""
|
||||
)
|
||||
self._arp_spoof_active = False
|
||||
self._spoof_gateway = None
|
||||
self._spoof_targets = []
|
||||
self._interface = config.get("network", {}).get(
|
||||
"primary_interface", "eth0"
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Main entry
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def execute(self, reason: str = "manual") -> None:
|
||||
"""Execute the full kill switch sequence. Does not return."""
|
||||
logger.critical("KILL SWITCH ACTIVATED — reason: %s", reason)
|
||||
|
||||
# Set boot flag first so reboot completes the wipe
|
||||
self._set_boot_flag()
|
||||
|
||||
self.bus.emit("KILL_SWITCH",
|
||||
{"reason": reason, "stage": "initiated"},
|
||||
source_module="kill_switch")
|
||||
|
||||
# Phase 1: Stop everything
|
||||
self._stop_all_processes()
|
||||
|
||||
# Phase 2: Corrective ARP
|
||||
self._send_corrective_arp()
|
||||
|
||||
# Phase 3: Destroy LUKS header
|
||||
self._destroy_luks()
|
||||
|
||||
# Phase 4: Shred sensitive files
|
||||
self._shred_sensitive_files()
|
||||
|
||||
# Phase 5: Zero-fill data partition
|
||||
self._zero_fill_partition()
|
||||
|
||||
# Phase 6: Clear RAM
|
||||
self._clear_ram()
|
||||
|
||||
# Phase 7: Remove boot flag and reboot
|
||||
self._clear_boot_flag()
|
||||
self._reboot()
|
||||
|
||||
def check_boot_flag(self) -> bool:
|
||||
"""Check if a wipe was interrupted and should resume.
|
||||
|
||||
Call this on boot. Returns True if wipe was resumed.
|
||||
"""
|
||||
if os.path.exists(WIPE_FLAG):
|
||||
logger.critical("Boot flag detected — resuming interrupted wipe")
|
||||
self.execute(reason="interrupted_resume")
|
||||
return True
|
||||
return False
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# ARP state (called by active modules to register spoof state)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def set_arp_spoof_state(self, active: bool, gateway: str = None,
|
||||
targets: list = None, interface: str = None) -> None:
|
||||
"""Register current ARP spoof state for corrective ARP on kill."""
|
||||
self._arp_spoof_active = active
|
||||
self._spoof_gateway = gateway
|
||||
self._spoof_targets = targets or []
|
||||
if interface:
|
||||
self._interface = interface
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Phase 1: Stop processes
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _stop_all_processes(self) -> None:
|
||||
"""Stop all modules and managed subprocesses."""
|
||||
logger.info("Kill switch phase 1: stopping all processes")
|
||||
|
||||
if self.engine:
|
||||
try:
|
||||
self.engine.stop_all(timeout=3.0)
|
||||
except Exception:
|
||||
logger.exception("Error stopping modules")
|
||||
|
||||
if self.tool_manager:
|
||||
try:
|
||||
self.tool_manager.stop_all()
|
||||
except Exception:
|
||||
logger.exception("Error stopping tools")
|
||||
|
||||
# Kill any remaining tool PIDs
|
||||
for pid in self.tool_manager.get_all_pids():
|
||||
try:
|
||||
os.kill(pid, signal.SIGKILL)
|
||||
except (OSError, ProcessLookupError):
|
||||
pass
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Phase 2: Corrective ARP
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _send_corrective_arp(self) -> None:
|
||||
"""Send corrective gratuitous ARPs if ARP spoof was active."""
|
||||
if not self._arp_spoof_active or not self._spoof_gateway:
|
||||
return
|
||||
|
||||
logger.info("Kill switch phase 2: sending corrective ARP")
|
||||
try:
|
||||
gw_ip = self._spoof_gateway
|
||||
gw_mac = self._resolve_mac(gw_ip)
|
||||
if not gw_mac:
|
||||
logger.warning("Could not resolve gateway MAC for corrective ARP")
|
||||
return
|
||||
|
||||
for target_ip in self._spoof_targets:
|
||||
target_mac = self._resolve_mac(target_ip)
|
||||
if target_mac:
|
||||
self._send_arp_reply(
|
||||
src_mac=gw_mac, src_ip=gw_ip,
|
||||
dst_mac=target_mac, dst_ip=target_ip,
|
||||
interface=self._interface
|
||||
)
|
||||
# Also restore gateway's ARP cache
|
||||
for target_ip in self._spoof_targets:
|
||||
target_mac = self._resolve_mac(target_ip)
|
||||
if target_mac:
|
||||
self._send_arp_reply(
|
||||
src_mac=target_mac, src_ip=target_ip,
|
||||
dst_mac=gw_mac, dst_ip=gw_ip,
|
||||
interface=self._interface
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("Corrective ARP failed")
|
||||
|
||||
@staticmethod
|
||||
def _resolve_mac(ip: str) -> Optional[str]:
|
||||
"""Resolve IP to MAC from /proc/net/arp."""
|
||||
try:
|
||||
with open("/proc/net/arp", "r") as f:
|
||||
for line in f:
|
||||
parts = line.split()
|
||||
if len(parts) >= 4 and parts[0] == ip:
|
||||
mac = parts[3]
|
||||
if mac != "00:00:00:00:00:00":
|
||||
return mac
|
||||
except (FileNotFoundError, PermissionError):
|
||||
pass
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _send_arp_reply(src_mac: str, src_ip: str, dst_mac: str, dst_ip: str,
|
||||
interface: str) -> None:
|
||||
"""Send a single ARP reply using raw socket."""
|
||||
try:
|
||||
sock = socket.socket(socket.AF_PACKET, socket.SOCK_RAW, socket.htons(0x0806))
|
||||
sock.bind((interface, 0))
|
||||
|
||||
src_mac_bytes = bytes.fromhex(src_mac.replace(":", ""))
|
||||
dst_mac_bytes = bytes.fromhex(dst_mac.replace(":", ""))
|
||||
src_ip_bytes = socket.inet_aton(src_ip)
|
||||
dst_ip_bytes = socket.inet_aton(dst_ip)
|
||||
|
||||
# Ethernet header + ARP reply
|
||||
frame = (
|
||||
dst_mac_bytes + # Dst MAC
|
||||
src_mac_bytes + # Src MAC
|
||||
b"\x08\x06" + # EtherType: ARP
|
||||
b"\x00\x01" + # Hardware type: Ethernet
|
||||
b"\x08\x00" + # Protocol type: IPv4
|
||||
b"\x06" + # Hardware size: 6
|
||||
b"\x04" + # Protocol size: 4
|
||||
b"\x00\x02" + # Opcode: reply
|
||||
src_mac_bytes + # Sender MAC
|
||||
src_ip_bytes + # Sender IP
|
||||
dst_mac_bytes + # Target MAC
|
||||
dst_ip_bytes # Target IP
|
||||
)
|
||||
|
||||
# Send 3 times for reliability
|
||||
for _ in range(3):
|
||||
sock.send(frame)
|
||||
time.sleep(0.1)
|
||||
|
||||
sock.close()
|
||||
except Exception as e:
|
||||
logger.warning("ARP reply send failed: %s", e)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Phase 3: LUKS header destruction
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _destroy_luks(self) -> None:
|
||||
"""Destroy the LUKS header, making encrypted data unrecoverable."""
|
||||
logger.info("Kill switch phase 3: destroying LUKS header")
|
||||
|
||||
if not self._luks_partition:
|
||||
logger.warning("No LUKS partition configured, skipping")
|
||||
return
|
||||
|
||||
# Close the LUKS device first
|
||||
try:
|
||||
subprocess.run(
|
||||
["cryptsetup", "close", self._luks_device.split("/")[-1]],
|
||||
timeout=5, capture_output=True
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Overwrite LUKS header (first 16MB covers all key slots)
|
||||
try:
|
||||
subprocess.run(
|
||||
["dd", "if=/dev/urandom", f"of={self._luks_partition}",
|
||||
"bs=1M", "count=16", "conv=notrunc"],
|
||||
timeout=30, capture_output=True, check=True
|
||||
)
|
||||
logger.info("LUKS header destroyed")
|
||||
except Exception as e:
|
||||
logger.error("LUKS header destruction failed: %s", e)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Phase 4: Shred sensitive files
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _shred_sensitive_files(self) -> None:
|
||||
"""Securely delete sensitive files and directories."""
|
||||
logger.info("Kill switch phase 4: shredding sensitive files")
|
||||
|
||||
base = Path(self._install_path)
|
||||
|
||||
# Shred specific files first
|
||||
for rel_path in SENSITIVE_FILES:
|
||||
fpath = base / rel_path
|
||||
if fpath.exists():
|
||||
self._shred_file(str(fpath))
|
||||
|
||||
# Then directories
|
||||
for rel_path in DEFAULT_SENSITIVE_PATHS:
|
||||
dpath = base / rel_path
|
||||
if dpath.exists():
|
||||
for fpath in dpath.rglob("*"):
|
||||
if fpath.is_file():
|
||||
self._shred_file(str(fpath))
|
||||
try:
|
||||
shutil.rmtree(str(dpath), ignore_errors=True)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# State database
|
||||
state_db = os.path.expanduser("~/.bigbrother/state.db")
|
||||
for f in [state_db, state_db + "-wal", state_db + "-shm"]:
|
||||
if os.path.exists(f):
|
||||
self._shred_file(f)
|
||||
|
||||
@staticmethod
|
||||
def _shred_file(path: str) -> None:
|
||||
"""Overwrite file with random data, then unlink."""
|
||||
try:
|
||||
size = os.path.getsize(path)
|
||||
with open(path, "r+b") as f:
|
||||
# 3 passes: random, zeros, random
|
||||
for pattern in [None, b"\x00", None]:
|
||||
f.seek(0)
|
||||
remaining = size
|
||||
while remaining > 0:
|
||||
chunk = min(remaining, 65536)
|
||||
if pattern is None:
|
||||
data = os.urandom(chunk)
|
||||
else:
|
||||
data = pattern * chunk
|
||||
f.write(data)
|
||||
remaining -= chunk
|
||||
f.flush()
|
||||
os.fsync(f.fileno())
|
||||
os.unlink(path)
|
||||
except Exception as e:
|
||||
# Fallback: just delete
|
||||
try:
|
||||
os.unlink(path)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Phase 5: Zero-fill
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _zero_fill_partition(self) -> None:
|
||||
"""Zero-fill the data partition."""
|
||||
logger.info("Kill switch phase 5: zero-filling data partition")
|
||||
|
||||
storage_dir = os.path.join(self._install_path, "storage")
|
||||
if not os.path.isdir(storage_dir):
|
||||
storage_dir = self._install_path
|
||||
|
||||
try:
|
||||
# Write zeros until disk full
|
||||
zero_file = os.path.join(storage_dir, ".wipe")
|
||||
subprocess.run(
|
||||
["dd", "if=/dev/zero", f"of={zero_file}",
|
||||
"bs=1M", "status=none"],
|
||||
timeout=600, capture_output=True
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
finally:
|
||||
try:
|
||||
os.unlink(os.path.join(storage_dir, ".wipe"))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Phase 6: Clear RAM
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def _clear_ram() -> None:
|
||||
"""Drop caches and clear as much RAM as possible."""
|
||||
logger.info("Kill switch phase 6: clearing RAM")
|
||||
try:
|
||||
# Drop page cache, dentries, inodes
|
||||
with open("/proc/sys/vm/drop_caches", "w") as f:
|
||||
f.write("3")
|
||||
except (PermissionError, FileNotFoundError):
|
||||
pass
|
||||
|
||||
# smem or sdmem would be better but may not be installed
|
||||
try:
|
||||
subprocess.run(["sdmem", "-f", "-ll"], timeout=60,
|
||||
capture_output=True)
|
||||
except (FileNotFoundError, subprocess.TimeoutExpired):
|
||||
pass
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Phase 7: Reboot
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def _reboot() -> None:
|
||||
"""Reboot the system."""
|
||||
logger.info("Kill switch phase 7: rebooting")
|
||||
try:
|
||||
subprocess.run(["reboot"], timeout=5)
|
||||
except Exception:
|
||||
try:
|
||||
os.system("reboot")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Boot flag
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def _set_boot_flag() -> None:
|
||||
"""Create boot flag to resume wipe if interrupted."""
|
||||
try:
|
||||
Path(WIPE_FLAG).parent.mkdir(parents=True, exist_ok=True)
|
||||
Path(WIPE_FLAG).touch()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def _clear_boot_flag() -> None:
|
||||
"""Remove boot flag after successful wipe."""
|
||||
try:
|
||||
os.unlink(WIPE_FLAG)
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Dead man's switch
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def start_dead_man_switch(self, interval_hours: int = 24) -> None:
|
||||
"""Start dead man's switch timer.
|
||||
|
||||
If not reset within interval_hours, trigger kill switch.
|
||||
Intended to be reset by operator check-in via C2.
|
||||
"""
|
||||
self._dms_interval = interval_hours * 3600
|
||||
self._dms_last_checkin = time.time()
|
||||
|
||||
import threading
|
||||
def _dms_loop():
|
||||
while True:
|
||||
time.sleep(60)
|
||||
elapsed = time.time() - self._dms_last_checkin
|
||||
if elapsed > self._dms_interval:
|
||||
logger.critical("Dead man's switch triggered (no check-in for %dh)",
|
||||
interval_hours)
|
||||
self.execute(reason="dead_man_switch")
|
||||
|
||||
t = threading.Thread(target=_dms_loop, daemon=True, name="bb-dms")
|
||||
t.start()
|
||||
|
||||
def reset_dead_man_switch(self) -> None:
|
||||
"""Reset the dead man's switch timer (operator check-in)."""
|
||||
self._dms_last_checkin = time.time()
|
||||
@@ -0,0 +1,329 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Per-module and per-subprocess resource monitoring.
|
||||
|
||||
Tracks RAM, CPU, disk usage, and thermal readings. Publishes
|
||||
RESOURCE_WARNING, RESOURCE_CRITICAL, and THERMAL_WARNING events.
|
||||
OOM-kills the lowest-priority module when memory exceeds tier limits.
|
||||
Sheds modules at 70C thermal threshold.
|
||||
"""
|
||||
|
||||
import os
|
||||
import time
|
||||
import logging
|
||||
import threading
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
from core.bus import EventBus
|
||||
from core.state import StateManager
|
||||
|
||||
logger = logging.getLogger("bb.resource_monitor")
|
||||
|
||||
# Thermal thresholds (Celsius)
|
||||
THERMAL_WARNING = 55
|
||||
THERMAL_CRITICAL = 70
|
||||
|
||||
# Disk usage threshold
|
||||
DISK_WARNING_PCT = 80
|
||||
DISK_CRITICAL_PCT = 90
|
||||
|
||||
# Check interval
|
||||
CHECK_INTERVAL = 15.0 # seconds
|
||||
|
||||
|
||||
class ResourceMonitor:
|
||||
"""Monitors system and per-process resources.
|
||||
|
||||
Integrates with Engine (module processes) and ToolManager (subprocess PIDs)
|
||||
to provide unified resource tracking.
|
||||
"""
|
||||
|
||||
def __init__(self, bus: EventBus, state: StateManager, config: dict,
|
||||
engine=None, tool_manager=None):
|
||||
self.bus = bus
|
||||
self.state = state
|
||||
self.config = config
|
||||
self.engine = engine
|
||||
self.tool_manager = tool_manager
|
||||
self._thread: Optional[threading.Thread] = None
|
||||
self._running = False
|
||||
|
||||
# Load tier limits
|
||||
tier = config.get("device", {}).get("platform", "generic")
|
||||
from core.engine import HARDWARE_TIERS, detect_hardware_tier
|
||||
if tier == "auto":
|
||||
tier = detect_hardware_tier()
|
||||
self._tier_limits = HARDWARE_TIERS.get(tier, HARDWARE_TIERS["generic"])
|
||||
self._tier = tier
|
||||
|
||||
# Thermal zone path (try common locations)
|
||||
self._thermal_path = self._find_thermal_zone()
|
||||
|
||||
# State tracking
|
||||
self._thermal_warned = False
|
||||
self._resource_warned = False
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Lifecycle
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def start(self) -> None:
|
||||
if self._running:
|
||||
return
|
||||
self._running = True
|
||||
self._thread = threading.Thread(
|
||||
target=self._monitor_loop, daemon=True, name="bb-resource-monitor"
|
||||
)
|
||||
self._thread.start()
|
||||
logger.info("ResourceMonitor started (tier=%s, thermal_zone=%s)",
|
||||
self._tier, self._thermal_path or "none")
|
||||
|
||||
def stop(self) -> None:
|
||||
self._running = False
|
||||
if self._thread and self._thread.is_alive():
|
||||
self._thread.join(timeout=5.0)
|
||||
logger.info("ResourceMonitor stopped")
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Monitor loop
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _monitor_loop(self) -> None:
|
||||
while self._running:
|
||||
try:
|
||||
self._check_resources()
|
||||
except Exception:
|
||||
logger.exception("Resource check error")
|
||||
time.sleep(CHECK_INTERVAL)
|
||||
|
||||
def _check_resources(self) -> None:
|
||||
"""Run all resource checks."""
|
||||
self._check_memory()
|
||||
self._check_disk()
|
||||
self._check_thermal()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Memory
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _check_memory(self) -> None:
|
||||
"""Check total memory usage against tier limits."""
|
||||
mem = self.get_system_memory()
|
||||
max_ram = self._tier_limits.get("max_ram_mb", 0)
|
||||
if max_ram <= 0:
|
||||
return
|
||||
|
||||
used_mb = mem["used_mb"]
|
||||
|
||||
if used_mb > max_ram * 0.95:
|
||||
# Critical — OOM kill lowest priority module
|
||||
logger.critical("Memory critical: %.0fMB / %dMB", used_mb, max_ram)
|
||||
self.bus.emit("RESOURCE_CRITICAL",
|
||||
{"resource": "memory", "used_mb": used_mb, "limit_mb": max_ram},
|
||||
source_module="resource_monitor")
|
||||
self._oom_kill()
|
||||
elif used_mb > max_ram * 0.80:
|
||||
if not self._resource_warned:
|
||||
logger.warning("Memory warning: %.0fMB / %dMB", used_mb, max_ram)
|
||||
self.bus.emit("RESOURCE_WARNING",
|
||||
{"resource": "memory", "used_mb": used_mb, "limit_mb": max_ram},
|
||||
source_module="resource_monitor")
|
||||
self._resource_warned = True
|
||||
else:
|
||||
self._resource_warned = False
|
||||
|
||||
def _oom_kill(self) -> None:
|
||||
"""Kill the lowest-priority (highest priority number) running module."""
|
||||
if self.engine is None:
|
||||
return
|
||||
|
||||
# Find lowest-priority running module
|
||||
worst_name = None
|
||||
worst_priority = -1000
|
||||
|
||||
status = self.engine.get_status()
|
||||
for name, info in status.items():
|
||||
if not info.get("running"):
|
||||
continue
|
||||
prio = info.get("priority", 0)
|
||||
if prio > worst_priority:
|
||||
worst_priority = prio
|
||||
worst_name = name
|
||||
|
||||
if worst_name:
|
||||
logger.warning("OOM killing module: %s (priority=%d)", worst_name, worst_priority)
|
||||
self.engine.stop(worst_name)
|
||||
self.bus.emit("MODULE_STOPPED",
|
||||
{"module": worst_name, "reason": "oom_kill"},
|
||||
source_module="resource_monitor")
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Disk
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _check_disk(self) -> None:
|
||||
"""Check disk usage on the data partition."""
|
||||
storage_path = self.config.get("storage_path", "/opt/.cache/bb/storage")
|
||||
try:
|
||||
st = os.statvfs(storage_path)
|
||||
except OSError:
|
||||
try:
|
||||
st = os.statvfs("/")
|
||||
except OSError:
|
||||
return
|
||||
|
||||
total = st.f_blocks * st.f_frsize
|
||||
free = st.f_bavail * st.f_frsize
|
||||
used_pct = ((total - free) / total * 100) if total > 0 else 0
|
||||
|
||||
if used_pct >= DISK_CRITICAL_PCT:
|
||||
self.bus.emit("RESOURCE_CRITICAL",
|
||||
{"resource": "disk", "used_pct": round(used_pct, 1),
|
||||
"path": storage_path},
|
||||
source_module="resource_monitor")
|
||||
elif used_pct >= DISK_WARNING_PCT:
|
||||
self.bus.emit("RESOURCE_WARNING",
|
||||
{"resource": "disk", "used_pct": round(used_pct, 1),
|
||||
"path": storage_path},
|
||||
source_module="resource_monitor")
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Thermal
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _check_thermal(self) -> None:
|
||||
"""Read CPU temperature and emit warnings / shed modules."""
|
||||
temp = self.get_temperature()
|
||||
if temp is None:
|
||||
return
|
||||
|
||||
if temp >= THERMAL_CRITICAL:
|
||||
logger.critical("Thermal critical: %.1fC — shedding modules", temp)
|
||||
self.bus.emit("THERMAL_WARNING",
|
||||
{"temperature": temp, "action": "shedding"},
|
||||
source_module="resource_monitor")
|
||||
self._thermal_shed()
|
||||
elif temp >= THERMAL_WARNING:
|
||||
if not self._thermal_warned:
|
||||
logger.warning("Thermal warning: %.1fC", temp)
|
||||
self.bus.emit("THERMAL_WARNING",
|
||||
{"temperature": temp, "action": "warning"},
|
||||
source_module="resource_monitor")
|
||||
self._thermal_warned = True
|
||||
else:
|
||||
self._thermal_warned = False
|
||||
|
||||
def _thermal_shed(self) -> None:
|
||||
"""Shed (stop) active modules to reduce heat — keep passive + core."""
|
||||
if self.engine is None:
|
||||
return
|
||||
|
||||
status = self.engine.get_status()
|
||||
# Stop active modules first (highest priority number = least important)
|
||||
shedded = []
|
||||
for name, info in sorted(status.items(),
|
||||
key=lambda x: x[1].get("priority", 0), reverse=True):
|
||||
if not info.get("running"):
|
||||
continue
|
||||
if info.get("type") == "active":
|
||||
self.engine.stop(name)
|
||||
shedded.append(name)
|
||||
|
||||
if shedded:
|
||||
logger.warning("Thermal shed: stopped %s", shedded)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Readings
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def get_system_memory() -> dict:
|
||||
"""Read system memory from /proc/meminfo."""
|
||||
info = {"total_mb": 0, "available_mb": 0, "used_mb": 0}
|
||||
try:
|
||||
with open("/proc/meminfo", "r") as f:
|
||||
for line in f:
|
||||
if line.startswith("MemTotal:"):
|
||||
info["total_mb"] = int(line.split()[1]) / 1024
|
||||
elif line.startswith("MemAvailable:"):
|
||||
info["available_mb"] = int(line.split()[1]) / 1024
|
||||
info["used_mb"] = info["total_mb"] - info["available_mb"]
|
||||
except (FileNotFoundError, PermissionError, ValueError):
|
||||
pass
|
||||
return {k: round(v, 1) for k, v in info.items()}
|
||||
|
||||
def get_temperature(self) -> Optional[float]:
|
||||
"""Read CPU temperature in Celsius."""
|
||||
if not self._thermal_path:
|
||||
return None
|
||||
try:
|
||||
with open(self._thermal_path, "r") as f:
|
||||
raw = f.read().strip()
|
||||
temp = int(raw) / 1000.0 # millidegrees
|
||||
return round(temp, 1)
|
||||
except (FileNotFoundError, PermissionError, ValueError):
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def get_cpu_usage() -> float:
|
||||
"""Get overall CPU usage percentage from /proc/stat (1-second sample)."""
|
||||
def _read():
|
||||
with open("/proc/stat", "r") as f:
|
||||
parts = f.readline().split()
|
||||
return [int(x) for x in parts[1:]]
|
||||
|
||||
try:
|
||||
t1 = _read()
|
||||
time.sleep(1.0)
|
||||
t2 = _read()
|
||||
d = [t2[i] - t1[i] for i in range(len(t1))]
|
||||
idle = d[3] + (d[4] if len(d) > 4 else 0)
|
||||
total = sum(d)
|
||||
return round((1.0 - idle / total) * 100, 1) if total > 0 else 0.0
|
||||
except Exception:
|
||||
return 0.0
|
||||
|
||||
def get_process_resources(self, pid: int) -> dict:
|
||||
"""Get RAM/CPU for a specific PID."""
|
||||
from core.tool_manager import ToolManager
|
||||
return ToolManager._read_proc_stats(pid)
|
||||
|
||||
def get_snapshot(self) -> dict:
|
||||
"""Full resource snapshot."""
|
||||
return {
|
||||
"system_memory": self.get_system_memory(),
|
||||
"temperature": self.get_temperature(),
|
||||
"tier": self._tier,
|
||||
"tier_limits": self._tier_limits,
|
||||
}
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def _find_thermal_zone() -> Optional[str]:
|
||||
"""Find the CPU thermal zone file."""
|
||||
base = Path("/sys/class/thermal")
|
||||
if not base.exists():
|
||||
return None
|
||||
|
||||
for zone in sorted(base.iterdir()):
|
||||
type_file = zone / "type"
|
||||
temp_file = zone / "temp"
|
||||
if not temp_file.exists():
|
||||
continue
|
||||
try:
|
||||
ztype = type_file.read_text().strip().lower()
|
||||
if any(k in ztype for k in ("cpu", "soc", "x86_pkg")):
|
||||
return str(temp_file)
|
||||
except (PermissionError, FileNotFoundError):
|
||||
continue
|
||||
|
||||
# Fallback: first thermal zone with a temp file
|
||||
for zone in sorted(base.iterdir()):
|
||||
temp_file = zone / "temp"
|
||||
if temp_file.exists():
|
||||
return str(temp_file)
|
||||
|
||||
return None
|
||||
@@ -0,0 +1,192 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Cron-like task scheduler with jitter.
|
||||
|
||||
Modules register periodic tasks with an interval and optional jitter range.
|
||||
Tasks execute in a thread pool. Used for PCAP rotation, health checks,
|
||||
baseline captures, data pruning, resource checks.
|
||||
"""
|
||||
|
||||
import time
|
||||
import random
|
||||
import logging
|
||||
import threading
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Callable, Optional
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
logger = logging.getLogger("bb.scheduler")
|
||||
|
||||
|
||||
@dataclass
|
||||
class ScheduledTask:
|
||||
"""A registered periodic task."""
|
||||
name: str
|
||||
callback: Callable
|
||||
interval: float # seconds between executions
|
||||
jitter: float = 0.0 # max random jitter added to interval (seconds)
|
||||
run_immediately: bool = False
|
||||
enabled: bool = True
|
||||
|
||||
# Runtime state
|
||||
last_run: float = 0.0
|
||||
next_run: float = 0.0
|
||||
run_count: int = 0
|
||||
error_count: int = 0
|
||||
_last_error: Optional[str] = None
|
||||
|
||||
|
||||
class Scheduler:
|
||||
"""Thread-safe periodic task scheduler with jitter support.
|
||||
|
||||
Tasks are checked every second and dispatched to a thread pool.
|
||||
Jitter is re-randomized on each scheduling cycle to avoid patterns.
|
||||
"""
|
||||
|
||||
def __init__(self, max_workers: int = 4):
|
||||
self._tasks: dict[str, ScheduledTask] = {}
|
||||
self._lock = threading.Lock()
|
||||
self._running = False
|
||||
self._thread: Optional[threading.Thread] = None
|
||||
self._pool = ThreadPoolExecutor(max_workers=max_workers,
|
||||
thread_name_prefix="bb-sched")
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Lifecycle
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def start(self) -> None:
|
||||
"""Start the scheduler loop."""
|
||||
if self._running:
|
||||
return
|
||||
self._running = True
|
||||
# Initialize next_run for tasks that want immediate execution
|
||||
now = time.time()
|
||||
with self._lock:
|
||||
for task in self._tasks.values():
|
||||
if task.run_immediately:
|
||||
task.next_run = now
|
||||
elif task.next_run == 0.0:
|
||||
task.next_run = now + task.interval + random.uniform(0, task.jitter)
|
||||
|
||||
self._thread = threading.Thread(
|
||||
target=self._scheduler_loop, daemon=True, name="bb-scheduler"
|
||||
)
|
||||
self._thread.start()
|
||||
logger.info("Scheduler started (%d tasks)", len(self._tasks))
|
||||
|
||||
def stop(self) -> None:
|
||||
"""Stop the scheduler and wait for running tasks to finish."""
|
||||
self._running = False
|
||||
if self._thread and self._thread.is_alive():
|
||||
self._thread.join(timeout=5.0)
|
||||
self._pool.shutdown(wait=False)
|
||||
logger.info("Scheduler stopped")
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Task management
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def register(self, name: str, callback: Callable, interval: float,
|
||||
jitter: float = 0.0, run_immediately: bool = False) -> None:
|
||||
"""Register a periodic task.
|
||||
|
||||
Args:
|
||||
name: Unique task name.
|
||||
callback: Callable (no arguments) to execute.
|
||||
interval: Base interval in seconds.
|
||||
jitter: Maximum random seconds added to interval.
|
||||
run_immediately: If True, run as soon as the scheduler starts.
|
||||
"""
|
||||
now = time.time()
|
||||
task = ScheduledTask(
|
||||
name=name,
|
||||
callback=callback,
|
||||
interval=interval,
|
||||
jitter=jitter,
|
||||
run_immediately=run_immediately,
|
||||
next_run=now if run_immediately else now + interval + random.uniform(0, jitter),
|
||||
)
|
||||
with self._lock:
|
||||
self._tasks[name] = task
|
||||
logger.debug("Registered task: %s (interval=%.1fs, jitter=%.1fs)",
|
||||
name, interval, jitter)
|
||||
|
||||
def unregister(self, name: str) -> None:
|
||||
"""Remove a task."""
|
||||
with self._lock:
|
||||
self._tasks.pop(name, None)
|
||||
|
||||
def enable(self, name: str) -> None:
|
||||
"""Enable a disabled task."""
|
||||
with self._lock:
|
||||
task = self._tasks.get(name)
|
||||
if task:
|
||||
task.enabled = True
|
||||
|
||||
def disable(self, name: str) -> None:
|
||||
"""Disable a task (stays registered but won't execute)."""
|
||||
with self._lock:
|
||||
task = self._tasks.get(name)
|
||||
if task:
|
||||
task.enabled = False
|
||||
|
||||
def run_now(self, name: str) -> None:
|
||||
"""Trigger immediate execution of a task."""
|
||||
with self._lock:
|
||||
task = self._tasks.get(name)
|
||||
if task:
|
||||
task.next_run = time.time()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Scheduler loop
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _scheduler_loop(self) -> None:
|
||||
"""Check tasks every second, dispatch due tasks to thread pool."""
|
||||
while self._running:
|
||||
now = time.time()
|
||||
with self._lock:
|
||||
tasks = list(self._tasks.values())
|
||||
|
||||
for task in tasks:
|
||||
if not task.enabled:
|
||||
continue
|
||||
if now >= task.next_run:
|
||||
self._pool.submit(self._execute_task, task)
|
||||
# Schedule next run with fresh jitter
|
||||
jitter = random.uniform(0, task.jitter) if task.jitter > 0 else 0
|
||||
task.next_run = now + task.interval + jitter
|
||||
|
||||
time.sleep(1.0)
|
||||
|
||||
def _execute_task(self, task: ScheduledTask) -> None:
|
||||
"""Execute a task and update bookkeeping."""
|
||||
try:
|
||||
task.callback()
|
||||
task.run_count += 1
|
||||
task.last_run = time.time()
|
||||
except Exception as e:
|
||||
task.error_count += 1
|
||||
task._last_error = str(e)
|
||||
logger.exception("Scheduler task %s failed", task.name)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Introspection
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def get_tasks(self) -> dict:
|
||||
"""Return info about all registered tasks."""
|
||||
with self._lock:
|
||||
return {
|
||||
name: {
|
||||
"interval": t.interval,
|
||||
"jitter": t.jitter,
|
||||
"enabled": t.enabled,
|
||||
"run_count": t.run_count,
|
||||
"error_count": t.error_count,
|
||||
"last_run": t.last_run,
|
||||
"next_run": t.next_run,
|
||||
"last_error": t._last_error,
|
||||
}
|
||||
for name, t in self._tasks.items()
|
||||
}
|
||||
+259
@@ -0,0 +1,259 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Persistent state manager backed by SQLite in WAL mode.
|
||||
|
||||
A dedicated writer thread coalesces writes with a 500ms flush interval
|
||||
to reduce SD card wear. Reads are served directly (WAL allows concurrent
|
||||
readers). Module status and generic key-value data are stored in separate
|
||||
tables.
|
||||
"""
|
||||
|
||||
import os
|
||||
import time
|
||||
import json
|
||||
import sqlite3
|
||||
import threading
|
||||
import logging
|
||||
from typing import Any, Optional
|
||||
from pathlib import Path
|
||||
|
||||
logger = logging.getLogger("bb.state")
|
||||
|
||||
_DEFAULT_DB = os.path.join(
|
||||
os.path.expanduser("~"), ".bigbrother", "state.db"
|
||||
)
|
||||
|
||||
|
||||
class StateManager:
|
||||
"""SQLite-backed persistent state with coalesced writes.
|
||||
|
||||
Thread-safe. The write queue is drained every 500ms by a background
|
||||
thread that batches all pending operations into a single transaction.
|
||||
"""
|
||||
|
||||
FLUSH_INTERVAL = 0.5 # seconds
|
||||
|
||||
def __init__(self, db_path: str = _DEFAULT_DB):
|
||||
self._db_path = db_path
|
||||
Path(db_path).parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Writer thread owns this connection
|
||||
self._write_conn: Optional[sqlite3.Connection] = None
|
||||
# Per-thread read connections via _local
|
||||
self._local = threading.local()
|
||||
|
||||
self._write_queue: list[tuple] = [] # (sql, params) tuples
|
||||
self._queue_lock = threading.Lock()
|
||||
self._writer_thread: Optional[threading.Thread] = None
|
||||
self._running = False
|
||||
|
||||
self._init_db()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Lifecycle
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def start(self) -> None:
|
||||
"""Start the background writer thread."""
|
||||
if self._running:
|
||||
return
|
||||
self._running = True
|
||||
self._writer_thread = threading.Thread(
|
||||
target=self._writer_loop, daemon=True, name="bb-state-writer"
|
||||
)
|
||||
self._writer_thread.start()
|
||||
logger.info("StateManager writer started (db=%s)", self._db_path)
|
||||
|
||||
def stop(self) -> None:
|
||||
"""Flush remaining writes and stop the writer thread."""
|
||||
if not self._running:
|
||||
return
|
||||
self._running = False
|
||||
if self._writer_thread and self._writer_thread.is_alive():
|
||||
self._writer_thread.join(timeout=5.0)
|
||||
# Final flush
|
||||
self._flush()
|
||||
if self._write_conn:
|
||||
self._write_conn.close()
|
||||
self._write_conn = None
|
||||
logger.info("StateManager stopped")
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Schema init
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _init_db(self) -> None:
|
||||
conn = sqlite3.connect(self._db_path)
|
||||
conn.execute("PRAGMA journal_mode=WAL")
|
||||
conn.execute("PRAGMA synchronous=NORMAL")
|
||||
conn.executescript("""
|
||||
CREATE TABLE IF NOT EXISTS kv_store (
|
||||
module TEXT NOT NULL,
|
||||
key TEXT NOT NULL,
|
||||
value TEXT,
|
||||
updated REAL NOT NULL,
|
||||
PRIMARY KEY (module, key)
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS module_status (
|
||||
module TEXT PRIMARY KEY,
|
||||
status TEXT NOT NULL DEFAULT 'stopped',
|
||||
pid INTEGER,
|
||||
started REAL,
|
||||
extra TEXT,
|
||||
updated REAL NOT NULL
|
||||
);
|
||||
""")
|
||||
conn.commit()
|
||||
conn.close()
|
||||
|
||||
def _get_read_conn(self) -> sqlite3.Connection:
|
||||
"""Return a per-thread read connection."""
|
||||
conn = getattr(self._local, "conn", None)
|
||||
if conn is None:
|
||||
conn = sqlite3.connect(self._db_path, check_same_thread=False)
|
||||
conn.execute("PRAGMA journal_mode=WAL")
|
||||
conn.execute("PRAGMA synchronous=NORMAL")
|
||||
conn.row_factory = sqlite3.Row
|
||||
self._local.conn = conn
|
||||
return conn
|
||||
|
||||
def _get_write_conn(self) -> sqlite3.Connection:
|
||||
if self._write_conn is None:
|
||||
self._write_conn = sqlite3.connect(self._db_path, check_same_thread=False)
|
||||
self._write_conn.execute("PRAGMA journal_mode=WAL")
|
||||
self._write_conn.execute("PRAGMA synchronous=NORMAL")
|
||||
return self._write_conn
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Key-value store
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def set(self, module: str, key: str, value: Any) -> None:
|
||||
"""Set a key-value pair (coalesced write)."""
|
||||
val_str = json.dumps(value) if not isinstance(value, str) else value
|
||||
sql = """INSERT INTO kv_store (module, key, value, updated)
|
||||
VALUES (?, ?, ?, ?)
|
||||
ON CONFLICT(module, key)
|
||||
DO UPDATE SET value=excluded.value, updated=excluded.updated"""
|
||||
self._enqueue(sql, (module, key, val_str, time.time()))
|
||||
|
||||
def get(self, module: str, key: str, default: Any = None) -> Any:
|
||||
"""Get a value by module + key (direct read, no queue delay)."""
|
||||
conn = self._get_read_conn()
|
||||
row = conn.execute(
|
||||
"SELECT value FROM kv_store WHERE module=? AND key=?",
|
||||
(module, key),
|
||||
).fetchone()
|
||||
if row is None:
|
||||
return default
|
||||
return row[0]
|
||||
|
||||
def get_all(self, module: str) -> dict:
|
||||
"""Get all key-value pairs for a module."""
|
||||
conn = self._get_read_conn()
|
||||
rows = conn.execute(
|
||||
"SELECT key, value FROM kv_store WHERE module=?", (module,)
|
||||
).fetchall()
|
||||
return {r[0]: r[1] for r in rows}
|
||||
|
||||
def delete(self, module: str, key: str) -> None:
|
||||
"""Delete a key-value pair."""
|
||||
self._enqueue("DELETE FROM kv_store WHERE module=? AND key=?", (module, key))
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Module status
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def set_module_status(self, module: str, status: str, pid: int = None,
|
||||
extra: dict = None) -> None:
|
||||
"""Update module runtime status."""
|
||||
extra_str = json.dumps(extra) if extra else None
|
||||
sql = """INSERT INTO module_status (module, status, pid, started, extra, updated)
|
||||
VALUES (?, ?, ?, ?, ?, ?)
|
||||
ON CONFLICT(module)
|
||||
DO UPDATE SET status=excluded.status, pid=excluded.pid,
|
||||
started=CASE WHEN excluded.status='running'
|
||||
THEN excluded.started
|
||||
ELSE module_status.started END,
|
||||
extra=excluded.extra, updated=excluded.updated"""
|
||||
started = time.time() if status == "running" else None
|
||||
self._enqueue(sql, (module, status, pid, started, extra_str, time.time()))
|
||||
|
||||
def get_module_status(self, module: str) -> dict:
|
||||
"""Get module runtime status."""
|
||||
conn = self._get_read_conn()
|
||||
row = conn.execute(
|
||||
"SELECT status, pid, started, extra, updated FROM module_status WHERE module=?",
|
||||
(module,),
|
||||
).fetchone()
|
||||
if row is None:
|
||||
return {"status": "unknown", "pid": None, "started": None, "extra": None}
|
||||
return {
|
||||
"status": row[0],
|
||||
"pid": row[1],
|
||||
"started": row[2],
|
||||
"extra": json.loads(row[3]) if row[3] else None,
|
||||
"updated": row[4],
|
||||
}
|
||||
|
||||
def get_all_module_status(self) -> dict:
|
||||
"""Get status for all modules."""
|
||||
conn = self._get_read_conn()
|
||||
rows = conn.execute(
|
||||
"SELECT module, status, pid, started, extra, updated FROM module_status"
|
||||
).fetchall()
|
||||
return {
|
||||
r[0]: {
|
||||
"status": r[1], "pid": r[2], "started": r[3],
|
||||
"extra": json.loads(r[4]) if r[4] else None,
|
||||
"updated": r[5],
|
||||
}
|
||||
for r in rows
|
||||
}
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Write queue
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _enqueue(self, sql: str, params: tuple) -> None:
|
||||
with self._queue_lock:
|
||||
self._write_queue.append((sql, params))
|
||||
|
||||
def _flush(self) -> None:
|
||||
"""Execute all queued writes in a single transaction."""
|
||||
with self._queue_lock:
|
||||
batch = list(self._write_queue)
|
||||
self._write_queue.clear()
|
||||
|
||||
if not batch:
|
||||
return
|
||||
|
||||
conn = self._get_write_conn()
|
||||
try:
|
||||
with conn:
|
||||
for sql, params in batch:
|
||||
conn.execute(sql, params)
|
||||
except Exception:
|
||||
logger.exception("StateManager flush failed (%d ops)", len(batch))
|
||||
|
||||
def _writer_loop(self) -> None:
|
||||
"""Background loop: flush write queue every FLUSH_INTERVAL seconds."""
|
||||
while self._running:
|
||||
time.sleep(self.FLUSH_INTERVAL)
|
||||
try:
|
||||
self._flush()
|
||||
except Exception:
|
||||
logger.exception("Writer loop error")
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Maintenance
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def checkpoint(self) -> None:
|
||||
"""Force a WAL checkpoint (periodic call to limit WAL size)."""
|
||||
conn = self._get_write_conn()
|
||||
conn.execute("PRAGMA wal_checkpoint(TRUNCATE)")
|
||||
|
||||
def flush_sync(self) -> None:
|
||||
"""Synchronously flush all pending writes (for shutdown paths)."""
|
||||
self._flush()
|
||||
@@ -0,0 +1,416 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Subprocess lifecycle manager for external tools.
|
||||
|
||||
Manages bettercap, tcpdump, Responder, mitmproxy, ntlmrelayx as
|
||||
supervised subprocesses with:
|
||||
- Start/stop/restart with configurable retry (max 3, exponential backoff)
|
||||
- stdout/stderr capture
|
||||
- PID tracking for kill_switch integration
|
||||
- Resource monitoring via /proc/PID/stat
|
||||
- Process disguise integration (prctl rename)
|
||||
- Health check callbacks
|
||||
- Graceful shutdown (SIGTERM -> 5s -> SIGKILL)
|
||||
- Crash callbacks (e.g., ARP restoration on bettercap crash)
|
||||
"""
|
||||
|
||||
import os
|
||||
import time
|
||||
import signal
|
||||
import logging
|
||||
import subprocess
|
||||
import threading
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Callable, Optional
|
||||
|
||||
from core.bus import EventBus
|
||||
|
||||
logger = logging.getLogger("bb.tool_manager")
|
||||
|
||||
|
||||
@dataclass
|
||||
class ManagedTool:
|
||||
"""Configuration and runtime state for a managed subprocess."""
|
||||
name: str
|
||||
binary: str
|
||||
args: list = field(default_factory=list)
|
||||
health_check: Optional[Callable] = None
|
||||
crash_callback: Optional[Callable] = None
|
||||
max_restarts: int = 3
|
||||
disguise_name: str = ""
|
||||
env: dict = field(default_factory=dict)
|
||||
cwd: Optional[str] = None
|
||||
|
||||
# Runtime state
|
||||
process: Optional[subprocess.Popen] = field(default=None, repr=False)
|
||||
pid: Optional[int] = None
|
||||
restart_count: int = 0
|
||||
started_at: Optional[float] = None
|
||||
_stdout_thread: Optional[threading.Thread] = field(default=None, repr=False)
|
||||
_stderr_thread: Optional[threading.Thread] = field(default=None, repr=False)
|
||||
stdout_lines: list = field(default_factory=list, repr=False)
|
||||
stderr_lines: list = field(default_factory=list, repr=False)
|
||||
_max_log_lines: int = 1000
|
||||
|
||||
|
||||
class ToolManager:
|
||||
"""Manages external tool subprocesses with supervision.
|
||||
|
||||
Usage:
|
||||
tm = ToolManager(bus)
|
||||
tm.register(name="bettercap", binary="/usr/local/bin/bettercap",
|
||||
args=["-api-rest-address", "127.0.0.1"],
|
||||
health_check=lambda: api.is_alive(),
|
||||
crash_callback=on_crash, max_restarts=3,
|
||||
disguise_name="networkd-dispatcher")
|
||||
tm.start("bettercap")
|
||||
tm.stop("bettercap")
|
||||
"""
|
||||
|
||||
SIGTERM_TIMEOUT = 5.0 # seconds to wait after SIGTERM before SIGKILL
|
||||
|
||||
def __init__(self, bus: EventBus):
|
||||
self.bus = bus
|
||||
self._tools: dict[str, ManagedTool] = {}
|
||||
self._lock = threading.Lock()
|
||||
self._monitor_thread: Optional[threading.Thread] = None
|
||||
self._running = False
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Lifecycle
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def start_monitoring(self) -> None:
|
||||
"""Start the background health/crash monitor."""
|
||||
if self._running:
|
||||
return
|
||||
self._running = True
|
||||
self._monitor_thread = threading.Thread(
|
||||
target=self._monitor_loop, daemon=True, name="bb-tool-monitor"
|
||||
)
|
||||
self._monitor_thread.start()
|
||||
|
||||
def stop_monitoring(self) -> None:
|
||||
"""Stop monitoring thread."""
|
||||
self._running = False
|
||||
if self._monitor_thread and self._monitor_thread.is_alive():
|
||||
self._monitor_thread.join(timeout=3.0)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Registration
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def register(self, name: str, binary: str, args: list = None,
|
||||
health_check: Callable = None, crash_callback: Callable = None,
|
||||
max_restarts: int = 3, disguise_name: str = "",
|
||||
env: dict = None, cwd: str = None) -> None:
|
||||
"""Register a tool for management."""
|
||||
with self._lock:
|
||||
self._tools[name] = ManagedTool(
|
||||
name=name,
|
||||
binary=binary,
|
||||
args=args or [],
|
||||
health_check=health_check,
|
||||
crash_callback=crash_callback,
|
||||
max_restarts=max_restarts,
|
||||
disguise_name=disguise_name,
|
||||
env=env or {},
|
||||
cwd=cwd,
|
||||
)
|
||||
logger.info("Registered tool: %s (%s)", name, binary)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Start / Stop / Restart
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def start(self, name: str) -> bool:
|
||||
"""Start a registered tool subprocess."""
|
||||
with self._lock:
|
||||
tool = self._tools.get(name)
|
||||
if tool is None:
|
||||
logger.error("Tool %s not registered", name)
|
||||
return False
|
||||
|
||||
if tool.process and tool.process.poll() is None:
|
||||
logger.warning("Tool %s already running (pid=%d)", name, tool.pid)
|
||||
return True
|
||||
|
||||
return self._start_tool(tool)
|
||||
|
||||
def _start_tool(self, tool: ManagedTool) -> bool:
|
||||
"""Internal: launch the subprocess. Caller must hold _lock."""
|
||||
cmd = [tool.binary] + tool.args
|
||||
env = dict(os.environ)
|
||||
env.update(tool.env)
|
||||
|
||||
try:
|
||||
proc = subprocess.Popen(
|
||||
cmd,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
env=env,
|
||||
cwd=tool.cwd,
|
||||
preexec_fn=self._make_preexec(tool.disguise_name) if tool.disguise_name else None,
|
||||
)
|
||||
except FileNotFoundError:
|
||||
logger.error("Binary not found: %s", tool.binary)
|
||||
return False
|
||||
except Exception as e:
|
||||
logger.error("Failed to start %s: %s", tool.name, e)
|
||||
return False
|
||||
|
||||
tool.process = proc
|
||||
tool.pid = proc.pid
|
||||
tool.started_at = time.time()
|
||||
tool.stdout_lines.clear()
|
||||
tool.stderr_lines.clear()
|
||||
|
||||
# Stdout/stderr reader threads
|
||||
tool._stdout_thread = threading.Thread(
|
||||
target=self._pipe_reader, args=(proc.stdout, tool.stdout_lines, tool._max_log_lines),
|
||||
daemon=True, name=f"bb-{tool.name}-stdout",
|
||||
)
|
||||
tool._stderr_thread = threading.Thread(
|
||||
target=self._pipe_reader, args=(proc.stderr, tool.stderr_lines, tool._max_log_lines),
|
||||
daemon=True, name=f"bb-{tool.name}-stderr",
|
||||
)
|
||||
tool._stdout_thread.start()
|
||||
tool._stderr_thread.start()
|
||||
|
||||
self.bus.emit("TOOL_RESTARTED" if tool.restart_count > 0 else "MODULE_STARTED",
|
||||
{"tool": tool.name, "pid": proc.pid, "restart_count": tool.restart_count},
|
||||
source_module="tool_manager")
|
||||
logger.info("Started tool %s (pid=%d)", tool.name, proc.pid)
|
||||
return True
|
||||
|
||||
def stop(self, name: str) -> bool:
|
||||
"""Stop a tool. SIGTERM -> 5s timeout -> SIGKILL."""
|
||||
with self._lock:
|
||||
tool = self._tools.get(name)
|
||||
if tool is None:
|
||||
return True
|
||||
return self._stop_tool(tool)
|
||||
|
||||
def _stop_tool(self, tool: ManagedTool) -> bool:
|
||||
"""Internal: stop subprocess. Caller must hold _lock."""
|
||||
proc = tool.process
|
||||
if proc is None or proc.poll() is not None:
|
||||
tool.process = None
|
||||
tool.pid = None
|
||||
return True
|
||||
|
||||
# Graceful: SIGTERM
|
||||
try:
|
||||
proc.terminate()
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
try:
|
||||
proc.wait(timeout=self.SIGTERM_TIMEOUT)
|
||||
except subprocess.TimeoutExpired:
|
||||
# Forceful: SIGKILL
|
||||
logger.warning("Tool %s did not stop after SIGTERM, sending SIGKILL", tool.name)
|
||||
try:
|
||||
proc.kill()
|
||||
proc.wait(timeout=2.0)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
tool.process = None
|
||||
tool.pid = None
|
||||
tool.started_at = None
|
||||
logger.info("Stopped tool %s", tool.name)
|
||||
return True
|
||||
|
||||
def restart(self, name: str) -> bool:
|
||||
"""Stop then start a tool, resetting restart count."""
|
||||
with self._lock:
|
||||
tool = self._tools.get(name)
|
||||
if tool is None:
|
||||
return False
|
||||
self._stop_tool(tool)
|
||||
tool.restart_count = 0
|
||||
return self._start_tool(tool)
|
||||
|
||||
def stop_all(self) -> None:
|
||||
"""Stop all managed tools."""
|
||||
with self._lock:
|
||||
for tool in self._tools.values():
|
||||
self._stop_tool(tool)
|
||||
self.stop_monitoring()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Status / resources
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def get_status(self, name: str) -> dict:
|
||||
"""Get tool status including resource usage from /proc."""
|
||||
with self._lock:
|
||||
tool = self._tools.get(name)
|
||||
if tool is None:
|
||||
return {"error": "not registered"}
|
||||
|
||||
alive = tool.process is not None and tool.process.poll() is None
|
||||
result = {
|
||||
"name": tool.name,
|
||||
"running": alive,
|
||||
"pid": tool.pid,
|
||||
"uptime": time.time() - tool.started_at if tool.started_at and alive else 0,
|
||||
"restart_count": tool.restart_count,
|
||||
}
|
||||
|
||||
if alive and tool.pid:
|
||||
resources = self._read_proc_stats(tool.pid)
|
||||
result.update(resources)
|
||||
|
||||
return result
|
||||
|
||||
def get_all_status(self) -> dict:
|
||||
"""Get status for all registered tools."""
|
||||
with self._lock:
|
||||
names = list(self._tools.keys())
|
||||
return {n: self.get_status(n) for n in names}
|
||||
|
||||
def get_all_pids(self) -> list[int]:
|
||||
"""Return PIDs of all running tools (for kill_switch)."""
|
||||
pids = []
|
||||
with self._lock:
|
||||
for tool in self._tools.values():
|
||||
if tool.pid and tool.process and tool.process.poll() is None:
|
||||
pids.append(tool.pid)
|
||||
return pids
|
||||
|
||||
@staticmethod
|
||||
def _read_proc_stats(pid: int) -> dict:
|
||||
"""Read RAM/CPU from /proc/PID/stat and /proc/PID/statm."""
|
||||
result = {"ram_mb": 0.0, "cpu_pct": 0.0}
|
||||
try:
|
||||
with open(f"/proc/{pid}/statm", "r") as f:
|
||||
parts = f.read().split()
|
||||
# statm: size resident shared text lib data dt (in pages)
|
||||
page_size = os.sysconf("SC_PAGE_SIZE")
|
||||
result["ram_mb"] = round(int(parts[1]) * page_size / (1024 * 1024), 1)
|
||||
except (FileNotFoundError, PermissionError, IndexError, ValueError):
|
||||
pass
|
||||
|
||||
try:
|
||||
with open(f"/proc/{pid}/stat", "r") as f:
|
||||
parts = f.read().split()
|
||||
# Fields 14,15 = utime, stime in clock ticks
|
||||
utime = int(parts[13])
|
||||
stime = int(parts[14])
|
||||
total_ticks = utime + stime
|
||||
hz = os.sysconf("SC_CLK_TCK")
|
||||
uptime_f = open("/proc/uptime", "r")
|
||||
system_uptime = float(uptime_f.read().split()[0])
|
||||
uptime_f.close()
|
||||
start_ticks = int(parts[21])
|
||||
proc_uptime = system_uptime - (start_ticks / hz)
|
||||
if proc_uptime > 0:
|
||||
result["cpu_pct"] = round((total_ticks / hz) / proc_uptime * 100, 1)
|
||||
except (FileNotFoundError, PermissionError, IndexError, ValueError, ZeroDivisionError):
|
||||
pass
|
||||
|
||||
return result
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Monitor loop
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _monitor_loop(self) -> None:
|
||||
"""Background loop: detect crashes, run health checks, auto-restart."""
|
||||
while self._running:
|
||||
time.sleep(5.0)
|
||||
with self._lock:
|
||||
tools = list(self._tools.values())
|
||||
|
||||
for tool in tools:
|
||||
if tool.process is None:
|
||||
continue
|
||||
|
||||
rc = tool.process.poll()
|
||||
if rc is not None:
|
||||
# Process exited
|
||||
logger.warning("Tool %s exited (rc=%s)", tool.name, rc)
|
||||
self.bus.emit("TOOL_CRASHED",
|
||||
{"tool": tool.name, "exit_code": rc,
|
||||
"restart_count": tool.restart_count},
|
||||
source_module="tool_manager")
|
||||
|
||||
# Crash callback (e.g., corrective ARP)
|
||||
if tool.crash_callback:
|
||||
try:
|
||||
tool.crash_callback()
|
||||
except Exception:
|
||||
logger.exception("Crash callback failed for %s", tool.name)
|
||||
|
||||
# Auto-restart with exponential backoff
|
||||
if tool.restart_count < tool.max_restarts:
|
||||
backoff = min(2 ** tool.restart_count, 30)
|
||||
logger.info("Restarting %s in %ds (attempt %d/%d)",
|
||||
tool.name, backoff, tool.restart_count + 1,
|
||||
tool.max_restarts)
|
||||
time.sleep(backoff)
|
||||
tool.restart_count += 1
|
||||
with self._lock:
|
||||
self._start_tool(tool)
|
||||
else:
|
||||
logger.error("Tool %s exceeded max restarts (%d), giving up",
|
||||
tool.name, tool.max_restarts)
|
||||
tool.process = None
|
||||
tool.pid = None
|
||||
|
||||
elif tool.health_check:
|
||||
# Process alive — run health check
|
||||
try:
|
||||
if not tool.health_check():
|
||||
logger.warning("Health check failed for %s", tool.name)
|
||||
except Exception:
|
||||
logger.exception("Health check error for %s", tool.name)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def _pipe_reader(pipe, line_buffer: list, max_lines: int) -> None:
|
||||
"""Read lines from a subprocess pipe into a bounded buffer."""
|
||||
try:
|
||||
for line in iter(pipe.readline, b""):
|
||||
try:
|
||||
decoded = line.decode("utf-8", errors="replace").rstrip()
|
||||
except Exception:
|
||||
decoded = repr(line)
|
||||
line_buffer.append(decoded)
|
||||
if len(line_buffer) > max_lines:
|
||||
line_buffer.pop(0)
|
||||
except Exception:
|
||||
pass
|
||||
finally:
|
||||
try:
|
||||
pipe.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def _make_preexec(disguise_name: str):
|
||||
"""Return a preexec_fn that renames the process via prctl."""
|
||||
def _preexec():
|
||||
try:
|
||||
import ctypes
|
||||
libc = ctypes.CDLL("libc.so.6", use_errno=True)
|
||||
PR_SET_NAME = 15
|
||||
name_bytes = disguise_name[:15].encode("utf-8")
|
||||
libc.prctl(PR_SET_NAME, name_bytes, 0, 0, 0)
|
||||
except Exception:
|
||||
pass
|
||||
return _preexec
|
||||
|
||||
def get_output(self, name: str, stream: str = "stderr",
|
||||
lines: int = 50) -> list[str]:
|
||||
"""Get recent stdout/stderr lines from a tool."""
|
||||
with self._lock:
|
||||
tool = self._tools.get(name)
|
||||
if tool is None:
|
||||
return []
|
||||
buf = tool.stderr_lines if stream == "stderr" else tool.stdout_lines
|
||||
return list(buf[-lines:])
|
||||
Reference in New Issue
Block a user