415f5faae9
Pylint complains that main() is too long and that's indeed true. Factor out single process call (together with all the prepartion) into separate function and keep only common code in main().
304 lines
11 KiB
Python
304 lines
11 KiB
Python
# -*- encoding: utf8 -*-
|
|
#
|
|
# The Qubes OS Project, http://www.qubes-os.org
|
|
#
|
|
# Copyright (C) 2017 Marek Marczykowski-Górecki
|
|
# <marmarek@invisiblethingslab.com>
|
|
#
|
|
# This program is free software; you can redistribute it and/or modify
|
|
# it under the terms of the GNU Lesser General Public License as published by
|
|
# the Free Software Foundation; either version 2.1 of the License, or
|
|
# (at your option) any later version.
|
|
#
|
|
# This program is distributed in the hope that it will be useful,
|
|
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
|
# GNU Lesser General Public License for more details.
|
|
#
|
|
# You should have received a copy of the GNU Lesser General Public License along
|
|
# with this program; if not, see <http://www.gnu.org/licenses/>.
|
|
|
|
''' qvm-run tool'''
|
|
import contextlib
|
|
import os
|
|
import signal
|
|
import subprocess
|
|
import sys
|
|
|
|
import multiprocessing
|
|
|
|
import select
|
|
|
|
import qubesadmin.tools
|
|
import qubesadmin.exc
|
|
|
|
parser = qubesadmin.tools.QubesArgumentParser()
|
|
|
|
parser.add_argument('--user', '-u', metavar='USER',
|
|
help='run command in a qube as USER (available only from dom0)')
|
|
|
|
parser.add_argument('--autostart', '--auto', '-a',
|
|
action='store_true', default=True,
|
|
help='option ignored, this is default')
|
|
|
|
parser.add_argument('--no-autostart', '--no-auto', '-n',
|
|
action='store_false', dest='autostart',
|
|
help='do not autostart qube')
|
|
|
|
parser.add_argument('--pass-io', '-p',
|
|
action='store_true', dest='passio', default=False,
|
|
help='pass stdio from remote program')
|
|
|
|
parser.add_argument('--localcmd', metavar='COMMAND',
|
|
help='with --pass-io, pass stdio to the given program')
|
|
|
|
parser.add_argument('--gui',
|
|
action='store_true', default=True,
|
|
help='run the command with GUI (default on)')
|
|
|
|
parser.add_argument('--no-gui', '--nogui',
|
|
action='store_false', dest='gui',
|
|
help='run the command without GUI')
|
|
|
|
parser.add_argument('--colour-output', '--color-output', metavar='COLOUR',
|
|
action='store', dest='color_output', default=None,
|
|
help='mark the qube output with given ANSI colour (ie. "31" for red)')
|
|
|
|
parser.add_argument('--colour-stderr', '--color-stderr', metavar='COLOUR',
|
|
action='store', dest='color_stderr', default=None,
|
|
help='mark the qube stderr with given ANSI colour (ie. "31" for red)')
|
|
|
|
parser.add_argument('--no-colour-output', '--no-color-output',
|
|
action='store_false', dest='color_output',
|
|
help='disable colouring the stdio')
|
|
|
|
parser.add_argument('--no-colour-stderr', '--no-color-stderr',
|
|
action='store_false', dest='color_stderr',
|
|
help='disable colouring the stderr')
|
|
|
|
parser.add_argument('--filter-escape-chars',
|
|
action='store_true', dest='filter_esc',
|
|
default=os.isatty(sys.stdout.fileno()),
|
|
help='filter terminal escape sequences (default if output is terminal)')
|
|
|
|
parser.add_argument('--no-filter-escape-chars',
|
|
action='store_false', dest='filter_esc',
|
|
help='do not filter terminal escape sequences; DANGEROUS when output is a'
|
|
' terminal emulator')
|
|
|
|
parser.add_argument('--service',
|
|
action='store_true', dest='service',
|
|
help='run a qrexec service (named by COMMAND) instead of shell command')
|
|
|
|
target_parser = parser.add_mutually_exclusive_group()
|
|
|
|
target_parser.add_argument('--dispvm', action='store', nargs='?',
|
|
const=True, metavar='BASE_APPVM',
|
|
help='start a service in new Disposable VM; '
|
|
'optionally specify base AppVM for DispVM')
|
|
target_parser.add_argument('VMNAME',
|
|
nargs='?',
|
|
action=qubesadmin.tools.VmNameAction)
|
|
|
|
# add those manually instead of vmname_args, because of mutually exclusive
|
|
# group with --dispvm; parsing is still handled by QubesArgumentParser
|
|
target_parser.add_argument('--all', action='store_true', dest='all_domains',
|
|
help='run command on all running qubes')
|
|
|
|
parser.add_argument('--exclude', action='append', default=[],
|
|
help='exclude the qube from --all')
|
|
|
|
parser.add_argument('cmd', metavar='COMMAND',
|
|
help='command or service to run')
|
|
|
|
def copy_stdin(stream):
|
|
'''Copy stdin to *stream*'''
|
|
# multiprocessing.Process have sys.stdin connected to /dev/null, use fd 0
|
|
# directly
|
|
while True:
|
|
try:
|
|
# select so this code works even if fd 0 is non-blocking
|
|
select.select([0], [], [])
|
|
data = os.read(0, 65536)
|
|
if data is None or data == b'':
|
|
break
|
|
stream.write(data)
|
|
stream.flush()
|
|
except KeyboardInterrupt:
|
|
break
|
|
stream.close()
|
|
|
|
def print_no_color(msg, file, color):
|
|
'''Print a *msg* to *file* without coloring it.
|
|
Namely reset to base color first, print a message, then restore color.
|
|
'''
|
|
if color:
|
|
print('\033[0m{}\033[0;{}m'.format(msg, color), file=file)
|
|
else:
|
|
print(msg, file=file)
|
|
|
|
|
|
def run_command_single(args, vm):
|
|
'''Handle a single VM to run the command in'''
|
|
run_kwargs = {}
|
|
if not args.passio:
|
|
run_kwargs['stdout'] = subprocess.DEVNULL
|
|
run_kwargs['stderr'] = subprocess.DEVNULL
|
|
elif args.localcmd:
|
|
run_kwargs['stdin'] = subprocess.PIPE
|
|
run_kwargs['stdout'] = subprocess.PIPE
|
|
run_kwargs['stderr'] = None
|
|
else:
|
|
# connect process output to stdout/err directly if --pass-io is given
|
|
run_kwargs['stdout'] = None
|
|
run_kwargs['stderr'] = None
|
|
if args.filter_esc:
|
|
run_kwargs['filter_esc'] = True
|
|
|
|
if isinstance(args.app, qubesadmin.app.QubesLocal) and \
|
|
not args.passio and \
|
|
not args.localcmd and \
|
|
args.service and \
|
|
not args.dispvm:
|
|
# wait=False works only in dom0; but it's still useful, to save on
|
|
# simultaneous vchan connections
|
|
run_kwargs['wait'] = False
|
|
|
|
copy_proc = None
|
|
local_proc = None
|
|
if args.service:
|
|
service = args.cmd
|
|
else:
|
|
service = 'qubes.VMShell'
|
|
if args.gui and args.dispvm:
|
|
service += '+WaitForSession'
|
|
proc = vm.run_service(service,
|
|
user=args.user,
|
|
**run_kwargs)
|
|
if not args.service:
|
|
proc.stdin.write(vm.prepare_input_for_vmshell(args.cmd))
|
|
proc.stdin.flush()
|
|
if args.localcmd:
|
|
local_proc = subprocess.Popen(args.localcmd,
|
|
shell=True,
|
|
stdout=proc.stdin,
|
|
stdin=proc.stdout)
|
|
# stdin is closed below
|
|
proc.stdout.close()
|
|
elif args.passio:
|
|
copy_proc = multiprocessing.Process(target=copy_stdin,
|
|
args=(proc.stdin,))
|
|
copy_proc.start()
|
|
# keep the copying process running
|
|
proc.stdin.close()
|
|
return proc, copy_proc, local_proc
|
|
|
|
|
|
def main(args=None, app=None):
|
|
'''Main function of qvm-run tool'''
|
|
args = parser.parse_args(args, app=app)
|
|
if args.color_output is None and args.filter_esc:
|
|
args.color_output = '31'
|
|
|
|
if args.color_stderr is None and os.isatty(sys.stderr.fileno()):
|
|
args.color_stderr = 31
|
|
|
|
if len(args.domains) > 1 and args.passio and not args.localcmd:
|
|
parser.error('--passio cannot be used when more than 1 qube is chosen '
|
|
'and no --localcmd is used')
|
|
if args.localcmd and not args.passio:
|
|
parser.error('--localcmd have no effect without --pass-io')
|
|
if args.color_output and not args.filter_esc:
|
|
parser.error('--color-output must be used with --filter-escape-chars')
|
|
|
|
retcode = 0
|
|
|
|
verbose = args.verbose - args.quiet
|
|
if args.passio:
|
|
verbose -= 1
|
|
|
|
# --all and --exclude are handled by QubesArgumentParser
|
|
domains = args.domains
|
|
dispvm = None
|
|
if args.dispvm:
|
|
if args.exclude:
|
|
parser.error('Cannot use --exclude with --dispvm')
|
|
dispvm = qubesadmin.vm.DispVM.from_appvm(args.app,
|
|
None if args.dispvm is True else args.dispvm)
|
|
domains = [dispvm]
|
|
elif args.all_domains:
|
|
# --all consider only running VMs
|
|
domains = [vm for vm in domains if vm.is_running()]
|
|
if args.color_output:
|
|
sys.stdout.write('\033[0;{}m'.format(args.color_output))
|
|
sys.stdout.flush()
|
|
if args.color_stderr:
|
|
sys.stderr.write('\033[0;{}m'.format(args.color_stderr))
|
|
sys.stderr.flush()
|
|
copy_proc = None
|
|
try:
|
|
procs = []
|
|
for vm in domains:
|
|
if not args.autostart and not vm.is_running():
|
|
if verbose > 0:
|
|
print_no_color('Qube \'{}\' not started'.format(vm.name),
|
|
file=sys.stderr, color=args.color_stderr)
|
|
retcode = max(retcode, 1)
|
|
continue
|
|
try:
|
|
if verbose > 0:
|
|
print_no_color(
|
|
'Running \'{}\' on {}'.format(args.cmd, vm.name),
|
|
file=sys.stderr, color=args.color_stderr)
|
|
if args.gui and not args.dispvm:
|
|
wait_session = vm.run_service('qubes.WaitForSession',
|
|
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
|
|
try:
|
|
wait_session.communicate(vm.default_user.encode())
|
|
except KeyboardInterrupt:
|
|
with contextlib.suppress(ProcessLookupError):
|
|
wait_session.send_signal(signal.SIGINT)
|
|
break
|
|
proc, copy_proc, local_proc = run_command_single(args, vm)
|
|
procs.append((vm, proc))
|
|
if local_proc:
|
|
procs.append((vm, local_proc))
|
|
except qubesadmin.exc.QubesException as e:
|
|
if args.color_output:
|
|
sys.stdout.write('\033[0m')
|
|
sys.stdout.flush()
|
|
vm.log.error(str(e))
|
|
return -1
|
|
try:
|
|
for vm, proc in procs:
|
|
this_retcode = proc.wait()
|
|
if this_retcode and args.verbose > 0:
|
|
print_no_color(
|
|
'{}: command failed with code: {}'.format(
|
|
vm.name, this_retcode),
|
|
file=sys.stderr, color=args.color_stderr)
|
|
retcode = max(retcode, proc.wait())
|
|
except KeyboardInterrupt:
|
|
for vm, proc in procs:
|
|
with contextlib.suppress(ProcessLookupError):
|
|
proc.send_signal(signal.SIGINT)
|
|
for vm, proc in procs:
|
|
retcode = max(retcode, proc.wait())
|
|
finally:
|
|
if dispvm:
|
|
dispvm.cleanup()
|
|
if args.color_output:
|
|
sys.stdout.write('\033[0m')
|
|
sys.stdout.flush()
|
|
if args.color_stderr:
|
|
sys.stderr.write('\033[0m')
|
|
sys.stderr.flush()
|
|
if copy_proc is not None:
|
|
copy_proc.terminate()
|
|
|
|
return retcode
|
|
|
|
|
|
if __name__ == '__main__':
|
|
sys.exit(main())
|