mmgen_pexpect.py 8.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279
  1. #!/usr/bin/env python
  2. # -*- coding: UTF-8 -*-
  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/mmgen_pexpect.py: pexpect implementation for MMGen test suites
  20. """
  21. from mmgen.common import *
  22. from mmgen.test import getrandstr,ok,init_coverage
  23. try:
  24. import pexpect
  25. from pexpect.popen_spawn import PopenSpawn
  26. except:
  27. die(2,red('Pexpect module is missing. Cannnot run test suite'))
  28. if opt.buf_keypress:
  29. send_delay = 0.3
  30. else:
  31. send_delay = 0
  32. os.environ['MMGEN_DISABLE_HOLD_PROTECT'] = '1'
  33. stderr_save = sys.stderr
  34. def errmsg(s): stderr_save.write(s+'\n')
  35. def errmsg_r(s): stderr_save.write(s)
  36. def my_send(p,t,delay=send_delay,s=False):
  37. if delay: time.sleep(delay)
  38. ret = p.send(t) # returns num bytes written
  39. if delay: time.sleep(delay)
  40. if opt.verbose:
  41. ls = (' ','')[bool(opt.debug or not s)]
  42. es = (' ','')[bool(s)]
  43. msg('{}SEND {}{}'.format(ls,es,yellow("'{}'"%t.replace('\n',r'\n'))))
  44. return ret
  45. def my_expect(p,s,t='',delay=send_delay,regex=False,nonl=False,silent=False):
  46. quo = ('',"'")[type(s) == str]
  47. if not silent:
  48. if opt.verbose: msg_r('EXPECT {}'.format(yellow(quo+str(s)+quo)))
  49. elif not opt.exact_output: msg_r('+')
  50. try:
  51. if s == '': ret = 0
  52. else:
  53. f = (p.expect_exact,p.expect)[bool(regex)]
  54. ret = f(s,timeout=(60,5)[bool(opt.debug_pexpect)])
  55. except pexpect.TIMEOUT:
  56. if opt.debug_pexpect: raise
  57. errmsg(red('\nERROR. Expect {}{}{} timed out. Exiting'.format(quo,s,quo)))
  58. sys.exit(1)
  59. debug_pexpect_msg(p)
  60. if opt.debug or (opt.verbose and type(s) != str):
  61. msg_r(' ==> {} '.format(ret))
  62. if ret == -1:
  63. errmsg('Error. Expect returned {}'.format(ret))
  64. sys.exit(1)
  65. else:
  66. if t == '':
  67. if not nonl and not silent: vmsg('')
  68. else:
  69. my_send(p,t,delay,s)
  70. return ret
  71. def debug_pexpect_msg(p):
  72. if opt.debug_pexpect:
  73. errmsg('\n{}{}{}'.format(red('BEFORE ['),p.before,red(']')))
  74. errmsg('{}{}{}'.format(red('MATCH ['),p.after,red(']')))
  75. class MMGenPexpect(object):
  76. NL = '\r\n'
  77. if g.platform == 'linux' and opt.popen_spawn:
  78. import atexit
  79. atexit.register(lambda: os.system('stty sane'))
  80. NL = '\n'
  81. def __init__(self,name,mmgen_cmd,cmd_args,desc,no_output=False,passthru_args=[],msg_only=False,no_msg=False):
  82. cmd_args = ['--{}{}'.format(k.replace('_','-'),
  83. '='+getattr(opt,k) if getattr(opt,k) != True else ''
  84. ) for k in passthru_args if getattr(opt,k)] \
  85. + ['--data-dir='+os.path.join('test','data_dir')] + cmd_args
  86. if g.platform == 'win': cmd,args = 'python',[mmgen_cmd]+cmd_args
  87. else: cmd,args = mmgen_cmd,cmd_args
  88. for i in args:
  89. if type(i) not in (str,unicode):
  90. m1 = 'Error: missing input files in cmd line?:'
  91. m2 = '\nName: {}\nCmd: {}\nCmd args: {}'
  92. die(2,(m1+m2).format(name,cmd,args))
  93. if opt.popen_spawn:
  94. args = [(a,"'{}'".format(a))[' ' in a] for a in args]
  95. cmd_str = '{} {}'.format(cmd,' '.join(args)).replace('\\','/')
  96. if opt.coverage:
  97. fs = 'python -m trace --count --coverdir={} --file={} {c}'
  98. cmd_str = fs.format(*init_coverage(),c=cmd_str)
  99. if opt.log:
  100. log_fd.write(cmd_str+'\n')
  101. if not no_msg:
  102. if opt.verbose or opt.print_cmdline or opt.exact_output:
  103. clr1,clr2,eol = ((green,cyan,'\n'),(nocolor,nocolor,' '))[bool(opt.print_cmdline)]
  104. sys.stderr.write(green('Testing: {}\n'.format(desc)))
  105. if not msg_only:
  106. sys.stderr.write(clr1('Executing {}{}'.format(clr2(cmd_str),eol)))
  107. else:
  108. m = 'Testing {}: '.format(desc)
  109. msg_r(m)
  110. if msg_only: return
  111. if opt.direct_exec:
  112. msg('')
  113. from subprocess import call,check_output
  114. f = (call,check_output)[bool(no_output)]
  115. ret = f([cmd] + args)
  116. if f == call and ret != 0:
  117. die(1,red('ERROR: process returned a non-zero exit status ({})'.format(ret)))
  118. else:
  119. if opt.traceback:
  120. cmd,args = g.traceback_cmd,[cmd]+args
  121. cmd_str = g.traceback_cmd + ' ' + cmd_str
  122. # Msg('\ncmd_str: {}'.format(cmd_str))
  123. if opt.popen_spawn:
  124. self.p = PopenSpawn(cmd_str)
  125. else:
  126. self.p = pexpect.spawn(cmd,args)
  127. if opt.exact_output: self.p.logfile = sys.stdout
  128. def ok(self,exit_val=0):
  129. ret = self.p.wait()
  130. # Msg('expect: {} got: {}'.format(exit_val,ret))
  131. if ret != exit_val and not opt.coverage:
  132. die(1,red('test.py: spawned program exited with value {}'.format(ret)))
  133. if opt.profile: return
  134. if opt.verbose or opt.exact_output:
  135. sys.stderr.write(green('OK\n'))
  136. else: msg(' OK')
  137. def cmp_or_die(self,s,t,skip_ok=False,exit_val=0):
  138. ret = self.p.wait()
  139. if ret != exit_val:
  140. rdie(1,'test.py: spawned program exited with value {}'.format(ret))
  141. if s == t:
  142. if not skip_ok: ok()
  143. else:
  144. fs = 'ERROR: recoded data:\n{}\ndiffers from original data:\n{}'
  145. rdie(3,fs.format(repr(t),repr(s)))
  146. def license(self):
  147. if 'MMGEN_NO_LICENSE' in os.environ: return
  148. p = "'w' for conditions and warranty info, or 'c' to continue: "
  149. my_expect(self.p,p,'c')
  150. def label(self,label='Test Label'):
  151. p = 'Enter a wallet label, or hit ENTER for no label: '
  152. my_expect(self.p,p,label+'\n')
  153. def usr_rand_out(self,saved=False):
  154. fs = 'Generating encryption key from OS random data plus {}user-supplied entropy'
  155. my_expect(self.p,fs.format(('','saved ')[saved]))
  156. def usr_rand(self,num_chars):
  157. if opt.usr_random:
  158. self.interactive()
  159. my_send(self.p,'\n')
  160. else:
  161. rand_chars = list(getrandstr(num_chars,no_space=True))
  162. my_expect(self.p,'symbols left: ','x')
  163. try:
  164. vmsg_r('SEND ')
  165. while self.p.expect('left: ',0.1) == 0:
  166. ch = rand_chars.pop(0)
  167. msg_r(yellow(ch)+' ' if opt.verbose else '+')
  168. self.p.send(ch)
  169. except:
  170. vmsg('EOT')
  171. my_expect(self.p,'ENTER to continue: ','\n')
  172. def passphrase_new(self,desc,passphrase):
  173. my_expect(self.p,'Enter passphrase for {}: '.format(desc),passphrase+'\n')
  174. my_expect(self.p,'Repeat passphrase: ',passphrase+'\n')
  175. def passphrase(self,desc,passphrase,pwtype=''):
  176. if pwtype: pwtype += ' '
  177. my_expect(self.p,
  178. 'Enter {}passphrase for {}.*?: '.format(pwtype,desc),
  179. passphrase+'\n',regex=True)
  180. def hash_preset(self,desc,preset=''):
  181. my_expect(self.p,'Enter hash preset for {}'.format(desc))
  182. my_expect(self.p,'or hit ENTER .*?:',str(preset)+'\n',regex=True)
  183. def written_to_file(self,desc,overwrite_unlikely=False,query='Overwrite? ',oo=False):
  184. s1 = '{} written to file '.format(desc)
  185. s2 = query + "Type uppercase 'YES' to confirm: "
  186. ret = my_expect(self.p,([s1,s2],s1)[overwrite_unlikely])
  187. if ret == 1:
  188. my_send(self.p,'YES\n')
  189. # if oo:
  190. outfile = self.expect_getend("Overwriting file '").rstrip("'")
  191. return outfile
  192. # else:
  193. # ret = my_expect(self.p,s1)
  194. self.expect(self.NL,nonl=True)
  195. outfile = self.p.before.strip().strip("'")
  196. if opt.debug_pexpect: rmsg('Outfile [{}]'.format(outfile))
  197. vmsg('{} file: {}'.format(desc,cyan(outfile.replace("'",''))))
  198. return outfile
  199. def no_overwrite(self):
  200. self.expect("Overwrite? Type uppercase 'YES' to confirm: ",'\n')
  201. self.expect('Exiting at user request')
  202. def tx_view(self,view=None):
  203. repl = { 'terse':'t', 'full':'v' }[view] if view else 'n'
  204. my_expect(self.p,r'View .*?transaction.*? \(y\)es, \(N\)o, pager \(v\)iew.*?: ',repl,regex=True)
  205. if repl in ('t','v'):
  206. my_expect(self.p,r'any key to continue: ','\n')
  207. def expect_getend(self,s,regex=False):
  208. ret = self.expect(s,regex=regex,nonl=True)
  209. debug_pexpect_msg(self.p)
  210. # end = self.readline().strip()
  211. # readline() of partial lines doesn't work with PopenSpawn, so do this instead:
  212. self.expect(self.NL,nonl=True,silent=True)
  213. debug_pexpect_msg(self.p)
  214. end = self.p.before
  215. vmsg(' ==> {}'.format(cyan(end)))
  216. return end
  217. def interactive(self):
  218. return self.p.interact() # interact() not available with popen_spawn
  219. def kill(self,signal):
  220. return self.p.kill(signal)
  221. def logfile(self,arg):
  222. self.p.logfile = arg
  223. def expect(self,*args,**kwargs):
  224. return my_expect(self.p,*args,**kwargs)
  225. def send(self,*args,**kwargs):
  226. return my_send(self.p,*args,**kwargs)
  227. # def readline(self):
  228. # return self.p.readline()
  229. # def readlines(self):
  230. # return [l.rstrip()+'\n' for l in self.p.readlines()]
  231. def read(self,n=None):
  232. return self.p.read(n)
  233. def close(self):
  234. if not opt.popen_spawn:
  235. self.p.close()