home-network-k3s-impl/platform/vuln-scanner/src/state.py

"""Persisted scanner state and the time helpers every module shares.

state.json holds what must survive a restart: the last inventory, one record per unit of scan
(a digest, or a by-tag target id), the database info, the last run, remediation and heartbeat.
Every write is temp-file + fsync + rename, so a crash leaves the previous file intact.
"""
import json
import logging
import os
import re
import tempfile
from datetime import datetime, timedelta, timezone

log = logging.getLogger("vuln-scanner.state")

STATE_VERSION = 1
RETENTION = timedelta(days=7)  # how long a scan record outlives its target in the inventory

_RFC3339 = re.compile(r"^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}(\.\d+)?(Z|[+-]\d{2}:\d{2})$")


def utcnow():
    return datetime.now(timezone.utc)


def fmt_ts(value):
    return value.astimezone(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")


def parse_ts(value):
    """A timezone-aware datetime for an RFC 3339 string, or None for anything else."""
    if not isinstance(value, str) or not _RFC3339.match(value):
        return None
    try:
        return datetime.fromisoformat(value)
    except ValueError:
        return None


def atomic_write(path, data):
    directory = os.path.dirname(path) or "."
    fd, tmp = tempfile.mkstemp(dir=directory, prefix=".tmp-")
    try:
        with os.fdopen(fd, "wb") as handle:
            handle.write(data)
            handle.flush()
            os.fsync(handle.fileno())
        os.replace(tmp, path)
    except BaseException:
        try:
            os.unlink(tmp)
        except OSError:
            pass
        raise
    dir_fd = os.open(directory, os.O_RDONLY)
    try:
        os.fsync(dir_fd)
    finally:
        os.close(dir_fd)


def new_state():
    return {
        "version": STATE_VERSION,
        "inventory": [],
        "inventoryAt": None,
        "inventoryError": None,
        "scans": {},
        "database": None,
        "run": None,
        "lastCompleteRunAt": None,
        "remediation": [],
        "triageHeartbeatAt": None,
        "loadError": None,
        # The scan key Trivy is running against right now, persisted BEFORE the scan starts: an OOM
        # kill takes the whole container, so nothing can record the failure afterwards.
        "inFlight": None,
    }


def new_record(now):
    return {
        "firstDiscoveredAt": fmt_ts(now),
        "lastInventoriedAt": fmt_ts(now),
        "lastAttemptAt": None,
        "lastSuccessAt": None,
        "outcome": None,
        "reason": None,
        "failedSince": None,
        "killedAt": None,
        "findings": [],
        "firstSeen": {},
    }


class Store:
    """state.json on the data volume."""

    def __init__(self, data_dir):
        self.data_dir = data_dir
        self.path = os.path.join(data_dir, "state.json")

    def load(self):
        try:
            with open(self.path, "rb") as handle:
                loaded = json.loads(handle.read())
            if not isinstance(loaded, dict) or loaded.get("version") != STATE_VERSION:
                raise ValueError("unrecognised state version")
        except FileNotFoundError:
            return new_state()
        except (OSError, ValueError) as exc:
            log.critical("state.json is unreadable (%s); starting empty and marking the run degraded", exc)
            try:
                os.replace(self.path, self.path + ".corrupt")
            except OSError:
                pass
            fresh = new_state()
            fresh["loadError"] = "state.json was unreadable at startup"
            return fresh
        merged = new_state()
        merged.update(loaded)
        return merged

    def save(self, state):
        atomic_write(self.path, json.dumps(state, sort_keys=True).encode("utf-8"))


def merge_inventory(state, targets, now):
    """Install a fresh inventory; create a record per new unit; forget long-gone units."""
    state["inventory"] = targets
    state["inventoryAt"] = fmt_ts(now)
    state["inventoryError"] = None
    stamp = fmt_ts(now)
    for target in targets:
        key = scan_key(target)
        if key is None:
            continue
        record = state["scans"].get(key)
        if record is None:
            record = state["scans"][key] = new_record(now)
        record["lastInventoriedAt"] = stamp
    cutoff = now - RETENTION
    for key in [k for k, r in state["scans"].items() if (parse_ts(r["lastInventoriedAt"]) or now) < cutoff]:
        del state["scans"][key]


def scan_key(target):
    """The unit of scan: the digest, else the by-tag target id. VMs and nodes have none."""
    if target["type"] != "image":
        return None
    return target["digest"] or target["id"]