plm-lottery/tests/unit/test_master_key_check.py

"""B-89 (SEC-49, SEC-50, SEC-09): the master key must be the one that produced the database."""

import hashlib
import os
import sqlite3
import stat

import pytest
from cryptography.fernet import Fernet
from embit import script

from app.config import ConfigError, settings
from app.wallet import hd
from app.wallet.keystore import encrypt_round_seed
from app.wallet.plm_network import PLM_MAINNET
from app.wallet.verify import load_sample, sample_problems, verify_master_key


def _database(tmp_path, monkeypatch, users=(), rounds=()):
    """A file DB with just the columns the check reads."""
    path = tmp_path / "check.db"
    db = sqlite3.connect(path)
    db.execute("CREATE TABLE users (derivation_index INTEGER, address TEXT)")
    db.execute(
        "CREATE TABLE rounds (id INTEGER, pool_address TEXT, server_seed_enc TEXT, server_seed TEXT, server_seed_hash TEXT)"
    )
    db.executemany("INSERT INTO users VALUES (?, ?)", users)
    db.executemany("INSERT INTO rounds VALUES (?, ?, ?, ?, ?)", rounds)
    db.commit()
    db.close()
    monkeypatch.setattr(settings, "database_url", f"sqlite+aiosqlite:///{path}")


def _addr(path):
    return script.p2wpkh(hd._load_account_key().derive(path).to_public()).address(network=PLM_MAINNET)


def test_a_fresh_instance_only_needs_the_key_to_decrypt(tmp_path, monkeypatch):
    monkeypatch.setattr(settings, "database_url", f"sqlite+aiosqlite:///{tmp_path}/absent.db")
    verify_master_key()


def test_the_matching_key_passes(tmp_path, monkeypatch):
    _database(
        tmp_path,
        monkeypatch,
        users=[(i, _addr(f"0/{i}")) for i in range(25)],
        rounds=[(1, _addr("1/0"), None, None, None), (2, _addr("2/2"), None, None, None)],
    )
    verify_master_key()


def test_another_master_key_is_refused_without_printing_it(tmp_path, monkeypatch):
    """SEC-49: the rows came from a different xprv (the two-keys incident)."""
    _database(tmp_path, monkeypatch, users=[(0, _addr("0/0")), (1, _addr("0/1"))])
    monkeypatch.setattr(settings, "master_key_path", str(tmp_path / "other.xprv.enc"))
    hd._account_key = None
    hd.generate_master_key()
    with pytest.raises(ConfigError, match="derivation index 0") as excinfo:
        verify_master_key()
    assert "xprv" not in str(excinfo.value).lower().replace("master key", "")


def test_only_the_edges_of_the_user_table_are_sampled(tmp_path, monkeypatch):
    users = [(i, _addr(f"0/{i}")) for i in range(40)]
    _database(tmp_path, monkeypatch, users=users)
    assert [i for i, _ in load_sample().users] == list(range(10)) + list(range(30, 40))


def test_a_round_pool_address_must_be_its_own_or_the_reserve(tmp_path, monkeypatch):
    _database(tmp_path, monkeypatch, rounds=[(5, _addr("2/9"), None, None, None)])
    with pytest.raises(ConfigError, match="round 5"):
        verify_master_key()


def test_a_wrong_encryption_key_is_refused_at_startup(tmp_path, monkeypatch):
    monkeypatch.setattr(settings, "xprv_encryption_key", Fernet.generate_key().decode())
    hd._account_key = None
    with pytest.raises(ConfigError, match="cannot be loaded"):
        verify_master_key()


def test_an_unrevealed_secret_must_decrypt_and_match_its_hash(tmp_path, monkeypatch):
    secret = "ab" * 32
    good = (1, "", encrypt_round_seed(secret), None, hashlib.sha256(bytes.fromhex(secret)).hexdigest())
    bad_hash = (2, "", encrypt_round_seed(secret), None, "00" * 32)
    revealed = (3, "", "not-a-token", secret, "00" * 32)  # already revealed: not checked
    _database(tmp_path, monkeypatch, rounds=[good, bad_hash, revealed])
    problems = sample_problems(hd._load_account_key(), load_sample(), settings.xprv_encryption_key)
    assert problems == ["secret of round 2 does not match its published hash"]


def test_the_key_file_is_written_0600(tmp_path, monkeypatch):
    path = tmp_path / "k.xprv.enc"
    monkeypatch.setattr(settings, "master_key_path", str(path))
    os.close(os.open(path, os.O_CREAT | os.O_WRONLY, 0o644))  # a leftover world-readable file
    hd.generate_master_key(overwrite=True)
    assert stat.S_IMODE(path.stat().st_mode) == 0o600


def test_the_import_script_refuses_an_xprv_that_does_not_match_the_database(tmp_path, monkeypatch):
    """SEC-05 / SEC-50: a Bitcoin-looking xprv passes the version check, the addresses do not."""
    from embit.bip32 import HDKey

    from app.wallet.verify import account_key

    _database(tmp_path, monkeypatch, users=[(0, _addr("0/0"))])
    other = HDKey.from_seed(b"\x09" * 32, version=PLM_MAINNET["xprv"]).to_base58(version=PLM_MAINNET["xprv"])
    assert sample_problems(account_key(other), load_sample())
    mine = hd._load_account_key()
    assert sample_problems(mine, load_sample()) == []