#!/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("sensor.state") _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()