contrib: testgen: use base58 methods from test framework

2 files changed, 10 insertions, 121 deletions

diff --git a/contrib/testgen/base58.py b/contrib/testgen/base58.py
deleted file mode 100644
index 87341ccf96..0000000000
--- a/contrib/testgen/base58.py
+++ /dev/null
@@ -1,115 +0,0 @@
-# Copyright (c) 2012-2020 The Bitcoin Core developers
-# Distributed under the MIT software license, see the accompanying
-# file COPYING or http://www.opensource.org/licenses/mit-license.php.
-'''
-Bitcoin base58 encoding and decoding.
-
-Based on https://bitcointalk.org/index.php?topic=1026.0 (public domain)
-'''
-import hashlib
-
-# for compatibility with following code...
-class SHA256:
-    new = hashlib.sha256
-
-if str != bytes:
-    # Python 3.x
-    def ord(c):
-        return c
-    def chr(n):
-        return bytes( (n,) )
-
-__b58chars = '123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz'
-__b58base = len(__b58chars)
-b58chars = __b58chars
-
-def b58encode(v):
-    """ encode v, which is a string of bytes, to base58.
-    """
-    long_value = 0
-    for (i, c) in enumerate(v[::-1]):
-        if isinstance(c, str):
-            c = ord(c)
-        long_value += (256**i) * c
-
-    result = ''
-    while long_value >= __b58base:
-        div, mod = divmod(long_value, __b58base)
-        result = __b58chars[mod] + result
-        long_value = div
-    result = __b58chars[long_value] + result
-
-    # Bitcoin does a little leading-zero-compression:
-    # leading 0-bytes in the input become leading-1s
-    nPad = 0
-    for c in v:
-        if c == 0:
-            nPad += 1
-        else:
-            break
-
-    return (__b58chars[0]*nPad) + result
-
-def b58decode(v, length = None):
-    """ decode v into a string of len bytes
-    """
-    long_value = 0
-    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:
-        div, mod = divmod(long_value, 256)
-        result = chr(mod) + result
-        long_value = div
-    result = chr(long_value) + result
-
-    nPad = 0
-    for c in v:
-        if c == __b58chars[0]:
-            nPad += 1
-            continue
-        break
-
-    result = bytes(nPad) + result
-    if length is not None and len(result) != length:
-        return None
-
-    return result
-
-def checksum(v):
-    """Return 32-bit checksum based on SHA256"""
-    return SHA256.new(SHA256.new(v).digest()).digest()[0:4]
-
-def b58encode_chk(v):
-    """b58encode a string, with 32-bit checksum"""
-    return b58encode(v + checksum(v))
-
-def b58decode_chk(v):
-    """decode a base58 string, check and remove checksum"""
-    result = b58decode(v)
-    if result is None:
-        return None
-    if result[-4:] == checksum(result[:-4]):
-        return result[:-4]
-    else:
-        return None
-
-def get_bcaddress_version(strAddress):
-    """ Returns None if strAddress is invalid.  Otherwise returns integer version of address. """
-    addr = b58decode_chk(strAddress)
-    if addr is None or len(addr)!=21:
-        return None
-    version = addr[0]
-    return ord(version)
-
-if __name__ == '__main__':
-    # Test case (from http://gitorious.org/bitcoin/python-base58.git)
-    assert get_bcaddress_version('15VjRaDX9zpbA8LVnbrCAFzrVzN7ixHNsC') == 0
-    _ohai = 'o hai'.encode('ascii')
-    _tmp = b58encode(_ohai)
-    assert _tmp == 'DYB3oMS'
-    assert b58decode(_tmp, 5) == _ohai
-    print("Tests passed")
diff --git a/contrib/testgen/gen_key_io_test_vectors.py b/contrib/testgen/gen_key_io_test_vectors.py
index c69225c326..8017fad461 100755
--- a/contrib/testgen/gen_key_io_test_vectors.py
+++ b/contrib/testgen/gen_key_io_test_vectors.py
@@ -10,14 +10,14 @@ Usage:
     ./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)