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,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()
|
||||
}
|
||||
Reference in New Issue
Block a user