You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

gen_base58_test_vectors.py 4.4KB

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