diff --git a/pybgl/functions/address.py b/pybgl/functions/address.py index 8156128..1308d05 100644 --- a/pybgl/functions/address.py +++ b/pybgl/functions/address.py @@ -14,6 +14,17 @@ base32charset_upcase) +def _split_segwit_address(address): + try: + prefix, payload = address.split("1", 1) + except ValueError: + return None + prefix = prefix.lower() + if prefix not in (MAINNET_SEGWIT_ADDRESS_PREFIX, TESTNET_SEGWIT_ADDRESS_PREFIX): + return None + return prefix, payload + + def public_key_to_address(pubkey, testnet=False, p2sh_p2wpkh=False, witness_version=0): """ Get address from public key/script hash. In case PUBKEY, P2PKH, P2PKH public key/script hash is SHA256+RIPEMD160, @@ -105,12 +116,11 @@ def address_to_hash(address, hex=True): """ if address[0] in ADDRESS_PREFIX_LIST: h = decode_base58(address)[1:-4] - elif address.split("1")[0] in (MAINNET_SEGWIT_ADDRESS_PREFIX, - TESTNET_SEGWIT_ADDRESS_PREFIX): - address = address.split("1")[1] - h = rebase_5_to_8(rebase_32_to_5(address)[1:-6], False) else: - return None + segwit = _split_segwit_address(address) + if segwit is None: + return None + h = rebase_5_to_8(rebase_32_to_5(segwit[1])[1:-6], False) return h.hex() if hex else h @@ -129,8 +139,7 @@ def address_type(address, num=False): TESTNET_ADDRESS_PREFIX, TESTNET_ADDRESS_PREFIX_2): t = 'P2PKH' - elif address[:3] in (MAINNET_SEGWIT_ADDRESS_PREFIX, - TESTNET_SEGWIT_ADDRESS_PREFIX): + elif _split_segwit_address(address) is not None: if len(address) == 43: t = 'P2WPKH' elif len(address) == 63: @@ -149,16 +158,17 @@ def address_net_type(address): :param address: address in base58 or bech32 format. :return: address network type in string format or None. """ + segwit = _split_segwit_address(address) if address[0] in (MAINNET_SCRIPT_ADDRESS_PREFIX, MAINNET_ADDRESS_PREFIX): return "mainnet" - elif address[:3] == MAINNET_SEGWIT_ADDRESS_PREFIX: + elif segwit is not None and segwit[0] == MAINNET_SEGWIT_ADDRESS_PREFIX: return "mainnet" elif address[0] in (TESTNET_SCRIPT_ADDRESS_PREFIX, TESTNET_ADDRESS_PREFIX, TESTNET_ADDRESS_PREFIX_2): return "testnet" - elif address[:4] == TESTNET_SEGWIT_ADDRESS_PREFIX: + elif segwit is not None and segwit[0] == TESTNET_SEGWIT_ADDRESS_PREFIX: return "testnet" return None @@ -189,8 +199,7 @@ def address_to_script(address, hex=False): address_to_hash(address, hex=False), OP_EQUALVERIFY, OP_CHECKSIG] - elif (address[:3] == MAINNET_SEGWIT_ADDRESS_PREFIX) \ - or (address[:4] == TESTNET_SEGWIT_ADDRESS_PREFIX): + elif _split_segwit_address(address) is not None: h = address_to_hash(address, hex=False) s = [OP_0, bytes([len(h)]), @@ -307,6 +316,8 @@ def is_address_valid(address, testnet=False): def get_witness_version(address): - address = address.split("1")[1] - h = rebase_32_to_5(address) + segwit = _split_segwit_address(address) + if segwit is None: + return None + h = rebase_32_to_5(segwit[1]) return h[0] diff --git a/pybgl/test/bech32_address_decode.py b/pybgl/test/bech32_address_decode.py new file mode 100644 index 0000000..d4befb7 --- /dev/null +++ b/pybgl/test/bech32_address_decode.py @@ -0,0 +1,25 @@ +import unittest + +from pybgl.functions.address import ( + address_net_type, + address_to_hash, + address_to_script, + address_type, + get_witness_version, +) + + +class bech32_address_decode(unittest.TestCase): + def test_uppercase_segwit_decode_matches_lowercase(self): + address = "bgl1qdzjn6rd7e84lt2m5d3yf9jcg42ncdje7vhp4rl" + uppercase_address = address.upper() + + self.assertEqual(address_to_hash(uppercase_address), address_to_hash(address)) + self.assertEqual(address_to_script(uppercase_address), address_to_script(address)) + self.assertEqual(address_type(uppercase_address), "P2WPKH") + self.assertEqual(address_net_type(uppercase_address), "mainnet") + self.assertEqual(get_witness_version(uppercase_address), get_witness_version(address)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests.py b/tests.py index 8d48c99..c04dd0b 100644 --- a/tests.py +++ b/tests.py @@ -7,6 +7,7 @@ # from pybgl.test.shamir import * from pybgl.test.hash_functions import * from pybgl.test.create_transaction import * +from pybgl.test.bech32_address_decode import * # testLoad = unittest.TestLoader() # suites = testLoad.loadTestsFromModule(pybgl.test) # @@ -15,4 +16,4 @@ # # import unittest -unittest.main() \ No newline at end of file +unittest.main()