|
| 1 | +#!/usr/bin/env python3 |
| 2 | +# Copyright (c) 2012-2018 The Bitcoin Core developers |
| 3 | +# Distributed under the MIT software license, see the accompanying |
| 4 | +# file COPYING or http://www.opensource.org/licenses/mit-license.php. |
| 5 | +''' |
| 6 | +Generate valid and invalid base58 address and private key test vectors. |
| 7 | +
|
| 8 | +Usage: |
| 9 | + PYTHONPATH=../../test/functional/test_framework ./gen_key_io_test_vectors.py valid 50 > ../../src/test/data/key_io_valid.json |
| 10 | + PYTHONPATH=../../test/functional/test_framework ./gen_key_io_test_vectors.py invalid 50 > ../../src/test/data/key_io_invalid.json |
| 11 | +''' |
| 12 | +# 2012 Wladimir J. van der Laan |
| 13 | +# Released under MIT License |
| 14 | +import os |
| 15 | +from itertools import islice |
| 16 | +from base58 import b58encode_chk, b58decode_chk, b58chars |
| 17 | +import random |
| 18 | +from binascii import b2a_hex |
| 19 | +from segwit_addr import bech32_encode, decode, convertbits, CHARSET |
| 20 | + |
| 21 | +# key types |
| 22 | +PUBKEY_ADDRESS = 0 |
| 23 | +SCRIPT_ADDRESS = 5 |
| 24 | +PUBKEY_ADDRESS_TEST = 111 |
| 25 | +SCRIPT_ADDRESS_TEST = 196 |
| 26 | +PUBKEY_ADDRESS_REGTEST = 111 |
| 27 | +SCRIPT_ADDRESS_REGTEST = 196 |
| 28 | +PRIVKEY = 128 |
| 29 | +PRIVKEY_TEST = 239 |
| 30 | +PRIVKEY_REGTEST = 239 |
| 31 | + |
| 32 | +# script |
| 33 | +OP_0 = 0x00 |
| 34 | +OP_1 = 0x51 |
| 35 | +OP_2 = 0x52 |
| 36 | +OP_16 = 0x60 |
| 37 | +OP_DUP = 0x76 |
| 38 | +OP_EQUAL = 0x87 |
| 39 | +OP_EQUALVERIFY = 0x88 |
| 40 | +OP_HASH160 = 0xa9 |
| 41 | +OP_CHECKSIG = 0xac |
| 42 | +pubkey_prefix = (OP_DUP, OP_HASH160, 20) |
| 43 | +pubkey_suffix = (OP_EQUALVERIFY, OP_CHECKSIG) |
| 44 | +script_prefix = (OP_HASH160, 20) |
| 45 | +script_suffix = (OP_EQUAL,) |
| 46 | +p2wpkh_prefix = (OP_0, 20) |
| 47 | +p2wsh_prefix = (OP_0, 32) |
| 48 | + |
| 49 | +metadata_keys = ['isPrivkey', 'chain', 'isCompressed', 'tryCaseFlip'] |
| 50 | +# templates for valid sequences |
| 51 | +templates = [ |
| 52 | + # prefix, payload_size, suffix, metadata, output_prefix, output_suffix |
| 53 | + # None = N/A |
| 54 | + ((PUBKEY_ADDRESS,), 20, (), (False, 'main', None, None), pubkey_prefix, pubkey_suffix), |
| 55 | + ((SCRIPT_ADDRESS,), 20, (), (False, 'main', None, None), script_prefix, script_suffix), |
| 56 | + ((PUBKEY_ADDRESS_TEST,), 20, (), (False, 'test', None, None), pubkey_prefix, pubkey_suffix), |
| 57 | + ((SCRIPT_ADDRESS_TEST,), 20, (), (False, 'test', None, None), script_prefix, script_suffix), |
| 58 | + ((PUBKEY_ADDRESS_REGTEST,), 20, (), (False, 'regtest', None, None), pubkey_prefix, pubkey_suffix), |
| 59 | + ((SCRIPT_ADDRESS_REGTEST,), 20, (), (False, 'regtest', None, None), script_prefix, script_suffix), |
| 60 | + ((PRIVKEY,), 32, (), (True, 'main', False, None), (), ()), |
| 61 | + ((PRIVKEY,), 32, (1,), (True, 'main', True, None), (), ()), |
| 62 | + ((PRIVKEY_TEST,), 32, (), (True, 'test', False, None), (), ()), |
| 63 | + ((PRIVKEY_TEST,), 32, (1,), (True, 'test', True, None), (), ()), |
| 64 | + ((PRIVKEY_REGTEST,), 32, (), (True, 'regtest', False, None), (), ()), |
| 65 | + ((PRIVKEY_REGTEST,), 32, (1,), (True, 'regtest', True, None), (), ()) |
| 66 | +] |
| 67 | +# templates for valid bech32 sequences |
| 68 | +bech32_templates = [ |
| 69 | + # hrp, version, witprog_size, metadata, output_prefix |
| 70 | + ('bc', 0, 20, (False, 'main', None, True), p2wpkh_prefix), |
| 71 | + ('bc', 0, 32, (False, 'main', None, True), p2wsh_prefix), |
| 72 | + ('bc', 1, 2, (False, 'main', None, True), (OP_1, 2)), |
| 73 | + ('tb', 0, 20, (False, 'test', None, True), p2wpkh_prefix), |
| 74 | + ('tb', 0, 32, (False, 'test', None, True), p2wsh_prefix), |
| 75 | + ('tb', 2, 16, (False, 'test', None, True), (OP_2, 16)), |
| 76 | + ('bcrt', 0, 20, (False, 'regtest', None, True), p2wpkh_prefix), |
| 77 | + ('bcrt', 0, 32, (False, 'regtest', None, True), p2wsh_prefix), |
| 78 | + ('bcrt', 16, 40, (False, 'regtest', None, True), (OP_16, 40)) |
| 79 | +] |
| 80 | +# templates for invalid bech32 sequences |
| 81 | +bech32_ng_templates = [ |
| 82 | + # hrp, version, witprog_size, invalid_bech32, invalid_checksum, invalid_char |
| 83 | + ('tc', 0, 20, False, False, False), |
| 84 | + ('tb', 17, 32, False, False, False), |
| 85 | + ('bcrt', 3, 1, False, False, False), |
| 86 | + ('bc', 15, 41, False, False, False), |
| 87 | + ('tb', 0, 16, False, False, False), |
| 88 | + ('bcrt', 0, 32, True, False, False), |
| 89 | + ('bc', 0, 16, True, False, False), |
| 90 | + ('tb', 0, 32, False, True, False), |
| 91 | + ('bcrt', 0, 20, False, False, True) |
| 92 | +] |
| 93 | + |
| 94 | +def is_valid(v): |
| 95 | + '''Check vector v for validity''' |
| 96 | + if len(set(v) - set(b58chars)) > 0: |
| 97 | + return is_valid_bech32(v) |
| 98 | + result = b58decode_chk(v) |
| 99 | + if result is None: |
| 100 | + return is_valid_bech32(v) |
| 101 | + for template in templates: |
| 102 | + prefix = bytearray(template[0]) |
| 103 | + suffix = bytearray(template[2]) |
| 104 | + if result.startswith(prefix) and result.endswith(suffix): |
| 105 | + if (len(result) - len(prefix) - len(suffix)) == template[1]: |
| 106 | + return True |
| 107 | + return is_valid_bech32(v) |
| 108 | + |
| 109 | +def is_valid_bech32(v): |
| 110 | + '''Check vector v for bech32 validity''' |
| 111 | + for hrp in ['bc', 'tb', 'bcrt']: |
| 112 | + if decode(hrp, v) != (None, None): |
| 113 | + return True |
| 114 | + return False |
| 115 | + |
| 116 | +def gen_valid_base58_vector(template): |
| 117 | + '''Generate valid base58 vector''' |
| 118 | + prefix = bytearray(template[0]) |
| 119 | + payload = bytearray(os.urandom(template[1])) |
| 120 | + suffix = bytearray(template[2]) |
| 121 | + dst_prefix = bytearray(template[4]) |
| 122 | + dst_suffix = bytearray(template[5]) |
| 123 | + rv = b58encode_chk(prefix + payload + suffix) |
| 124 | + return rv, dst_prefix + payload + dst_suffix |
| 125 | + |
| 126 | +def gen_valid_bech32_vector(template): |
| 127 | + '''Generate valid bech32 vector''' |
| 128 | + hrp = template[0] |
| 129 | + witver = template[1] |
| 130 | + witprog = bytearray(os.urandom(template[2])) |
| 131 | + dst_prefix = bytearray(template[4]) |
| 132 | + rv = bech32_encode(hrp, [witver] + convertbits(witprog, 8, 5)) |
| 133 | + return rv, dst_prefix + witprog |
| 134 | + |
| 135 | +def gen_valid_vectors(): |
| 136 | + '''Generate valid test vectors''' |
| 137 | + glist = [gen_valid_base58_vector, gen_valid_bech32_vector] |
| 138 | + tlist = [templates, bech32_templates] |
| 139 | + while True: |
| 140 | + for template, valid_vector_generator in [(t, g) for g, l in zip(glist, tlist) for t in l]: |
| 141 | + rv, payload = valid_vector_generator(template) |
| 142 | + assert is_valid(rv) |
| 143 | + metadata = {x: y for x, y in zip(metadata_keys,template[3]) if y is not None} |
| 144 | + hexrepr = b2a_hex(payload) |
| 145 | + if isinstance(hexrepr, bytes): |
| 146 | + hexrepr = hexrepr.decode('utf8') |
| 147 | + yield (rv, hexrepr, metadata) |
| 148 | + |
| 149 | +def gen_invalid_base58_vector(template): |
| 150 | + '''Generate possibly invalid vector''' |
| 151 | + # kinds of invalid vectors: |
| 152 | + # invalid prefix |
| 153 | + # invalid payload length |
| 154 | + # invalid (randomized) suffix (add random data) |
| 155 | + # corrupt checksum |
| 156 | + corrupt_prefix = randbool(0.2) |
| 157 | + randomize_payload_size = randbool(0.2) |
| 158 | + corrupt_suffix = randbool(0.2) |
| 159 | + |
| 160 | + if corrupt_prefix: |
| 161 | + prefix = os.urandom(1) |
| 162 | + else: |
| 163 | + prefix = bytearray(template[0]) |
| 164 | + |
| 165 | + if randomize_payload_size: |
| 166 | + payload = os.urandom(max(int(random.expovariate(0.5)), 50)) |
| 167 | + else: |
| 168 | + payload = os.urandom(template[1]) |
| 169 | + |
| 170 | + if corrupt_suffix: |
| 171 | + suffix = os.urandom(len(template[2])) |
| 172 | + else: |
| 173 | + suffix = bytearray(template[2]) |
| 174 | + |
| 175 | + val = b58encode_chk(prefix + payload + suffix) |
| 176 | + if random.randint(0,10)<1: # line corruption |
| 177 | + if randbool(): # add random character to end |
| 178 | + val += random.choice(b58chars) |
| 179 | + else: # replace random character in the middle |
| 180 | + n = random.randint(0, len(val)) |
| 181 | + val = val[0:n] + random.choice(b58chars) + val[n+1:] |
| 182 | + |
| 183 | + return val |
| 184 | + |
| 185 | +def gen_invalid_bech32_vector(template): |
| 186 | + '''Generate possibly invalid bech32 vector''' |
| 187 | + no_data = randbool(0.1) |
| 188 | + to_upper = randbool(0.1) |
| 189 | + hrp = template[0] |
| 190 | + witver = template[1] |
| 191 | + witprog = bytearray(os.urandom(template[2])) |
| 192 | + |
| 193 | + if no_data: |
| 194 | + rv = bech32_encode(hrp, []) |
| 195 | + else: |
| 196 | + data = [witver] + convertbits(witprog, 8, 5) |
| 197 | + if template[3] and not no_data: |
| 198 | + if template[2] % 5 in {2, 4}: |
| 199 | + data[-1] |= 1 |
| 200 | + else: |
| 201 | + data.append(0) |
| 202 | + rv = bech32_encode(hrp, data) |
| 203 | + |
| 204 | + if template[4]: |
| 205 | + i = len(rv) - random.randrange(1, 7) |
| 206 | + rv = rv[:i] + random.choice(CHARSET.replace(rv[i], '')) + rv[i + 1:] |
| 207 | + if template[5]: |
| 208 | + i = len(hrp) + 1 + random.randrange(0, len(rv) - len(hrp) - 4) |
| 209 | + rv = rv[:i] + rv[i:i + 4].upper() + rv[i + 4:] |
| 210 | + |
| 211 | + if to_upper: |
| 212 | + rv = rv.swapcase() |
| 213 | + |
| 214 | + return rv |
| 215 | + |
| 216 | +def randbool(p = 0.5): |
| 217 | + '''Return True with P(p)''' |
| 218 | + return random.random() < p |
| 219 | + |
| 220 | +def gen_invalid_vectors(): |
| 221 | + '''Generate invalid test vectors''' |
| 222 | + # start with some manual edge-cases |
| 223 | + yield "", |
| 224 | + yield "x", |
| 225 | + glist = [gen_invalid_base58_vector, gen_invalid_bech32_vector] |
| 226 | + tlist = [templates, bech32_ng_templates] |
| 227 | + while True: |
| 228 | + for template, invalid_vector_generator in [(t, g) for g, l in zip(glist, tlist) for t in l]: |
| 229 | + val = invalid_vector_generator(template) |
| 230 | + if not is_valid(val): |
| 231 | + yield val, |
| 232 | + |
| 233 | +if __name__ == '__main__': |
| 234 | + import sys |
| 235 | + import json |
| 236 | + iters = {'valid':gen_valid_vectors, 'invalid':gen_invalid_vectors} |
| 237 | + try: |
| 238 | + uiter = iters[sys.argv[1]] |
| 239 | + except IndexError: |
| 240 | + uiter = gen_valid_vectors |
| 241 | + try: |
| 242 | + count = int(sys.argv[2]) |
| 243 | + except IndexError: |
| 244 | + count = 0 |
| 245 | + |
| 246 | + data = list(islice(uiter(), count)) |
| 247 | + json.dump(data, sys.stdout, sort_keys=True, indent=4) |
| 248 | + sys.stdout.write('\n') |
| 249 | + |
0 commit comments