session5/complete/helper.py

from unittest import TestCase, TestSuite, TextTestRunner

import hashlib


SIGHASH_ALL = 1
SIGHASH_NONE = 2
SIGHASH_SINGLE = 3
BASE58_ALPHABET = "123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz"


def run(test):
    suite = TestSuite()
    suite.addTest(test)
    TextTestRunner().run(suite)


def bytes_to_str(b, encoding="ascii"):
    """Returns a string version of the bytes"""
    # use the bytes.decode(encoding) method
    return b.decode(encoding)


def str_to_bytes(s, encoding="ascii"):
    """Returns a bytes version of the string"""
    # use the string.encode(encoding) method
    return s.encode(encoding)


def little_endian_to_int(b):
    """little_endian_to_int takes byte sequence as a little-endian number.
    Returns an integer"""
    # use the int.from_bytes(b, <endianness>) method
    return int.from_bytes(b, "little")


def int_to_little_endian(n, length):
    """endian_to_little_endian takes an integer and returns the little-endian
    byte sequence of length"""
    # use the int.to_bytes(length, <endianness>) method
    return n.to_bytes(length, "little")


def hash160(s):
    return hashlib.new("ripemd160", hashlib.sha256(s).digest()).digest()


def hash256(s):
    return hashlib.sha256(hashlib.sha256(s).digest()).digest()


def encode_base58(s):
    # determine how many 0 bytes (b'\x00') s starts with
    count = 0
    for c in s:
        if c == 0:
            count += 1
        else:
            break
    # convert from binary to hex, then hex to integer
    num = int(s.hex(), 16)
    result = ""
    prefix = "1" * count
    while num > 0:
        num, mod = divmod(num, 58)
        result = BASE58_ALPHABET[mod] + result
    return prefix + result


def encode_base58_checksum(raw):
    """Takes bytes and turns it into base58 encoding with checksum"""
    # checksum is the first 4 bytes of the hash256
    checksum = hash256(raw)[:4]
    # encode_base58 on the raw and the checksum
    return encode_base58(raw + checksum)


def decode_base58(s):
    num = 0
    for c in s:
        num *= 58
        num += BASE58_ALPHABET.index(c)
    combined = num.to_bytes(25, byteorder="big")
    checksum = combined[-4:]
    if hash256(combined[:-4])[:4] != checksum:
        raise RuntimeError(f"bad address: {checksum} {hash256(combined)[:4]}")
    return combined[1:-4]


def read_varint(s):
    """read_varint reads a variable integer from a stream"""
    i = s.read(1)[0]
    if i == 0xFD:
        # 0xfd means the next two bytes are the number
        return little_endian_to_int(s.read(2))
    elif i == 0xFE:
        # 0xfe means the next four bytes are the number
        return little_endian_to_int(s.read(4))
    elif i == 0xFF:
        # 0xff means the next eight bytes are the number
        return little_endian_to_int(s.read(8))
    else:
        # anything else is just the integer
        return i


def encode_varint(i):
    """encodes an integer as a varint"""
    if i < 0xFD:
        return bytes([i])
    elif i < 0x10000:
        return b"\xfd" + int_to_little_endian(i, 2)
    elif i < 0x100000000:
        return b"\xfe" + int_to_little_endian(i, 4)
    elif i < 0x10000000000000000:
        return b"\xff" + int_to_little_endian(i, 8)
    else:
        raise RuntimeError(f"integer too large: {i}")


def h160_to_p2pkh_address(h160, network="mainnet"):
    """Takes a byte sequence hash160 and returns a p2pkh address string"""
    # p2pkh has a prefix of b'\x00' for mainnet, b'\x6f' for testnet/signet
    if network in ("testnet", "signet"):
        prefix = b"\x6f"
    else:
        prefix = b"\x00"
    # return the encode_base58_checksum the prefix and h160
    return encode_base58_checksum(prefix + h160)


def h160_to_p2sh_address(h160, network="mainnet"):
    """Takes a byte sequence hash160 and returns a p2sh address string"""
    # p2sh has a prefix of b'\x05' for mainnet, b'\xc4' for testnet/signet
    if network in ("testnet", "signet"):
        prefix = b"\xc4"
    else:
        prefix = b"\x05"
    # return the encode_base58_checksum the prefix and h160
    return encode_base58_checksum(prefix + h160)


class HelperTest(TestCase):
    def test_bytes(self):
        b = b"hello world"
        s = "hello world"
        self.assertEqual(b, str_to_bytes(s))
        self.assertEqual(s, bytes_to_str(b))

    def test_little_endian_to_int(self):
        h = bytes.fromhex("99c3980000000000")
        want = 10011545
        self.assertEqual(little_endian_to_int(h), want)
        h = bytes.fromhex("a135ef0100000000")
        want = 32454049
        self.assertEqual(little_endian_to_int(h), want)

    def test_int_to_little_endian(self):
        n = 1
        want = b"\x01\x00\x00\x00"
        self.assertEqual(int_to_little_endian(n, 4), want)
        n = 10011545
        want = b"\x99\xc3\x98\x00\x00\x00\x00\x00"
        self.assertEqual(int_to_little_endian(n, 8), want)

    def test_base58(self):
        addr = "mnrVtF8DWjMu839VW3rBfgYaAfKk8983Xf"
        h160 = decode_base58(addr).hex()
        want = "507b27411ccf7f16f10297de6cef3f291623eddf"
        self.assertEqual(h160, want)
        got = encode_base58_checksum(b"\x6f" + bytes.fromhex(h160))
        self.assertEqual(got, addr)

    def test_encode_base58_checksum(self):
        raw = bytes.fromhex("005dedfbf9ea599dd4e3ca6a80b333c472fd0b3f69")
        want = "19ZewH8Kk1PDbSNdJ97FP4EiCjTRaZMZQA"
        self.assertEqual(encode_base58_checksum(raw), want)

    def test_p2pkh_address(self):
        h160 = bytes.fromhex("74d691da1574e6b3c192ecfb52cc8984ee7b6c56")
        want = "1BenRpVUFK65JFWcQSuHnJKzc4M8ZP8Eqa"
        self.assertEqual(h160_to_p2pkh_address(h160, network="mainnet"), want)
        want = "mrAjisaT4LXL5MzE81sfcDYKU3wqWSvf9q"
        self.assertEqual(h160_to_p2pkh_address(h160, network="signet"), want)

    def test_p2sh_address(self):
        h160 = bytes.fromhex("74d691da1574e6b3c192ecfb52cc8984ee7b6c56")
        want = "3CLoMMyuoDQTPRD3XYZtCvgvkadrAdvdXh"
        self.assertEqual(h160_to_p2sh_address(h160, network="mainnet"), want)
        want = "2N3u1R6uwQfuobCqbCgBkpsgBxvr1tZpe7B"
        self.assertEqual(h160_to_p2sh_address(h160, network="signet"), want)