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