gentest.py 5.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184
  1. #!/usr/bin/env python
  2. #
  3. # mmgen = Multi-Mode GENerator, command-line Bitcoin cold storage solution
  4. # Copyright (C)2013-2017 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
  29. rounds = 100
  30. opts_data = lambda: {
  31. 'desc': "Test address generation in various ways",
  32. 'usage':'[options] [spec] [rounds | dump file]',
  33. 'options': """
  34. -h, --help Print this help message
  35. --, --longhelp Print help message for long options (common options)
  36. -q, --quiet Produce quieter output
  37. -s, --segwit Generate Segwit (P2SH-P2WPKH) addresses
  38. -v, --verbose Produce more verbose output
  39. """,
  40. 'notes': """
  41. Tests:
  42. A/B: {prog} a:b [rounds] (compare output of two key generators)
  43. Speed: {prog} a [rounds] (test speed of one key generator)
  44. Compare: {prog} a <dump file> (compare output of a key generator against wallet dump)
  45. where a and b are one of:
  46. '1' - native Python ecdsa library (very slow)
  47. '2' - bitcoincore.org's secp256k1 library (default from v0.8.6)
  48. EXAMPLES:
  49. {prog} 1:2 100
  50. (compare output of native Python ECDSA with secp256k1 library, 100 rounds)
  51. {prog} 2 1000
  52. (test speed of secp256k1 library address generation, 1000 rounds)
  53. {prog} 2 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. sys.argv = [sys.argv[0]] + ['--skip-cfg-file'] + sys.argv[1:]
  58. cmd_args = opts.init(opts_data,add_opts=['exact_output'])
  59. if not 1 <= len(cmd_args) <= 2: opts.usage()
  60. urounds,fh = None,None
  61. dump = []
  62. if len(cmd_args) == 2:
  63. try:
  64. urounds = int(cmd_args[1])
  65. assert urounds > 0
  66. except:
  67. try:
  68. fh = open(cmd_args[1])
  69. except:
  70. die(1,"Second argument must be filename or positive integer")
  71. else:
  72. for line in fh.readlines():
  73. if 'addr=' in line:
  74. x,addr = line.split('addr=')
  75. dump.append([x.split()[0],addr.split()[0]])
  76. if urounds: rounds = urounds
  77. a,b = None,None
  78. try:
  79. a,b = cmd_args[0].split(':')
  80. except:
  81. try:
  82. a = cmd_args[0]
  83. a = int(a)
  84. assert 1 <= a <= len(g.key_generators)
  85. except:
  86. die(1,"First argument must be one or two generator IDs, colon separated")
  87. else:
  88. try:
  89. a,b = int(a),int(b)
  90. for i in (a,b): assert 1 <= i <= len(g.key_generators)
  91. assert a != b
  92. except:
  93. die(1,"%s: invalid generator IDs" % cmd_args[0])
  94. def match_error(sec,wif,a_addr,b_addr,a,b):
  95. m = ['','py-ecdsa','secp256k1','dump']
  96. qmsg_r(red('\nERROR: Addresses do not match!'))
  97. die(3,"""
  98. sec key : {}
  99. WIF key : {}
  100. {a:10}: {}
  101. {b:10}: {}
  102. """.format(sec,wif,a_addr,b_addr,pnm=g.proj_name,a=m[a],b=m[b]).rstrip())
  103. # Begin execution
  104. mmtype = ('L','S')[bool(opt.segwit)]
  105. compressed = True
  106. from mmgen.addr import KeyGenerator,AddrGenerator
  107. from mmgen.obj import PrivKey
  108. ag = AddrGenerator(('p2pkh','segwit')[bool(opt.segwit)])
  109. if a and b:
  110. m = "Comparing address generators '{}' and '{}'"
  111. qmsg(green(m.format(g.key_generators[a-1],g.key_generators[b-1])))
  112. last_t = time.time()
  113. kg_a = KeyGenerator(a)
  114. kg_b = KeyGenerator(b)
  115. for i in range(rounds):
  116. if time.time() - last_t >= 0.1:
  117. qmsg_r('\rRound %s/%s ' % (i+1,rounds))
  118. last_t = time.time()
  119. sec = PrivKey(os.urandom(32),compressed)
  120. a_addr = ag.to_addr(kg_a.to_pubhex(sec))
  121. b_addr = ag.to_addr(kg_b.to_pubhex(sec))
  122. vmsg('\nkey: %s\naddr: %s\n' % (sec.wif,a_addr))
  123. if a_addr != b_addr:
  124. match_error(sec,sec.wif,a_addr,b_addr,a,b)
  125. if not opt.segwit:
  126. compressed = not compressed
  127. qmsg_r('\rRound %s/%s ' % (i+1,rounds))
  128. qmsg(green(('\n','')[bool(opt.verbose)] + 'OK'))
  129. elif a and not fh:
  130. m = "Testing speed of address generator '{}'"
  131. qmsg(green(m.format(g.key_generators[a-1])))
  132. from struct import pack,unpack
  133. seed = os.urandom(28)
  134. print 'Incrementing key with each round'
  135. print 'Starting key:', hexlify(seed+pack('I',0))
  136. import time
  137. start = last_t = time.time()
  138. kg = KeyGenerator(a)
  139. for i in range(rounds):
  140. if time.time() - last_t >= 0.1:
  141. qmsg_r('\rRound %s/%s ' % (i+1,rounds))
  142. last_t = time.time()
  143. sec = PrivKey(seed+pack('I',i),compressed)
  144. a_addr = ag.to_addr(kg.to_pubhex(sec))
  145. vmsg('\nkey: %s\naddr: %s\n' % (sec.wif,a_addr))
  146. if not opt.segwit:
  147. compressed = not compressed
  148. qmsg_r('\rRound %s/%s ' % (i+1,rounds))
  149. qmsg('\n{} addresses generated in {:.2f} seconds'.format(rounds,time.time()-start))
  150. elif a and dump:
  151. m = "Comparing output of address generator '{}' against wallet dump '{}'"
  152. qmsg(green(m.format(g.key_generators[a-1],cmd_args[1])))
  153. kg = KeyGenerator(a)
  154. for n,[wif,a_addr] in enumerate(dump,1):
  155. qmsg_r('\rKey %s/%s ' % (n,len(dump)))
  156. try:
  157. sec = PrivKey(wif=wif)
  158. except:
  159. die(2,'\nInvalid {}net WIF address in dump file: {}'.format(('main','test')[g.testnet],wif))
  160. compressed = wif[0] != ('5','9')[g.testnet]
  161. b_addr = ag.to_addr(kg.to_pubhex(sec))
  162. if a_addr != b_addr:
  163. match_error(sec,wif,a_addr,b_addr,3,a)
  164. qmsg(green(('\n','')[bool(opt.verbose)] + 'OK'))