gentest.py 11 KB

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