gentest.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305
  1. #!/usr/bin/env python
  2. #
  3. # mmgen = Multi-Mode GENerator, command-line Bitcoin cold storage solution
  4. # Copyright (C)2013-2018 The MMGen Project <mmgen@tuta.io>
  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: Cryptocoin 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. os.environ['MMGEN_TEST_SUITE'] = '1'
  26. from binascii import hexlify
  27. # Import these _after_ local path's been added to sys.path
  28. from mmgen.common import *
  29. from mmgen.obj import MMGenAddrType
  30. rounds = 100
  31. opts_data = lambda: {
  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. -a, --all Test all supported coins for external generator 'ext'
  37. --, --longhelp Print help message for long options (common options)
  38. -q, --quiet Produce quieter output
  39. -t, --type=t Specify address type (valid options: 'compressed','segwit','zcash_z')
  40. -v, --verbose Produce more verbose output
  41. """,
  42. 'notes': """
  43. Tests:
  44. A/B: {prog} a:b [rounds] (compare output of two key generators)
  45. Speed: {prog} a [rounds] (test speed of one key generator)
  46. Compare: {prog} a <dump file> (compare output of a key generator against wallet dump)
  47. where a and b are one of:
  48. '1' - native Python ecdsa library (very slow)
  49. '2' - bitcoincore.org's secp256k1 library (default from v0.8.6)
  50. EXAMPLES:
  51. {prog} 1:2 100
  52. (compare output of native Python ECDSA with secp256k1 library, 100 rounds)
  53. {prog} 2:ext 100
  54. (compare output of secp256k1 library with external library (see below), 100 rounds)
  55. {prog} 2 1000
  56. (test speed of secp256k1 library address generation, 1000 rounds)
  57. {prog} 2 my.dump
  58. (compare addrs generated with secp256k1 library to {dn} wallet dump)
  59. External libraries required for the 'ext' generator:
  60. + pyethereum (for ETH,ETC) https://github.com/ethereum/pyethereum
  61. + zcash-mini (for zcash_z addresses) https://github.com/FiloSottile/zcash-mini
  62. + pycoin (for supported coins) https://github.com/richardkiss/pycoin
  63. + keyconv (for all other coins) https://github.com/exploitagency/vanitygen-plus
  64. ('keyconv' generates uncompressed addresses only)
  65. """.format(prog='gentest.py',pnm=g.proj_name,snum=rounds,dn=g.proto.daemon_name)
  66. }
  67. sys.argv = [sys.argv[0]] + ['--skip-cfg-file'] + sys.argv[1:]
  68. cmd_args = opts.init(opts_data,add_opts=['exact_output'])
  69. if not 1 <= len(cmd_args) <= 2: opts.usage()
  70. addr_type = MMGenAddrType(opt.type or g.proto.dfl_mmtype)
  71. def pyethereum_sec2addr(sec):
  72. return sec,eth.privtoaddr(sec).encode('hex')
  73. def keyconv_sec2addr(sec):
  74. p = sp.Popen(['keyconv','-C',g.coin,sec.wif],stderr=sp.PIPE,stdout=sp.PIPE)
  75. o = p.stdout.read().splitlines()
  76. # print p.stderr.read()
  77. return o[1].split()[1],o[0].split()[1]
  78. def zcash_mini_sec2addr(sec):
  79. p = sp.Popen(['zcash-mini','-key','-simple'],stderr=sp.PIPE,stdin=sp.PIPE,stdout=sp.PIPE)
  80. p.stdin.write(sec.wif+'\n')
  81. o = p.stdout.read().split()
  82. return sec.wif,o[0],o[-1]
  83. def pycoin_sec2addr(sec):
  84. coin = ci.external_tests['testnet']['pycoin'][g.coin] if g.testnet else g.coin
  85. key = pcku.parse_key(sec,PREFIX_TRANSFORMS,coin)
  86. if key is None: die(1,"can't parse {}".format(sec))
  87. o = pcku.create_output(sec,key)[0]
  88. suf = ('_uncompressed','')[addr_type.compressed]
  89. wif = o['wif{}'.format(suf)]
  90. addr = o['p2sh_segwit' if addr_type.name == 'segwit' else '{}_address{}'.format(coin,suf)]
  91. return wif,addr
  92. # pycoin/networks/all.py pycoin/networks/legacy_networks.py
  93. def init_external_prog():
  94. global b,b_desc,ext_lib,ext_sec2addr,sp,eth,pcku,PREFIX_TRANSFORMS,addr_type
  95. def test_support(k):
  96. if b == k: return True
  97. if b != 'ext' and b != k: return False
  98. if g.coin in ci.external_tests['mainnet'][k] and not g.testnet: return True
  99. if g.coin in ci.external_tests['testnet'][k]: return True
  100. return False
  101. if b == 'zcash_mini' or addr_type.name == 'zcash_z':
  102. import subprocess as sp
  103. from mmgen.protocol import init_coin
  104. ext_sec2addr = zcash_mini_sec2addr
  105. ext_lib = 'zcash_mini'
  106. init_coin('zec')
  107. addr_type = MMGenAddrType('Z')
  108. elif test_support('pyethereum'):
  109. try:
  110. import ethereum.utils as eth
  111. except:
  112. raise ImportError,"Unable to import 'pyethereum' module. Is pyethereum installed?"
  113. ext_sec2addr = pyethereum_sec2addr
  114. ext_lib = 'pyethereum'
  115. elif test_support('pycoin'):
  116. try:
  117. import pycoin.cmds.ku as pcku
  118. except:
  119. raise ImportError,"Unable to import module 'ku'. Is pycoin installed?"
  120. PREFIX_TRANSFORMS = pcku.prefix_transforms_for_network(g.coin)
  121. ext_sec2addr = pycoin_sec2addr
  122. ext_lib = 'pycoin'
  123. elif test_support('keyconv'):
  124. import subprocess as sp
  125. ext_sec2addr = keyconv_sec2addr
  126. ext_lib = 'keyconv'
  127. else:
  128. m = '{}: coin supported by MMGen but unsupported by gentest.py for {}'
  129. raise ValueError,m.format(g.coin,('mainnet','testnet')[g.testnet])
  130. b_desc = ext_lib
  131. b = 'ext'
  132. def match_error(sec,wif,a_addr,b_addr,a,b):
  133. qmsg_r(red('\nERROR: Values do not match!'))
  134. die(3,"""
  135. sec key : {}
  136. WIF key : {}
  137. {a:10}: {}
  138. {b:10}: {}
  139. """.format(sec,wif,a_addr,b_addr,pnm=g.proj_name,a=kg_a.desc,b=b_desc).rstrip())
  140. def compare_test():
  141. for k in ('segwit','compressed'):
  142. if addr_type.name == k and g.coin not in ci.external_tests_segwit_compressed[k]:
  143. msg('{} testing not supported for coin {}'.format(addr_type.name.capitalize(),g.coin))
  144. return
  145. if 'ext_lib' in globals():
  146. if g.coin not in ci.external_tests[('mainnet','testnet')[g.testnet]][ext_lib]:
  147. msg("Coin '{}' incompatible with external generator '{}'".format(g.coin,ext_lib))
  148. return
  149. m = "Comparing address generators '{}' and '{}' for coin {}"
  150. last_t = time.time()
  151. qmsg(green(m.format(kg_a.desc,(ext_lib if b == 'ext' else kg_b.desc),g.coin)))
  152. for i in range(rounds):
  153. if opt.verbose or time.time() - last_t >= 0.1:
  154. qmsg_r('\rRound %s/%s ' % (i+1,rounds))
  155. last_t = time.time()
  156. sec = PrivKey(os.urandom(32),compressed=addr_type.compressed,pubkey_type=addr_type.pubkey_type)
  157. ph = kg_a.to_pubhex(sec)
  158. a_addr = ag.to_addr(ph)
  159. if addr_type.name == 'zcash_z':
  160. a_vk = ag.to_viewkey(ph)
  161. if b == 'ext':
  162. if addr_type.name == 'zcash_z':
  163. b_wif,b_addr,b_vk = ext_sec2addr(sec)
  164. if b_vk != a_vk:
  165. match_error(sec,sec.wif,a_vk,b_vk,a,b)
  166. else:
  167. b_wif,b_addr = ext_sec2addr(sec)
  168. if b_wif != sec.wif:
  169. match_error(sec,sec.wif,sec.wif,b_wif,a,b)
  170. else:
  171. b_addr = ag.to_addr(kg_b.to_pubhex(sec))
  172. vmsg('\nkey: %s\naddr: %s\n' % (sec.wif,a_addr))
  173. if a_addr != b_addr:
  174. match_error(sec,sec.wif,a_addr,b_addr,a,ext_lib if b == 'ext' else b)
  175. qmsg_r('\rRound %s/%s ' % (i+1,rounds))
  176. qmsg(green(('\n','')[bool(opt.verbose)] + 'OK'))
  177. def speed_test():
  178. m = "Testing speed of address generator '{}' for coin {}"
  179. qmsg(green(m.format(kg_a.desc,g.coin)))
  180. from struct import pack,unpack
  181. seed = os.urandom(28)
  182. print 'Incrementing key with each round'
  183. print 'Starting key:', hexlify(seed+pack('I',0))
  184. import time
  185. start = last_t = time.time()
  186. for i in range(rounds):
  187. if time.time() - last_t >= 0.1:
  188. qmsg_r('\rRound %s/%s ' % (i+1,rounds))
  189. last_t = time.time()
  190. sec = PrivKey(seed+pack('I',i),compressed=addr_type.compressed,pubkey_type=addr_type.pubkey_type)
  191. a_addr = ag.to_addr(kg_a.to_pubhex(sec))
  192. vmsg('\nkey: %s\naddr: %s\n' % (sec.wif,a_addr))
  193. qmsg_r('\rRound %s/%s ' % (i+1,rounds))
  194. qmsg('\n{} addresses generated in {:.2f} seconds'.format(rounds,time.time()-start))
  195. def dump_test():
  196. m = "Comparing output of address generator '{}' against wallet dump '{}'"
  197. qmsg(green(m.format(kg_a.desc,cmd_args[1])))
  198. for n,[wif,a_addr] in enumerate(dump,1):
  199. qmsg_r('\rKey %s/%s ' % (n,len(dump)))
  200. try:
  201. sec = PrivKey(wif=wif)
  202. except:
  203. die(2,'\nInvalid {}net WIF address in dump file: {}'.format(('main','test')[g.testnet],wif))
  204. b_addr = ag.to_addr(kg_a.to_pubhex(sec))
  205. vmsg('\nwif: {}\naddr: {}\n'.format(wif,b_addr))
  206. if a_addr != b_addr:
  207. match_error(sec,wif,a_addr,b_addr,3,a)
  208. qmsg(green(('\n','')[bool(opt.verbose)] + 'OK'))
  209. from mmgen.altcoin import CoinInfo as ci
  210. urounds,fh = None,None
  211. dump = []
  212. if len(cmd_args) == 2:
  213. try:
  214. urounds = int(cmd_args[1])
  215. assert urounds > 0
  216. except:
  217. try:
  218. fh = open(cmd_args[1])
  219. except:
  220. die(1,'Second argument must be filename or positive integer')
  221. else:
  222. for line in fh.readlines():
  223. if 'addr=' in line:
  224. x,addr = line.split('addr=')
  225. dump.append([x.split()[0],addr.split()[0]])
  226. if urounds: rounds = urounds
  227. a,b = None,None
  228. b_desc = 'unknown'
  229. try:
  230. a,b = cmd_args[0].split(':')
  231. except:
  232. try:
  233. a = cmd_args[0]
  234. a = int(a)
  235. assert 1 <= a <= len(g.key_generators)
  236. except:
  237. die(1,'First argument must be one or two generator IDs, colon separated')
  238. else:
  239. try:
  240. a = int(a)
  241. assert 1 <= a <= len(g.key_generators),'{}: invalid key generator'.format(a)
  242. if b in ('ext','pyethereum','pycoin','keyconv','zcash_mini'):
  243. init_external_prog()
  244. else:
  245. b = int(b)
  246. assert 1 <= b <= len(g.key_generators),'{}: invalid key generator'.format(b)
  247. assert a != b,'Key generators are the same!'
  248. except Exception as e:
  249. die(1,'{}\n{}: invalid generator argument'.format(e[0],cmd_args[0]))
  250. from mmgen.addr import KeyGenerator,AddrGenerator
  251. from mmgen.obj import PrivKey
  252. kg_a = KeyGenerator(addr_type,a)
  253. ag = AddrGenerator(addr_type)
  254. if a and b:
  255. if opt.all:
  256. from mmgen.protocol import init_coin,init_genonly_altcoins,CoinProtocol
  257. init_genonly_altcoins('btc',trust_level=0)
  258. mmgen_supported = CoinProtocol.get_valid_coins(upcase=True)
  259. for coin in ci.external_tests[('mainnet','testnet')[g.testnet]][ext_lib]:
  260. if coin not in mmgen_supported: continue
  261. init_coin(coin)
  262. tmp_addr_type = addr_type if addr_type in g.proto.mmtypes else MMGenAddrType(g.proto.dfl_mmtype)
  263. kg_a = KeyGenerator(tmp_addr_type,a)
  264. ag = AddrGenerator(tmp_addr_type)
  265. compare_test()
  266. else:
  267. if b != 'ext':
  268. kg_b = KeyGenerator(addr_type,b)
  269. b_desc = kg_b.desc
  270. compare_test()
  271. elif a and not fh:
  272. speed_test()
  273. elif a and dump:
  274. b_desc = 'dump'
  275. dump_test()
  276. else:
  277. die(2,'Illegal invocation')