FabledSteward/steward/backup/check.py

"""Checking a backup: what it would restore, compared with what is in place,
without restoring anything (#5823).

The backup's rows come out of pg_restore as COPY text, with no database
involved; the rows in place go out through COPY in the same form. They are
compared table by table on their primary keys, over the columns both sides
have. Nothing is written.

What is compared is what a restore would replace: every table whose rows
the backup carries and this Steward keeps (contents.py). Columns that are
Steward's record of what it last saw (contents.ACTIVITY) and settings it
writes as it runs (contents.STATE_SETTINGS) are left out of the comparison:
they change by the minute and say nothing about whether the backup is good.

A sealed value (crypto's ``enc:v1:`` tokens) is opened with its own side's
key and compared as plaintext, so a secret sealed twice still matches. It
is never shown: the report says only that it differs. A backup whose key
doesn't open its own sealed values is reported, since its secrets would not
come back.
"""
from __future__ import annotations

import asyncio
import hashlib
import io
import json
import re
import tempfile
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path

from sqlalchemy import text
from sqlalchemy.ext.asyncio import create_async_engine

from .contents import (ACTIVITY, ACTIVITY_EVERYWHERE, AREAS, EXCLUDED, HAND_SET_ONLY,
                       STATE_SETTINGS)
from .dump import read_backup
from .pgdata import async_url as _async_url
from .pgdata import copy_rows, restore_script, script_tables

# Names shown for a row, the first a table has; else its primary key.
LABEL_COLUMNS = ("name", "username", "key", "title", "widget_key", "label", "mac")
# A setting's two values are shown when they fit; longer ones are cut.
SHOWN_CHARS = 120
_TOKEN = re.compile(r"enc:v1:[A-Za-z0-9_\-=]+")


@dataclass
class RowDiff:
    label: str
    columns: list[str]
    secret: bool = False
    backup_value: str | None = None      # shown only for a setting that isn't secret
    here_value: str | None = None


@dataclass
class TableCheck:
    table: str
    same: int = 0
    differ: list[RowDiff] = field(default_factory=list)
    only_backup: list[str] = field(default_factory=list)
    only_here: list[str] = field(default_factory=list)
    note: str = ""                       # a table one side lacks, columns one side lacks

    @property
    def changes(self) -> int:
        return len(self.differ) + len(self.only_backup) + len(self.only_here)


@dataclass
class AreaCheck:
    name: str
    tables: list[TableCheck]

    @property
    def same(self) -> int:
        return sum(t.same for t in self.tables)

    @property
    def changes(self) -> int:
        return sum(t.changes for t in self.tables)


@dataclass
class CheckReport:
    name: str
    manifest: dict
    schema: str                          # "same", "older" or "newer"
    backup_heads: list[str]
    here_heads: list[str]
    key_same: bool
    sealed: int                          # sealed values in the backup
    unopened: int                        # of those, how many its own key couldn't open
    areas: list[AreaCheck]
    checked_at: str = ""

    @property
    def made_at(self) -> datetime | None:
        try:
            return datetime.fromisoformat(self.manifest.get("created_at") or "")
        except ValueError:
            return None

    @property
    def rows(self) -> int:
        return sum(a.same + a.changes for a in self.areas)

    @property
    def changes(self) -> int:
        return sum(a.changes for a in self.areas)

    @property
    def ok(self) -> bool:
        """The backup reads back whole: its key opens its secrets. Whether it
        matches what is in place is `changes`; a difference isn't a fault."""
        return self.unopened == 0

    def summary(self) -> str:
        if self.unopened:
            return (f"The backup's key doesn't open {self.unopened} of its {self.sealed} sealed "
                    f"value{'s' if self.sealed != 1 else ''}: those secrets wouldn't come back.")
        if not self.changes:
            tables = sum(len(a.tables) for a in self.areas)
            return f"Matches what is in place: {self.rows} rows across {tables} table{'s' if tables != 1 else ''}."
        parts = [f"{a.changes} in {a.name}" for a in self.areas if a.changes]
        return (f"{self.changes} difference{'s' if self.changes != 1 else ''} from what is in place: "
                + ", ".join(parts) + ".")

    def as_entry(self) -> dict:
        """What a backup's history row keeps of its own check."""
        return {"ok": self.ok, "changes": self.changes, "rows": self.rows, "summary": self.summary()}


# ── Comparing ───────────────────────────────────────────────────────────────

def _fernet(key: str):
    from steward.core.crypto import _make_fernet
    return _make_fernet(key) if key else None


class _Opener:
    """Opens one side's sealed values, counting what it couldn't."""

    def __init__(self, key: str):
        self.fernet = _fernet(key)
        self.seen = 0
        self.failed = 0

    def open(self, value: str) -> tuple[str, bool]:
        """`value` with each sealed token replaced by a digest of its
        plaintext (never the plaintext itself), and whether it held one."""
        if "enc:v1:" not in value:
            return value, False

        def plain(m: re.Match) -> str:
            self.seen += 1
            try:
                opened = self.fernet.decrypt(m.group(0)[len("enc:v1:"):].encode("ascii"))
            except Exception:
                self.failed += 1
                return "\x00unopened:" + m.group(0)
            return "\x00sealed:" + hashlib.sha256(opened).hexdigest()
        return _TOKEN.sub(plain, value), True


def _offset(v: str) -> str:
    return re.sub(r"([+-]\d\d)$", r"\1:00", v)


def normalise(value: str | None, data_type: str) -> object:
    """A value as compared: an instant for a timestamp with zone (the two
    sides may print it in different zones), parsed JSON, a float."""
    if value is None:
        return None
    try:
        if data_type == "timestamp with time zone":
            return datetime.fromisoformat(_offset(value)).astimezone(timezone.utc)
        if data_type in ("json", "jsonb"):
            return json.dumps(json.loads(value), sort_keys=True)
        if data_type in ("double precision", "real", "numeric"):
            return float(value)
    except ValueError:
        pass
    return value


def _shown(value: str | None) -> str | None:
    if value is None:
        return None
    return value if len(value) <= SHOWN_CHARS else value[:SHOWN_CHARS - 1] + "…"


@dataclass
class _Side:
    columns: list[str]
    rows: list[list[str | None]]


def compare_table(table: str, backup: _Side, here: _Side, *, types: dict[str, str],
                  key: list[str], openers: tuple[_Opener, _Opener],
                  labels: dict[str, dict[str, str]]) -> TableCheck:
    """One table's rows, backup against here."""
    result = TableCheck(table)
    skip = ACTIVITY.get(table, frozenset()) | ACTIVITY_EVERYWHERE
    common = [c for c in backup.columns if c in here.columns and c not in skip]
    lacking = sorted((set(backup.columns) ^ set(here.columns)) - skip)
    if lacking:
        result.note = "Columns only one side has: " + ", ".join(lacking) + "."
    key = [c for c in key if c in common] or common
    hand_set = HAND_SET_ONLY.get(table)

    def index(side: _Side, opener: _Opener) -> dict[tuple, dict[str, object]]:
        at = {c: side.columns.index(c) for c in common}
        out = {}
        for raw in side.rows:
            row: dict[str, object] = {}
            sealed: set[str] = set()
            for c in common:
                v = raw[at[c]]
                if isinstance(v, str):
                    v, was = opener.open(v)
                    if was:
                        sealed.add(c)
                row[c] = normalise(v, types.get(c, ""))
            if table == "app_settings" and row.get("key") in STATE_SETTINGS:
                continue
            if hand_set and not any(raw[side.columns.index(c)] for c in hand_set if c in side.columns):
                continue
            row["\x00sealed"] = sealed
            row["\x00raw"] = {c: raw[at[c]] for c in common}
            out[tuple(str(row[c]) for c in key)] = row
        return out

    theirs, ours = index(backup, openers[0]), index(here, openers[1])

    def label(row: dict) -> str:
        raw = row["\x00raw"]
        if table == "host_agent_registrations" and raw.get("host_id"):
            return labels.get("hosts", {}).get(raw["host_id"], raw["host_id"])
        if table == "ansible_target_groups":
            return (f"{labels.get('ansible_targets', {}).get(raw.get('target_id'), raw.get('target_id'))} in "
                    f"{labels.get('ansible_groups', {}).get(raw.get('group_id'), raw.get('group_id'))}")
        for c in LABEL_COLUMNS:
            if raw.get(c):
                return str(raw[c])
        return " / ".join(str(raw.get(c)) for c in key)

    for k in sorted(theirs.keys() | ours.keys()):
        b, h = theirs.get(k), ours.get(k)
        if h is None:
            result.only_backup.append(label(b))
        elif b is None:
            result.only_here.append(label(h))
        else:
            cols = [c for c in common if b[c] != h[c]]
            if not cols:
                result.same += 1
                continue
            secret = bool((b["\x00sealed"] | h["\x00sealed"]) & set(cols))
            diff = RowDiff(label(b), cols, secret)
            if table == "app_settings" and not secret:
                diff.backup_value = _shown(b["\x00raw"].get("value_json"))
                diff.here_value = _shown(h["\x00raw"].get("value_json"))
            result.differ.append(diff)
    return result


def _labels(tables: dict[str, _Side]) -> dict[str, dict[str, str]]:
    """id → name for the tables other rows are named through."""
    out: dict[str, dict[str, str]] = {}
    for t in ("hosts", "ansible_targets", "ansible_groups"):
        side = tables.get(t)
        if side and "id" in side.columns and "name" in side.columns:
            i, n = side.columns.index("id"), side.columns.index("name")
            out[t] = {r[i]: r[n] for r in side.rows}
    return out


# ── Reading both sides ──────────────────────────────────────────────────────

async def _backup_rows(dump: Path) -> dict[str, tuple[list[str], list[list[str | None]]]]:
    """Every table's rows in the dump, as pg_restore would load them."""
    return script_tables(await restore_script(dump))


async def _here(database_url: str, tables: list[str]):
    """The rows in place for `tables`, with each column's type, each table's
    primary key and the migration heads."""
    engine = create_async_engine(_async_url(database_url))
    try:
        async with engine.connect() as conn:
            types: dict[str, dict[str, str]] = {}
            order: dict[str, list[str]] = {}
            for t, c, dt in (await conn.execute(text(
                    "SELECT table_name, column_name, data_type FROM information_schema.columns "
                    "WHERE table_schema = 'public' AND table_name = ANY(:t) "
                    "ORDER BY table_name, ordinal_position"), {"t": tables})).all():
                types.setdefault(t, {})[c] = dt
                order.setdefault(t, []).append(c)
            keys: dict[str, list[str]] = {}
            for t, c in (await conn.execute(text(
                    "SELECT cl.relname, a.attname FROM pg_index i "
                    "JOIN pg_class cl ON cl.oid = i.indrelid "
                    "JOIN pg_namespace ns ON ns.oid = cl.relnamespace "
                    "JOIN pg_attribute a ON a.attrelid = cl.oid AND a.attnum = ANY(i.indkey) "
                    "WHERE i.indisprimary AND ns.nspname = 'public' AND cl.relname = ANY(:t) "
                    "ORDER BY cl.relname, a.attnum"), {"t": tables})).all():
                keys.setdefault(t, []).append(c)
            heads = sorted((await conn.execute(text("SELECT version_num FROM alembic_version"))).scalars())
            raw = (await conn.get_raw_connection()).driver_connection
            rows: dict[str, _Side] = {}
            for t in tables:
                if t not in order:
                    continue
                buf = io.BytesIO()
                await raw.copy_from_table(t, columns=order[t], output=buf, format="text")
                rows[t] = _Side(order[t], copy_rows(buf.getvalue()))
            return rows, types, keys, heads
    finally:
        await engine.dispose()


def known_revisions(plugin_dirs: list[Path] | None = None) -> set[str]:
    """Every migration revision this Steward has, core and plugins: a backup
    head not among them was made by a newer Steward."""
    import steward
    roots = [Path(steward.__file__).parent / "migrations" / "versions"]
    for d in plugin_dirs or [Path(steward.__file__).parent.parent / "plugins"]:
        if Path(d).is_dir():
            roots += [p / "migrations" / "versions" for p in Path(d).iterdir()]
    found = set()
    for root in roots:
        if root.is_dir():
            for f in root.glob("*.py"):
                m = re.search(r'^revision\s*(?::\s*str\s*)?=\s*["\']([^"\']+)["\']',
                              f.read_text(errors="replace"), re.M)
                if m:
                    found.add(m.group(1))
    return found


def schema_of(backup_heads: list[str], here_heads: list[str], known: set[str]) -> str:
    if sorted(backup_heads) == sorted(here_heads):
        return "same"
    return "newer" if any(h not in known for h in backup_heads) else "older"


async def check_backup(path: Path, database_url: str, here_key: str, *,
                       plugin_dirs: list[Path] | None = None, name: str = "") -> CheckReport:
    """Compare the backup at `path` with the database at `database_url`.
    Raises BackupError when the file isn't a backup or can't be read."""
    with tempfile.TemporaryDirectory(prefix="steward-check-") as tmp:
        info, dump, backup_key = await asyncio.to_thread(read_backup, path, Path(tmp))
        theirs = await _backup_rows(dump)
    tables = sorted(t for t in theirs if t not in EXCLUDED)
    here, types, keys, heads = await _here(database_url, tables)
    openers = (_Opener(backup_key), _Opener(here_key))
    backup_sides = {t: _Side(*theirs[t]) for t in tables}
    labels = _labels(backup_sides)
    for t, names in _labels(here).items():           # a name in place wins
        labels[t] = {**labels.get(t, {}), **names}

    checks: dict[str, TableCheck] = {}
    for t in tables:
        if t not in here:
            checks[t] = TableCheck(t, only_backup=[f"{len(theirs[t][1])} rows"],
                                   note="This Steward has no such table.")
            continue
        checks[t] = compare_table(t, backup_sides[t], here[t], types=types.get(t, {}),
                                  key=keys.get(t, []), openers=openers, labels=labels)
    areas = []
    placed: set[str] = set()
    for area, members in AREAS:
        found = [checks[t] for t in members if t in checks]
        placed |= set(members)
        if found:
            areas.append(AreaCheck(area, found))
    rest = [checks[t] for t in tables if t not in placed]
    if rest:
        areas.append(AreaCheck("Other", rest))
    backup_heads = sorted(info.get("migration_heads") or [])
    return CheckReport(
        name=name or path.name, manifest=info,
        schema=schema_of(backup_heads, heads, known_revisions(plugin_dirs)),
        backup_heads=backup_heads, here_heads=heads, key_same=backup_key == here_key,
        sealed=openers[0].seen, unopened=openers[0].failed, areas=areas,
        checked_at=datetime.now(timezone.utc).isoformat())