gentest.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301
  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: 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
  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. ext_sec2addr = zcash_mini_sec2addr
  104. ext_lib = 'zcash_mini'
  105. elif test_support('pyethereum'):
  106. try:
  107. import ethereum.utils as eth
  108. except:
  109. raise ImportError,"Unable to import 'pyethereum' module. Is pyethereum installed?"
  110. ext_sec2addr = pyethereum_sec2addr
  111. ext_lib = 'pyethereum'
  112. elif test_support('pycoin'):
  113. try:
  114. import pycoin.cmds.ku as pcku
  115. except:
  116. raise ImportError,"Unable to import module 'ku'. Is pycoin installed?"
  117. PREFIX_TRANSFORMS = pcku.prefix_transforms_for_network(g.coin)
  118. ext_sec2addr = pycoin_sec2addr
  119. ext_lib = 'pycoin'
  120. elif test_support('keyconv'):
  121. import subprocess as sp
  122. ext_sec2addr = keyconv_sec2addr
  123. ext_lib = 'keyconv'
  124. else:
  125. m = '{}: coin supported by MMGen but unsupported by gentest.py for {}'
  126. raise ValueError,m.format(g.coin,('mainnet','testnet')[g.testnet])
  127. b_desc = ext_lib
  128. b = 'ext'
  129. def match_error(sec,wif,a_addr,b_addr,a,b):
  130. qmsg_r(red('\nERROR: Values do not match!'))
  131. die(3,"""
  132. sec key : {}
  133. WIF key : {}
  134. {a:10}: {}
  135. {b:10}: {}
  136. """.format(sec,wif,a_addr,b_addr,pnm=g.proj_name,a=kg_a.desc,b=b_desc).rstrip())
  137. def compare_test():
  138. for k in ('segwit','compressed'):
  139. if addr_type.name == k and g.coin not in ci.external_tests_segwit_compressed[k]:
  140. msg('{} testing not supported for coin {}'.format(addr_type.name.capitalize(),g.coin))
  141. return
  142. if 'ext_lib' in globals():
  143. if g.coin not in ci.external_tests[('mainnet','testnet')[g.testnet]][ext_lib]:
  144. msg("Coin '{}' incompatible with external generator '{}'".format(g.coin,ext_lib))
  145. return
  146. m = "Comparing address generators '{}' and '{}' for coin {}"
  147. last_t = time.time()
  148. qmsg(green(m.format(kg_a.desc,(ext_lib if b == 'ext' else kg_b.desc),g.coin)))
  149. for i in range(rounds):
  150. if opt.verbose or time.time() - last_t >= 0.1:
  151. qmsg_r('\rRound %s/%s ' % (i+1,rounds))
  152. last_t = time.time()
  153. sec = PrivKey(os.urandom(32),compressed=addr_type.compressed,pubkey_type=addr_type.pubkey_type)
  154. ph = kg_a.to_pubhex(sec)
  155. a_addr = ag.to_addr(ph)
  156. if addr_type.name == 'zcash_z':
  157. a_vk = ag.to_viewkey(ph)
  158. if b == 'ext':
  159. if addr_type.name == 'zcash_z':
  160. b_wif,b_addr,b_vk = ext_sec2addr(sec)
  161. if b_vk != a_vk:
  162. match_error(sec,sec.wif,a_vk,b_vk,a,b)
  163. else:
  164. b_wif,b_addr = ext_sec2addr(sec)
  165. if b_wif != sec.wif:
  166. match_error(sec,sec.wif,sec.wif,b_wif,a,b)
  167. else:
  168. b_addr = ag.to_addr(kg_b.to_pubhex(sec))
  169. vmsg('\nkey: %s\naddr: %s\n' % (sec.wif,a_addr))
  170. if a_addr != b_addr:
  171. match_error(sec,sec.wif,a_addr,b_addr,a,ext_lib if b == 'ext' else b)
  172. qmsg_r('\rRound %s/%s ' % (i+1,rounds))
  173. qmsg(green(('\n','')[bool(opt.verbose)] + 'OK'))
  174. def speed_test():
  175. m = "Testing speed of address generator '{}' for coin {}"
  176. qmsg(green(m.format(kg_a.desc,g.coin)))
  177. from struct import pack,unpack
  178. seed = os.urandom(28)
  179. print 'Incrementing key with each round'
  180. print 'Starting key:', hexlify(seed+pack('I',0))
  181. import time
  182. start = last_t = time.time()
  183. for i in range(rounds):
  184. if time.time() - last_t >= 0.1:
  185. qmsg_r('\rRound %s/%s ' % (i+1,rounds))
  186. last_t = time.time()
  187. sec = PrivKey(seed+pack('I',i),compressed=addr_type.compressed,pubkey_type=addr_type.pubkey_type)
  188. a_addr = ag.to_addr(kg_a.to_pubhex(sec))
  189. vmsg('\nkey: %s\naddr: %s\n' % (sec.wif,a_addr))
  190. qmsg_r('\rRound %s/%s ' % (i+1,rounds))
  191. qmsg('\n{} addresses generated in {:.2f} seconds'.format(rounds,time.time()-start))
  192. def dump_test():
  193. m = "Comparing output of address generator '{}' against wallet dump '{}'"
  194. qmsg(green(m.format(kg_a.desc,cmd_args[1])))
  195. for n,[wif,a_addr] in enumerate(dump,1):
  196. qmsg_r('\rKey %s/%s ' % (n,len(dump)))
  197. try:
  198. sec = PrivKey(wif=wif)
  199. except:
  200. die(2,'\nInvalid {}net WIF address in dump file: {}'.format(('main','test')[g.testnet],wif))
  201. b_addr = ag.to_addr(kg_a.to_pubhex(sec))
  202. if a_addr != b_addr:
  203. match_error(sec,wif,a_addr,b_addr,3,a)
  204. qmsg(green(('\n','')[bool(opt.verbose)] + 'OK'))
  205. from mmgen.altcoin import CoinInfo as ci
  206. urounds,fh = None,None
  207. dump = []
  208. if len(cmd_args) == 2:
  209. try:
  210. urounds = int(cmd_args[1])
  211. assert urounds > 0
  212. except:
  213. try:
  214. fh = open(cmd_args[1])
  215. except:
  216. die(1,'Second argument must be filename or positive integer')
  217. else:
  218. for line in fh.readlines():
  219. if 'addr=' in line:
  220. x,addr = line.split('addr=')
  221. dump.append([x.split()[0],addr.split()[0]])
  222. if urounds: rounds = urounds
  223. a,b = None,None
  224. b_desc = 'unknown'
  225. try:
  226. a,b = cmd_args[0].split(':')
  227. except:
  228. try:
  229. a = cmd_args[0]
  230. a = int(a)
  231. assert 1 <= a <= len(g.key_generators)
  232. except:
  233. die(1,'First argument must be one or two generator IDs, colon separated')
  234. else:
  235. try:
  236. a = int(a)
  237. assert 1 <= a <= len(g.key_generators),'{}: invalid key generator'.format(a)
  238. if b in ('ext','pyethereum','pycoin','keyconv','zcash_mini'):
  239. init_external_prog()
  240. else:
  241. b = int(b)
  242. assert 1 <= b <= len(g.key_generators),'{}: invalid key generator'.format(b)
  243. assert a != b,'Key generators are the same!'
  244. except Exception as e:
  245. die(1,'{}\n{}: invalid generator argument'.format(e[0],cmd_args[0]))
  246. from mmgen.addr import KeyGenerator,AddrGenerator
  247. from mmgen.obj import PrivKey
  248. kg_a = KeyGenerator(addr_type,a)
  249. ag = AddrGenerator(addr_type)
  250. if a and b:
  251. if opt.all:
  252. from mmgen.protocol import init_coin,init_genonly_altcoins
  253. init_genonly_altcoins('btc',trust_level=0)
  254. mmgen_supported = [e[1] for e in ci.coin_constants['mainnet']]
  255. for coin in ci.external_tests[('mainnet','testnet')[g.testnet]][ext_lib]:
  256. if coin not in mmgen_supported and ext_lib != 'pyethereum': continue
  257. init_coin(coin)
  258. tmp_addr_type = addr_type if addr_type in g.proto.mmtypes else MMGenAddrType(g.proto.dfl_mmtype)
  259. kg_a = KeyGenerator(tmp_addr_type,a)
  260. ag = AddrGenerator(tmp_addr_type)
  261. compare_test()
  262. else:
  263. if b != 'ext':
  264. kg_b = KeyGenerator(addr_type,b)
  265. b_desc = kg_b.desc
  266. compare_test()
  267. elif a and not fh:
  268. speed_test()
  269. elif a and dump:
  270. b_desc = 'dump'
  271. dump_test()
  272. else:
  273. die(2,'Illegal invocation')