contrib: testgen: use base58 methods from test framework

Author: theStack

contrib/testgen/base58.py

@@ -10,14 +10,14 @@
     ./gen_key_io_test_vectors.py invalid 70 > ../../src/test/data/key_io_invalid.json
 '''
 
-import os
 from itertools import islice
-from base58 import b58encode_chk, b58decode_chk, b58chars
+import os
 import random
 import sys
 
 sys.path.append(os.path.join(os.path.dirname(__file__), '../../test/functional'))
 
+from test_framework.address import base58_to_byte, byte_to_base58, b58chars  # noqa: E402
 from test_framework.script import OP_0, OP_1, OP_2, OP_3, OP_16, OP_DUP, OP_EQUAL, OP_EQUALVERIFY, OP_HASH160, OP_CHECKSIG  # noqa: E402
 from test_framework.segwit_addr import bech32_encode, decode_segwit_address, convertbits, CHARSET, Encoding  # noqa: E402
 
@@ -108,8 +108,10 @@ def is_valid(v):
     '''Check vector v for validity'''
     if len(set(v) - set(b58chars)) > 0:
         return is_valid_bech32(v)
-    result = b58decode_chk(v)
-    if result is None:
+    try:
+        payload, version = base58_to_byte(v)
+        result = bytes([version]) + payload
+    except AssertionError:  # thrown if checksum doesn't match
         return is_valid_bech32(v)
     for template in templates:
         prefix = bytearray(template[0])
@@ -133,7 +135,8 @@ def gen_valid_base58_vector(template):
     suffix = bytearray(template[2])
     dst_prefix = bytearray(template[4])
     dst_suffix = bytearray(template[5])
-    rv = b58encode_chk(prefix + payload + suffix)
+    assert len(prefix) == 1
+    rv = byte_to_base58(payload + suffix, prefix[0])
     return rv, dst_prefix + payload + dst_suffix
 
 def gen_valid_bech32_vector(template):
@@ -184,7 +187,8 @@ def gen_invalid_base58_vector(template):
     else:
         suffix = bytearray(template[2])
 
-    val = b58encode_chk(prefix + payload + suffix)
+    assert len(prefix) == 1
+    val = byte_to_base58(payload + suffix, prefix[0])
     if random.randint(0,10)<1: # line corruption
         if randbool(): # add random character to end
             val += random.choice(b58chars)