FabledSteward/steward/backup/restoring.py

"""Restoring mode (#5816): what a running Steward does while a backup goes back
in, from Settings → Backups.

1. Read the backup and refuse it up front if it can't go in: not a backup,
   or a different secret key while STEWARD_SECRET_KEY pins this one.
2. Take a safety backup of Steward as it stands: kept on /data/backups and
   sent to the bucket when one is set. A restore can then be undone the same
   way. Without it nothing goes further.
3. While it runs, Steward is restoring: the scheduler holds its ticks, agent
   ingest answers 503 (agents keep their samples and send them again), and
   every page says what is happening.
4. In place for a backup of this schema (restore.restore_in_place), or a
   full replace the person agreed to for an older one.
5. The backup's key is put in place, the restore audited, and the process
   exits with RESTART_EXIT_CODE. Swarm, compose or systemd starts it again,
   and settings, the key and plugins load fresh. The person who restored
   signs in again: users came back from the backup.

A restore that fails in place has changed nothing; Steward carries on and
says why. A full replace that fails part-way leaves an empty database, so
Steward restarts into first-run setup, where the safety backup can be
uploaded.
"""
from __future__ import annotations

import asyncio
import json
import logging
import os
import re
import secrets
import shutil
import signal
import time
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path

logger = logging.getLogger(__name__)

SAFETY_DIR = Path("/data/backups")
# Safety backups kept on /data, newest first; older ones are on the bucket.
SAFETY_KEEP = 5
# Not 0: systemd's Restart=on-failure restarts only a failed exit.
RESTART_EXIT_CODE = 3
# Long enough for the restoring page's next poll to say it is done.
EXIT_DELAY_SECONDS = 5
# A backup chosen for restore waits here between its check and the
# confirmation (#5817); one not confirmed within the hour is removed.
STAGED_SECONDS = 3600


@dataclass
class Restoring:
    active: bool = False
    name: str = ""
    by: str = ""
    since: datetime | None = None
    step: str = ""
    error: str = ""
    restart_requested: bool = False


STATE = Restoring()
_tasks: set[asyncio.Task] = set()


def is_restoring() -> bool:
    return STATE.active


def start_restore(app, path: Path, name: str, *, username: str, replace_all: bool = False,
                  owned_dir: Path | None = None, safety: bool = True) -> bool:
    """Start restoring `path` in the background; False if a restore is already
    under way. `owned_dir`, when given, is removed once it is over. `safety`
    False skips the safety backup: first-run setup has nothing to keep."""
    if STATE.active or STATE.restart_requested:
        return False
    STATE.active, STATE.name, STATE.by = True, name, username
    STATE.since, STATE.step, STATE.error = datetime.now(timezone.utc), "Reading the backup", ""
    task = asyncio.get_running_loop().create_task(
        _run(app, path, name, username=username, replace_all=replace_all, owned_dir=owned_dir,
             safety=safety))
    _tasks.add(task)
    task.add_done_callback(_tasks.discard)
    return True


async def _run(app, path: Path, name: str, *, username: str, replace_all: bool,
               owned_dir: Path | None, safety: bool = True) -> None:
    from steward.core.audit import log_audit

    from .dump import BackupError
    from .s3 import S3Error
    wiped = False
    try:
        key = await asyncio.to_thread(_backup_key, path)
        pinned = os.environ.get("STEWARD_SECRET_KEY")
        if pinned and pinned != key:
            raise BackupError("STEWARD_SECRET_KEY is set and differs from the backup's key; "
                              "set it to the backup's key first, or its passwords and keys can't be read.")
        safety_name = ""
        if safety:
            STATE.step = "Saving a safety backup of Steward as it stands"
            safety_name = await _safety_backup(app)
        STATE.step = "Restoring"
        plugin_dirs = [Path(d) for d in app.config["PLUGIN_DIRS"]]
        if replace_all:
            from .restore import restore
            wiped = True
            done = await restore(path, app.config["DATABASE_URL"], wipe=True, plugin_dirs=plugin_dirs)
            detail = {"mode": "replaced everything"}
        else:
            from .restore import restore_in_place
            done = await restore_in_place(path, app.config["DATABASE_URL"], plugin_dirs=plugin_dirs)
            detail = {"mode": "in place", "tables": len(done.tables), "rows": done.rows,
                      "put_right": done.fixed}
        if done.secret_key != app.config["SECRET_KEY"] and not pinned:
            _write_key(done.secret_key)
        await log_audit(app, None, username, "backup.restored", entity_type="backup",
                        entity_id=name, detail=detail | ({"safety_backup": safety_name} if safety else {}))
        logger.warning("Restored %s (%s); restarting", name, detail["mode"])
        STATE.step = "Restored. Steward is restarting; sign in again when it is back"
        _restart_soon()
    except (BackupError, S3Error, OSError) as exc:
        _failed(app, name, username, str(exc), wiped)
    except Exception as exc:
        logger.exception("Restore of %s failed", name)
        _failed(app, name, username, f"Unexpected error: {exc}", wiped)
    finally:
        if owned_dir is not None:
            shutil.rmtree(owned_dir, ignore_errors=True)


def _failed(app, name: str, username: str, error: str, wiped: bool) -> None:
    from steward.core.audit import log_audit
    logger.error("Restore of %s failed: %s", name, error)
    STATE.error = error
    if wiped:
        # The database was dropped before the restore failed: start again
        # into first-run setup, where the safety backup can be uploaded.
        STATE.step = ("The restore failed after the database was emptied; Steward is restarting "
                      "into first-run setup, where the safety backup can be uploaded")
        _restart_soon()
        return
    STATE.active = False
    task = asyncio.get_running_loop().create_task(log_audit(
        app, None, username, "backup.restore_failed", entity_type="backup", entity_id=name,
        detail={"error": error}))
    _tasks.add(task)
    task.add_done_callback(_tasks.discard)


def _backup_key(path: Path) -> str:
    import tarfile

    from .dump import KEY_NAME, BackupError
    try:
        with tarfile.open(path, "r") as tar:
            return tar.extractfile(KEY_NAME).read().decode().strip()
    except (tarfile.TarError, OSError, KeyError, AttributeError):
        raise BackupError(f"{path.name} isn't a Steward backup.") from None


async def _safety_backup(app) -> str:
    """A backup of Steward as it stands, before anything is replaced: on
    /data/backups always, and on the bucket when one is set. Its run is in
    Recent backups like any other."""
    from steward.core.settings import get_all_settings

    from .dump import make_backup
    from .job import _record, migration_heads, target_from
    from .s3 import put_file
    SAFETY_DIR.mkdir(parents=True, exist_ok=True)
    url = app.config["DATABASE_URL"]
    made = await make_backup(url, app.config["SECRET_KEY"], await migration_heads(url), SAFETY_DIR)
    async with app.db_sessionmaker() as db:
        target = target_from(await get_all_settings(db))
    entry = {"at": datetime.now(timezone.utc).isoformat(), "ok": True, "by": "before restore",
             "name": made.name, "size": made.size, "secs": 0.0, "error": ""}
    if target is not None:
        await put_file(target, made.name, made.path)
    else:
        entry["error"] = "Kept on /data/backups only: no bucket is set up."
    await _record(app, entry, None, "")
    _prune()
    return made.name


def _prune() -> None:
    kept = sorted(SAFETY_DIR.glob("steward-*.tar"), reverse=True)
    for old in kept[SAFETY_KEEP:]:
        old.unlink(missing_ok=True)


def _write_key(key: str) -> None:
    from steward.config import _SECRET_KEY_FILE
    _SECRET_KEY_FILE.parent.mkdir(parents=True, exist_ok=True)
    _SECRET_KEY_FILE.write_text(key)
    _SECRET_KEY_FILE.chmod(0o600)


def _restart_soon() -> None:
    STATE.restart_requested = True
    loop = asyncio.get_running_loop()
    loop.call_later(EXIT_DELAY_SECONDS, os.kill, os.getpid(), signal.SIGTERM)


# ── Chosen, checked, waiting for a yes (#5817) ─────────────────────────────

@dataclass
class Staged:
    token: str
    dir: Path
    name: str = ""
    schema: str = ""
    ok: bool = False

    @property
    def path(self) -> Path:
        return self.dir / "backup.tar"

    def save(self) -> None:
        (self.dir / "staged.json").write_text(json.dumps(
            {"name": self.name, "schema": self.schema, "ok": self.ok}))


def _staging() -> Path:
    return SAFETY_DIR / "incoming"


def stage() -> Staged:
    """A fresh place to hold one backup until its restore is confirmed. Old
    ones not confirmed are removed first."""
    root = _staging()
    root.mkdir(parents=True, exist_ok=True)
    cutoff = time.time() - STAGED_SECONDS
    for old in root.iterdir():
        if old.stat().st_mtime < cutoff:
            shutil.rmtree(old, ignore_errors=True)
    token = secrets.token_hex(16)
    (root / token).mkdir(mode=0o700)
    return Staged(token, root / token)


def staged(token: str) -> Staged | None:
    """The backup held under `token`, or None when there is none (or it was
    removed for its age)."""
    if not re.fullmatch(r"[0-9a-f]{32}", token or ""):
        return None
    d = _staging() / token
    try:
        meta = json.loads((d / "staged.json").read_text())
    except (OSError, ValueError):
        return None
    if not (d / "backup.tar").is_file():
        return None
    return Staged(token, d, meta.get("name", ""), meta.get("schema", ""), bool(meta.get("ok")))


# ── While it runs ───────────────────────────────────────────────────────────

async def gate():
    """before_request: while restoring, nothing reads or writes the database.
    Agents and scripts get a 503 to retry; a page says what is happening."""
    from quart import Response, g, render_template, request
    if not (STATE.active or STATE.restart_requested) or request.path.startswith("/static/"):
        return None
    g.database_down = True        # the page's own reads (the sidebar's alert count) wait
    headers = {"Retry-After": "30"}
    wants_page = "text/html" in request.headers.get("Accept", "")
    if request.method != "GET" or request.headers.get("HX-Request") or not wants_page:
        return Response("Steward is restoring a backup; try again shortly.", status=503,
                        headers=headers, content_type="text/plain; charset=utf-8")
    page = await render_template("errors/restoring.html", state=STATE)
    return Response(page, status=503, headers=headers)