gentest.py 5.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173
  1. #!/usr/bin/env python
  2. #
  3. # mmgen = Multi-Mode GENerator, command-line Bitcoin cold storage solution
  4. # Copyright (C)2013-2016 Philemon <mmgen-py@yandex.com>
  5. #
  6. # This program is free software: you can redistribute it and/or modify
  7. # it under the terms of the GNU General Public License as published by
  8. # the Free Software Foundation, either version 3 of the License, or
  9. # (at your option) any later version.
  10. #
  11. # This program is distributed in the hope that it will be useful,
  12. # but WITHOUT ANY WARRANTY; without even the implied warranty of
  13. # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
  14. # GNU General Public License for more details.
  15. #
  16. # You should have received a copy of the GNU General Public License
  17. # along with this program. If not, see <http://www.gnu.org/licenses/>.
  18. """
  19. test/gentest.py: Bitcoin key/address generation tests for the MMGen suite
  20. """
  21. import sys,os
  22. pn = os.path.dirname(sys.argv[0])
  23. os.chdir(os.path.join(pn,os.pardir))
  24. sys.path.__setitem__(0,os.path.abspath(os.curdir))
  25. from binascii import hexlify
  26. # Import these _after_ local path's been added to sys.path
  27. from mmgen.common import *
  28. from mmgen.bitcoin import hex2wif,privnum2addr
  29. start_mscolor()
  30. rounds = 100
  31. opts_data = {
  32. 'desc': "Test address generation in various ways",
  33. 'usage':'[options] [spec] [rounds | dump file]',
  34. 'options': """
  35. -h, --help Print this help message
  36. -v, --verbose Produce more verbose output
  37. -q, --quiet Produce quieter output
  38. """,
  39. 'notes': """
  40. Tests:
  41. A/B: {prog} a:b [rounds] (compare output of two key generators)
  42. Speed: {prog} a [rounds] (test speed of one key generator)
  43. Compare: {prog} a <dump file> (compare output of a key generator against wallet dump)
  44. where a and b are one of:
  45. '1' - native Python ecdsa library (very slow)
  46. '2' - 'keyconv' utility from the 'vanitygen' package (old default)
  47. '3' - bitcoincore.org's secp256k1 library (default from v0.8.6)
  48. EXAMPLES:
  49. {prog} 2:3 1000
  50. (compare output of 'keyconv' with secp256k1 library, 1000 rounds)
  51. {prog} 3 1000
  52. (test speed of secp256k1 library address generation, 1000 rounds)
  53. {prog} 3 my.dump
  54. (compare addrs generated with secp256k1 library to bitcoind wallet dump)
  55. """.format(prog='gentest.py',pnm=g.proj_name,snum=rounds)
  56. }
  57. cmd_args = opts.init(opts_data,add_opts=['exact_output'])
  58. if not 1 <= len(cmd_args) <= 2: opts.usage()
  59. urounds,fh = None,None
  60. dump = []
  61. if len(cmd_args) == 2:
  62. try:
  63. urounds = int(cmd_args[1])
  64. assert urounds > 0
  65. except:
  66. try:
  67. fh = open(cmd_args[1])
  68. except:
  69. die(1,"Second argument must be filename or positive integer")
  70. else:
  71. for line in fh.readlines():
  72. if 'addr=' in line:
  73. x,addr = line.split('addr=')
  74. dump.append([x.split()[0],addr.split()[0]])
  75. if urounds: rounds = urounds
  76. a,b = None,None
  77. try:
  78. a,b = cmd_args[0].split(':')
  79. except:
  80. try:
  81. a = cmd_args[0]
  82. a = int(a)
  83. assert 1 <= a <= len(g.key_generators)
  84. except:
  85. die(1,"First argument must be one or two generator IDs, colon separated")
  86. else:
  87. try:
  88. a,b = int(a),int(b)
  89. for i in a,b: assert 1 <= i <= len(g.key_generators)
  90. assert a != b
  91. except:
  92. die(1,"%s: invalid generator IDs" % cmd_args[0])
  93. def match_error(sec,wif,a_addr,b_addr,a,b):
  94. m = ['','py-ecdsa','keyconv','secp256k1','dump']
  95. msg_r(red('\nERROR: Addresses do not match!'))
  96. die(3,"""
  97. sec key : {}
  98. WIF key : {}
  99. {a:10}: {}
  100. {b:10}: {}
  101. """.format(sec,wif,a_addr,b_addr,pnm=g.proj_name,a=m[a],b=m[b]).rstrip())
  102. if a and b:
  103. m = "Comparing address generators '{}' and '{}'"
  104. msg(green(m.format(g.key_generators[a-1],g.key_generators[b-1])))
  105. from mmgen.addr import get_privhex2addr_f
  106. gen_a = get_privhex2addr_f(selector=a)
  107. gen_b = get_privhex2addr_f(selector=b)
  108. compressed = False
  109. for i in range(1,rounds+1):
  110. msg_r('\rRound %s/%s ' % (i,rounds))
  111. sec = hexlify(os.urandom(32))
  112. wif = hex2wif(sec,compressed=compressed)
  113. a_addr = gen_a(sec,compressed)
  114. b_addr = gen_b(sec,compressed)
  115. vmsg('\nkey: %s\naddr: %s\n' % (wif,a_addr))
  116. if a_addr != b_addr:
  117. match_error(sec,wif,a_addr,b_addr,a,b)
  118. if a != 2 and b != 2:
  119. compressed = not compressed
  120. msg(green(('\n','')[bool(opt.verbose)] + 'OK'))
  121. elif a and not fh:
  122. m = "Testing speed of address generator '{}'"
  123. msg(green(m.format(g.key_generators[a-1])))
  124. from mmgen.addr import get_privhex2addr_f
  125. gen_a = get_privhex2addr_f(selector=a)
  126. import time
  127. start = time.time()
  128. from struct import pack,unpack
  129. seed = os.urandom(28)
  130. print 'Incrementing key with each round'
  131. print 'Starting key:', hexlify(seed+pack('I',0))
  132. compressed = False
  133. for i in range(rounds):
  134. if not opt.quiet: msg_r('\rRound %s/%s ' % (i+1,rounds))
  135. sec = hexlify(seed+pack('I',i))
  136. wif = hex2wif(sec,compressed=compressed)
  137. a_addr = gen_a(sec,compressed)
  138. vmsg('\nkey: %s\naddr: %s\n' % (wif,a_addr))
  139. if a != 2:
  140. compressed = not compressed
  141. elapsed = int(time.time() - start)
  142. if not opt.quiet: msg('')
  143. msg('%s addresses generated in %s second%s' % (rounds,elapsed,('s','')[elapsed==1]))
  144. elif a and dump:
  145. m = "Comparing output of address generator '{}' against wallet dump '{}'"
  146. msg(green(m.format(g.key_generators[a-1],cmd_args[1])))
  147. if a == 2:
  148. msg("NOTE: for compressed addresses, 'python-ecdsa' generator will be used")
  149. from mmgen.addr import get_privhex2addr_f
  150. gen_a = get_privhex2addr_f(selector=a)
  151. from mmgen.bitcoin import wif2hex
  152. for n,[wif,a_addr] in enumerate(dump,1):
  153. msg_r('\rKey %s/%s ' % (n,len(dump)))
  154. sec = wif2hex(wif)
  155. compressed = wif[0] != ('5','9')[g.testnet]
  156. b_addr = gen_a(sec,compressed)
  157. if a_addr != b_addr:
  158. match_error(sec,wif,a_addr,b_addr,1 if compressed and a==2 else a,4)
  159. msg(green(('\n','')[bool(opt.verbose)] + 'OK'))