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