aboutsummaryrefslogtreecommitdiff
path: root/contrib/testgen/gen_base58_test_vectors.py
blob: c76de5ce750e3d328a83c5d407e74bf901bb9d85 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
#!/usr/bin/env python3
# Copyright (c) 2012-2018 The Bitcoin Core developers
# Distributed under the MIT software license, see the accompanying
# file COPYING or http://www.opensource.org/licenses/mit-license.php.
'''
Generate valid and invalid base58 address and private key test vectors.

Usage:
    gen_base58_test_vectors.py valid 50 > ../../src/test/data/base58_keys_valid.json
    gen_base58_test_vectors.py invalid 50 > ../../src/test/data/base58_keys_invalid.json
'''
# 2012 Wladimir J. van der Laan
# Released under MIT License
import os
from itertools import islice
from base58 import b58encode_chk, b58decode_chk, b58chars
import random
from binascii import b2a_hex

# key types
PUBKEY_ADDRESS = 0
SCRIPT_ADDRESS = 5
PUBKEY_ADDRESS_TEST = 111
SCRIPT_ADDRESS_TEST = 196
PRIVKEY = 128
PRIVKEY_TEST = 239

metadata_keys = ['isPrivkey', 'isTestnet', 'addrType', 'isCompressed']
# templates for valid sequences
templates = [
  # prefix, payload_size, suffix, metadata
  #                                  None = N/A
  ((PUBKEY_ADDRESS,),      20, (),   (False, False, 'pubkey', None)),
  ((SCRIPT_ADDRESS,),      20, (),   (False, False, 'script',  None)),
  ((PUBKEY_ADDRESS_TEST,), 20, (),   (False, True,  'pubkey', None)),
  ((SCRIPT_ADDRESS_TEST,), 20, (),   (False, True,  'script',  None)),
  ((PRIVKEY,),             32, (),   (True,  False, None,  False)),
  ((PRIVKEY,),             32, (1,), (True,  False, None,  True)),
  ((PRIVKEY_TEST,),        32, (),   (True,  True,  None,  False)),
  ((PRIVKEY_TEST,),        32, (1,), (True,  True,  None,  True))
]

def is_valid(v):
    '''Check vector v for validity'''
    result = b58decode_chk(v)
    if result is None:
        return False
    for template in templates:
        prefix = bytearray(template[0])
        suffix = bytearray(template[2])
        if result.startswith(prefix) and result.endswith(suffix):
            if (len(result) - len(prefix) - len(suffix)) == template[1]:
                return True
    return False

def gen_valid_vectors():
    '''Generate valid test vectors'''
    while True:
        for template in templates:
            prefix = bytearray(template[0])
            payload = bytearray(os.urandom(template[1]))
            suffix = bytearray(template[2])
            rv = b58encode_chk(prefix + payload + suffix)
            assert is_valid(rv)
            metadata = {x: y for x, y in zip(metadata_keys,template[3]) if y is not None}
            hexrepr = b2a_hex(payload)
            if isinstance(hexrepr, bytes):
                hexrepr = hexrepr.decode('utf8')
            yield (rv, hexrepr, metadata)

def gen_invalid_vector(template, corrupt_prefix, randomize_payload_size, corrupt_suffix):
    '''Generate possibly invalid vector'''
    if corrupt_prefix:
        prefix = os.urandom(1)
    else:
        prefix = bytearray(template[0])

    if randomize_payload_size:
        payload = os.urandom(max(int(random.expovariate(0.5)), 50))
    else:
        payload = os.urandom(template[1])

    if corrupt_suffix:
        suffix = os.urandom(len(template[2]))
    else:
        suffix = bytearray(template[2])

    return b58encode_chk(prefix + payload + suffix)

def randbool(p = 0.5):
    '''Return True with P(p)'''
    return random.random() < p

def gen_invalid_vectors():
    '''Generate invalid test vectors'''
    # start with some manual edge-cases
    yield "",
    yield "x",
    while True:
        # kinds of invalid vectors:
        #   invalid prefix
        #   invalid payload length
        #   invalid (randomized) suffix (add random data)
        #   corrupt checksum
        for template in templates:
            val = gen_invalid_vector(template, randbool(0.2), randbool(0.2), randbool(0.2))
            if random.randint(0,10)<1: # line corruption
                if randbool(): # add random character to end
                    val += random.choice(b58chars)
                else: # replace random character in the middle
                    n = random.randint(0, len(val))
                    val = val[0:n] + random.choice(b58chars) + val[n+1:]
            if not is_valid(val):
                yield val,

if __name__ == '__main__':
    import sys
    import json
    iters = {'valid':gen_valid_vectors, 'invalid':gen_invalid_vectors}
    try:
        uiter = iters[sys.argv[1]]
    except IndexError:
        uiter = gen_valid_vectors
    try:
        count = int(sys.argv[2])
    except IndexError:
        count = 0

    data = list(islice(uiter(), count))
    json.dump(data, sys.stdout, sort_keys=True, indent=4)
    sys.stdout.write('\n')