qvm-template-upgrade

#!/usr/bin/python3
#
# Qubes OS Template Upgrade Script
# Supports Fedora and Debian Templates
# https://www.kennethrrosen.cloud
#
# Copyright (C) 2024 by Kenneth R. Rosen
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License; 
# as published by the Free Software Foundation; either version 2
# 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 General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.

import argparse
import subprocess
import sys
import qubesadmin
from qubesadmin.tools import qvm_clone, qvm_start, qvm_run, qvm_shutdown

def message(text):
    """
    Print a message in bold.
    """
    print(f"\033[1m{text}\033[0m")

def get_template_type(template):
    """
    Get the type of the template (Debian or Fedora).
    """
    try:
        vm = app.domains[template]
        vm_type = vm.features.get('os-distribution', None)
        if not vm_type:
            vm_type = vm.run("grep ^ID= /etc/os-release | cut -d'=' -f2", 
                              capture_output=True, text=True, check=True).stdout.strip().strip('"')
        if not vm_type or vm_type is None:
            print(f"Error: Could not determine template type for {template}. It might be EOL.", 
                  file=sys.stderr)
            sys.exit(1)
        return vm_type
    except subprocess.CalledProcessError:
        print(f"Error: Could not determine template type for {template}.", file=sys.stderr)
        sys.exit(1)

def get_template_version(template):
    """
    Get the version of the template.
    """
    try:
        vm = app.domains[template]
        version = vm.features.get('os-version', None)
        if not version:
            version = vm.run("grep -oP '\\d+' /etc/fedora-release",
                              capture_output=True, text=True).stdout.strip()
        if not version:
            version = vm.run("grep ^VERSION_CODENAME= /etc/os-release | cut -d'=' -f2",
                              capture_output=True, text=True).stdout.strip('"')
        if not version:
            version = vm.run("grep -o '.*' /etc/debian_version",
                              capture_output=True, text=True).stdout.strip()
        if not version or version is None:
            print(f"Error: Could not determine template version for {template}. It might be EOL.",
                  file=sys.stderr)
            sys.exit(1)
        return version
    except qubesadmin.exc.QubesDaemonAccessError:
        print(f"Error: Could not determine template version for {template}.", file=sys.stderr)
        sys.exit(1)

def next_debian_version(current_version):
    """
    Map current Debian version to the next version.
    """
    version_map = {
        "buster": "bullseye",
        "bullseye": "bookworm",
        "bookworm": "trixie",
        "trixie": "fornextversion"  # Update with the next version as necessary
    }
    return version_map.get(current_version, None)

def change_qvm_features(template):
    """
    Change QVM features for the specified template.
    """
    qvm_start.run(['--skip-if-running', template], check=True)
    qvm_run.run(['-u', 'root', template, 'qvm-features-request --commit'], check=True)

def upgrade_debian_template(template, clone, new_template_name, old_name, new_name):
    """
    Upgrade a Debian template.
    """
    max_retries = 3
    retry_count = 0
    success = False

    if clone and not new_template_name:
        message("Error: New template name required when cloning.")
        sys.exit(1)
    elif clone:
        message(f"Cloning {template} to {new_template_name}...")
        if not qvm_clone.run([template, new_template_name]).returncode == 0:
            message("Failed to clone template. Exiting.")
            sys.exit(1)
        template = new_template_name

    message(f"Upgrading {template} from {old_name} to {new_name}. Patience...")
    qvm_start.run(['--skip-if-running', template], check=True)

    message("Updating APT repositories...")
    qvm_run.run(f"if [ -n \"$(find /etc/apt/sources.list.d/ -name '*.list')\" ]; then sed -i 's/{old_name}/{new_name}/g' /etc/apt/sources.list.d/*.list; fi", check=True)
    qvm_run.run(f"if [ -n \"$(find /etc/apt/sources.list.d/ -name '*.sources')\" ]; then sed -i 's/{old_name}/{new_name}/g' /etc/apt/sources.list.d/*.sources; fi", check=True)

    while retry_count < max_retries:
        if qvm_run.run(['-p', '-u', 'root', template,
                        "DEBIAN_FRONTEND=noninteractive apt-get update && DEBIAN_FRONTEND=noninteractive apt-get -o Dpkg::Options::='--force-confnew' --fix-missing full-upgrade -y && DEBIAN_FRONTEND=noninteractive apt-get autoremove -y && apt-get clean"]).returncode == 0:
            success = True
            break
        retry_count += 1
        message(f"Attempt {retry_count} failed. Retrying...")

    if not success:
        message(f"Upgrade failed after {max_retries} attempts. Exiting.")
        sys.exit(1)

    message(f"Shutting down {template}...")
    change_qvm_features(template)
    qvm_shutdown.run(['--wait', template], check=True)

    message(f"Upgrade to {new_name} completed successfully for {template}.")

def upgrade_fedora_template(template, clone, new_template_name, current_num):
    """
    Upgrade a Fedora template.
    """
    new_num = current_num + 1
    max_retries = 3
    retry_count = 0
    success = False

    if clone:
        message(f"Cloning {template} to {new_template_name}...")
        qvm_clone.run([template, new_template_name], check=True)
        template = new_template_name
    else:
        new_template_name = template

    message(f"Upgrading Fedora {current_num} to Fedora {new_num}. Patience...")
    qvm_start.run(['--skip-if-running', template], check=True)

    while retry_count < max_retries:
        if qvm_run.run(['-p', '-u', 'root', template,
                        f"dnf clean all && dnf --releasever={new_num} distro-sync --best --allowerasing -y && dnf update -y && dnf upgrade -y"]).returncode == 0:
            success = True
            break
        retry_count += 1
        message(f"Attempt {retry_count} failed. Retrying...")

    if not success:
        message(f"Upgrade failed after {max_retries} attempts. Exiting.")
        sys.exit(1)

    message(f"Shutting down {template}...")
    change_qvm_features(template)
    qvm_shutdown.run(['--wait', template], check=True)

    message(f"Upgrade completed successfully for {template}.")

def main():
    """
    Main function to parse arguments and initiate the upgrade process.
    """
    parser = argparse.ArgumentParser(description="Qubes OS Template Upgrade Script")
    parser.add_argument('template', help='Name of the template to upgrade')
    parser.add_argument('-c', '--clone', action='store_true', help='Clone the template before upgrading')
    parser.add_argument('-N', '--new-template', help='New template name (required if cloning)')

    args = parser.parse_args()

    if args.clone and not args.new_template:
        parser.error("--new-template is required if --clone is specified")

    global app
    app = qubesadmin.Qubes()

    template = args.template
    clone = args.clone
    new_template_name = args.new_template

    message("Determining template type and version...")
    template_type = get_template_type(template)
    template_version = get_template_version(template)

    message(f"Template type: {template_type}")
    message(f"Template version: {template_version}")

    if not template_type or template_type is None:
        print(f"Error: Could not determine template type for {template}.", file=sys.stderr)
        sys.exit(1)

    if not template_version or template_version is None:
        print(f"Error: Could not determine template version for {template}.", file=sys.stderr)
        sys.exit(1)

    if template_type == "debian":
        new_name = next_debian_version(template_version)
        if new_name is None:
            print("Error: Unknown Debian version.", file=sys.stderr)
            sys.exit(1)
        upgrade_debian_template(template, clone, new_template_name, template_version, new_name)
    elif template_type == "fedora":
        upgrade_fedora_template(template, clone, new_template_name, int(template_version))
    else:
        print("Error: Unsupported template type. Use 'debian' or 'fedora'.", file=sys.stderr)
        sys.exit(1)

if __name__ == '__main__':
    main()