qvm_run.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349
  1. # -*- encoding: utf-8 -*-
  2. #
  3. # The Qubes OS Project, http://www.qubes-os.org
  4. #
  5. # Copyright (C) 2017 Marek Marczykowski-Górecki
  6. # <marmarek@invisiblethingslab.com>
  7. #
  8. # This program is free software; you can redistribute it and/or modify
  9. # it under the terms of the GNU Lesser General Public License as published by
  10. # the Free Software Foundation; either version 2.1 of the License, or
  11. # (at your option) any later version.
  12. #
  13. # This program is distributed in the hope that it will be useful,
  14. # but WITHOUT ANY WARRANTY; without even the implied warranty of
  15. # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
  16. # GNU Lesser General Public License for more details.
  17. #
  18. # You should have received a copy of the GNU Lesser General Public License along
  19. # with this program; if not, see <http://www.gnu.org/licenses/>.
  20. ''' qvm-run tool'''
  21. import contextlib
  22. import os
  23. import shlex
  24. import signal
  25. import subprocess
  26. import sys
  27. import multiprocessing
  28. import select
  29. import qubesadmin.tools
  30. import qubesadmin.exc
  31. import qubesadmin.utils
  32. parser = qubesadmin.tools.QubesArgumentParser()
  33. parser.add_argument('--user', '-u', metavar='USER',
  34. help='run command in a qube as USER (available only from dom0)')
  35. parser.add_argument('--autostart', '--auto', '-a',
  36. action='store_true', default=True,
  37. help='option ignored, this is default')
  38. parser.add_argument('--no-autostart', '--no-auto', '-n',
  39. action='store_false', dest='autostart',
  40. help='do not autostart/unpause qube')
  41. parser.add_argument('--pass-io', '-p',
  42. action='store_true', dest='passio', default=False,
  43. help='pass stdio from remote program')
  44. parser.add_argument('--localcmd', metavar='COMMAND',
  45. help='with --pass-io, pass stdio to the given program')
  46. parser.add_argument('--gui',
  47. action='store_true', default=True,
  48. help='run the command with GUI (default on)')
  49. parser.add_argument('--no-gui', '--nogui',
  50. action='store_false', dest='gui',
  51. help='run the command without GUI')
  52. parser.add_argument('--colour-output', '--color-output', metavar='COLOUR',
  53. action='store', dest='color_output', default=None,
  54. help='mark the qube output with given ANSI colour (ie. "31" for red)')
  55. parser.add_argument('--colour-stderr', '--color-stderr', metavar='COLOUR',
  56. action='store', dest='color_stderr', default=None,
  57. help='mark the qube stderr with given ANSI colour (ie. "31" for red)')
  58. parser.add_argument('--no-colour-output', '--no-color-output',
  59. action='store_false', dest='color_output',
  60. help='disable colouring the stdio')
  61. parser.add_argument('--no-colour-stderr', '--no-color-stderr',
  62. action='store_false', dest='color_stderr',
  63. help='disable colouring the stderr')
  64. parser.add_argument('--filter-escape-chars',
  65. action='store_true', dest='filter_esc',
  66. default=os.isatty(sys.stdout.fileno()),
  67. help='filter terminal escape sequences (default if output is terminal)')
  68. parser.add_argument('--no-filter-escape-chars',
  69. action='store_false', dest='filter_esc',
  70. help='do not filter terminal escape sequences; DANGEROUS when output is a'
  71. ' terminal emulator')
  72. parser.add_argument('--service',
  73. action='store_true', dest='service',
  74. help='run a qrexec service (named by COMMAND) instead of shell command')
  75. parser.add_argument('--no-shell', action='store_true',
  76. help='treat COMMAND as a simple executable, not a shell command')
  77. target_parser = parser.add_mutually_exclusive_group()
  78. target_parser.add_argument('--dispvm', action='store', nargs='?',
  79. const=True, metavar='BASE_APPVM',
  80. help='start a service in new Disposable VM; '
  81. 'optionally specify base AppVM for DispVM')
  82. target_parser.add_argument('VMNAME',
  83. nargs='?',
  84. action=qubesadmin.tools.VmNameAction)
  85. # add those manually instead of vmname_args, because of mutually exclusive
  86. # group with --dispvm; parsing is still handled by QubesArgumentParser
  87. target_parser.add_argument('--all', action='store_true', dest='all_domains',
  88. help='run command on all running qubes')
  89. parser.add_argument('--exclude', action='append', default=[],
  90. help='exclude the qube from --all')
  91. parser.add_argument('cmd', metavar='COMMAND',
  92. help='command or service to run')
  93. parser.add_argument('cmd_args', nargs='*', metavar='ARG',
  94. help='command arguments (implies --no-shell)')
  95. def copy_stdin(stream):
  96. '''Copy stdin to *stream*'''
  97. # multiprocessing.Process have sys.stdin connected to /dev/null, use fd 0
  98. # directly
  99. while True:
  100. try:
  101. # select so this code works even if fd 0 is non-blocking
  102. select.select([0], [], [])
  103. data = os.read(0, 65536)
  104. if data is None or data == b'':
  105. break
  106. stream.write(data)
  107. stream.flush()
  108. except KeyboardInterrupt:
  109. break
  110. stream.close()
  111. def print_no_color(msg, file, color):
  112. '''Print a *msg* to *file* without coloring it.
  113. Namely reset to base color first, print a message, then restore color.
  114. '''
  115. if color:
  116. print('\033[0m{}\033[0;{}m'.format(msg, color), file=file)
  117. else:
  118. print(msg, file=file)
  119. def run_command_single(args, vm):
  120. '''Handle a single VM to run the command in'''
  121. run_kwargs = {}
  122. if not args.passio:
  123. run_kwargs['stdout'] = subprocess.DEVNULL
  124. run_kwargs['stderr'] = subprocess.DEVNULL
  125. elif args.localcmd:
  126. run_kwargs['stdin'] = subprocess.PIPE
  127. run_kwargs['stdout'] = subprocess.PIPE
  128. run_kwargs['stderr'] = None
  129. else:
  130. # connect process output to stdout/err directly if --pass-io is given
  131. run_kwargs['stdout'] = None
  132. run_kwargs['stderr'] = None
  133. if args.filter_esc:
  134. run_kwargs['filter_esc'] = True
  135. if isinstance(args.app, qubesadmin.app.QubesLocal) and \
  136. not args.passio and \
  137. not args.localcmd and \
  138. args.service and \
  139. not args.dispvm:
  140. # wait=False works only in dom0; but it's still useful, to save on
  141. # simultaneous vchan connections
  142. run_kwargs['wait'] = False
  143. use_exec = len(args.cmd_args) > 0 or args.no_shell
  144. copy_proc = None
  145. local_proc = None
  146. shell_cmd = None
  147. if args.service:
  148. service = args.cmd
  149. elif use_exec:
  150. all_args = [args.cmd] + args.cmd_args
  151. if vm.features.check_with_template('vmexec', False):
  152. service = 'qubes.VMExec'
  153. if args.gui and args.dispvm:
  154. service = 'qubes.VMExecGUI'
  155. service += '+' + qubesadmin.utils.encode_for_vmexec(all_args)
  156. else:
  157. service = 'qubes.VMShell'
  158. if args.gui and args.dispvm:
  159. service += '+WaitForSession'
  160. shell_cmd = ' '.join(shlex.quote(arg) for arg in all_args)
  161. else:
  162. service = 'qubes.VMShell'
  163. if args.gui and args.dispvm:
  164. service += '+WaitForSession'
  165. shell_cmd = args.cmd
  166. proc = vm.run_service(service,
  167. user=args.user,
  168. **run_kwargs)
  169. if shell_cmd:
  170. proc.stdin.write(vm.prepare_input_for_vmshell(shell_cmd))
  171. proc.stdin.flush()
  172. if args.localcmd:
  173. local_proc = subprocess.Popen(args.localcmd,
  174. shell=True,
  175. stdout=proc.stdin,
  176. stdin=proc.stdout)
  177. # stdin is closed below
  178. proc.stdout.close()
  179. elif args.passio:
  180. copy_proc = multiprocessing.Process(target=copy_stdin,
  181. args=(proc.stdin,))
  182. copy_proc.start()
  183. # keep the copying process running
  184. proc.stdin.close()
  185. return proc, copy_proc, local_proc
  186. def main(args=None, app=None):
  187. '''Main function of qvm-run tool'''
  188. args = parser.parse_args(args, app=app)
  189. if args.passio:
  190. if args.color_output is None and args.filter_esc:
  191. args.color_output = 31
  192. if args.color_stderr is None and os.isatty(sys.stderr.fileno()):
  193. args.color_stderr = 31
  194. if len(args.domains) > 1 and args.passio and not args.localcmd:
  195. parser.error('--passio cannot be used when more than 1 qube is chosen '
  196. 'and no --localcmd is used')
  197. if args.localcmd and not args.passio:
  198. parser.error('--localcmd have no effect without --pass-io')
  199. if args.color_output and not args.filter_esc:
  200. parser.error('--color-output must be used with --filter-escape-chars')
  201. if args.service and args.no_shell:
  202. parser.error('--no-shell does not apply to --service')
  203. retcode = 0
  204. verbose = args.verbose - args.quiet
  205. if args.passio:
  206. verbose -= 1
  207. # --all and --exclude are handled by QubesArgumentParser
  208. domains = args.domains
  209. dispvm = None
  210. if args.dispvm:
  211. if args.exclude:
  212. parser.error('Cannot use --exclude with --dispvm')
  213. dispvm = qubesadmin.vm.DispVM.from_appvm(args.app,
  214. None if args.dispvm is True else args.dispvm)
  215. domains = [dispvm]
  216. elif args.all_domains:
  217. # --all consider only running VMs
  218. domains = [vm for vm in domains if vm.is_running()]
  219. if args.color_output:
  220. sys.stdout.write('\033[0;{}m'.format(args.color_output))
  221. sys.stdout.flush()
  222. if args.color_stderr:
  223. sys.stderr.write('\033[0;{}m'.format(args.color_stderr))
  224. sys.stderr.flush()
  225. copy_proc = None
  226. try:
  227. procs = []
  228. for vm in domains:
  229. if not args.autostart and not vm.is_running():
  230. if verbose > 0:
  231. print_no_color('Qube \'{}\' not started'.format(vm.name),
  232. file=sys.stderr, color=args.color_stderr)
  233. retcode = max(retcode, 1)
  234. continue
  235. if vm.is_paused():
  236. if not args.autostart:
  237. if verbose > 0:
  238. print_no_color(
  239. 'Qube \'{}\' is paused'.format(vm.name),
  240. file=sys.stderr, color=args.color_stderr)
  241. retcode = max(retcode, 1)
  242. continue
  243. try:
  244. vm.unpause()
  245. except qubesadmin.exc.QubesException:
  246. if verbose > 0:
  247. print_no_color(
  248. 'Qube \'{}\' cannot be unpaused'.format(
  249. vm.name),
  250. file=sys.stderr, color=args.color_stderr)
  251. retcode = max(retcode, 1)
  252. continue
  253. try:
  254. if verbose > 0:
  255. print_no_color(
  256. 'Running \'{}\' on {}'.format(args.cmd, vm.name),
  257. file=sys.stderr, color=args.color_stderr)
  258. if args.gui and not args.dispvm:
  259. wait_session = vm.run_service('qubes.WaitForSession',
  260. stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
  261. try:
  262. wait_session.communicate(vm.default_user.encode())
  263. except KeyboardInterrupt:
  264. with contextlib.suppress(ProcessLookupError):
  265. wait_session.send_signal(signal.SIGINT)
  266. break
  267. proc, copy_proc, local_proc = run_command_single(args, vm)
  268. procs.append((vm, proc))
  269. if local_proc:
  270. procs.append((vm, local_proc))
  271. except qubesadmin.exc.QubesException as e:
  272. if args.color_output:
  273. sys.stdout.write('\033[0m')
  274. sys.stdout.flush()
  275. vm.log.error(str(e))
  276. return -1
  277. try:
  278. for vm, proc in procs:
  279. this_retcode = proc.wait()
  280. if this_retcode and verbose > 0:
  281. print_no_color(
  282. '{}: command failed with code: {}'.format(
  283. vm.name, this_retcode),
  284. file=sys.stderr, color=args.color_stderr)
  285. retcode = max(retcode, proc.wait())
  286. except KeyboardInterrupt:
  287. for vm, proc in procs:
  288. with contextlib.suppress(ProcessLookupError):
  289. proc.send_signal(signal.SIGINT)
  290. for vm, proc in procs:
  291. retcode = max(retcode, proc.wait())
  292. finally:
  293. if dispvm:
  294. dispvm.cleanup()
  295. if args.color_output:
  296. sys.stdout.write('\033[0m')
  297. sys.stdout.flush()
  298. if args.color_stderr:
  299. sys.stderr.write('\033[0m')
  300. sys.stderr.flush()
  301. if copy_proc is not None:
  302. copy_proc.terminate()
  303. return retcode
  304. if __name__ == '__main__':
  305. sys.exit(main())