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.2KB

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