ts_shared.py 8.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300
  1. #!/usr/bin/env python3
  2. #
  3. # mmgen = Multi-Mode GENerator, command-line Bitcoin cold storage solution
  4. # Copyright (C)2013-2023 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.test_py_d.ts_shared: Shared methods for the test.py test suite
  20. """
  21. import os
  22. from mmgen.util import ymsg,get_extension
  23. from mmgen.wallet import get_wallet_cls
  24. from ..include.common import cmp_or_die,strip_ansi_escapes,joinpath
  25. from .common import ref_bw_file,ref_bw_hash_preset,ref_dir
  26. class TestSuiteShared:
  27. 'shared methods for the test.py test suite'
  28. def txcreate_ui_common(
  29. self,
  30. t,
  31. caller = None,
  32. menu = [],
  33. inputs = '1',
  34. file_desc = 'Unsigned transaction',
  35. input_sels_prompt = 'to spend',
  36. bad_input_sels = False,
  37. interactive_fee = '',
  38. fee_desc = 'transaction fee',
  39. fee_info_pat = None,
  40. add_comment = '',
  41. view = 't',
  42. save = True,
  43. tweaks = [],
  44. used_chg_addr_resp = None,
  45. auto_chg_addr = None):
  46. txdo = (caller or self.test_name)[:4] == 'txdo'
  47. expect_pat = r'\[q\]uit menu, .*?:.'
  48. delete_pat = r'Enter account number .*:.'
  49. confirm_pat = r'Is this what you want.*:.'
  50. if used_chg_addr_resp is not None:
  51. t.expect('reuse harms your privacy.*:.*',used_chg_addr_resp,regex=True)
  52. if auto_chg_addr is not None:
  53. e1 = 'Choose a change address:.*Enter a number> '
  54. e2 = fr'Using .*{auto_chg_addr}.* as.*address'
  55. res = t.expect([e1,e2],regex=True)
  56. if res == 0:
  57. choice = [s.split(')')[0].lstrip() for s in t.p.match[0].split('\n') if auto_chg_addr in s][0]
  58. t.send(f'{choice}\n')
  59. t.expect(e2,regex=True)
  60. t.send('y')
  61. pat = expect_pat
  62. for choice in menu + ['q']:
  63. t.expect(pat,choice,regex=True)
  64. if self.proto.base_proto == 'Ethereum':
  65. pat = confirm_pat if pat == delete_pat else delete_pat if choice == 'D' else expect_pat
  66. if bad_input_sels:
  67. for r in ('x','3-1','9999'):
  68. t.expect(input_sels_prompt+': ',r+'\n')
  69. t.expect(input_sels_prompt+': ',inputs+'\n')
  70. have_est_fee = t.expect([f'{fee_desc}: ','OK? (Y/n): ']) == 1
  71. if have_est_fee and not interactive_fee:
  72. t.send('y')
  73. else:
  74. if have_est_fee:
  75. t.send('n')
  76. t.expect(f'{fee_desc}: ',interactive_fee+'\n')
  77. else:
  78. t.send(interactive_fee+'\n')
  79. if fee_info_pat:
  80. t.expect(fee_info_pat,regex=True)
  81. t.expect('OK? (Y/n): ','y')
  82. t.expect('(Y/n): ','\n') # chg amt OK prompt
  83. if 'confirm_non_mmgen' in tweaks:
  84. t.expect('Continue? (Y/n)','\n')
  85. t.do_comment(add_comment)
  86. t.view_tx(view)
  87. if not txdo:
  88. t.expect('(y/N): ',('n','y')[save])
  89. t.written_to_file(file_desc)
  90. return t
  91. def txsign_ui_common(
  92. self,
  93. t,
  94. caller = None,
  95. view = 't',
  96. add_comment = '',
  97. file_desc = 'Signed transaction',
  98. ni = False,
  99. save = True,
  100. do_passwd = False,
  101. has_label = False):
  102. txdo = (caller or self.test_name)[:4] == 'txdo'
  103. if do_passwd:
  104. t.passphrase('MMGen wallet',self.wpasswd)
  105. if not ni and not txdo:
  106. t.view_tx(view)
  107. t.do_comment(add_comment,has_label=has_label)
  108. t.expect('(Y/n): ',('n','y')[save])
  109. t.written_to_file(file_desc)
  110. return t
  111. def txsend_ui_common(
  112. self,
  113. t,
  114. caller = None,
  115. view = 'n',
  116. add_comment = '',
  117. file_desc = 'Sent transaction',
  118. confirm_send = True,
  119. bogus_send = True,
  120. quiet = False,
  121. has_label = False):
  122. txdo = (caller or self.test_name)[:4] == 'txdo'
  123. if not txdo:
  124. t.license() # MMGEN_NO_LICENSE is set, so does nothing
  125. t.view_tx(view)
  126. t.do_comment(add_comment,has_label=has_label)
  127. self._do_confirm_send(t,quiet=quiet,confirm_send=confirm_send)
  128. if bogus_send:
  129. txid = ''
  130. t.expect('BOGUS transaction NOT sent')
  131. else:
  132. txid = strip_ansi_escapes(t.expect_getend('Transaction sent: '))
  133. assert len(txid) == 64, f'{txid!r}: Incorrect txid length!'
  134. t.written_to_file(file_desc)
  135. return txid
  136. def txsign_end(self,t,tnum=None,has_label=False):
  137. t.expect('Signing transaction')
  138. t.do_comment(False,has_label=has_label)
  139. t.expect(r'Save signed transaction.*?\? \(Y/n\): ','y',regex=True)
  140. t.written_to_file('Signed transaction' + (' #' + tnum if tnum else ''), oo=True)
  141. return t
  142. def txsign(
  143. self,
  144. wf,
  145. txfile,
  146. pf = '',
  147. bumpf = '',
  148. save = True,
  149. has_label = False,
  150. extra_opts = [],
  151. extra_desc = '',
  152. view = 'n',
  153. dfl_wallet = False):
  154. opts = extra_opts + ['-d',self.tmpdir,txfile] + ([wf] if wf else [])
  155. t = self.spawn('mmgen-txsign', opts, extra_desc)
  156. t.license()
  157. t.view_tx(view)
  158. wcls = get_wallet_cls( ext = 'mmdat' if dfl_wallet else get_extension(wf) )
  159. if wcls.enc and wcls.type != 'brain':
  160. t.passphrase(wcls.desc,self.wpasswd)
  161. if save:
  162. self.txsign_end(t,has_label=has_label)
  163. else:
  164. t.do_comment(False,has_label=has_label)
  165. t.expect('Save signed transaction? (Y/n): ','n')
  166. t.expect('not saved')
  167. t.req_exit_val = 1
  168. return t
  169. def ref_brain_chk(self,bw_file=ref_bw_file):
  170. wf = joinpath(ref_dir,bw_file)
  171. add_args = [f'-l{self.seed_len}', f'-p{ref_bw_hash_preset}']
  172. return self.walletchk(wf,pf=None,add_args=add_args,sid=self.ref_bw_seed_id)
  173. def walletchk(
  174. self,
  175. wf,
  176. pf,
  177. wcls = None,
  178. add_args = [],
  179. sid = None,
  180. extra_desc = '',
  181. dfl_wallet = False):
  182. hp = self.hash_preset if hasattr(self,'hash_preset') else '1'
  183. wcls = wcls or get_wallet_cls(ext=get_extension(wf))
  184. t = self.spawn('mmgen-walletchk',
  185. ([] if dfl_wallet else ['-i',wcls.fmt_codes[0]])
  186. + add_args + ['-p',hp]
  187. + ([wf] if wf else []),
  188. extra_desc=extra_desc)
  189. if wcls.type != 'incog_hidden':
  190. t.expect(f"Getting {wcls.desc} from file '")
  191. if wcls.enc and wcls.type != 'brain':
  192. t.passphrase(wcls.desc,self.wpasswd)
  193. t.expect(['Passphrase is OK', 'Passphrase.* are correct'],regex=True)
  194. chk = t.expect_getend(f'Valid {wcls.desc} for Seed ID ')[:8]
  195. if sid:
  196. cmp_or_die(chk,sid)
  197. return t
  198. def addrgen(
  199. self,
  200. wf,
  201. pf = None,
  202. check_ref = False,
  203. ftype = 'addr',
  204. id_str = None,
  205. extra_args = [],
  206. mmtype = None,
  207. stdout = False,
  208. dfl_wallet = False):
  209. passgen = ftype[:4] == 'pass'
  210. if not mmtype and not passgen:
  211. mmtype = self.segwit_mmtype
  212. cmd_pfx = (ftype,'pass')[passgen]
  213. t = self.spawn(
  214. f'mmgen-{cmd_pfx}gen',
  215. ['-d',self.tmpdir] + extra_args +
  216. ([],['--type='+str(mmtype)])[bool(mmtype)] +
  217. ([],['--stdout'])[stdout] +
  218. ([],[wf])[bool(wf)] +
  219. ([],[id_str])[bool(id_str)] +
  220. [getattr(self,f'{cmd_pfx}_idx_list')],
  221. extra_desc=f'({mmtype})' if mmtype in ('segwit','bech32') else '')
  222. t.license()
  223. wcls = get_wallet_cls( ext = 'mmdat' if dfl_wallet else get_extension(wf) )
  224. t.passphrase(wcls.desc,self.wpasswd)
  225. t.expect('Passphrase is OK')
  226. desc = ('address','password')[passgen]
  227. chk = t.expect_getend(rf'Checksum for {desc} data .*?: ',regex=True)
  228. if passgen:
  229. t.expect('Encrypt password list? (y/N): ','N')
  230. t.read() if stdout else t.written_to_file(('Addresses','Password list')[passgen])
  231. if check_ref:
  232. chk_ref = (
  233. self.chk_data[self.test_name] if passgen else
  234. self.chk_data[self.test_name][self.fork][self.proto.testnet])
  235. cmp_or_die(chk,chk_ref,desc=f'{ftype}list data checksum')
  236. return t
  237. def keyaddrgen(self,wf,pf=None,check_ref=False,mmtype=None):
  238. if not mmtype:
  239. mmtype = self.segwit_mmtype
  240. args = ['-d',self.tmpdir,self.usr_rand_arg,wf,self.addr_idx_list]
  241. t = self.spawn('mmgen-keygen',
  242. ([],['--type='+str(mmtype)])[bool(mmtype)] + args,
  243. extra_desc=f'({mmtype})' if mmtype in ('segwit','bech32') else '')
  244. t.license()
  245. wcls = get_wallet_cls(ext=get_extension(wf))
  246. t.passphrase(wcls.desc,self.wpasswd)
  247. chk = t.expect_getend(r'Checksum for key-address data .*?: ',regex=True)
  248. if check_ref:
  249. chk_ref = self.chk_data[self.test_name][self.fork][self.proto.testnet]
  250. cmp_or_die(chk,chk_ref,desc='key-address list data checksum')
  251. t.expect('Encrypt key list? (y/N): ','y')
  252. t.usr_rand(self.usr_rand_chars)
  253. t.hash_preset('new key-address list','1')
  254. t.passphrase_new('new key-address list',self.kapasswd)
  255. t.written_to_file('Encrypted secret keys',oo=True)
  256. return t
  257. def _do_confirm_send(self,t,quiet=False,confirm_send=True,sure=True):
  258. if sure:
  259. t.expect('Are you sure you want to broadcast this')
  260. m = ('YES, I REALLY WANT TO DO THIS','YES')[quiet]
  261. t.expect(f'{m!r} to confirm: ',('',m)[confirm_send]+'\n')