Cobra
ffd384f64b
Replace all sensor.* logger namespaces with __name__ (generic module identifiers instead of discoverable 'sensor.*' prefixes). Change hardcoded 'bb' API user to 'admin' in config and code defaults. Change hardcoded relay_user 'bb' to 'operator' — prevents network profiling from exposing tool identity via SSH config. Fixes #457, #458, #459
320 lines
12 KiB
Python
320 lines
12 KiB
Python
#!/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(__name__)
|
|
|
|
_DEFAULT_DB = os.path.join(
|
|
os.path.expanduser("~"), ".implant", "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="sensor-state-writer"
|
|
)
|
|
self._writer_thread.start()
|
|
logger.info("StateManager writer started (db=%s)", self._db_path)
|
|
|
|
def stop(self, secure_wipe: bool = False) -> None:
|
|
"""Flush remaining writes and stop the writer thread.
|
|
|
|
Args:
|
|
secure_wipe: If True, securely overwrite and delete the state database.
|
|
"""
|
|
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
|
|
|
|
# Secure wipe if requested
|
|
if secure_wipe:
|
|
from utils.crypto import secure_wipe_file
|
|
# SQLite creates multiple files in WAL mode: db, db-wal, db-shm
|
|
for suffix in ["", "-wal", "-shm"]:
|
|
filepath = self._db_path + suffix
|
|
if os.path.isfile(filepath):
|
|
try:
|
|
if secure_wipe_file(filepath, passes=3):
|
|
logger.info("Securely wiped %s", filepath)
|
|
else:
|
|
logger.warning("Secure wipe failed for %s (used regular delete)", filepath)
|
|
except Exception as e:
|
|
logger.error("Failed to wipe %s: %s", filepath, e)
|
|
|
|
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, timeout=30)
|
|
self._write_conn.execute("PRAGMA journal_mode=WAL")
|
|
self._write_conn.execute("PRAGMA synchronous=NORMAL")
|
|
self._write_conn.execute("PRAGMA busy_timeout=30000")
|
|
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
|
|
}
|
|
|
|
def reset_module_status(self) -> None:
|
|
"""Mark all running modules as stopped.
|
|
|
|
Called at engine startup to clear stale PIDs from a previous run.
|
|
Without this, the watchdog picks up old PIDs and reports dead modules
|
|
that were never started in the current session.
|
|
"""
|
|
conn = self._get_write_conn()
|
|
try:
|
|
conn.execute("BEGIN IMMEDIATE")
|
|
conn.execute(
|
|
"UPDATE module_status SET status='stopped', pid=NULL, updated=? "
|
|
"WHERE status='running'",
|
|
(time.time(),),
|
|
)
|
|
conn.commit()
|
|
except Exception:
|
|
try:
|
|
conn.rollback()
|
|
except Exception:
|
|
pass
|
|
|
|
# ------------------------------------------------------------------
|
|
# 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()
|
|
for attempt in range(6):
|
|
try:
|
|
conn.execute("BEGIN IMMEDIATE")
|
|
for sql, params in batch:
|
|
conn.execute(sql, params)
|
|
conn.execute("COMMIT")
|
|
return
|
|
except sqlite3.OperationalError as e:
|
|
try:
|
|
conn.execute("ROLLBACK")
|
|
except Exception:
|
|
pass
|
|
if attempt < 5:
|
|
time.sleep(0.5 * (attempt + 1))
|
|
else:
|
|
logger.error("StateManager flush failed after 6 attempts (%d ops): %s", len(batch), e)
|
|
except Exception:
|
|
try:
|
|
conn.execute("ROLLBACK")
|
|
except Exception:
|
|
pass
|
|
logger.exception("StateManager flush failed (%d ops)", len(batch))
|
|
return
|
|
|
|
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()
|