stackchain-dashboard/scripts/rotate_private_state.py

#!/usr/bin/env python3
"""Rewrap shared private-state envelopes under the configured active key."""

from __future__ import annotations

import json
import os
import sqlite3
import sys
from dataclasses import dataclass
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from src.private_state import connect_private_sqlite  # noqa: E402
from src.state_encryption import (  # noqa: E402
    PrivateStateCipher,
    PrivateStateEncryptionError,
    private_state_encryption_config,
)


@dataclass(frozen=True)
class Field:
    column: str
    binding: str
    alias: str | None = None
    plaintext_string: bool = False


@dataclass(frozen=True)
class Table:
    name: str
    context_columns: tuple[str, ...]
    fields: tuple[Field, ...]


@dataclass(frozen=True)
class Store:
    name: str
    identity: str
    env: str
    filename: str
    tables: tuple[Table, ...]


STORES = (
    Store("today", "today", "STACKCHAIN_TODAY_DB", "today.sqlite3", (
        Table("today_plans", ("login",), (Field("ids", "plan:{login}"),)),
        Table("tomorrow_plans", ("login",), (Field("payload", "tomorrow:{login}"),)),
        Table("tomorrow_promotions", ("login", "promotion_id"), (Field("result", "tomorrow-promotion:{login}:{promotion_id}"),)),
        Table("week_plans", ("login",), (Field("payload", "week:{login}"),)),
        Table("week_promotions", ("login", "promotion_id"), (Field("result", "week-promotion:{login}:{promotion_id}"),)),
        Table("week_reschedules", ("login", "operation_id"), (Field("result", "week-reschedule:{login}:{operation_id}"),)),
        Table("today_sessions", ("login",), (Field("device_id", "session:{login}"),)),
        Table("today_recaps", ("login",), (
            Field("session_id", "recap-id:{login}", "session_id", True),
            Field("items", "recap-items:{login}:{session_id}"),
        )),
        Table("today_time_logs", ("login",), (
            Field("session_id", "time-log-session:{login}", "session_id", True),
            Field("identity", "time-log-identity:{login}:{session_id}", "identity", True),
            Field("actual_minutes", "time-log-payload:{login}:{session_id}:{identity}"),
        )),
    )),
    Store("later", "later", "STACKCHAIN_LATER_DB", "later.sqlite3", (
        Table("later_plans", ("login",), (Field("records", "plan:{login}"),)),
        Table("later_item_revisions", ("login",), (Field("item_id", "item-revision:{login}", plaintext_string=True),)),
    )),
    Store("live-snapshot", "live-snapshot", "STACKCHAIN_LIVE_SNAPSHOT_DB", "live-snapshot.sqlite3", (
        Table("live_snapshot", ("generation",), (Field("value_json", "{generation}"),)),
    )),
    Store("available-issue-snapshot", "available-issue-snapshot", "STACKCHAIN_AVAILABLE_ISSUE_SNAPSHOT_DB", "available-issue-snapshot.sqlite3", (
        Table("available_issue_snapshot", (), (Field("items_json", "singleton"),)),
    )),
    Store("saved-searches", "saved-searches", "STACKCHAIN_SAVED_SEARCH_DB", "saved-searches.sqlite3", (
        Table("saved_searches", ("login",), (Field("views", "views:{login}"),)),
    )),
    Store(
        "completed-filed-reviews",
        "completed-filed-reviews",
        "STACKCHAIN_COMPLETED_FILED_REVIEW_DB",
        "completed-filed-reviews.sqlite3",
        (
            Table(
                "completed_filed_review_collections",
                ("login",),
                (Field("receipts", "receipts:{login}"),),
            ),
        ),
    ),
    Store("security-events", "security-events", "STACKCHAIN_SECURITY_EVENT_DB", "security-events.sqlite3", (
        Table("security_events", ("id",), (Field("payload", "event:{id}"),)),
    )),
    Store("idempotency-ledger", "idempotency-ledger", "STACKCHAIN_IDEMPOTENCY_DB", "idempotency.sqlite3", (
        Table("idempotency_operations", ("key",), (
            Field("fingerprint", "{key}:fingerprint"),
            Field("response_json", "{key}:response"),
        )),
    )),
)


def _tables(connection: sqlite3.Connection) -> set[str]:
    return {
        row[0]
        for row in connection.execute(
            "SELECT name FROM sqlite_master WHERE type = 'table'"
        )
    }


def rotate_store(path: Path, spec: Store, config) -> dict[str, int]:
    counts = {"current": 0, "failed": 0, "migrated": 0, "total": 0}
    if not path.exists():
        return counts
    cipher = PrivateStateCipher(config, store=spec.identity)
    with connect_private_sqlite(path, timeout=1.0) as connection:
        connection.row_factory = sqlite3.Row
        available = _tables(connection)
        for table in spec.tables:
            if table.name not in available:
                continue
            columns = (*table.context_columns, *(field.column for field in table.fields))
            query = f"SELECT rowid AS _rotation_rowid, {', '.join(columns)} FROM {table.name}"
            for row in connection.execute(query).fetchall():
                context = {name: row[name] for name in table.context_columns}
                updates: dict[str, str] = {}
                row_failed = False
                row_counts = {"current": 0, "migrated": 0, "total": 0}
                for field in table.fields:
                    payload = row[field.column]
                    if payload is None:
                        continue
                    row_counts["total"] += 1
                    try:
                        binding = field.binding.format(**context)
                        if isinstance(payload, str) and not payload.startswith(("v1:", "v2:")) and field.plaintext_string:
                            value, stale = payload, True
                        else:
                            value, stale = cipher.open(str(payload), binding=binding)
                        if field.alias:
                            context[field.alias] = value
                        if stale:
                            updates[field.column] = cipher.seal(value, binding=binding)
                            row_counts["migrated"] += 1
                        else:
                            row_counts["current"] += 1
                    except (KeyError, PrivateStateEncryptionError, ValueError):
                        row_failed = True
                        break
                counts["total"] += row_counts["total"]
                if row_failed:
                    counts["failed"] += 1
                    continue
                if updates:
                    assignments = ", ".join(f"{name} = ?" for name in updates)
                    connection.execute(
                        f"UPDATE {table.name} SET {assignments} WHERE rowid = ?",
                        (*updates.values(), row["_rotation_rowid"]),
                    )
                counts["migrated"] += row_counts["migrated"]
                counts["current"] += row_counts["current"]
    return counts


def main() -> int:
    try:
        config = private_state_encryption_config()
        if isinstance(config, bytes):
            raise PrivateStateEncryptionError("rotation requires a keyring")
        state = Path(os.getenv("STACKCHAIN_STATE_DIR", ".stackchain-state"))
        report = {
            spec.name: rotate_store(
                Path(os.getenv(spec.env, str(state / spec.filename))), spec, config
            )
            for spec in STORES
        }
    except (OSError, sqlite3.Error, PrivateStateEncryptionError):
        print(json.dumps({"error": "Private-state rotation configuration is unavailable"}, sort_keys=True))
        return 2
    print(json.dumps(report, sort_keys=True))
    return 1 if any(item["failed"] for item in report.values()) else 0


if __name__ == "__main__":
    raise SystemExit(main())