contrib/testgen/gen_base58_test_vectors.py

   1 #!/usr/bin/env python
   2 # Copyright (c) 2012-2016 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 address and private key test vectors.
   7
   8 Usage:
   9     gen_base58_test_vectors.py valid 50 > ../../src/test/data/base58_keys_valid.json
  10     gen_base58_test_vectors.py invalid 50 > ../../src/test/data/base58_keys_invalid.json
  11 '''
  12 # 2012 Wladimir J. van der Laan
  13 # Released under MIT License
  14 import os
  15 from itertools import islice
  16 from base58 import b58encode_chk, b58decode_chk, b58chars
  17 import random
  18 from binascii import b2a_hex
  19
  20 # key types
  21 PUBKEY_ADDRESS = 0
  22 SCRIPT_ADDRESS = 5
  23 PUBKEY_ADDRESS_TEST = 111
  24 SCRIPT_ADDRESS_TEST = 196
  25 PRIVKEY = 128
  26 PRIVKEY_TEST = 239
  27
  28 metadata_keys = ['isPrivkey', 'isTestnet', 'addrType', 'isCompressed']
  29 # templates for valid sequences
  30 templates = [
  31   # prefix, payload_size, suffix, metadata
  32   #                                  None = N/A
  33   ((PUBKEY_ADDRESS,),      20, (),   (False, False, 'pubkey', None)),
  34   ((SCRIPT_ADDRESS,),      20, (),   (False, False, 'script',  None)),
  35   ((PUBKEY_ADDRESS_TEST,), 20, (),   (False, True,  'pubkey', None)),
  36   ((SCRIPT_ADDRESS_TEST,), 20, (),   (False, True,  'script',  None)),
  37   ((PRIVKEY,),             32, (),   (True,  False, None,  False)),
  38   ((PRIVKEY,),             32, (1,), (True,  False, None,  True)),
  39   ((PRIVKEY_TEST,),        32, (),   (True,  True,  None,  False)),
  40   ((PRIVKEY_TEST,),        32, (1,), (True,  True,  None,  True))
  41 ]
  42
  43 def is_valid(v):
  44     '''Check vector v for validity'''
  45     result = b58decode_chk(v)
  46     if result is None:
  47         return False
  48     for template in templates:
  49         prefix = str(bytearray(template[0]))
  50         suffix = str(bytearray(template[2]))
  51         if result.startswith(prefix) and result.endswith(suffix):
  52             if (len(result) - len(prefix) - len(suffix)) == template[1]:
  53                 return True
  54     return False
  55
  56 def gen_valid_vectors():
  57     '''Generate valid test vectors'''
  58     while True:
  59         for template in templates:
  60             prefix = str(bytearray(template[0]))
  61             payload = os.urandom(template[1])
  62             suffix = str(bytearray(template[2]))
  63             rv = b58encode_chk(prefix + payload + suffix)
  64             assert is_valid(rv)
  65             metadata = dict([(x,y) for (x,y) in zip(metadata_keys,template[3]) if y is not None])
  66             yield (rv, b2a_hex(payload), metadata)
  67
  68 def gen_invalid_vector(template, corrupt_prefix, randomize_payload_size, corrupt_suffix):
  69     '''Generate possibly invalid vector'''
  70     if corrupt_prefix:
  71         prefix = os.urandom(1)
  72     else:
  73         prefix = str(bytearray(template[0]))
  74
  75     if randomize_payload_size:
  76         payload = os.urandom(max(int(random.expovariate(0.5)), 50))
  77     else:
  78         payload = os.urandom(template[1])
  79
  80     if corrupt_suffix:
  81         suffix = os.urandom(len(template[2]))
  82     else:
  83         suffix = str(bytearray(template[2]))
  84
  85     return b58encode_chk(prefix + payload + suffix)
  86
  87 def randbool(p = 0.5):
  88     '''Return True with P(p)'''
  89     return random.random() < p
  90
  91 def gen_invalid_vectors():
  92     '''Generate invalid test vectors'''
  93     # start with some manual edge-cases
  94     yield "",
  95     yield "x",
  96     while True:
  97         # kinds of invalid vectors:
  98         #   invalid prefix
  99         #   invalid payload length
 100         #   invalid (randomized) suffix (add random data)
 101         #   corrupt checksum
 102         for template in templates:
 103             val = gen_invalid_vector(template, randbool(0.2), randbool(0.2), randbool(0.2))
 104             if random.randint(0,10)<1: # line corruption
 105                 if randbool(): # add random character to end
 106                     val += random.choice(b58chars)
 107                 else: # replace random character in the middle
 108                     n = random.randint(0, len(val))
 109                     val = val[0:n] + random.choice(b58chars) + val[n+1:]
 110             if not is_valid(val):
 111                 yield val,
 112
 113 if __name__ == '__main__':
 114     import sys, json
 115     iters = {'valid':gen_valid_vectors, 'invalid':gen_invalid_vectors}
 116     try:
 117         uiter = iters[sys.argv[1]]
 118     except IndexError:
 119         uiter = gen_valid_vectors
 120     try:
 121         count = int(sys.argv[2])
 122     except IndexError:
 123         count = 0
 124
 125     data = list(islice(uiter(), count))
 126     json.dump(data, sys.stdout, sort_keys=True, indent=4)
 127     sys.stdout.write('\n')
 128