Make base58 python contrib code work with python3

2 files changed, 27 insertions, 20 deletions

diff --git a/contrib/testgen/base58.py b/contrib/testgen/base58.py
index 816d40b49c..0dbb79a707 100644
--- a/contrib/testgen/base58.py
+++ b/contrib/testgen/base58.py
@@ -28,7 +28,9 @@ def b58encode(v):
     """
     long_value = 0
     for (i, c) in enumerate(v[::-1]):
-        long_value += (256**i) * ord(c)
+        if isinstance(c, str):
+            c = ord(c)
+        long_value += (256**i) * c
 
     result = ''
     while long_value >= __b58base:
@@ -41,7 +43,7 @@ def b58encode(v):
     # leading 0-bytes in the input become leading-1s
     nPad = 0
     for c in v:
-        if c == '\0': nPad += 1
+        if c == 0: nPad += 1
         else: break
 
     return (__b58chars[0]*nPad) + result
@@ -50,8 +52,10 @@ def b58decode(v, length = None):
     """ decode v into a string of len bytes
     """
     long_value = 0
-    for (i, c) in enumerate(v[::-1]):
-        long_value += __b58chars.find(c) * (__b58base**i)
+    for i, c in enumerate(v[::-1]):
+        pos = __b58chars.find(c)
+        assert pos != -1
+        long_value += pos * (__b58base**i)
 
     result = bytes()
     while long_value >= 256:
@@ -62,10 +66,12 @@ def b58decode(v, length = None):
 
     nPad = 0
     for c in v:
-        if c == __b58chars[0]: nPad += 1
-        else: break
+        if c == __b58chars[0]:
+            nPad += 1
+            continue
+        break
 
-    result = chr(0)*nPad + result
+    result = bytes(nPad) + result
     if length is not None and len(result) != length:
         return None
 
diff --git a/contrib/testgen/gen_base58_test_vectors.py b/contrib/testgen/gen_base58_test_vectors.py
index af6d7e1112..4351605786 100755
--- a/contrib/testgen/gen_base58_test_vectors.py
+++ b/contrib/testgen/gen_base58_test_vectors.py
@@ -1,15 +1,13 @@
-#!/usr/bin/env python2
+#!/usr/bin/env python3
 # Copyright (c) 2012-2017 The Bitcoin Core developers
 # Distributed under the MIT software license, see the accompanying
 # file COPYING or http://www.opensource.org/licenses/mit-license.php.
 '''
 Generate valid and invalid base58 address and private key test vectors.
 
-Usage: 
+Usage:
     gen_base58_test_vectors.py valid 50 > ../../src/test/data/base58_keys_valid.json
     gen_base58_test_vectors.py invalid 50 > ../../src/test/data/base58_keys_invalid.json
-
-Note that this script is Python2 only, and will fail in Python3
 '''
 # 2012 Wladimir J. van der Laan
 # Released under MIT License
@@ -48,8 +46,8 @@ def is_valid(v):
     if result is None:
         return False
     for template in templates:
-        prefix = str(bytearray(template[0]))
-        suffix = str(bytearray(template[2]))
+        prefix = bytearray(template[0])
+        suffix = bytearray(template[2])
         if result.startswith(prefix) and result.endswith(suffix):
             if (len(result) - len(prefix) - len(suffix)) == template[1]:
                 return True
@@ -59,20 +57,23 @@ def gen_valid_vectors():
     '''Generate valid test vectors'''
     while True:
         for template in templates:
-            prefix = str(bytearray(template[0]))
-            payload = os.urandom(template[1]) 
-            suffix = str(bytearray(template[2]))
+            prefix = bytearray(template[0])
+            payload = bytearray(os.urandom(template[1]))
+            suffix = bytearray(template[2])
             rv = b58encode_chk(prefix + payload + suffix)
             assert is_valid(rv)
-            metadata = dict([(x,y) for (x,y) in zip(metadata_keys,template[3]) if y is not None])
-            yield (rv, b2a_hex(payload), metadata)
+            metadata = {x: y for x, y in zip(metadata_keys,template[3]) if y is not None}
+            hexrepr = b2a_hex(payload)
+            if isinstance(hexrepr, bytes):
+                hexrepr = hexrepr.decode('utf8')
+            yield (rv, hexrepr, metadata)
 
 def gen_invalid_vector(template, corrupt_prefix, randomize_payload_size, corrupt_suffix):
     '''Generate possibly invalid vector'''
     if corrupt_prefix:
         prefix = os.urandom(1)
     else:
-        prefix = str(bytearray(template[0]))
+        prefix = bytearray(template[0])
     
     if randomize_payload_size:
         payload = os.urandom(max(int(random.expovariate(0.5)), 50))
@@ -82,7 +83,7 @@ def gen_invalid_vector(template, corrupt_prefix, randomize_payload_size, corrupt
     if corrupt_suffix:
         suffix = os.urandom(len(template[2]))
     else:
-        suffix = str(bytearray(template[2]))
+        suffix = bytearray(template[2])
 
     return b58encode_chk(prefix + payload + suffix)