test_stacks_wire.py

"""
Unit tests for Stacks wire format encoding.

Run with: python -m pytest test_stacks_wire.py -v
"""

import pytest
from stacks_wire import (
    StacksNetwork, StacksOpcode,
    LeaderBlockCommit, LeaderKeyRegister, UserBurnSupport,
    PreStx, StackStx, TransferStx, DelegateStx,
    create_op_return_script, get_burn_address,
    BURN_COMMITMENT_WINDOW,
    # C32 encoding functions
    C32_ALPHABET,
    _c32_encode,
    _c32_checksum,
    c32check_encode,
    hash160_to_stx_address,
    pubkey_to_stx_address,
    STX_ADDRESS_VERSION_MAINNET_SINGLESIG,
    STX_ADDRESS_VERSION_MAINNET_MULTISIG,
    STX_ADDRESS_VERSION_TESTNET_SINGLESIG,
    STX_ADDRESS_VERSION_TESTNET_MULTISIG,
)


class TestStacksNetwork:
    def test_mainnet_magic(self):
        assert StacksNetwork.MAINNET.value == b'X2'

    def test_testnet_magic(self):
        assert StacksNetwork.TESTNET.value == b'T2'


class TestStacksOpcode:
    def test_leader_block_commit_opcode(self):
        assert StacksOpcode.LEADER_BLOCK_COMMIT.value == ord('[')
        assert chr(StacksOpcode.LEADER_BLOCK_COMMIT.value) == '['

    def test_leader_key_register_opcode(self):
        assert StacksOpcode.LEADER_KEY_REGISTER.value == ord('^')

    def test_user_burn_support_opcode(self):
        assert StacksOpcode.USER_BURN_SUPPORT.value == ord('_')

    def test_pre_stx_opcode(self):
        assert StacksOpcode.PRE_STX.value == ord('p')

    def test_stack_stx_opcode(self):
        assert StacksOpcode.STACK_STX.value == ord('x')

    def test_transfer_stx_opcode(self):
        assert StacksOpcode.TRANSFER_STX.value == ord('$')

    def test_delegate_stx_opcode(self):
        assert StacksOpcode.DELEGATE_STX.value == ord('#')


class TestLeaderBlockCommit:
    def test_encode_length(self):
        """Leader block commit should encode to exactly 80 bytes."""
        commit = LeaderBlockCommit(
            block_hash=bytes(32),
            new_seed=bytes(32),
            parent_block=100,
            parent_txoff=5,
            key_block=90,
            key_txoff=2,
            burn_parent_modulus=3
        )
        payload = commit.encode(StacksNetwork.MAINNET)
        assert len(payload) == 80

    def test_encode_magic_and_opcode(self):
        """Verify magic bytes and opcode are correct."""
        commit = LeaderBlockCommit(
            block_hash=bytes(32),
            new_seed=bytes(32),
            parent_block=0,
            parent_txoff=0,
            key_block=0,
            key_txoff=0,
            burn_parent_modulus=0
        )
        payload = commit.encode(StacksNetwork.MAINNET)
        assert payload[:2] == b'X2'
        assert payload[2] == ord('[')

    def test_encode_testnet(self):
        """Verify testnet magic bytes."""
        commit = LeaderBlockCommit(
            block_hash=bytes(32),
            new_seed=bytes(32),
            parent_block=0,
            parent_txoff=0,
            key_block=0,
            key_txoff=0,
            burn_parent_modulus=0
        )
        payload = commit.encode(StacksNetwork.TESTNET)
        assert payload[:2] == b'T2'

    def test_burn_parent_modulus(self):
        """Verify burn_parent_modulus is encoded with modulo."""
        commit = LeaderBlockCommit(
            block_hash=bytes(32),
            new_seed=bytes(32),
            parent_block=0,
            parent_txoff=0,
            key_block=0,
            key_txoff=0,
            burn_parent_modulus=10  # Should become 10 % 6 = 4
        )
        payload = commit.encode()
        assert payload[-1] == 4

    def test_invalid_block_hash_length(self):
        """Should raise error for invalid block hash length."""
        with pytest.raises(ValueError, match="block_hash must be 32 bytes"):
            commit = LeaderBlockCommit(
                block_hash=bytes(16),  # Wrong length
                new_seed=bytes(32),
                parent_block=0,
                parent_txoff=0,
                key_block=0,
                key_txoff=0,
                burn_parent_modulus=0
            )
            commit.encode()


class TestLeaderKeyRegister:
    def test_encode_length(self):
        """Leader key register should encode to exactly 80 bytes."""
        register = LeaderKeyRegister(
            consensus_hash=bytes(20),
            proving_public_key=bytes(32),
            memo=b''
        )
        payload = register.encode(StacksNetwork.MAINNET)
        assert len(payload) == 80

    def test_encode_magic_and_opcode(self):
        """Verify magic bytes and opcode."""
        register = LeaderKeyRegister(
            consensus_hash=bytes(20),
            proving_public_key=bytes(32),
            memo=b''
        )
        payload = register.encode(StacksNetwork.MAINNET)
        assert payload[:2] == b'X2'
        assert payload[2] == ord('^')

    def test_memo_padding(self):
        """Memo should be zero-padded to 25 bytes."""
        register = LeaderKeyRegister(
            consensus_hash=bytes(20),
            proving_public_key=bytes(32),
            memo=b'test'
        )
        payload = register.encode()
        # Memo starts at byte 55
        memo_section = payload[55:80]
        assert memo_section[:4] == b'test'
        assert memo_section[4:] == bytes(21)  # Zero-padded

    def test_memo_too_long(self):
        """Should raise error for memo > 25 bytes."""
        with pytest.raises(ValueError, match="memo must be at most 25 bytes"):
            register = LeaderKeyRegister(
                consensus_hash=bytes(20),
                proving_public_key=bytes(32),
                memo=bytes(30)
            )
            register.encode()


class TestUserBurnSupport:
    def test_encode_length(self):
        """User burn support should encode to exactly 80 bytes."""
        support = UserBurnSupport(
            consensus_hash=bytes(20),
            proving_public_key=bytes(32),
            block_hash_160=bytes(20),
            key_block=100,
            key_vtxindex=5
        )
        payload = support.encode(StacksNetwork.MAINNET)
        assert len(payload) == 80

    def test_encode_magic_and_opcode(self):
        """Verify magic bytes and opcode."""
        support = UserBurnSupport(
            consensus_hash=bytes(20),
            proving_public_key=bytes(32),
            block_hash_160=bytes(20),
            key_block=0,
            key_vtxindex=0
        )
        payload = support.encode(StacksNetwork.MAINNET)
        assert payload[:2] == b'X2'
        assert payload[2] == ord('_')

    def test_key_block_24bit(self):
        """key_block should be encoded as 24-bit (3 bytes)."""
        support = UserBurnSupport(
            consensus_hash=bytes(20),
            proving_public_key=bytes(32),
            block_hash_160=bytes(20),
            key_block=0x123456,
            key_vtxindex=0
        )
        payload = support.encode()
        # key_block at bytes 75-77
        assert payload[75:78] == bytes.fromhex('123456')


class TestPreStx:
    def test_encode_length(self):
        """Pre-STX should encode to exactly 3 bytes."""
        pre_stx = PreStx()
        payload = pre_stx.encode(StacksNetwork.MAINNET)
        assert len(payload) == 3

    def test_encode_content(self):
        """Verify Pre-STX encoding."""
        pre_stx = PreStx()
        payload = pre_stx.encode(StacksNetwork.MAINNET)
        assert payload == b'X2p'


class TestStackStx:
    def test_encode_length(self):
        """Stack STX should encode to exactly 20 bytes."""
        stack = StackStx(stacked_ustx=1000000, num_cycles=1)
        payload = stack.encode(StacksNetwork.MAINNET)
        assert len(payload) == 20

    def test_encode_magic_and_opcode(self):
        """Verify magic bytes and opcode."""
        stack = StackStx(stacked_ustx=1000000, num_cycles=1)
        payload = stack.encode(StacksNetwork.MAINNET)
        assert payload[:2] == b'X2'
        assert payload[2] == ord('x')

    def test_amount_encoding(self):
        """Verify amount is encoded as u128 big-endian."""
        stack = StackStx(stacked_ustx=0x0102030405060708090a0b0c0d0e0f10, num_cycles=1)
        payload = stack.encode()
        # Amount at bytes 3-18
        assert payload[3:19] == bytes.fromhex('0102030405060708090a0b0c0d0e0f10')

    def test_cycles_range(self):
        """num_cycles must be 1-12."""
        with pytest.raises(ValueError, match="num_cycles must be between 1 and 12"):
            stack = StackStx(stacked_ustx=1000, num_cycles=0)
            stack.encode()

        with pytest.raises(ValueError, match="num_cycles must be between 1 and 12"):
            stack = StackStx(stacked_ustx=1000, num_cycles=13)
            stack.encode()


class TestTransferStx:
    def test_encode_length(self):
        """Transfer STX should encode to exactly 19 bytes."""
        transfer = TransferStx(amount_ustx=1000000)
        payload = transfer.encode(StacksNetwork.MAINNET)
        assert len(payload) == 19

    def test_encode_magic_and_opcode(self):
        """Verify magic bytes and opcode."""
        transfer = TransferStx(amount_ustx=1000000)
        payload = transfer.encode(StacksNetwork.MAINNET)
        assert payload[:2] == b'X2'
        assert payload[2] == ord('$')


class TestDelegateStx:
    def test_encode_length(self):
        """Delegate STX should encode to exactly 19 bytes."""
        delegate = DelegateStx(amount_ustx=1000000)
        payload = delegate.encode(StacksNetwork.MAINNET)
        assert len(payload) == 19

    def test_encode_magic_and_opcode(self):
        """Verify magic bytes and opcode."""
        delegate = DelegateStx(amount_ustx=1000000)
        payload = delegate.encode(StacksNetwork.MAINNET)
        assert payload[:2] == b'X2'
        assert payload[2] == ord('#')


class TestOpReturnScript:
    def test_small_payload(self):
        """Small payloads use direct push."""
        payload = bytes(50)
        script = create_op_return_script(payload)
        assert script[0] == 0x6a  # OP_RETURN
        assert script[1] == 50   # Length
        assert script[2:] == payload

    def test_medium_payload(self):
        """Payloads 76-80 bytes use OP_PUSHDATA1."""
        payload = bytes(80)
        script = create_op_return_script(payload)
        assert script[0] == 0x6a  # OP_RETURN
        assert script[1] == 0x4c  # OP_PUSHDATA1
        assert script[2] == 80   # Length
        assert script[3:] == payload

    def test_payload_too_large(self):
        """Should reject payloads > 80 bytes."""
        with pytest.raises(ValueError, match="exceeds maximum"):
            create_op_return_script(bytes(81))


class TestBurnAddress:
    def test_mainnet_burn_address(self):
        """Verify mainnet burn address."""
        assert get_burn_address(StacksNetwork.MAINNET) == "1111111111111111111114oLvT2"

    def test_testnet_burn_address(self):
        """Verify testnet burn address."""
        assert get_burn_address(StacksNetwork.TESTNET) == "mfWxJ45yp2SFn7UciZyNpvDKrzbhyfKrY8"


class TestC32Alphabet:
    """Test C32 alphabet matches the Crockford base32 variant used by Stacks."""

    def test_alphabet_length(self):
        """C32 alphabet should have 32 characters."""
        assert len(C32_ALPHABET) == 32

    def test_alphabet_content(self):
        """Verify exact alphabet (excludes I, L, O, U)."""
        assert C32_ALPHABET == "0123456789ABCDEFGHJKMNPQRSTVWXYZ"

    def test_excluded_characters(self):
        """Verify confusing characters are excluded."""
        assert 'I' not in C32_ALPHABET
        assert 'L' not in C32_ALPHABET
        assert 'O' not in C32_ALPHABET
        assert 'U' not in C32_ALPHABET


class TestC32Encode:
    """
    Test C32 encoding using test vectors from the Hiro reference implementation.
    Reference: https://github.com/hirosystems/stacks-encoding-native-js/blob/main/src/address/c32.rs
    """

    def test_encode_empty(self):
        """Empty input should return empty string."""
        assert _c32_encode(b'') == ""

    def test_encode_all_zeros(self):
        """All zero bytes should encode to all '0' characters."""
        assert _c32_encode(bytes(20)) == "0" * 20

    def test_encode_simple_values(self):
        """Test simple encoding cases from reference."""
        # Reference test vectors
        test_cases = [
            ("01", "1"),
            ("22", "12"),
            ("0001", "01"),
            ("000001", "001"),
            ("00000001", "0001"),
            ("10", "G"),
            ("0100", "80"),
            ("1000", "400"),
            ("010000", "2000"),
            ("100000", "10000"),
            ("01000000", "G0000"),
            ("10000000", "800000"),
            ("0100000000", "4000000"),
        ]
        for hex_input, expected in test_cases:
            data = bytes.fromhex(hex_input)
            result = _c32_encode(data)
            assert result == expected, f"Failed for {hex_input}: got {result}, expected {expected}"

    def test_encode_hash160(self):
        """Test encoding a typical hash160 value."""
        data = bytes.fromhex("a46ff88886c2ef9762d970b4d2c63678835bd39d")
        expected = "MHQZH246RBQSERPSE2TD5HHPF21NQMWX"
        assert _c32_encode(data) == expected

    def test_encode_leading_zeros_preserved(self):
        """Leading zero bytes should be preserved as '0' characters."""
        # 19 leading zero bytes + 0x01
        data = bytes.fromhex("0000000000000000000000000000000000000001")
        result = _c32_encode(data)
        assert result == "00000000000000000001"

    def test_encode_high_byte_value(self):
        """Test with high byte values."""
        data = bytes.fromhex("1000000000000000000000000000000000000001")
        result = _c32_encode(data)
        assert result == "20000000000000000000000000000001"


class TestC32Checksum:
    """Test C32 checksum calculation (double SHA256)."""

    def test_checksum_length(self):
        """Checksum should always be 4 bytes."""
        data = bytes.fromhex("a46ff88886c2ef9762d970b4d2c63678835bd39d")
        checksum = _c32_checksum(22, data)
        assert len(checksum) == 4

    def test_checksum_deterministic(self):
        """Same input should produce same checksum."""
        data = bytes.fromhex("a46ff88886c2ef9762d970b4d2c63678835bd39d")
        checksum1 = _c32_checksum(22, data)
        checksum2 = _c32_checksum(22, data)
        assert checksum1 == checksum2

    def test_checksum_version_matters(self):
        """Different versions should produce different checksums."""
        data = bytes.fromhex("a46ff88886c2ef9762d970b4d2c63678835bd39d")
        checksum_v22 = _c32_checksum(22, data)
        checksum_v26 = _c32_checksum(26, data)
        assert checksum_v22 != checksum_v26


class TestC32CheckEncode:
    """Test c32check encoding (version + data + checksum)."""

    def test_version_char_mainnet_singlesig(self):
        """Version 22 should produce 'P' prefix (after 'S')."""
        data = bytes(20)
        result = c32check_encode(22, data)
        assert result[0] == 'P'

    def test_version_char_mainnet_multisig(self):
        """Version 20 should produce 'M' prefix (after 'S')."""
        data = bytes(20)
        result = c32check_encode(20, data)
        assert result[0] == 'M'

    def test_version_char_testnet_singlesig(self):
        """Version 26 should produce 'T' prefix (after 'S')."""
        data = bytes(20)
        result = c32check_encode(26, data)
        assert result[0] == 'T'

    def test_version_char_testnet_multisig(self):
        """Version 21 should produce 'N' prefix (after 'S')."""
        data = bytes(20)
        result = c32check_encode(21, data)
        assert result[0] == 'N'


class TestHash160ToStxAddress:
    """
    Test hash160 to STX address conversion using official test vectors.
    Reference: https://github.com/hirosystems/stacks-encoding-native-js/blob/main/src/address/c32.rs
    """

    # Test vectors from the Hiro reference implementation
    HASH160_TEST_VECTORS = [
        "a46ff88886c2ef9762d970b4d2c63678835bd39d",
        "0000000000000000000000000000000000000000",
        "0000000000000000000000000000000000000001",
        "1000000000000000000000000000000000000001",
        "1000000000000000000000000000000000000000",
    ]

    # Expected addresses for version 22 (mainnet singlesig, 'SP' prefix)
    MAINNET_SINGLESIG_EXPECTED = [
        "SP2J6ZY48GV1EZ5V2V5RB9MP66SW86PYKKNRV9EJ7",
        "SP000000000000000000002Q6VF78",
        "SP00000000000000000005JA84HQ",
        "SP80000000000000000000000000000004R0CMNV",
        "SP800000000000000000000000000000033H8YKK",
    ]

    # Expected addresses for version 26 (testnet singlesig, 'ST' prefix)
    TESTNET_SINGLESIG_EXPECTED = [
        "ST2J6ZY48GV1EZ5V2V5RB9MP66SW86PYKKQYAC0RQ",
        "ST000000000000000000002AMW42H",
        "ST000000000000000000042DB08Y",
        "ST80000000000000000000000000000006BYJ4R4",
        "ST80000000000000000000000000000002YBNPV3",
    ]

    # Expected addresses for version 20 (mainnet multisig, 'SM' prefix)
    MAINNET_MULTISIG_EXPECTED = [
        "SM2J6ZY48GV1EZ5V2V5RB9MP66SW86PYKKQVX8X0G",
        "SM0000000000000000000062QV6X",
        "SM00000000000000000005VR75B2",
        "SM80000000000000000000000000000004WBEWKC",
        "SM80000000000000000000000000000000JGSYGV",
    ]

    # Expected addresses for version 21 (testnet multisig, 'SN' prefix)
    TESTNET_MULTISIG_EXPECTED = [
        "SN2J6ZY48GV1EZ5V2V5RB9MP66SW86PYKKP6D2ZK9",
        "SN000000000000000000003YDHWKJ",
        "SN00000000000000000005341MC8",
        "SN800000000000000000000000000000066KZWY0",
        "SN800000000000000000000000000000006H75AK",
    ]

    def test_mainnet_singlesig_addresses(self):
        """Test mainnet singlesig address generation (SP prefix)."""
        for i, hex_hash in enumerate(self.HASH160_TEST_VECTORS):
            hash160 = bytes.fromhex(hex_hash)
            result = hash160_to_stx_address(hash160, StacksNetwork.MAINNET, is_multisig=False)
            expected = self.MAINNET_SINGLESIG_EXPECTED[i]
            assert result == expected, f"Mismatch for {hex_hash}: got {result}, expected {expected}"

    def test_testnet_singlesig_addresses(self):
        """Test testnet singlesig address generation (ST prefix)."""
        for i, hex_hash in enumerate(self.HASH160_TEST_VECTORS):
            hash160 = bytes.fromhex(hex_hash)
            result = hash160_to_stx_address(hash160, StacksNetwork.TESTNET, is_multisig=False)
            expected = self.TESTNET_SINGLESIG_EXPECTED[i]
            assert result == expected, f"Mismatch for {hex_hash}: got {result}, expected {expected}"

    def test_mainnet_multisig_addresses(self):
        """Test mainnet multisig address generation (SM prefix)."""
        for i, hex_hash in enumerate(self.HASH160_TEST_VECTORS):
            hash160 = bytes.fromhex(hex_hash)
            result = hash160_to_stx_address(hash160, StacksNetwork.MAINNET, is_multisig=True)
            expected = self.MAINNET_MULTISIG_EXPECTED[i]
            assert result == expected, f"Mismatch for {hex_hash}: got {result}, expected {expected}"

    def test_testnet_multisig_addresses(self):
        """Test testnet multisig address generation (SN prefix)."""
        for i, hex_hash in enumerate(self.HASH160_TEST_VECTORS):
            hash160 = bytes.fromhex(hex_hash)
            result = hash160_to_stx_address(hash160, StacksNetwork.TESTNET, is_multisig=True)
            expected = self.TESTNET_MULTISIG_EXPECTED[i]
            assert result == expected, f"Mismatch for {hex_hash}: got {result}, expected {expected}"

    def test_address_starts_with_s(self):
        """All STX addresses should start with 'S'."""
        hash160 = bytes.fromhex("a46ff88886c2ef9762d970b4d2c63678835bd39d")
        for network in [StacksNetwork.MAINNET, StacksNetwork.TESTNET]:
            for is_multisig in [False, True]:
                result = hash160_to_stx_address(hash160, network, is_multisig)
                assert result.startswith('S')

    def test_invalid_hash160_length(self):
        """Should reject hash160 that is not 20 bytes."""
        with pytest.raises(ValueError, match="hash160 must be 20 bytes"):
            hash160_to_stx_address(bytes(19), StacksNetwork.MAINNET)

        with pytest.raises(ValueError, match="hash160 must be 20 bytes"):
            hash160_to_stx_address(bytes(21), StacksNetwork.MAINNET)


class TestAddressVersionConstants:
    """Test address version byte constants."""

    def test_mainnet_singlesig_version(self):
        assert STX_ADDRESS_VERSION_MAINNET_SINGLESIG == 22

    def test_mainnet_multisig_version(self):
        assert STX_ADDRESS_VERSION_MAINNET_MULTISIG == 20

    def test_testnet_singlesig_version(self):
        assert STX_ADDRESS_VERSION_TESTNET_SINGLESIG == 26

    def test_testnet_multisig_version(self):
        assert STX_ADDRESS_VERSION_TESTNET_MULTISIG == 21

    def test_version_to_prefix_mapping(self):
        """Verify version bytes map to correct C32 prefix characters."""
        # Version 22 -> 'P' (index 22 in C32_ALPHABET)
        assert C32_ALPHABET[22] == 'P'
        # Version 20 -> 'M' (index 20 in C32_ALPHABET)
        assert C32_ALPHABET[20] == 'M'
        # Version 26 -> 'T' (index 26 in C32_ALPHABET)
        assert C32_ALPHABET[26] == 'T'
        # Version 21 -> 'N' (index 21 in C32_ALPHABET)
        assert C32_ALPHABET[21] == 'N'


class TestPubkeyToStxAddress:
    """Test public key to STX address conversion."""

    def test_compressed_pubkey(self):
        """Test conversion of a compressed public key."""
        # Example compressed public key (33 bytes)
        pubkey = bytes.fromhex(
            "0279be667ef9dcbbac55a06295ce870b07029bfcdb2dce28d959f2815b16f81798"
        )
        result = pubkey_to_stx_address(pubkey, StacksNetwork.MAINNET)
        # Should produce a valid SP address
        assert result.startswith("SP")
        assert len(result) > 10

    def test_uncompressed_pubkey(self):
        """Test conversion of an uncompressed public key."""
        # Example uncompressed public key (65 bytes)
        pubkey = bytes.fromhex(
            "04"
            "79be667ef9dcbbac55a06295ce870b07029bfcdb2dce28d959f2815b16f81798"
            "483ada7726a3c4655da4fbfc0e1108a8fd17b448a68554199c47d08ffb10d4b8"
        )
        result = pubkey_to_stx_address(pubkey, StacksNetwork.MAINNET)
        # Should produce a valid SP address (different from compressed)
        assert result.startswith("SP")
        assert len(result) > 10

    def test_testnet_pubkey(self):
        """Test pubkey conversion for testnet."""
        pubkey = bytes.fromhex(
            "0279be667ef9dcbbac55a06295ce870b07029bfcdb2dce28d959f2815b16f81798"
        )
        result = pubkey_to_stx_address(pubkey, StacksNetwork.TESTNET)
        # Should produce a valid ST address
        assert result.startswith("ST")


if __name__ == '__main__':
    pytest.main([__file__, '-v'])