"""Pre-computed state store for vacuum-walld. Collects system state at startup and on demand. Handlers read from the state instead of invoking subprocesses on every request. The collectors themselves live in ``daemon/collectors/`` (they make read-only ``sudo`` queries and so do not belong in ``lib/``). Importing that package registers them here as a side effect; registration must happen before the first ``populate()``/``poll()`` call. """ import logging from copy import deepcopy from datetime import UTC, datetime from pathlib import Path from typing import Any, ClassVar logger = logging.getLogger(__name__) PROJECT_DIR = Path(__file__).resolve().parent.parent _DEFAULT_POLL_INTERVALS: dict[str, int] = { "firewall": 30, "wireguard": 10, "dnsmasq": 10, "networkd": 10, "system": 1, # nginx/acme state derives from config files and rendered artifacts on # disk; poll so drift (manual edits, out-of-band applies) is re-collected. "nginx": 60, "acme": 300, } # --------------------------------------------------------------------------- # State store # --------------------------------------------------------------------------- class State: """In-memory state store keyed by subsystem name. Each subsystem's value is a dict collected from the corresponding ``collect_*`` function. A value of ``None`` means the subsystem has not been populated yet or the last collection failed. Attributes: SUBSYSTEMS: Ordered list of subsystem names. _data: Dict mapping subsystem names to their state data. """ SUBSYSTEMS: ClassVar[list[str]] = [ "firewall", "dnsmasq", "nginx", "acme", "wireguard", "networkd", "system", ] def __init__(self) -> None: """Initialize the state store with empty subsystem slots.""" self._data: dict[str, dict[str, Any] | None] = { name: None for name in self.SUBSYSTEMS } self._versions: dict[str, int] = {name: 0 for name in self.SUBSYSTEMS} self._last_broadcast: dict[str, int] | None = None def bump(self, subsystem: str) -> None: """Increment the version counter for *subsystem*. Args: subsystem: Subsystem name. """ if subsystem in self._versions: self._versions[subsystem] += 1 def get_versions(self) -> dict[str, int]: """Return a shallow copy of all subsystem versions. Returns: Dict mapping subsystem names to their current version integers. """ return dict(self._versions) def get_updated_versions(self) -> dict[str, int]: """Return versions that changed since the last broadcast. After calling, ``_last_broadcast`` is updated to match current versions. Returns: Dict of subsystems whose versions changed, or empty dict. """ if self._last_broadcast is None: self._last_broadcast = dict(self._versions) return {} updated: dict[str, int] = {} for name, v in self._versions.items(): if v != self._last_broadcast.get(name, 0): updated[name] = v self._last_broadcast[name] = v return updated def get(self, subsystem: str) -> dict[str, Any] | None: """Get state data for *subsystem*. Args: subsystem: Subsystem name. Returns: State dict, or ``None`` if not populated. """ return self._data.get(subsystem) def get_snapshot(self) -> dict[str, dict[str, Any] | None]: """Return all subsystem state dicts. Used for the initial WS snapshot on connect. Returns: Dict mapping every subsystem name to its state data (``None`` when not populated or the last collection failed). """ return {name: self._data.get(name) for name in self.SUBSYSTEMS} def set(self, subsystem: str, data: dict[str, Any] | None) -> None: """Set state data for *subsystem*. Args: subsystem: Subsystem name. data: State data, or ``None`` to clear. """ self._data[subsystem] = data def populate(self, subsystems: list[str] | None = None) -> None: """Collect state for *subsystems* (all if ``None``). Args: subsystems: List of subsystem names to collect. Collects all subsystems when ``None``. """ targets = subsystems or self.SUBSYSTEMS for name in targets: collector = _COLLECTORS.get(name) if collector is None: continue try: self._data[name] = collector() except Exception: logger.warning( "State collection failed for %s, clearing state", name, exc_info=True, ) self._data[name] = None def is_populated(self) -> bool: """Check whether all subsystem states have been populated. Returns: ``True`` if every subsystem has non-``None`` state data. """ return all(v is not None for v in self._data.values()) def poll(self, subsystem: str) -> tuple[bool, bool]: """Run the collector for *subsystem* and compare against current state. The polled equivalent of ``populate()`` — same try/except safety, but with two-layer diff before storing. Args: subsystem: Subsystem name to poll. Returns: ``(structural_changed, volatile_changed)``. ``(False, False)`` on collector failure (no broadcast on failure to avoid noisy ticks). """ collector = _COLLECTORS.get(subsystem) if collector is None: return (False, False) vol = _VOLATILE.get(subsystem, frozenset()) try: new_data = collector() except Exception: logger.warning( "Poll collection failed for %s", subsystem, exc_info=True, ) return (False, False) old_data = self._data.get(subsystem) structural, volatile = _diff_layers(old_data, new_data, vol) self._data[subsystem] = new_data return (structural, volatile) def poll_all( self, intervals: dict[str, int] | None = None, ) -> dict[str, tuple[bool, bool]]: """Poll all subsystems that have a polling interval configured. Args: intervals: Subsystems to poll keyed by name. Defaults to ``_DEFAULT_POLL_INTERVALS``. Returns: Dict of ``{subsystem: (structural, volatile)}`` for each polled subsystem. """ targets = intervals or _DEFAULT_POLL_INTERVALS results: dict[str, tuple[bool, bool]] = {} for name in targets: results[name] = self.poll(name) return results # Singleton state = State() # --------------------------------------------------------------------------- # Collector registry # --------------------------------------------------------------------------- _COLLECTORS: dict[str, Any] = {} _VOLATILE: dict[str, frozenset[str]] = {} def register_collector(subsystem: str, fn: Any) -> Any: """Register *fn* as the state collector for *subsystem*. Args: subsystem: Subsystem name to register for. fn: Collector function to register. Returns: The *fn* function (for decorator usage). """ _COLLECTORS[subsystem] = fn return fn def register_volatile(subsystem: str, keys: frozenset[str]) -> None: """Register volatile field paths for *subsystem*. Args: subsystem: Subsystem name. keys: Frozenset of dot-separated volatile field paths (e.g. ``status.peers[].transfer_received``). """ _VOLATILE[subsystem] = keys # --------------------------------------------------------------------------- # Two-layer diff # --------------------------------------------------------------------------- def _strip_volatile( data: dict[str, Any], volatile: frozenset[str], pop_keys: frozenset[str] | None = None, ) -> dict[str, Any]: """Return a copy of *data* with volatile fields zeroed out. For list-of-dicts fields (``key[].subkey``), strips the volatile sub-keys from each dict in the list. For scalar/dict fields, sets them to ``None``. Args: data: State dict to process. volatile: Frozenset of dot-separated volatile field paths. pop_keys: Optional keys to remove from the root dict before stripping. Returns: A new dict with volatile fields replaced by ``None``. """ stripped = deepcopy(data) if pop_keys: for k in pop_keys: stripped.pop(k, None) for vpath in volatile: _strip_volatile_path(stripped, vpath) return stripped def _strip_volatile_item(item: dict[str, Any], keys: list[str]) -> None: """Recursively strip volatile keys from *item*, handling nested ``[]`` markers.""" for i, k in enumerate(keys): if "[]" in k: base_key = k.replace("[]", "") rest = keys[i + 1 :] target = item.get(base_key, []) if isinstance(target, list): for t in target: if isinstance(t, dict): _strip_volatile_item(t, rest) elif isinstance(target, dict): for v in target.values(): if isinstance(v, dict): _strip_volatile_item(v, rest) return elif i == len(keys) - 1: item[k] = None return else: if isinstance(item, dict) and k in item: item = item[k] else: return def _strip_volatile_path(stripped: dict[str, Any], vpath: str) -> None: """Strip a single volatile path from *stripped*, supporting nested ``[]`` markers.""" list_marker = vpath.index("[]") if "[]" in vpath else -1 if list_marker == -1: segments = vpath.split(".") parent = stripped for i, seg in enumerate(segments): if i == len(segments) - 1: if isinstance(parent, dict) and seg in parent: parent[seg] = None else: if isinstance(parent, dict) and seg in parent: parent = parent[seg] else: return return # Split into prefix and item keys. prefix = vpath[:list_marker].split(".") item_keys = ( vpath[list_marker + 3 :].split(".") if list_marker + 3 < len(vpath) else [] ) parent = stripped for seg in prefix: if isinstance(parent, dict) and seg in parent: parent = parent[seg] else: return if isinstance(parent, list): items = parent elif isinstance(parent, dict): items = list(parent.values()) else: return for item in items: if isinstance(item, dict) and item_keys: _strip_volatile_item(item, item_keys) def _diff_layers( old: dict[str, Any] | None, new: dict[str, Any], volatile: frozenset[str], ) -> tuple[bool, bool]: """Compare *old* and *new* state using two-layer diff. The two-layer strategy distinguishes between: 1. Structural changes (config, topology) → triggers full client re-fetch 2. Volatile changes (byte counters, timestamps) → triggers lightweight tick If structural data changed, volatile is suppressed (False) because the structural change already triggers a full re-fetch, making the volatile signal redundant. Args: old: Previous state data, or ``None`` if not yet populated. new: New state data from collector. volatile: Frozenset of volatile field paths. Returns: ``(structural_changed, volatile_changed)``. """ if old is None: return (True, True) # Structural diff: compare with volatile fields zeroed out, plus timestamp # removed. If these differ, the configuration or topology has changed. pop_keys = frozenset(("timestamp",)) old_struct = _strip_volatile(old, volatile, pop_keys) new_struct = _strip_volatile(new, volatile, pop_keys) structural = old_struct != new_struct # Volatile diff: compare without timestamp # Volatile diff: only relevant if structural is unchanged. Compare full # data (minus timestamp). If this differs, only volatile fields changed # (e.g. WireGuard transfer counters), and a lightweight tick suffices. volatile_changed = False if not structural: old_no_ts = {k: v for k, v in old.items() if k != "timestamp"} new_no_ts = {k: v for k, v in new.items() if k != "timestamp"} volatile_changed = old_no_ts != new_no_ts return (structural, volatile_changed) def _now_iso() -> str: """Return the current UTC time as an ISO 8601 string.""" return datetime.now(UTC).isoformat() __all__ = [ "PROJECT_DIR", "_COLLECTORS", "_DEFAULT_POLL_INTERVALS", "_VOLATILE", "State", "_diff_layers", "_now_iso", "_strip_volatile", "register_collector", "register_volatile", "state", ]