#!/usr/bin/python3

# SPDX-FileCopyrightText: 2015 Jason Mehring <https://nrgaway@gmail.com>
# SPDX-FileCopyrightText: 2016 Bahtiar Gadimov <https:/bahtiar.gadimov.de>
# SPDX-FileCopyrightText: 2016 - 2019 Marek Marczykowsk-Gorecki
#                                            <marmarek@invisiblethingslab.com>
# SPDX-FileCopyrightText: 2017 unman <unman@thirdeyesecurity.org>
# SPDX-FileCopyrightText: 2020 Frederic Pierret <https://github.com/fepitre>
# SPDX-FileCopyrightText: 2021 ctrlaltf24
# SPDX-FileCopyrightText: 2023 Benjamin Grande M. S. <ben.grande.b@gmail.com
#
# SPDX-License-Identifier: GPL-2.0-only

"""
Wrapper argound qubesctl to run the update states more easily.
With qubesctl, you specify the salt update states to run:
    qubesctl state.sls update.qubes-dom0
    qubesctl --skip-dom0 --templates state.sls update.qubes-vm
    qubesctl --skip-dom0 --standalones state.sls update.qubes-vm
    qubesctl --skip-dom0 --targets=VMNAME,VMNAME2 state.sls update.qubes-vm
While with qvm-update, it detects the salt update state for you:
    qvm-update --dom0
    qvm-update --templates
    qvm-update --standalones
    qvm-update --targets=VMNAME,VMNAME2
To do a full system update, qubesctl requires a very extensive command:
    qubesctl state.sls update.qubes-dom0
    targets="$(qvm-ls --no-spinner --raw-daata --fields NAME,CLASS | grep -v -e AppVM -e DispVM | cut -d "|" -f 1 | tr "\n" ",")"
    qubesctl --skip-dom0 --targets="$targets" state.sls update.qubes-vm
On the other side, qvm-update is plain simple:
    qvm-update --all
Note: Full system update: Dom0, TemplateVMs and StandaloneVMs.
Why not use the qvm-update-gui? Its maximum concurrency is fixed at 1.
Developed for R4.1. Can the R4.2 be concurrent?
"""

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('--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('--all', action='store_true',
                       help='Target all VMs (TemplateVMs, StandaloneVMs, '
                            'Dom0)')
    group.add_argument('--dom0', action='store_true',
                       help='Target Dom0')
    parser.add_argument('--templates', action='store_true',
                       help='Target all TemplatesVMs')
    parser.add_argument('--standalones', action='store_true',
                       help='Target all StandaloneVMs')
    args = parser.parse_args(args)

    args.command = ['state.sls']

    if args.dom0:
        args.command.append('update.qubes-dom0')
        try:
            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()

    targets = []
    if args.dom0:
        # already handled
        return 0
    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.all:
        targets = [vm for vm in app.domains.values()
                    if not vm.klass == 'AppVM' and not vm.klass == 'DispVM']
    elif args.targets:
        names = args.targets.split(',')
        targets = [vm for vm in app.domains.values()
                    if vm.name in names
                    and not vm.klass == 'AppVM' and not vm.klass == 'DispVM']

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

    exit_code = ''
    vms_to_go = targets
    if 'dom0' in vms_to_go:
        args.command.append('update.qubes-vm')
        vms_to_go = [vm for vm in targets
                    if not vm.name == 'dom0']
        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()
        args.command.remove('update.qubes-vm')
        args.command.append('update.qubes-dom0')
        try:
            subprocess.check_call(['qubesctl', '--dom0-only'] + args.command)
        except subprocess.CalledProcessError:
            print("DOM0 configuration failed, not continuing", file=sys.stderr)
            return 1
    else:
        args.command.append('update.qubes-vm')
        # remove dom0 - already handled
        targets = [vm for vm in targets if not vm.name == 'dom0']
        runner = qubessalt.ManageVMRunner(app, targets, args.command,
                                          show_output=args.show_output,
                                          force_color=args.force_color,
                                          max_concurrency=args.max_concurrency)
        exit_code = runner.run()

    return exit_code


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())