plm-lottery/app/wallet/verify.py

"""B-89 (SEC-49, SEC-50): is the master key the one that produced what the database holds?

A different xprv, or a wrong encryption key, used to surface only at the first
registration or payout — and a different xprv never did: it quietly moved the pool
address. The check compares a fixed sample of stored addresses with what the key derives,
reading the database with plain `sqlite3` (read-only), so the startup check and the
key scripts share it without an event loop."""

import hashlib
import logging
import sqlite3
from typing import NamedTuple

from cryptography.fernet import Fernet
from embit import script
from embit.bip32 import HDKey

from app.config import ConfigError, settings
from app.wallet.plm_network import ACCOUNT_PATH, PLM_MAINNET

logger = logging.getLogger(__name__)

_SAMPLE_SIZE = 10


class Sample(NamedTuple):
    users: list[tuple[int, str]]  # (derivation_index, address)
    rounds: list[tuple[int, str]]  # (id, pool_address) of the last rounds with one
    seeds: list[tuple[int, str, str]]  # (id, server_seed_enc, server_seed_hash), unrevealed


def database_file(url: str | None = None) -> str | None:
    """The SQLite file `DATABASE_URL` points at, or None for anything else."""
    url = url or settings.database_url
    prefix = "sqlite+aiosqlite:///"
    return url[len(prefix):] if url.startswith(prefix) else None


def load_sample(path: str | None = None) -> Sample:
    """The lowest and highest `_SAMPLE_SIZE` users, the last rounds that have a pool
    address and every round whose secret is still unrevealed. A missing file or missing
    tables (a fresh instance) is an empty sample: there is nothing to compare."""
    path = path or database_file()
    empty = Sample([], [], [])
    if path is None:
        return empty
    try:
        db = sqlite3.connect(f"file:{path}?mode=ro", uri=True)
    except sqlite3.Error:
        return empty
    try:
        users = {
            index: address
            for order in ("ASC", "DESC")
            for index, address in db.execute(
                f"SELECT derivation_index, address FROM users ORDER BY derivation_index {order} LIMIT ?",
                (_SAMPLE_SIZE,),
            )
        }
        rounds = db.execute(
            "SELECT id, pool_address FROM rounds WHERE pool_address != '' ORDER BY id DESC LIMIT ?",
            (_SAMPLE_SIZE,),
        ).fetchall()
        seeds = db.execute(
            "SELECT id, server_seed_enc, server_seed_hash FROM rounds "
            "WHERE server_seed_enc IS NOT NULL AND server_seed IS NULL"
        ).fetchall()
        return Sample(sorted(users.items()), rounds, seeds)
    except sqlite3.OperationalError:  # no such table
        return empty
    finally:
        db.close()


def account_key(xprv_base58: str) -> HDKey:
    return HDKey.from_base58(xprv_base58).derive(ACCOUNT_PATH)


def _address(account: HDKey, path: str) -> str:
    return script.p2wpkh(account.derive(path).to_public()).address(network=PLM_MAINNET)


def sample_problems(account: HDKey, sample: Sample, encryption_key: str | None = None) -> list[str]:
    """One line per check that failed; never any key material. With `encryption_key`
    the unrevealed round secrets are checked as well (they depend on it, not on the xprv)."""
    problems = []
    for index, stored in sample.users:
        if _address(account, f"0/{index}") != stored:
            problems.append(f"user address at derivation index {index} differs from the one derived")
    if sample.rounds:
        reserve = _address(account, "1/0")
        for round_id, stored in sample.rounds:
            if stored not in (reserve, _address(account, f"2/{round_id}")):
                problems.append(f"pool address of round {round_id} is neither its own nor the reserve address")
    if encryption_key is not None:
        fernet = Fernet(encryption_key)
        for round_id, token, committed in sample.seeds:
            try:
                secret = fernet.decrypt(token.encode()).decode()
            except Exception:
                problems.append(f"secret of round {round_id} does not decrypt")
                continue
            if hashlib.sha256(bytes.fromhex(secret)).hexdigest() != committed:
                problems.append(f"secret of round {round_id} does not match its published hash")
    return problems


def verify_master_key() -> None:
    """Called from the lifespan, before any task starts: decrypt the key and compare the
    sample. Raises `ConfigError` (the process exits) on the first sign of a mismatch."""
    from app.wallet import hd

    try:
        account = hd._load_account_key()
    except Exception as exc:
        raise ConfigError(
            f"the master key at {settings.master_key_path} cannot be loaded ({type(exc).__name__}): "
            "check MASTER_KEY_PATH and XPRV_ENCRYPTION_KEY_FILE"
        ) from exc
    problems = sample_problems(account, load_sample(), settings.xprv_encryption_key)
    if problems:
        raise ConfigError(
            "the master key does not match the database — refusing to serve: " + "; ".join(problems)
        )
    logger.info("Master key verified against the database sample")