firegex/backend/application/snapshots.py

"""Configuration snapshots and compensation. No authentication state is copied."""
import copy
from application import services as svc, firewall as fw
from modules.services.models import STATUS
from modules.firewall.models import FirewallSettings

METRICS = {"blocked", "filtering_since", "over_limit_hits", "over_limit_first", "over_limit_last"}


def configuration(value):
    if isinstance(value, dict):
        return {k: configuration(v) for k, v in value.items() if k not in METRICS}
    if isinstance(value, list):
        return [configuration(v) for v in value]
    if isinstance(value, bool):
        return int(value)
    return value


def public(value):
    if isinstance(value, dict):
        return {k: ("[redacted]" if v else None) if k in {"tls_key", "tls_cert", "password", "access_token"}
                else public(v) for k, v in value.items()}
    if isinstance(value, list):
        return [public(v) for v in value]
    return value


def service_snapshot(service_id):
    svc.firewall.flush_blocks()
    rows = svc.db.query("SELECT * FROM services WHERE service_id=?", service_id)
    result = {"scope": "service", "service_id": service_id, "service": rows[0] if rows else None,
              "addresses": [], "filters": [], "regexes": [], "functions": [], "code": {}}
    if not rows:
        return result
    result["addresses"] = svc.db.query("SELECT * FROM service_addresses WHERE service_id=? ORDER BY address_id", service_id)
    result["filters"] = svc.db.query("SELECT * FROM filters WHERE service_id=? ORDER BY position", service_id)
    for row in result["filters"]:
        fid = row["filter_id"]
        result["regexes"] += svc.db.query("SELECT * FROM regexes WHERE filter_id=? ORDER BY regex_id", fid)
        result["functions"] += svc.db.query("SELECT * FROM pyfilters WHERE filter_id=? ORDER BY position", fid)
        if row["kind"] == "pyfilter":
            result["code"][fid] = svc.read_code(fid)
    return result


def firewall_snapshot():
    policy = fw.firewall.policy
    return {"scope": "firewall", "rules": svc_copy_rules(), "policy": getattr(policy, "value", policy),
            "settings": fw.firewall.settings.model_dump(mode="json"), "enabled": fw.firewall.enabled}


def svc_copy_rules():
    return [{k: v for k, v in r.items() if k != "rule_id"}
            for r in fw.db.query("SELECT * FROM rules ORDER BY rule_id")]


def snapshot(scope, service_id=None):
    if scope == "service":
        return service_snapshot(service_id)
    if scope == "firewall":
        return firewall_snapshot()
    return {"scope": "instance", "firewall": firewall_snapshot(), "services": [
        service_snapshot(r["service_id"]) for r in svc.db.query("SELECT service_id FROM services ORDER BY service_id")]}


def insert_row(table, row):
    # Tables and columns only come from our own snapshots/prepared candidates.
    columns = list(row)
    return (f'INSERT INTO "{table}" ({", ".join(chr(34)+k+chr(34) for k in columns)}) '
            f'VALUES ({", ".join("?" for _ in columns)})', *[row[k] for k in columns])


def write_service_data(value):
    """Replace this service's rows, preserving live counters for identities that survive."""
    value = copy.deepcopy(value)
    sid = value["service_id"]
    current = service_snapshot(sid)
    tables = {"addresses": "service_addresses", "filters": "filters", "regexes": "regexes", "functions": "pyfilters"}
    for key in ("filters", "regexes", "functions"):
        identity = lambda r: (r.get("filter_id"), r.get("regex_id"), r.get("name") if key == "functions" else None)
        old = {identity(r): r for r in current[key]}
        for row in value[key]:
            if identity(row) in old:
                for metric in METRICS & row.keys():
                    row[metric] = old[identity(row)].get(metric, row[metric])
    if current["service"] and value["service"]:
        for metric in METRICS & value["service"].keys():
            value["service"][metric] = current["service"].get(metric)
    queries = [
        ("DELETE FROM regexes WHERE filter_id IN (SELECT filter_id FROM filters WHERE service_id=?)", sid),
        ("DELETE FROM pyfilters WHERE filter_id IN (SELECT filter_id FROM filters WHERE service_id=?)", sid),
        ("DELETE FROM filters WHERE service_id=?", sid),
        ("DELETE FROM service_addresses WHERE service_id=?", sid),
        ("DELETE FROM services WHERE service_id=?", sid),
    ]
    if value["service"]:
        queries.append(insert_row("services", value["service"]))
        for key, table in tables.items():
            queries.extend(insert_row(table, row) for row in value[key])
    svc.db.queries(queries)
    for fid in current["code"].keys() - value["code"].keys():
        svc.clear_code(fid)
    for fid, source in value["code"].items():
        svc.write_code(fid, source)


async def restore(value):
    if value["scope"] == "instance":
        wanted = {s["service_id"] for s in value["services"]}
        for row in svc.db.query("SELECT service_id FROM services"):
            if row["service_id"] not in wanted:
                await svc.delete_service(row["service_id"])
        for service in value["services"]:
            await restore(service)
        await restore(value["firewall"])
    elif value["scope"] == "firewall":
        fw._write_rules(value["rules"])
        fw.firewall.policy = value["policy"]
        fw.firewall.settings = FirewallSettings(**value["settings"])
        fw.firewall.enabled = value["enabled"]
        await fw.firewall.reload()
        await fw.refresh_frontend()
    else:
        sid = value["service_id"]
        await svc.firewall.remove(sid, persist=False)
        write_service_data(value)
        if value["service"]:
            await svc.firewall.reload()
            manager = svc.firewall.get(sid)
            if value["service"]["status"] == STATUS.ACTIVE and not manager.active:
                raise RuntimeError("Restored configuration could not start its datapath")
        await svc.refresh_frontend()