qubesctl

#!/usr/bin/python3
"""
Directly call a salt command in the modules, does not require a running salt
minion to run.
"""
from __future__ import print_function
import argparse

import sys
import subprocess
import qubessalt
import qubesadmin
import qubesadmin.vm


def main(args=None):  # pylint: disable=missing-docstring
    parser = argparse.ArgumentParser()
    parser.add_argument('--show-output', action='store_true',
                        help='Show output of management commands')
    parser.add_argument('--force-color', action='store_true',
                        help='Force color output, allow control characters '
                             'from VM, UNSAFE')
    parser.add_argument('--skip-dom0', action='store_true',
                        help='Skip dom0 configuration (VM creation etc)')
    parser.add_argument('--max-concurrency', action='store',
                        help='Maximum number of VMs configured simultaneously '
                             '(default: %(default)d)',
                        type=int, default=4)
    group = parser.add_mutually_exclusive_group()
    group.add_argument('--targets', action='store',
                       help='Comma separated list of VMs to target')
    group.add_argument('--templates', action='store_true',
                       help='Target all TemplatesVMs')
    group.add_argument('--standalones', action='store_true',
                       help='Target all StandaloneVMs')
    group.add_argument('--app', action='store_true',
                       help='Target all AppVMs')
    group.add_argument('--all', action='store_true',
                       help='Target all non-disposable VMs (TemplateVMs and '
                            'AppVMs)')
    parser.add_argument('command',
                        help='Salt command to execute (for example: '
                             'state.highstate)',
                        nargs=argparse.REMAINDER)
    args = parser.parse_args(args)

    if not args.skip_dom0:
        try:
            # TODO handle args.show_output - if false, log to some file
            subprocess.check_call(['qubesctl', '--dom0-only'] + args.command)
        except subprocess.CalledProcessError:
            print("DOM0 configuration failed, not continuing", file=sys.stderr)
            return 1

    app = qubesadmin.Qubes()
    # Load VM list only after dom0 salt call - some new VMs might be created

    targets = []
    if args.templates:
        targets += [vm for vm in app.domains.values()
                    if vm.klass == 'TemplateVM']
    if args.standalones:
        targets += [vm for vm in app.domains.values()
                    if vm.klass == 'StandaloneVM']
    if args.app:
        targets += [vm for vm in app.domains.values()
                    if vm.klass == 'AppVM']
    if args.all:
        # all but DispVMs
        targets = [vm for vm in app.domains.values()
                   if not vm.klass == 'DispVM']
    elif args.targets:
        names = args.targets.split(',')
        targets = [vm for vm in app.domains.values() if vm.name in names]

    # remove dom0 - already handled
    targets = [vm for vm in targets if vm.name != 'dom0']

    if args.show_output and args.force_color:
        args.command.insert(0, '--force-color')

    # templates first
    vms_to_go = [vm for vm in targets
                 if vm.klass == 'TemplateVM']
    runner = qubessalt.ManageVMRunner(app, vms_to_go, args.command,
                                      show_output=args.show_output,
                                      force_color=args.force_color,
                                      max_concurrency=args.max_concurrency)
    exit_code = runner.run()
    # then non-templates (AppVMs)
    vms_to_go = [vm for vm in targets
                 if not vm.klass == 'TemplateVM']
    runner = qubessalt.ManageVMRunner(app, vms_to_go, args.command,
                                      show_output=args.show_output,
                                      force_color=args.force_color,
                                      max_concurrency=args.max_concurrency)
    return max(exit_code, runner.run())


if __name__ == '__main__':
    # --dom0-only is a passthrough to salt-call
    if len(sys.argv) > 1 and sys.argv[1] == '--dom0-only':
        try:
            import qubes.mgmt.patches
        except ImportError:
            pass
        from salt.scripts import salt_call

        sys.argv[1] = '--local'
        salt_call()
    else:
        sys.exit(main())