stacks_wire.py

"""
Stacks Transaction Wire Format Encoding

Implements Bitcoin wire formats for Stacks blockchain operations as defined in:
- SIP-001: Burn Election (https://github.com/stacksgov/sips/blob/main/sips/sip-001/sip-001-burn-election.md)
- SIP-007: Stacking Consensus

Wire format structures use OP_RETURN outputs with specific byte layouts.
"""

import struct
from enum import Enum
from dataclasses import dataclass
from typing import Optional


class StacksNetwork(Enum):
    """Network identifiers for Stacks magic bytes."""
    MAINNET = b'X2'
    TESTNET = b'T2'
    REGTEST = b'id'  # Legacy/regtest magic


class StacksOpcode(Enum):
    """
    Stacks burnchain operation opcodes (ASCII encoded).

    These single-byte opcodes identify the transaction type in the OP_RETURN payload.
    """
    LEADER_BLOCK_COMMIT = ord('[')      # 0x5B
    LEADER_KEY_REGISTER = ord('^')      # 0x5E
    USER_BURN_SUPPORT = ord('_')        # 0x5F
    PRE_STX = ord('p')                  # 0x70
    STACK_STX = ord('x')                # 0x78
    TRANSFER_STX = ord('$')             # 0x24
    DELEGATE_STX = ord('#')             # 0x23


# Constants
BURN_COMMITMENT_WINDOW = 6
MAX_OP_RETURN_SIZE = 80
CONSENSUS_HASH_LENGTH = 20
VRF_PUBLIC_KEY_LENGTH = 32
BLOCK_HASH_LENGTH = 32
SEED_LENGTH = 32


@dataclass
class LeaderBlockCommit:
    """
    Leader Block Commit transaction (SIP-001).

    Commits a leader's burn amount and chain tip preference.
    Requires two Bitcoin outputs:
    1. OP_RETURN with encoded data
    2. Burn output (sends BTC to burn address)

    Wire format (80 bytes):
    - 0-1:   magic (2 bytes)
    - 2:     opcode '[' (1 byte)
    - 3-34:  block_hash (32 bytes) - Stacks block header hash
    - 35-66: new_seed (32 bytes) - Next VRF seed
    - 67-70: parent_block (4 bytes) - Parent burn block height
    - 71-72: parent_txoff (2 bytes) - Parent transaction index
    - 73-76: key_block (4 bytes) - VRF key registration block height
    - 77-78: key_txoff (2 bytes) - VRF key registration tx index
    - 79:    burn_parent_modulus (1 byte) - burn_block_height % BURN_COMMITMENT_WINDOW
    """
    block_hash: bytes           # 32 bytes - Stacks anchored block header hash
    new_seed: bytes             # 32 bytes - VRF seed for sortition
    parent_block: int           # 4 bytes - Parent burn block height
    parent_txoff: int           # 2 bytes - Parent transaction offset
    key_block: int              # 4 bytes - Block height of VRF key registration
    key_txoff: int              # 2 bytes - Tx index of VRF key registration
    burn_parent_modulus: int    # 1 byte - Linkage for same-miner commits

    def encode(self, network: StacksNetwork = StacksNetwork.MAINNET) -> bytes:
        """Encode to OP_RETURN payload bytes."""
        if len(self.block_hash) != BLOCK_HASH_LENGTH:
            raise ValueError(f"block_hash must be {BLOCK_HASH_LENGTH} bytes")
        if len(self.new_seed) != SEED_LENGTH:
            raise ValueError(f"new_seed must be {SEED_LENGTH} bytes")

        payload = bytearray()
        payload.extend(network.value)                           # 2 bytes: magic
        payload.append(StacksOpcode.LEADER_BLOCK_COMMIT.value)  # 1 byte: opcode
        payload.extend(self.block_hash)                         # 32 bytes
        payload.extend(self.new_seed)                           # 32 bytes
        payload.extend(struct.pack('>I', self.parent_block))    # 4 bytes big-endian
        payload.extend(struct.pack('>H', self.parent_txoff))    # 2 bytes big-endian
        payload.extend(struct.pack('>I', self.key_block))       # 4 bytes big-endian
        payload.extend(struct.pack('>H', self.key_txoff))       # 2 bytes big-endian
        payload.append(self.burn_parent_modulus % BURN_COMMITMENT_WINDOW)  # 1 byte

        assert len(payload) == 80, f"Payload must be 80 bytes, got {len(payload)}"
        return bytes(payload)


@dataclass
class LeaderKeyRegister:
    """
    Leader VRF Key Registration transaction (SIP-001).

    Registers a cryptographic proving key used in sortition.
    Must be confirmed before submitting block commits.

    Wire format (80 bytes):
    - 0-1:   magic (2 bytes)
    - 2:     opcode '^' (1 byte)
    - 3-22:  consensus_hash (20 bytes) - Current burnchain consensus hash
    - 23-54: proving_public_key (32 bytes) - VRF verification key
    - 55-79: memo (25 bytes) - Arbitrary data, zero-padded
    """
    consensus_hash: bytes       # 20 bytes - Current Stacks consensus hash
    proving_public_key: bytes   # 32 bytes - VRF public key
    memo: bytes = b''           # Up to 25 bytes - Optional memo field

    def encode(self, network: StacksNetwork = StacksNetwork.MAINNET) -> bytes:
        """Encode to OP_RETURN payload bytes."""
        if len(self.consensus_hash) != CONSENSUS_HASH_LENGTH:
            raise ValueError(f"consensus_hash must be {CONSENSUS_HASH_LENGTH} bytes")
        if len(self.proving_public_key) != VRF_PUBLIC_KEY_LENGTH:
            raise ValueError(f"proving_public_key must be {VRF_PUBLIC_KEY_LENGTH} bytes")
        if len(self.memo) > 25:
            raise ValueError("memo must be at most 25 bytes")

        payload = bytearray()
        payload.extend(network.value)                           # 2 bytes: magic
        payload.append(StacksOpcode.LEADER_KEY_REGISTER.value)  # 1 byte: opcode
        payload.extend(self.consensus_hash)                     # 20 bytes
        payload.extend(self.proving_public_key)                 # 32 bytes
        payload.extend(self.memo.ljust(25, b'\x00'))            # 25 bytes, zero-padded

        assert len(payload) == 80, f"Payload must be 80 bytes, got {len(payload)}"
        return bytes(payload)


@dataclass
class UserBurnSupport:
    """
    User Support Burns transaction (SIP-001).

    Allows users to contribute burns to preferred chain tips.

    Wire format (80 bytes):
    - 0-1:   magic (2 bytes)
    - 2:     opcode '_' (1 byte)
    - 3-22:  consensus_hash (20 bytes) - Truncated consensus hash
    - 23-54: proving_public_key (32 bytes) - Supported miner's VRF key
    - 55-74: block_hash_160 (20 bytes) - Truncated block hash (HASH160)
    - 75-77: key_block (3 bytes) - VRF registration block (24-bit)
    - 78-79: key_vtxindex (2 bytes) - VRF registration tx index
    """
    consensus_hash: bytes       # 20 bytes
    proving_public_key: bytes   # 32 bytes - Miner's VRF key to support
    block_hash_160: bytes       # 20 bytes - HASH160 of supported block
    key_block: int              # 3 bytes (24-bit) - Miner's key registration block
    key_vtxindex: int           # 2 bytes

    def encode(self, network: StacksNetwork = StacksNetwork.MAINNET) -> bytes:
        """Encode to OP_RETURN payload bytes."""
        if len(self.consensus_hash) != CONSENSUS_HASH_LENGTH:
            raise ValueError(f"consensus_hash must be {CONSENSUS_HASH_LENGTH} bytes")
        if len(self.proving_public_key) != VRF_PUBLIC_KEY_LENGTH:
            raise ValueError(f"proving_public_key must be {VRF_PUBLIC_KEY_LENGTH} bytes")
        if len(self.block_hash_160) != 20:
            raise ValueError("block_hash_160 must be 20 bytes")

        payload = bytearray()
        payload.extend(network.value)                         # 2 bytes: magic
        payload.append(StacksOpcode.USER_BURN_SUPPORT.value)  # 1 byte: opcode
        payload.extend(self.consensus_hash)                   # 20 bytes
        payload.extend(self.proving_public_key)               # 32 bytes
        payload.extend(self.block_hash_160)                   # 20 bytes
        # key_block as 3 bytes big-endian (24-bit)
        payload.extend(struct.pack('>I', self.key_block)[1:]) # 3 bytes
        payload.extend(struct.pack('>H', self.key_vtxindex))  # 2 bytes

        assert len(payload) == 80, f"Payload must be 80 bytes, got {len(payload)}"
        return bytes(payload)


@dataclass
class PreStx:
    """
    Pre-STX Authorization transaction (SIP-007).

    Pre-authorizes a Stacks address for subsequent STX operations.
    The first input's address becomes the authorized Stacks address.

    Wire format (3 bytes minimum):
    - 0-1: magic (2 bytes)
    - 2:   opcode 'p' (1 byte)

    Note: This is a minimal OP_RETURN. The authorization comes from
    the spending address in the first input.
    """

    def encode(self, network: StacksNetwork = StacksNetwork.MAINNET) -> bytes:
        """Encode to OP_RETURN payload bytes."""
        payload = bytearray()
        payload.extend(network.value)               # 2 bytes: magic
        payload.append(StacksOpcode.PRE_STX.value)  # 1 byte: opcode
        return bytes(payload)


@dataclass
class StackStx:
    """
    Stack STX transaction (SIP-007).

    Locks STX tokens for Proof-of-Transfer (PoX) participation.

    Wire format (20 bytes):
    - 0-1:   magic (2 bytes)
    - 2:     opcode 'x' (1 byte)
    - 3-18:  stacked_ustx (16 bytes) - Amount in micro-STX (u128 big-endian)
    - 19:    num_cycles (1 byte) - Number of reward cycles to lock
    """
    stacked_ustx: int   # u128 - Amount of micro-STX to lock
    num_cycles: int     # u8 - Number of reward cycles (1-12)

    def encode(self, network: StacksNetwork = StacksNetwork.MAINNET) -> bytes:
        """Encode to OP_RETURN payload bytes."""
        if self.num_cycles < 1 or self.num_cycles > 12:
            raise ValueError("num_cycles must be between 1 and 12")
        if self.stacked_ustx < 0 or self.stacked_ustx >= 2**128:
            raise ValueError("stacked_ustx must be a valid u128")

        payload = bytearray()
        payload.extend(network.value)                # 2 bytes: magic
        payload.append(StacksOpcode.STACK_STX.value) # 1 byte: opcode
        # u128 big-endian (16 bytes)
        payload.extend(self.stacked_ustx.to_bytes(16, byteorder='big'))
        payload.append(self.num_cycles)              # 1 byte

        assert len(payload) == 20, f"Payload must be 20 bytes, got {len(payload)}"
        return bytes(payload)


@dataclass
class TransferStx:
    """
    Transfer STX transaction (SIP-007).

    Transfers STX between addresses via Bitcoin transaction.

    Wire format (19 bytes):
    - 0-1:   magic (2 bytes)
    - 2:     opcode '$' (1 byte)
    - 3-18:  amount_ustx (16 bytes) - Amount in micro-STX (u128 big-endian)

    The recipient is encoded in the second output (after OP_RETURN).
    """
    amount_ustx: int    # u128 - Amount of micro-STX to transfer

    def encode(self, network: StacksNetwork = StacksNetwork.MAINNET) -> bytes:
        """Encode to OP_RETURN payload bytes."""
        if self.amount_ustx < 0 or self.amount_ustx >= 2**128:
            raise ValueError("amount_ustx must be a valid u128")

        payload = bytearray()
        payload.extend(network.value)                    # 2 bytes: magic
        payload.append(StacksOpcode.TRANSFER_STX.value)  # 1 byte: opcode
        # u128 big-endian (16 bytes)
        payload.extend(self.amount_ustx.to_bytes(16, byteorder='big'))

        assert len(payload) == 19, f"Payload must be 19 bytes, got {len(payload)}"
        return bytes(payload)


@dataclass
class DelegateStx:
    """
    Delegate STX transaction (SIP-007).

    Delegates STX to a pool operator for stacking.

    Wire format (19 bytes):
    - 0-1:   magic (2 bytes)
    - 2:     opcode '#' (1 byte)
    - 3-18:  amount_ustx (16 bytes) - Amount in micro-STX (u128 big-endian)

    The delegate address is encoded in the second output.
    """
    amount_ustx: int    # u128 - Amount of micro-STX to delegate

    def encode(self, network: StacksNetwork = StacksNetwork.MAINNET) -> bytes:
        """Encode to OP_RETURN payload bytes."""
        if self.amount_ustx < 0 or self.amount_ustx >= 2**128:
            raise ValueError("amount_ustx must be a valid u128")

        payload = bytearray()
        payload.extend(network.value)                    # 2 bytes: magic
        payload.append(StacksOpcode.DELEGATE_STX.value)  # 1 byte: opcode
        # u128 big-endian (16 bytes)
        payload.extend(self.amount_ustx.to_bytes(16, byteorder='big'))

        assert len(payload) == 19, f"Payload must be 19 bytes, got {len(payload)}"
        return bytes(payload)


def create_op_return_script(payload: bytes) -> bytes:
    """
    Create a Bitcoin OP_RETURN script from payload data.

    Script format: OP_RETURN <push_data>

    Args:
        payload: The data to embed (max 80 bytes)

    Returns:
        Complete scriptPubKey bytes
    """
    if len(payload) > MAX_OP_RETURN_SIZE:
        raise ValueError(f"Payload exceeds maximum OP_RETURN size of {MAX_OP_RETURN_SIZE} bytes")

    OP_RETURN = 0x6a

    script = bytearray()
    script.append(OP_RETURN)

    # Push data opcode
    if len(payload) <= 75:
        script.append(len(payload))  # Direct push
    elif len(payload) <= 255:
        script.append(0x4c)  # OP_PUSHDATA1
        script.append(len(payload))
    else:
        raise ValueError("Payload too large for OP_RETURN")

    script.extend(payload)
    return bytes(script)


# Canonical burn address for Stacks operations
STACKS_BURN_ADDRESS_MAINNET = "1111111111111111111114oLvT2"
STACKS_BURN_ADDRESS_TESTNET = "mfWxJ45yp2SFn7UciZyNpvDKrzbhyfKrY8"


def get_burn_address(network: StacksNetwork) -> str:
    """Get the canonical burn address for the given network."""
    if network == StacksNetwork.MAINNET:
        return STACKS_BURN_ADDRESS_MAINNET
    else:
        return STACKS_BURN_ADDRESS_TESTNET


# =============================================================================
# C32Check Encoding for STX Addresses
# =============================================================================

# C32 alphabet (Crockford base32 variant)
C32_ALPHABET = "0123456789ABCDEFGHJKMNPQRSTVWXYZ"

# STX address version bytes
STX_ADDRESS_VERSION_MAINNET_SINGLESIG = 22   # 0x16 -> "SP"
STX_ADDRESS_VERSION_MAINNET_MULTISIG = 20    # 0x14 -> "SM"
STX_ADDRESS_VERSION_TESTNET_SINGLESIG = 26   # 0x1a -> "ST"
STX_ADDRESS_VERSION_TESTNET_MULTISIG = 21    # 0x15 -> "SN"


def _c32_encode(data: bytes) -> str:
    """
    Encode bytes using c32 (Crockford base32 variant).
    """
    if len(data) == 0:
        return ""

    # Convert bytes to a single large integer
    num = int.from_bytes(data, 'big')

    if num == 0:
        return C32_ALPHABET[0] * len(data)

    result = []
    while num > 0:
        result.append(C32_ALPHABET[num % 32])
        num //= 32

    # Add leading zeros
    for byte in data:
        if byte == 0:
            result.append(C32_ALPHABET[0])
        else:
            break

    return ''.join(reversed(result))


def _c32_checksum(version: int, data: bytes) -> bytes:
    """
    Calculate c32check checksum.

    The checksum is the first 4 bytes of SHA256(SHA256(version + data)).
    """
    import hashlib
    payload = bytes([version]) + data
    hash1 = hashlib.sha256(payload).digest()
    hash2 = hashlib.sha256(hash1).digest()
    return hash2[:4]


def _c32_version_char(version: int) -> str:
    """Get the c32 character for a version byte."""
    return C32_ALPHABET[version]


def c32check_encode(version: int, data: bytes) -> str:
    """
    Encode data with c32check (version byte + data + checksum).

    Args:
        version: Version byte (e.g., 22 for mainnet singlesig)
        data: The hash160 (20 bytes) to encode

    Returns:
        c32check encoded string (without 'S' prefix)
    """
    checksum = _c32_checksum(version, data)
    # Encode version separately, then data+checksum together
    version_char = _c32_version_char(version)
    encoded_data = _c32_encode(data + checksum)
    return version_char + encoded_data


def hash160_to_stx_address(hash160: bytes, network: StacksNetwork = StacksNetwork.MAINNET,
                           is_multisig: bool = False) -> str:
    """
    Convert a hash160 (public key hash) to a STX address.

    Args:
        hash160: 20-byte public key hash (RIPEMD160(SHA256(pubkey)))
        network: Stacks network (MAINNET or TESTNET)
        is_multisig: Whether this is a multisig address

    Returns:
        STX address string (e.g., "SP2J6ZY48GV1EZ5V2V5RB9MP66SW86PYKKNRV9EJ7")
    """
    if len(hash160) != 20:
        raise ValueError("hash160 must be 20 bytes")

    if network == StacksNetwork.MAINNET:
        version = STX_ADDRESS_VERSION_MAINNET_MULTISIG if is_multisig else STX_ADDRESS_VERSION_MAINNET_SINGLESIG
    else:
        version = STX_ADDRESS_VERSION_TESTNET_MULTISIG if is_multisig else STX_ADDRESS_VERSION_TESTNET_SINGLESIG

    c32_encoded = c32check_encode(version, hash160)
    return "S" + c32_encoded


def bitcoin_address_to_stx_address(btc_address: str, stx_network: StacksNetwork = StacksNetwork.MAINNET) -> str:
    """
    Convert a Bitcoin address to its corresponding STX address.

    Both addresses share the same underlying public key hash but use different encodings.
    Supports:
    - Legacy P2PKH (starts with 1 or m/n for testnet)
    - P2SH (starts with 3 or 2 for testnet)
    - Native SegWit P2WPKH (starts with bc1q or tb1q)

    Note: P2WSH (bc1q with 32-byte program) and P2TR (bc1p) are not supported
    as they use different hash constructions.

    Args:
        btc_address: Bitcoin address
        stx_network: Target Stacks network

    Returns:
        Corresponding STX address

    Raises:
        ValueError: If address type is not supported
    """
    from electrum import constants

    # Check if it's a bech32/bech32m address (native SegWit)
    if btc_address.lower().startswith(('bc1', 'tb1', 'bcrt1')):
        from electrum.segwit_addr import decode_segwit_address

        # Determine the HRP (human-readable part) based on address prefix
        if btc_address.lower().startswith('bc1'):
            hrp = 'bc'
        elif btc_address.lower().startswith('tb1'):
            hrp = 'tb'
        else:
            hrp = 'bcrt'

        result = decode_segwit_address(hrp, btc_address)
        if result is None:
            raise ValueError("Invalid bech32 address")

        witver, witprog = result

        # P2WPKH: witness version 0, 20-byte program (the pubkey hash)
        if witver == 0 and len(witprog) == 20:
            hash160 = bytes(witprog)
            return hash160_to_stx_address(hash160, stx_network, is_multisig=False)

        # P2WSH: witness version 0, 32-byte program (script hash) - not directly convertible
        elif witver == 0 and len(witprog) == 32:
            raise ValueError("P2WSH addresses cannot be converted to STX addresses (different hash)")

        # P2TR: witness version 1 (Taproot) - not directly convertible
        elif witver == 1:
            raise ValueError("Taproot (P2TR) addresses cannot be converted to STX addresses")

        else:
            raise ValueError(f"Unsupported witness version {witver}")

    # Legacy base58 address (P2PKH or P2SH)
    else:
        from electrum.bitcoin import b58_address_to_hash160

        try:
            addr_type, hash160 = b58_address_to_hash160(btc_address)
        except Exception as e:
            raise ValueError(f"Cannot decode Bitcoin address: {e}")

        # Determine if multisig based on address type
        # P2PKH (legacy) -> singlesig, P2SH -> could be multisig
        is_multisig = (addr_type == constants.net.ADDRTYPE_P2SH)

        return hash160_to_stx_address(hash160, stx_network, is_multisig)


def pubkey_to_stx_address(pubkey: bytes, network: StacksNetwork = StacksNetwork.MAINNET) -> str:
    """
    Convert a public key to its STX address.

    Args:
        pubkey: Compressed or uncompressed public key bytes
        network: Stacks network

    Returns:
        STX address string
    """
    import hashlib

    # hash160 = RIPEMD160(SHA256(pubkey))
    sha256_hash = hashlib.sha256(pubkey).digest()
    ripemd160 = hashlib.new('ripemd160')
    ripemd160.update(sha256_hash)
    hash160 = ripemd160.digest()

    return hash160_to_stx_address(hash160, network, is_multisig=False)