awg-dash/backend/awg_api/store/state.py

"""Снимки состояния и восстановление после отката (rollback, issue #20).

Снимок — это JSON-копия состояния (режим, расписание GeoIP, правила,
исключения) на момент перед применением. Снимок сохраняется ДО мутации,
поэтому даже при неудачном apply (внешний iptables упал) есть источник для
восстановления.
"""

from __future__ import annotations

import json
import sqlite3
from dataclasses import dataclass
from typing import Any

from . import audit
from .clock import utcnow
from .errors import StorageError
from .models import Entity


@dataclass(frozen=True, slots=True)
class StateSnapshot:
    """Сохранённый снимок состояния."""

    id: int
    entity: str
    entity_id: str
    created_at: str
    payload: dict[str, Any]
    reason: str


def capture_state(conn: sqlite3.Connection, *, include_exceptions: bool = True) -> dict[str, Any]:
    """Снять текущее состояние в виде JSON-совместимого словаря."""
    mode_row = conn.execute("SELECT mode FROM routing_mode WHERE id = 1").fetchone()
    geoip_row = conn.execute(
        "SELECT enabled, cron, provider FROM geoip_schedule WHERE id = 1"
    ).fetchone()
    rules = [
        {
            "name": str(row["name"]),
            "kind": str(row["kind"]),
            "match": str(row["match"]),
            "target": str(row["target"]),
            "priority": int(row["priority"]),
            "enabled": bool(row["enabled"]),
        }
        for row in conn.execute(
            "SELECT name, kind, match, target, priority, enabled FROM routing_rules"
            " ORDER BY priority, id"
        ).fetchall()
    ]
    payload: dict[str, Any] = {
        "mode": "off" if mode_row is None else str(mode_row["mode"]),
        "geoip": {
            "enabled": True if geoip_row is None else bool(geoip_row["enabled"]),
            "cron": "0 4 * * *" if geoip_row is None else str(geoip_row["cron"]),
            "provider": "ripe" if geoip_row is None else str(geoip_row["provider"]),
        },
        "rules": rules,
    }
    if include_exceptions:
        payload["exceptions"] = [
            {
                "match": str(row["match"]),
                "reason": str(row["reason"]),
                "enabled": bool(row["enabled"]),
            }
            for row in conn.execute(
                "SELECT match, reason, enabled FROM rule_exceptions ORDER BY id"
            ).fetchall()
        ]
    return payload


def save_snapshot(
    conn: sqlite3.Connection,
    payload: dict[str, Any],
    *,
    entity: str = Entity.APPLY.value,
    entity_id: str = "state",
    reason: str = "",
) -> int:
    """Сохранить снимок состояния и вернуть идентификатор."""
    cursor = conn.execute(
        "INSERT INTO state_snapshots (entity, entity_id, created_at, payload, reason)"
        " VALUES (?, ?, ?, ?, ?)",
        (entity, entity_id, utcnow(), json.dumps(payload, sort_keys=True), reason),
    )
    return int(cursor.lastrowid or 0)


def load_snapshot(conn: sqlite3.Connection, snapshot_id: int) -> StateSnapshot:
    """Прочитать снимок по id."""
    row = conn.execute("SELECT * FROM state_snapshots WHERE id = ?", (snapshot_id,)).fetchone()
    if row is None:
        raise StorageError(f"снимок состояния {snapshot_id} не найден")
    return StateSnapshot(
        id=int(row["id"]),
        entity=str(row["entity"]),
        entity_id=str(row["entity_id"]),
        created_at=str(row["created_at"]),
        payload=json.loads(str(row["payload"])),
        reason=str(row["reason"]),
    )


def restore_state(conn: sqlite3.Connection, payload: dict[str, Any]) -> None:
    """Полностью заменить изменяемые таблицы состоянием из снимка."""
    conn.execute("DELETE FROM routing_rules")
    for rule in payload.get("rules", []):
        conn.execute(
            "INSERT INTO routing_rules (name, kind, match, target, priority, enabled,"
            " created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
            (
                rule["name"],
                rule["kind"],
                rule["match"],
                rule["target"],
                int(rule["priority"]),
                1 if rule["enabled"] else 0,
                utcnow(),
                utcnow(),
            ),
        )
    if "exceptions" in payload:
        conn.execute("DELETE FROM rule_exceptions")
        for exception in payload["exceptions"]:
            conn.execute(
                "INSERT INTO rule_exceptions (match, reason, enabled, created_at, updated_at)"
                " VALUES (?, ?, ?, ?, ?)",
                (
                    exception["match"],
                    exception.get("reason", ""),
                    1 if exception["enabled"] else 0,
                    utcnow(),
                    utcnow(),
                ),
            )
    geoip = payload.get("geoip", {})
    conn.execute(
        "UPDATE routing_mode SET mode = ?, updated_at = ? WHERE id = 1",
        (payload.get("mode", "off"), utcnow()),
    )
    conn.execute(
        "UPDATE geoip_schedule SET enabled = ?, cron = ?, provider = ?,"
        " updated_at = ? WHERE id = 1",
        (
            1 if geoip.get("enabled", True) else 0,
            geoip.get("cron", "0 4 * * *"),
            geoip.get("provider", "ripe"),
            utcnow(),
        ),
    )


def rollback_to_snapshot(
    conn: sqlite3.Connection, snapshot_id: int, *, actor: str, reason: str = ""
) -> StateSnapshot:
    """Восстановить состояние из снимка и записать это в журнал аудита."""
    snapshot = load_snapshot(conn, snapshot_id)
    audit.record(
        conn,
        actor=actor,
        action="rollback",
        entity=snapshot.entity,
        entity_id=snapshot.entity_id,
        old_value=capture_state(conn),
        new_value=snapshot.payload,
        detail=reason or f"rollback к снимку #{snapshot.id}",
    )
    restore_state(conn, snapshot.payload)
    return snapshot