plm-lottery/app/wallet/psbt_builder.py

from dataclasses import dataclass

from embit import script
from embit.bip32 import HDKey
from embit.finalizer import finalize_psbt
from embit.psbt import PSBT
from embit.transaction import Transaction, TransactionInput, TransactionOutput

# Standard P2WPKH size estimates (vbytes): 10.5-byte overhead (version+counts+locktime+
# segwit marker/flag), ~68 vbytes per input, ~31 vbytes per output. Used to size the fee
# before signing (fee only needs to be "minimized ~1 sat/vB", not maximally precise).
_TX_OVERHEAD_VBYTES = 11
_P2WPKH_INPUT_VBYTES = 68
_P2WPKH_OUTPUT_VBYTES = 31

# BIP125 opt-in RBF: any sequence < 0xfffffffe signals replaceability. Set on every
# input we create so a stuck tx can later be fee-bumped (tx/broadcast.py, stage 9).
RBF_SEQUENCE = 0xFFFFFFFD


class InsufficientFundsError(Exception):
    pass


@dataclass
class Utxo:
    txid: str
    vout: int
    amount_sats: int


@dataclass
class BuiltTransaction:
    raw_hex: str
    txid: str
    fee_sats: int
    recipient_sats: int
    change_sats: int
    spent_utxos: list[Utxo]


def estimate_vsize(n_inputs: int, n_outputs: int) -> int:
    return _TX_OVERHEAD_VBYTES + n_inputs * _P2WPKH_INPUT_VBYTES + n_outputs * _P2WPKH_OUTPUT_VBYTES


def select_utxos(utxos: list[Utxo], target_sats: int) -> tuple[list[Utxo], int]:
    """Greedily select UTXOs (largest first, to minimize input count) covering
    target_sats — the amount deducted from the sender's balance. The fee is paid
    out of target_sats (see build_signed_transaction), not added on top of it."""
    ordered = sorted(utxos, key=lambda u: u.amount_sats, reverse=True)
    selected: list[Utxo] = []
    total = 0
    for utxo in ordered:
        selected.append(utxo)
        total += utxo.amount_sats
        if total >= target_sats:
            return selected, total
    raise InsufficientFundsError("not enough confirmed balance to cover amount")


def build_signed_transaction(
    *,
    signing_key: HDKey,
    from_script: script.Script,
    utxos: list[Utxo],
    to_address: str,
    amount_sats: int,
    change_address: str,
    fee_rate_sat_vb: int,
) -> BuiltTransaction:
    """Build, sign and finalize a single-recipient P2WPKH transaction with change
    back to change_address.

    `amount_sats` is deducted from the sender's balance in full: the recipient
    receives `amount_sats - fee`, change = total_in - amount_sats. This matches the
    spec everywhere a single-recipient tx is used (bet, withdrawal): "fee deducted
    from the amount being moved", not paid on top by the sender.
    """
    selected, total_in = select_utxos(utxos, amount_sats)
    fee = estimate_vsize(len(selected), 2) * fee_rate_sat_vb
    recipient_amount = amount_sats - fee
    if recipient_amount <= 0:
        raise InsufficientFundsError("amount too small to cover the network fee")
    change = total_in - amount_sats

    # TransactionInput.txid is natural/display byte order (as in tx_hash from Electrum);
    # embit reverses it internally when serializing to wire format.
    vin = [TransactionInput(bytes.fromhex(u.txid), u.vout, sequence=RBF_SEQUENCE) for u in selected]
    vout = [TransactionOutput(recipient_amount, script.Script.from_address(to_address))]
    if change > 0:
        vout.append(TransactionOutput(change, script.Script.from_address(change_address)))

    tx = Transaction(vin=vin, vout=vout)
    psbt = PSBT(tx)
    for i, utxo in enumerate(selected):
        psbt.inputs[i].witness_utxo = TransactionOutput(utxo.amount_sats, from_script)

    signed_count = psbt.sign_with(signing_key)
    if signed_count != len(selected):
        raise RuntimeError(f"expected {len(selected)} signatures, got {signed_count}")

    final_tx = finalize_psbt(psbt)
    if final_tx is None:
        raise RuntimeError("failed to finalize PSBT")

    raw = final_tx.serialize()
    return BuiltTransaction(
        raw_hex=raw.hex(),
        txid=final_tx.txid().hex(),
        fee_sats=fee,
        recipient_sats=recipient_amount,
        change_sats=change,
        spent_utxos=selected,
    )


@dataclass
class PayoutTransaction:
    raw_hex: str
    txid: str
    fee_sats: int
    winner_sats: int
    commission_sats: int
    change_sats: int
    spent_utxos: list[Utxo]


def build_payout_transaction(
    *,
    signing_key: HDKey,
    from_script: script.Script,
    utxos: list[Utxo],
    winner_address: str,
    winner_share_sats: int,
    fee_address: str,
    commission_sats: int,
    change_address: str,
    fee_rate_sat_vb: int,
) -> PayoutTransaction:
    """Build, sign and finalize the round payout: pool -> winner + fee address,
    with change back to the pool itself. Per spec, only the winner's share
    absorbs the tx fee — the commission (fee_address) output is untouched."""
    target = winner_share_sats + commission_sats
    selected, total_in = select_utxos(utxos, target)
    fee = estimate_vsize(len(selected), 3) * fee_rate_sat_vb  # winner + commission + pool change
    winner_amount = winner_share_sats - fee
    if winner_amount <= 0:
        raise InsufficientFundsError("winner share too small to cover the network fee")
    change = total_in - target

    vin = [TransactionInput(bytes.fromhex(u.txid), u.vout, sequence=RBF_SEQUENCE) for u in selected]
    vout = [
        TransactionOutput(winner_amount, script.Script.from_address(winner_address)),
        TransactionOutput(commission_sats, script.Script.from_address(fee_address)),
    ]
    if change > 0:
        vout.append(TransactionOutput(change, script.Script.from_address(change_address)))

    tx = Transaction(vin=vin, vout=vout)
    psbt = PSBT(tx)
    for i, utxo in enumerate(selected):
        psbt.inputs[i].witness_utxo = TransactionOutput(utxo.amount_sats, from_script)

    signed_count = psbt.sign_with(signing_key)
    if signed_count != len(selected):
        raise RuntimeError(f"expected {len(selected)} signatures, got {signed_count}")

    final_tx = finalize_psbt(psbt)
    if final_tx is None:
        raise RuntimeError("failed to finalize PSBT")

    raw = final_tx.serialize()
    return PayoutTransaction(
        raw_hex=raw.hex(),
        txid=final_tx.txid().hex(),
        fee_sats=fee,
        winner_sats=winner_amount,
        commission_sats=commission_sats,
        change_sats=change,
        spent_utxos=selected,
    )