pexpect.py 7.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252
  1. #!/usr/bin/env python3
  2. #
  3. # mmgen = Multi-Mode GENerator, command-line Bitcoin cold storage solution
  4. # Copyright (C)2013-2024 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.include.pexpect: pexpect implementation for MMGen test suites
  20. """
  21. import sys,time
  22. from mmgen.color import red,yellow,green,cyan
  23. from mmgen.util import msg,msg_r,rmsg,die
  24. from .common import cfg,vmsg,vmsg_r,getrandstr,strip_ansi_escapes
  25. try:
  26. import pexpect
  27. from pexpect.popen_spawn import PopenSpawn
  28. except ImportError as e:
  29. die(2,red(f'Pexpect module is missing. Cannnot run test suite ({e!r})'))
  30. def debug_pexpect_msg(p):
  31. msg('\n{}{}{}'.format( red('BEFORE ['), p.before, red(']') ))
  32. msg('{}{}{}'.format( red('MATCH ['), p.after, red(']') ))
  33. NL = '\n'
  34. class MMGenPexpect:
  35. def __init__(
  36. self,
  37. args,
  38. no_output = False,
  39. spawn_env = None,
  40. pexpect_spawn = False,
  41. send_delay = None,
  42. timeout = None,
  43. direct_exec = False ):
  44. self.pexpect_spawn = pexpect_spawn
  45. self.send_delay = send_delay
  46. self.skip_ok = False
  47. self.sent_value = None
  48. self.spawn_env = spawn_env
  49. if direct_exec or cfg.direct_exec:
  50. from subprocess import Popen,DEVNULL
  51. redir = DEVNULL if (no_output or not cfg.exact_output) else None
  52. self.ep = Popen([args[0]] + args[1:], stderr=redir, env=spawn_env )
  53. else:
  54. timeout = int(timeout or cfg.pexpect_timeout or 0) or (60,5)[bool(cfg.debug_pexpect)]
  55. if pexpect_spawn:
  56. self.p = pexpect.spawn(args[0],args[1:],encoding='utf8',timeout=timeout,env=spawn_env)
  57. else:
  58. self.p = PopenSpawn(args,encoding='utf8',timeout=timeout,env=spawn_env)
  59. if cfg.exact_output:
  60. self.p.logfile = sys.stdout
  61. def do_decrypt_ka_data(
  62. self,
  63. pw,
  64. desc = 'key-address data',
  65. check = True,
  66. have_yes_opt = False):
  67. self.passphrase(desc,pw)
  68. if not have_yes_opt:
  69. self.expect('Check key-to-address validity? (y/N): ',('n','y')[check])
  70. def view_tx(self,view):
  71. self.expect(r'View.* transaction.*\? .*: ',view,regex=True)
  72. if view not in 'n\n':
  73. self.expect('to continue: ','\n')
  74. def do_comment(self,add_comment,has_label=False):
  75. p = ('Add a comment to transaction','Edit transaction comment')[has_label]
  76. self.expect(f'{p}? (y/N): ',('n','y')[bool(add_comment)])
  77. if add_comment:
  78. self.expect('Comment: ',add_comment+'\n')
  79. def ok(self, exit_val=None):
  80. if not self.pexpect_spawn:
  81. self.p.sendeof()
  82. self.p.read()
  83. ret = self.p.wait()
  84. if ret != (exit_val or 0) and not cfg.coverage:
  85. die( 'TestSuiteSpawnedScriptException', f'Spawned script exited with value {ret}' )
  86. if cfg.profile:
  87. return
  88. if not self.skip_ok:
  89. m = 'OK\n' if ret == 0 else f'OK[{ret}]\n'
  90. sys.stderr.write( green(m) if cfg.exact_output or cfg.verbose else ' '+m )
  91. return self
  92. def license(self):
  93. if self.spawn_env.get('MMGEN_NO_LICENSE'):
  94. return
  95. self.expect("'w' for conditions and warranty info, or 'c' to continue: ",'c')
  96. def label(self,label='Test Label (UTF-8) α'):
  97. self.expect('Enter a wallet label, or hit ENTER for no label: ',label+'\n')
  98. def usr_rand(self,num_chars):
  99. if cfg.usr_random:
  100. self.interactive()
  101. self.send('\n')
  102. else:
  103. rand_chars = list(getrandstr(num_chars,no_space=True))
  104. vmsg_r('SEND ')
  105. while rand_chars:
  106. ch = rand_chars.pop(0)
  107. msg_r(yellow(ch)+' ' if cfg.verbose else '+')
  108. self.expect('left: ',ch,delay=0.005)
  109. self.expect('ENTER to continue: ','\n')
  110. def passphrase_new(self,desc,passphrase):
  111. self.expect(f'Enter passphrase for {desc}: ',passphrase+'\n')
  112. self.expect('Repeat passphrase: ',passphrase+'\n')
  113. def passphrase(self,desc,passphrase,pwtype=''):
  114. if pwtype:
  115. pwtype += ' '
  116. self.expect(f'Enter {pwtype}passphrase for {desc}.*?: ',passphrase+'\n',regex=True)
  117. def hash_preset(self,desc,preset=''):
  118. self.expect(f'Enter hash preset for {desc}')
  119. self.expect('or hit ENTER .*?:',str(preset)+'\n',regex=True)
  120. def written_to_file(self,desc,overwrite_unlikely=False,query='Overwrite? '):
  121. s1 = f'{desc} written to file '
  122. s2 = query + "Type uppercase 'YES' to confirm: "
  123. ret = self.expect(([s1,s2],s1)[overwrite_unlikely])
  124. if ret == 1:
  125. self.send('YES\n')
  126. return self.expect_getend("Overwriting file '").rstrip("'")
  127. self.expect(NL,nonl=True)
  128. outfile = self.p.before.strip().strip("'")
  129. if cfg.debug_pexpect:
  130. rmsg(f'Outfile [{outfile}]')
  131. vmsg('{} file: {}'.format( desc, cyan(outfile.replace('"',"")) ))
  132. return outfile
  133. def hincog_create(self,hincog_bytes):
  134. ret = self.expect(['Create? (Y/n): ',"'YES' to confirm: "])
  135. if ret == 0:
  136. self.send('\n')
  137. self.expect('Enter file size: ',str(hincog_bytes)+'\n')
  138. else:
  139. self.send('YES\n')
  140. return ret
  141. def no_overwrite(self):
  142. self.expect("Overwrite? Type uppercase 'YES' to confirm: ",'\n')
  143. self.expect('Exiting at user request')
  144. def expect_getend(self,s,regex=False):
  145. self.expect(s,regex=regex,nonl=True)
  146. if cfg.debug_pexpect:
  147. debug_pexpect_msg(self.p)
  148. # readline() of partial lines doesn't work with PopenSpawn, so do this instead:
  149. self.expect(NL,nonl=True,silent=True)
  150. if cfg.debug_pexpect:
  151. debug_pexpect_msg(self.p)
  152. end = self.p.before.rstrip()
  153. if not cfg.debug:
  154. vmsg(f' ==> {cyan(end)}')
  155. return end
  156. def interactive(self):
  157. return self.p.interact() # interact() not available with popen_spawn
  158. def kill(self,signal):
  159. return self.p.kill(signal)
  160. def match_expect_list(self,expect_list,greedy=False):
  161. allrep = '.*' if greedy else '.*?'
  162. expect = (
  163. r'(\b|\s)' +
  164. fr'\s{allrep}\s'.join(s.replace(r'.',r'\.').replace(' ',r'\s+') for s in expect_list) +
  165. r'(\b|\s)' )
  166. import re
  167. m = re.search( expect, self.read(strip_color=True), re.DOTALL )
  168. assert m, f'No match found for regular expression {expect!r}'
  169. return m
  170. def expect(self,s,t='',delay=None,regex=False,nonl=False,silent=False):
  171. if not silent:
  172. if cfg.verbose:
  173. msg_r('EXPECT ' + yellow(str(s)))
  174. elif not cfg.exact_output:
  175. msg_r('+')
  176. try:
  177. ret = (self.p.expect_exact,self.p.expect)[bool(regex)](s) if s else 0
  178. except pexpect.TIMEOUT as e:
  179. if cfg.debug_pexpect:
  180. raise
  181. m1 = f'\nERROR. Expect {s!r} timed out. Exiting\n'
  182. m2 = f'before: [{self.p.before}]\n'
  183. m3 = f'sent value: [{self.sent_value}]' if self.sent_value is not None else ''
  184. raise pexpect.TIMEOUT(m1+m2+m3) from e
  185. if cfg.debug_pexpect:
  186. debug_pexpect_msg(self.p)
  187. if cfg.verbose and not isinstance(s,str):
  188. msg_r(f' ==> {ret} ')
  189. if ret == -1:
  190. die(4,f'Error. Expect returned {ret}')
  191. else:
  192. if t:
  193. self.send(t,delay,s)
  194. else:
  195. if not nonl and not silent:
  196. vmsg('')
  197. return ret
  198. def send(self,t,delay=None,s=False):
  199. delay = delay or self.send_delay
  200. if delay:
  201. time.sleep(delay)
  202. ret = self.p.send(t) # returns num bytes written
  203. self.sent_value = t if ret else None
  204. if cfg.demo and delay:
  205. time.sleep(delay)
  206. if cfg.verbose:
  207. ls = '' if cfg.debug or not s else ' '
  208. es = '' if s else ' '
  209. yt = yellow('{!r}'.format( t.replace('\n',r'\n') ))
  210. msg(f'{ls}SEND {es}{yt}')
  211. return ret
  212. def read(self,n=-1,strip_color=False):
  213. return strip_ansi_escapes(self.p.read(n)).replace('\r','') if strip_color else self.p.read(n)
  214. def close(self):
  215. if self.pexpect_spawn:
  216. self.p.close()