fix: decode uppercase segwit addresses

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]

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()

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()
+unittest.main()