btc_tools.py

# -*- coding: utf-8 -*-
"""
Python Bitcoin Utils CLI — key generation, address derivation, raw transaction builder.
Provides low-level Bitcoin primitives: ECDSA key pairs, all address formats
(P2PKH, P2SH-P2WPKH, P2WPKH, P2TR), and a minimal raw transaction constructor.
"""
import hashlib
import hmac
import os
import struct
from dataclasses import dataclass
from typing import Optional

from config import BtcConfig, load_config, NETWORKS


SECP256K1_ORDER = 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEBAAEDCE6AF48A03BBFD25E8CD0364141
SECP256K1_GX = 0x79BE667EF9DCBBAC55A06295CE870B07029BFCDB2DCE28D959F2815B16F81798
SECP256K1_GY = 0x483ADA7726A3C4655DA4FBFC0E1108A8FD17B448A68554199C47D08FFB10D4B8
SECP256K1_P = 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEFFFFFC2F

BASE58_ALPHABET = "123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz"
BECH32_CHARSET = "qpzry9x8gf2tvdw0s3jn54khce6mua7l"


def _sha256(data: bytes) -> bytes:
    return hashlib.sha256(data).digest()


def _hash160(data: bytes) -> bytes:
    return hashlib.new("ripemd160", _sha256(data)).digest()


def _double_sha256(data: bytes) -> bytes:
    return _sha256(_sha256(data))


def _base58_encode(payload: bytes) -> str:
    n = int.from_bytes(payload, "big")
    result = []
    while n > 0:
        n, r = divmod(n, 58)
        result.append(BASE58_ALPHABET[r])
    for byte in payload:
        if byte == 0:
            result.append(BASE58_ALPHABET[0])
        else:
            break
    return "".join(reversed(result))


def _base58check_encode(version: int, payload: bytes) -> str:
    versioned = bytes([version]) + payload
    checksum = _double_sha256(versioned)[:4]
    return _base58_encode(versioned + checksum)


def _bech32_polymod(values: list[int]) -> int:
    gen = [0x3b6a57b2, 0x26508e6d, 0x1ea119fa, 0x3d4233dd, 0x2a1462b3]
    chk = 1
    for v in values:
        b = chk >> 25
        chk = ((chk & 0x1ffffff) << 5) ^ v
        for i in range(5):
            chk ^= gen[i] if ((b >> i) & 1) else 0
    return chk


def _bech32_hrp_expand(hrp: str) -> list[int]:
    return [ord(c) >> 5 for c in hrp] + [0] + [ord(c) & 31 for c in hrp]


def _bech32_create_checksum(hrp: str, data: list[int]) -> list[int]:
    values = _bech32_hrp_expand(hrp) + data
    polymod = _bech32_polymod(values + [0, 0, 0, 0, 0, 0]) ^ 1
    return [(polymod >> 5 * (5 - i)) & 31 for i in range(6)]


def _convertbits(data: bytes, frombits: int, tobits: int, pad: bool = True) -> list[int]:
    acc = 0
    bits = 0
    ret = []
    maxv = (1 << tobits) - 1
    for value in data:
        acc = (acc << frombits) | value
        bits += frombits
        while bits >= tobits:
            bits -= tobits
            ret.append((acc >> bits) & maxv)
    if pad and bits:
        ret.append((acc << (tobits - bits)) & maxv)
    return ret


def bech32_encode(hrp: str, witver: int, witprog: bytes) -> str:
    """Encode a witness program into a bech32/bech32m address."""
    data = [witver] + _convertbits(witprog, 8, 5)
    checksum = _bech32_create_checksum(hrp, data)
    return hrp + "1" + "".join(BECH32_CHARSET[d] for d in data + checksum)


@dataclass
class KeyPair:
    private_key: bytes
    public_key_compressed: bytes
    public_key_uncompressed: bytes
    wif: str
    network: str

    @property
    def private_hex(self) -> str:
        return self.private_key.hex()

    @property
    def public_hex(self) -> str:
        return self.public_key_compressed.hex()


def _point_add(px: int, py: int, qx: int, qy: int) -> tuple[int, int]:
    """Elliptic curve point addition on secp256k1."""
    if px == 0 and py == 0:
        return qx, qy
    if qx == 0 and qy == 0:
        return px, py
    if px == qx and py != qy:
        return 0, 0
    if px == qx:
        lam = (3 * px * px) * pow(2 * py, SECP256K1_P - 2, SECP256K1_P) % SECP256K1_P
    else:
        lam = (qy - py) * pow(qx - px, SECP256K1_P - 2, SECP256K1_P) % SECP256K1_P
    rx = (lam * lam - px - qx) % SECP256K1_P
    ry = (lam * (px - rx) - py) % SECP256K1_P
    return rx, ry


def _scalar_mult(k: int, px: int, py: int) -> tuple[int, int]:
    """Double-and-add scalar multiplication."""
    rx, ry = 0, 0
    while k > 0:
        if k & 1:
            rx, ry = _point_add(rx, ry, px, py)
        px, py = _point_add(px, py, px, py)
        k >>= 1
    return rx, ry


def generate_keypair(cfg: Optional[BtcConfig] = None) -> KeyPair:
    """Generate a new ECDSA secp256k1 key pair."""
    if cfg is None:
        cfg = load_config()

    while True:
        priv_bytes = os.urandom(32)
        priv_int = int.from_bytes(priv_bytes, "big")
        if 0 < priv_int < SECP256K1_ORDER:
            break

    pub_x, pub_y = _scalar_mult(priv_int, SECP256K1_GX, SECP256K1_GY)
    prefix = b"\x02" if pub_y % 2 == 0 else b"\x03"
    compressed = prefix + pub_x.to_bytes(32, "big")
    uncompressed = b"\x04" + pub_x.to_bytes(32, "big") + pub_y.to_bytes(32, "big")

    net_params = NETWORKS.get(cfg.network, NETWORKS["testnet"])
    wif_version = net_params["wif_prefix"]
    wif = _base58check_encode(wif_version, priv_bytes + b"\x01")

    return KeyPair(
        private_key=priv_bytes,
        public_key_compressed=compressed,
        public_key_uncompressed=uncompressed,
        wif=wif,
        network=cfg.network,
    )


@dataclass
class DerivedAddresses:
    p2pkh: str
    p2sh_p2wpkh: str
    p2wpkh: str
    p2tr: str
    network: str


def derive_addresses(pubkey_compressed: bytes,
                     cfg: Optional[BtcConfig] = None) -> DerivedAddresses:
    """Derive all standard address formats from a compressed public key."""
    if cfg is None:
        cfg = load_config()

    net = NETWORKS.get(cfg.network, NETWORKS["testnet"])
    h160 = _hash160(pubkey_compressed)

    p2pkh = _base58check_encode(net["prefix_p2pkh"], h160)

    witness_program = b"\x00\x14" + h160
    script_hash = _hash160(witness_program)
    p2sh = _base58check_encode(net["prefix_p2sh"], script_hash)

    p2wpkh = bech32_encode(net["bech32_hrp"], 0, h160)

    x_only = pubkey_compressed[1:]
    tweaked = _sha256(x_only)
    p2tr = bech32_encode(net["bech32_hrp"], 1, tweaked)

    return DerivedAddresses(
        p2pkh=p2pkh,
        p2sh_p2wpkh=p2sh,
        p2wpkh=p2wpkh,
        p2tr=p2tr,
        network=cfg.network,
    )


@dataclass
class TxInput:
    txid: str
    vout: int
    script_sig: bytes = b""
    sequence: int = 0xFFFFFFFD

    def serialize(self) -> bytes:
        txid_bytes = bytes.fromhex(self.txid)[::-1]
        result = txid_bytes
        result += struct.pack("<I", self.vout)
        result += _varint(len(self.script_sig))
        result += self.script_sig
        result += struct.pack("<I", self.sequence)
        return result


@dataclass
class TxOutput:
    value_sat: int
    script_pubkey: bytes

    def serialize(self) -> bytes:
        result = struct.pack("<q", self.value_sat)
        result += _varint(len(self.script_pubkey))
        result += self.script_pubkey
        return result


def _varint(n: int) -> bytes:
    if n < 0xFD:
        return struct.pack("<B", n)
    if n <= 0xFFFF:
        return b"\xfd" + struct.pack("<H", n)
    if n <= 0xFFFFFFFF:
        return b"\xfe" + struct.pack("<I", n)
    return b"\xff" + struct.pack("<Q", n)


@dataclass
class RawTransaction:
    version: int = 2
    inputs: list = None
    outputs: list = None
    locktime: int = 0

    def __post_init__(self):
        if self.inputs is None:
            self.inputs = []
        if self.outputs is None:
            self.outputs = []

    def add_input(self, txid: str, vout: int, sequence: int = 0xFFFFFFFD) -> None:
        self.inputs.append(TxInput(txid=txid, vout=vout, sequence=sequence))

    def add_output(self, value_sat: int, script_pubkey: bytes) -> None:
        self.outputs.append(TxOutput(value_sat=value_sat, script_pubkey=script_pubkey))

    def serialize_unsigned(self) -> bytes:
        """Serialize the transaction without witness data (for signing)."""
        result = struct.pack("<I", self.version)
        result += _varint(len(self.inputs))
        for inp in self.inputs:
            result += inp.serialize()
        result += _varint(len(self.outputs))
        for out in self.outputs:
            result += out.serialize()
        result += struct.pack("<I", self.locktime)
        return result

    def txid(self) -> str:
        """Compute the transaction ID (double SHA-256, reversed)."""
        raw = self.serialize_unsigned()
        return _double_sha256(raw)[::-1].hex()

    @property
    def virtual_size(self) -> int:
        raw = self.serialize_unsigned()
        return len(raw)

    def fee_estimate(self, fee_rate_sat_vb: int) -> int:
        return self.virtual_size * fee_rate_sat_vb


def p2pkh_script(address_hash160: bytes) -> bytes:
    """Build a P2PKH scriptPubKey: OP_DUP OP_HASH160 <hash> OP_EQUALVERIFY OP_CHECKSIG."""
    return b"\x76\xa9\x14" + address_hash160 + b"\x88\xac"


def p2wpkh_script(pubkey_hash: bytes) -> bytes:
    """Build a P2WPKH scriptPubKey: OP_0 <20-byte-hash>."""
    return b"\x00\x14" + pubkey_hash