gen_key_io_test_vectors.py raw
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
266