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