pexpect.py 7.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219
  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/pexpect.py: pexpect implementation for MMGen test suites
  20. """
  21. import sys,os,time
  22. from mmgen.globalvars import g
  23. from mmgen.opts import opt
  24. from mmgen.util import msg,msg_r,vmsg,vmsg_r,rmsg,red,yellow,green,cyan,die,rdie
  25. from test.common import getrandstr
  26. try:
  27. import pexpect
  28. from pexpect.popen_spawn import PopenSpawn
  29. except:
  30. die(2,red('Pexpect module is missing. Cannnot run test suite'))
  31. def debug_pexpect_msg(p):
  32. if opt.debug_pexpect:
  33. msg('\n{}{}{}'.format(red('BEFORE ['),p.before,red(']')))
  34. msg('{}{}{}'.format(red('MATCH ['),p.after,red(']')))
  35. NL = '\n'
  36. class MMGenPexpect(object):
  37. def __init__(self,args,no_output=False):
  38. if opt.direct_exec:
  39. msg('')
  40. from subprocess import call,check_output
  41. f = (call,check_output)[bool(no_output)]
  42. ret = f([args[0]] + args[1:])
  43. if f == call and ret != 0:
  44. die(1,red('ERROR: process returned a non-zero exit status ({})'.format(ret)))
  45. else:
  46. if opt.pexpect_spawn:
  47. self.p = pexpect.spawn(args[0],args[1:],encoding='utf8')
  48. self.p.delaybeforesend = 0
  49. else:
  50. self.p = PopenSpawn(args,encoding='utf8')
  51. # self.p.delaybeforesend = 0 # TODO: try this here too
  52. if opt.exact_output: self.p.logfile = sys.stdout
  53. self.req_exit_val = 0
  54. self.skip_ok = False
  55. self.timeout = int(opt.pexpect_timeout or 0) or (60,5)[bool(opt.debug_pexpect)]
  56. self.sent_value = None
  57. def do_decrypt_ka_data(self,hp,pw,desc='key-address data',check=True,have_yes_opt=False):
  58. # self.hash_preset(desc,hp)
  59. self.passphrase(desc,pw)
  60. if not have_yes_opt:
  61. self.expect('Check key-to-address validity? (y/N): ',('n','y')[check])
  62. def view_tx(self,view):
  63. self.expect('View.* transaction.*\? .*: ',view,regex=True)
  64. if view not in 'n\n':
  65. self.expect('to continue: ','\n')
  66. def do_comment(self,add_comment,has_label=False):
  67. p = ('Add a comment to transaction','Edit transaction comment')[has_label]
  68. self.expect('{}? (y/N): '.format(p),('n','y')[bool(add_comment)])
  69. if add_comment:
  70. self.expect('Comment: ',add_comment+'\n')
  71. def ok(self):
  72. ret = self.p.wait()
  73. if ret != self.req_exit_val and not opt.coverage:
  74. die(1,red('test.py: spawned program exited with value {}'.format(ret)))
  75. if opt.profile: return
  76. if not self.skip_ok:
  77. sys.stderr.write(green('OK\n') if opt.exact_output or opt.verbose else (' OK\n'))
  78. return self
  79. def license(self):
  80. if 'MMGEN_NO_LICENSE' in os.environ: return
  81. self.expect("'w' for conditions and warranty info, or 'c' to continue: ",'c')
  82. def label(self,label='Test Label (UTF-8) α'):
  83. self.expect('Enter a wallet label, or hit ENTER for no label: ',label+'\n')
  84. def usr_rand_out(self,saved=False):
  85. fs = 'Generating encryption key from OS random data plus {}user-supplied entropy'
  86. self.expect(fs.format(('','saved ')[saved]))
  87. def usr_rand(self,num_chars):
  88. if opt.usr_random:
  89. self.interactive()
  90. self.send('\n')
  91. else:
  92. rand_chars = list(getrandstr(num_chars,no_space=True))
  93. vmsg_r('SEND ')
  94. while rand_chars:
  95. ch = rand_chars.pop(0)
  96. msg_r(yellow(ch)+' ' if opt.verbose else '+')
  97. ret = self.expect('left: ',ch,delay=0.005)
  98. self.expect('ENTER to continue: ','\n')
  99. def passphrase_new(self,desc,passphrase):
  100. self.expect('Enter passphrase for {}: '.format(desc),passphrase+'\n')
  101. self.expect('Repeat passphrase: ',passphrase+'\n')
  102. def passphrase(self,desc,passphrase,pwtype=''):
  103. if pwtype: pwtype += ' '
  104. self.expect('Enter {}passphrase for {}.*?: '.format(pwtype,desc),passphrase+'\n',regex=True)
  105. def hash_preset(self,desc,preset=''):
  106. self.expect('Enter hash preset for {}'.format(desc))
  107. self.expect('or hit ENTER .*?:',str(preset)+'\n',regex=True)
  108. def written_to_file(self,desc,overwrite_unlikely=False,query='Overwrite? ',oo=False):
  109. s1 = '{} written to file '.format(desc)
  110. s2 = query + "Type uppercase 'YES' to confirm: "
  111. ret = self.expect(([s1,s2],s1)[overwrite_unlikely])
  112. if ret == 1:
  113. self.send('YES\n')
  114. return self.expect_getend("Overwriting file '").rstrip("'")
  115. self.expect(NL,nonl=True)
  116. outfile = self.p.before.strip().strip("'")
  117. if opt.debug_pexpect:
  118. rmsg('Outfile [{}]'.format(outfile))
  119. vmsg('{} file: {}'.format(desc,cyan(outfile.replace("'",''))))
  120. return outfile
  121. def no_overwrite(self):
  122. self.expect("Overwrite? Type uppercase 'YES' to confirm: ",'\n')
  123. self.expect('Exiting at user request')
  124. def expect_getend(self,s,regex=False):
  125. ret = self.expect(s,regex=regex,nonl=True)
  126. debug_pexpect_msg(self.p)
  127. # readline() of partial lines doesn't work with PopenSpawn, so do this instead:
  128. self.expect(NL,nonl=True,silent=True)
  129. debug_pexpect_msg(self.p)
  130. end = self.p.before.rstrip()
  131. if not g.debug:
  132. vmsg(' ==> {}'.format(cyan(end)))
  133. return end
  134. def interactive(self):
  135. return self.p.interact() # interact() not available with popen_spawn
  136. def kill(self,signal):
  137. return self.p.kill(signal)
  138. def expect(self,s,t='',delay=None,regex=False,nonl=False,silent=False):
  139. delay = delay or (0,0.3)[bool(opt.buf_keypress)]
  140. if not silent:
  141. if opt.verbose:
  142. msg_r('EXPECT ' + yellow(str(s)))
  143. elif not opt.exact_output: msg_r('+')
  144. try:
  145. if s == '':
  146. ret = 0
  147. else:
  148. f = (self.p.expect_exact,self.p.expect)[bool(regex)]
  149. ret = f(s,self.timeout)
  150. except pexpect.TIMEOUT:
  151. if opt.debug_pexpect: raise
  152. m1 = red('\nERROR. Expect {!r} timed out. Exiting\n'.format(s))
  153. m2 = 'before: [{}]\n'.format(self.p.before)
  154. m3 = 'sent value: [{}]'.format(self.sent_value) if self.sent_value != None else ''
  155. rdie(1,m1+m2+m3)
  156. debug_pexpect_msg(self.p)
  157. if opt.verbose and type(s) != str:
  158. msg_r(' ==> {} '.format(ret))
  159. if ret == -1:
  160. rdie(1,'Error. Expect returned {}'.format(ret))
  161. else:
  162. if t == '':
  163. if not nonl and not silent: vmsg('')
  164. else:
  165. self.send(t,delay,s)
  166. return ret
  167. def send(self,t,delay=None,s=False):
  168. self.sent_value = None
  169. delay = delay or (0,0.3)[bool(opt.buf_keypress)]
  170. if delay: time.sleep(delay)
  171. ret = self.p.send(t) # returns num bytes written
  172. if ret:
  173. self.sent_value = t
  174. if delay: time.sleep(delay)
  175. if opt.verbose:
  176. ls = (' ','')[bool(opt.debug or not s)]
  177. es = (' ','')[bool(s)]
  178. msg('{}SEND {}{}'.format(ls,es,yellow("'{}'".format(t.replace('\n',r'\n')))))
  179. return ret
  180. def read(self,n=-1):
  181. return self.p.read(n)
  182. def close(self):
  183. if opt.pexpect_spawn:
  184. self.p.close()