P2WPKH-P2SH (SegWit) Wallet Address Generation

gen.py

"""
Implementation of "Pay to Witness Public Key Hash nested in BIP16 Pay to Script Hash" (P2WPKH-P2SH) address generation.
Described in BIP 141 (SegWit) https://github.com/bitcoin/bips/blob/master/bip-0141.mediawiki#P2WPKH_nested_in_BIP16_P2SH

I.e.: Public key -> SegWit address
"""

import hashlib
from hashlib import sha256

# MIT License
# pip install git+https://github.com/prusnak/bip32utils
from bip32utils import Base58


def pk_to_p2wpkh_as_p2sh_addr(pk_as_hex_string, testnet=False):
    pk_bytes = bytes.fromhex(pk_as_hex_string)
    assert len(pk_bytes) == 33 and (pk_bytes.startswith(b"\x02") or pk_bytes.startswith(b"\x03")), \
        "Only compressed public keys are compatible with p2sh-p2wpkh addresses. " \
        "See https://github.com/bitcoin/bips/blob/master/bip-0049.mediawiki."
    pk_hash = hash160_bytes(pk_bytes)
    push_20 = bytes.fromhex("0014")
    script_sig = push_20 + pk_hash
    address_bytes = hash160_bytes(script_sig)
    prefix = b"\xc4" if testnet else b"\x05"
    address = Base58.check_encode(prefix + address_bytes)
    return address


def hash160_bytes(byte_input):
    return hashlib.new('ripemd160', sha256(byte_input).digest()).digest()


# Test vector from https://github.com/bitcoin/bips/blob/master/bip-0049.mediawiki
public_key_hex = "03a1af804ac108a8a51782198c2d034b28bf90c8803f5a53f76276fa69a4eae77f"

testnet_address = pk_to_p2wpkh_as_p2sh_addr(public_key_hex, testnet=True)
print(testnet_address)
assert testnet_address == "2Mww8dCYPUpKHofjgcXcBCEGmniw9CoaiD2", "Generated address matches BIP44 test vector"

mainnet_address = pk_to_p2wpkh_as_p2sh_addr(public_key_hex)
print(mainnet_address)
assert mainnet_address == "36NvZTcMsMowbt78wPzJaHHWaNiyR73Y4g", "Generated address matches /my/ test vector"