|
15 | 15 | from itertools import islice
|
16 | 16 | from base58 import b58encode_chk, b58decode_chk, b58chars
|
17 | 17 | import random
|
18 |
| -from segwit_addr import bech32_encode, decode_segwit_address, convertbits, CHARSET |
| 18 | +from segwit_addr import bech32_encode, decode_segwit_address, convertbits, CHARSET, Encoding |
19 | 19 |
|
20 | 20 | # key types
|
21 | 21 | PUBKEY_ADDRESS = 0
|
|
32 | 32 | OP_0 = 0x00
|
33 | 33 | OP_1 = 0x51
|
34 | 34 | OP_2 = 0x52
|
| 35 | +OP_3 = 0x53 |
35 | 36 | OP_16 = 0x60
|
36 | 37 | OP_DUP = 0x76
|
37 | 38 | OP_EQUAL = 0x87
|
|
44 | 45 | script_suffix = (OP_EQUAL,)
|
45 | 46 | p2wpkh_prefix = (OP_0, 20)
|
46 | 47 | p2wsh_prefix = (OP_0, 32)
|
| 48 | +p2tr_prefix = (OP_1, 32) |
47 | 49 |
|
48 | 50 | metadata_keys = ['isPrivkey', 'chain', 'isCompressed', 'tryCaseFlip']
|
49 | 51 | # templates for valid sequences
|
|
65 | 67 | ]
|
66 | 68 | # templates for valid bech32 sequences
|
67 | 69 | bech32_templates = [
|
68 |
| - # hrp, version, witprog_size, metadata, output_prefix |
69 |
| - ('bc', 0, 20, (False, 'main', None, True), p2wpkh_prefix), |
70 |
| - ('bc', 0, 32, (False, 'main', None, True), p2wsh_prefix), |
71 |
| - ('bc', 1, 2, (False, 'main', None, True), (OP_1, 2)), |
72 |
| - ('tb', 0, 20, (False, 'test', None, True), p2wpkh_prefix), |
73 |
| - ('tb', 0, 32, (False, 'test', None, True), p2wsh_prefix), |
74 |
| - ('tb', 2, 16, (False, 'test', None, True), (OP_2, 16)), |
75 |
| - ('bcrt', 0, 20, (False, 'regtest', None, True), p2wpkh_prefix), |
76 |
| - ('bcrt', 0, 32, (False, 'regtest', None, True), p2wsh_prefix), |
77 |
| - ('bcrt', 16, 40, (False, 'regtest', None, True), (OP_16, 40)) |
| 70 | + # hrp, version, witprog_size, metadata, encoding, output_prefix |
| 71 | + ('bc', 0, 20, (False, 'main', None, True), Encoding.BECH32, p2wpkh_prefix), |
| 72 | + ('bc', 0, 32, (False, 'main', None, True), Encoding.BECH32, p2wsh_prefix), |
| 73 | + ('bc', 1, 32, (False, 'main', None, True), Encoding.BECH32M, p2tr_prefix), |
| 74 | + ('bc', 2, 2, (False, 'main', None, True), Encoding.BECH32M, (OP_2, 2)), |
| 75 | + ('tb', 0, 20, (False, 'test', None, True), Encoding.BECH32, p2wpkh_prefix), |
| 76 | + ('tb', 0, 32, (False, 'test', None, True), Encoding.BECH32, p2wsh_prefix), |
| 77 | + ('tb', 1, 32, (False, 'test', None, True), Encoding.BECH32M, p2tr_prefix), |
| 78 | + ('tb', 3, 16, (False, 'test', None, True), Encoding.BECH32M, (OP_3, 16)), |
| 79 | + ('bcrt', 0, 20, (False, 'regtest', None, True), Encoding.BECH32, p2wpkh_prefix), |
| 80 | + ('bcrt', 0, 32, (False, 'regtest', None, True), Encoding.BECH32, p2wsh_prefix), |
| 81 | + ('bcrt', 1, 32, (False, 'regtest', None, True), Encoding.BECH32M, p2tr_prefix), |
| 82 | + ('bcrt', 16, 40, (False, 'regtest', None, True), Encoding.BECH32M, (OP_16, 40)) |
78 | 83 | ]
|
79 | 84 | # templates for invalid bech32 sequences
|
80 | 85 | bech32_ng_templates = [
|
81 |
| - # hrp, version, witprog_size, invalid_bech32, invalid_checksum, invalid_char |
82 |
| - ('tc', 0, 20, False, False, False), |
83 |
| - ('tb', 17, 32, False, False, False), |
84 |
| - ('bcrt', 3, 1, False, False, False), |
85 |
| - ('bc', 15, 41, False, False, False), |
86 |
| - ('tb', 0, 16, False, False, False), |
87 |
| - ('bcrt', 0, 32, True, False, False), |
88 |
| - ('bc', 0, 16, True, False, False), |
89 |
| - ('tb', 0, 32, False, True, False), |
90 |
| - ('bcrt', 0, 20, False, False, True) |
| 86 | + # hrp, version, witprog_size, encoding, invalid_bech32, invalid_checksum, invalid_char |
| 87 | + ('tc', 0, 20, Encoding.BECH32, False, False, False), |
| 88 | + ('bt', 1, 32, Encoding.BECH32M, False, False, False), |
| 89 | + ('tb', 17, 32, Encoding.BECH32M, False, False, False), |
| 90 | + ('bcrt', 3, 1, Encoding.BECH32M, False, False, False), |
| 91 | + ('bc', 15, 41, Encoding.BECH32M, False, False, False), |
| 92 | + ('tb', 0, 16, Encoding.BECH32, False, False, False), |
| 93 | + ('bcrt', 0, 32, Encoding.BECH32, True, False, False), |
| 94 | + ('bc', 0, 16, Encoding.BECH32, True, False, False), |
| 95 | + ('tb', 0, 32, Encoding.BECH32, False, True, False), |
| 96 | + ('bcrt', 0, 20, Encoding.BECH32, False, False, True), |
| 97 | + ('bc', 0, 20, Encoding.BECH32M, False, False, False), |
| 98 | + ('tb', 0, 32, Encoding.BECH32M, False, False, False), |
| 99 | + ('bcrt', 0, 20, Encoding.BECH32M, False, False, False), |
| 100 | + ('bc', 1, 32, Encoding.BECH32, False, False, False), |
| 101 | + ('tb', 2, 16, Encoding.BECH32, False, False, False), |
| 102 | + ('bcrt', 16, 20, Encoding.BECH32, False, False, False), |
91 | 103 | ]
|
92 | 104 |
|
93 | 105 | def is_valid(v):
|
@@ -127,8 +139,9 @@ def gen_valid_bech32_vector(template):
|
127 | 139 | hrp = template[0]
|
128 | 140 | witver = template[1]
|
129 | 141 | witprog = bytearray(os.urandom(template[2]))
|
130 |
| - dst_prefix = bytearray(template[4]) |
131 |
| - rv = bech32_encode(hrp, [witver] + convertbits(witprog, 8, 5)) |
| 142 | + encoding = template[4] |
| 143 | + dst_prefix = bytearray(template[5]) |
| 144 | + rv = bech32_encode(hrp, [witver] + convertbits(witprog, 8, 5), encoding) |
132 | 145 | return rv, dst_prefix + witprog
|
133 | 146 |
|
134 | 147 | def gen_valid_vectors():
|
@@ -186,22 +199,23 @@ def gen_invalid_bech32_vector(template):
|
186 | 199 | hrp = template[0]
|
187 | 200 | witver = template[1]
|
188 | 201 | witprog = bytearray(os.urandom(template[2]))
|
| 202 | + encoding = template[3] |
189 | 203 |
|
190 | 204 | if no_data:
|
191 |
| - rv = bech32_encode(hrp, []) |
| 205 | + rv = bech32_encode(hrp, [], encoding) |
192 | 206 | else:
|
193 | 207 | data = [witver] + convertbits(witprog, 8, 5)
|
194 |
| - if template[3] and not no_data: |
| 208 | + if template[4] and not no_data: |
195 | 209 | if template[2] % 5 in {2, 4}:
|
196 | 210 | data[-1] |= 1
|
197 | 211 | else:
|
198 | 212 | data.append(0)
|
199 |
| - rv = bech32_encode(hrp, data) |
| 213 | + rv = bech32_encode(hrp, data, encoding) |
200 | 214 |
|
201 |
| - if template[4]: |
| 215 | + if template[5]: |
202 | 216 | i = len(rv) - random.randrange(1, 7)
|
203 | 217 | rv = rv[:i] + random.choice(CHARSET.replace(rv[i], '')) + rv[i + 1:]
|
204 |
| - if template[5]: |
| 218 | + if template[6]: |
205 | 219 | i = len(hrp) + 1 + random.randrange(0, len(rv) - len(hrp) - 4)
|
206 | 220 | rv = rv[:i] + rv[i:i + 4].upper() + rv[i + 4:]
|
207 | 221 |
|
|
0 commit comments