plm-lottery/app/wallet/hd.py

import os

from embit import script
from embit.bip32 import HDKey
from embit.ec import PrivateKey

from app.config import settings
from app.wallet.keystore import decrypt_xprv, encrypt_xprv
from app.wallet.plm_network import ACCOUNT_PATH, PLM_MAINNET

_account_key: HDKey | None = None


def generate_master_key(overwrite: bool = False) -> None:
    """One-time ops bootstrap: create a random master seed, encrypt it, write it to
    disk. Not exposed via any API endpoint — run manually before first launch."""
    if os.path.exists(settings.master_key_path) and not overwrite:
        raise FileExistsError(f"{settings.master_key_path} already exists")
    root = HDKey.from_seed(os.urandom(32), version=PLM_MAINNET["xprv"])
    with open(settings.master_key_path, "wb") as f:
        f.write(encrypt_xprv(root.to_base58(version=PLM_MAINNET["xprv"])))


def _load_account_key() -> HDKey:
    global _account_key
    if _account_key is None:
        with open(settings.master_key_path, "rb") as f:
            token = f.read()
        root = HDKey.from_base58(decrypt_xprv(token))
        _account_key = root.derive(ACCOUNT_PATH)
    return _account_key


def derive_user_key(derivation_index: int) -> HDKey:
    return _load_account_key().derive(f"0/{derivation_index}")


def derive_user_address(derivation_index: int) -> str:
    pub = derive_user_key(derivation_index).to_public()
    return script.p2wpkh(pub).address(network=PLM_MAINNET)


def derive_user_wif(derivation_index: int) -> str:
    """Exports a user's raw private key (WIF) for manual server-side intervention
    (e.g. sweeping funds back to a user, or out, if something gets stuck). This is
    a custodial system — the server already holds the master key this is derived
    from — but callers must still treat the result as a live secret: log access,
    never persist it, never return it over an unauthenticated channel."""
    key = derive_user_key(derivation_index)
    return PrivateKey(key.secret, compressed=True, network=PLM_MAINNET).wif(network=PLM_MAINNET)


def derive_pool_key(round_id: int | None = None) -> HDKey:
    """A pool key. With a `round_id`, that round's own key on branch 2 of the account
    (`m/84'/746'/0'/2/<round_id>`, B-80, RND-35): its bets are paid to it and its
    payout is signed with it, so one round's money and history never mix with another's.
    Without one, the **reserve** key at branch 1, index 0 — the single pool address every
    round used before B-80, still the pool of the rounds opened back then and the
    destination of the payout's bump reserve. Branch 0 is users. Not a spec requirement
    of the coin, an implementation choice to keep every key under the one encrypted
    master key rather than in separate secrets."""
    if round_id is None:
        return _load_account_key().derive("1/0")
    if round_id < 1:
        raise ValueError(f"invalid round id {round_id}")
    return _load_account_key().derive(f"2/{round_id}")


def derive_pool_address(round_id: int | None = None) -> str:
    """The pool address of a round (B-80), or the reserve address without a `round_id`."""
    pub = derive_pool_key(round_id).to_public()
    return script.p2wpkh(pub).address(network=PLM_MAINNET)


def derive_round_pool_key(round_id: int, pool_address: str) -> HDKey:
    """The key that controls `pool_address`, the address stored on round `round_id`.

    A round opened before B-80 stores the reserve address and is signed with the reserve
    key; every other round with its own. Checking the derived address against the stored
    one means a key can never be used for money it does not control, and a mismatch (a
    different master key, a corrupted row) fails loudly instead of building a transaction
    that spends someone else's coins."""
    for key in (derive_pool_key(round_id), derive_pool_key()):
        if script.p2wpkh(key.to_public()).address(network=PLM_MAINNET) == pool_address:
            return key
    raise ValueError(f"round {round_id}: no pool key controls the stored address {pool_address!r}")