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 qubesadmin
import qubesadmin.tools.qvm_clone as qvm_clone
import qubesadmin.tools.qvm_start as qvm_start
import qubesadmin.tools.qvm_run as qvm_run
import qubesadmin.tools.qvm_shutdown as qvm_shutdown
import argparse
import subprocess
import sys

def message(text):
    print(f"\033[1m{text}\033[0m")

def get_template_type(template):
    try:
        vm = app.domains[template]
        type = vm.features.get('os-distribution', None)
        if not type:
            type = subprocess.run(
                ['qvm-run', '-p', template, "cat /etc/os-release | grep ^ID="],
                capture_output=True, text=True, check=True
            ).stdout.split('=')[1].strip()
        if not type or type == "unknown":
            print(f"Error: Could not determine template type for {template}. It might be EOL.", file=sys.stderr)
            sys.exit(1)
        return 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):
    try:
        vm = app.domains[template]
        version = vm.features.get('os-version', None)
        if not version:
            version = subprocess.run(
                ['qvm-run', '-p', template, "grep -oP '\\d+' /etc/fedora-release"],
                capture_output=True, text=True
            ).stdout.strip()
        if not version:
            version = subprocess.run(
                ['qvm-run', '-p', template, "grep ^VERSION_CODENAME= /etc/os-release | cut -d'=' -f2"],
                capture_output=True, text=True
            ).stdout.strip()
        if not version:
            version = subprocess.run(
                ['qvm-run', '-p', template, "cat /etc/debian_version"],
                capture_output=True, text=True
            ).stdout.strip()
        if not version or version == "unknown":
            print(f"Error: Could not determine template version for {template}. It might be EOL.", file=sys.stderr)
            sys.exit(1)
        return version
    except subprocess.CalledProcessError:
        print(f"Error: Could not determine template version for {template}.", file=sys.stderr)
        sys.exit(1)

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

def change_qvm_features(template):
    subprocess.run(['qvm-start', '--skip-if-running', template], check=True)
    subprocess.run(['qvm-run', '-u', 'root', template, 'qvm-features-request --commit'], check=True)

def upgrade_debian_template(template, clone, new_template_name, old_name, new_name):
    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 subprocess.run(['qvm-clone', 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...")
    subprocess.run(['qvm-start', '--skip-if-running', template], check=True)

    message("Updating APT repositories...")
    subprocess.run(['qvm-run', '-u', 'root', template, f"sed -i 's/{old_name}/{new_name}/g' /etc/apt/sources.list /etc/apt/sources.list.d/*.list"], check=True)
    subprocess.run(['qvm-run', '-u', 'root', template, f"if ls /etc/apt/sources.list.d/*.sources 1> /dev/null 2>&1; then sed -i 's/{old_name}/{new_name}/g' /etc/apt/sources.list.d/*.sources; fi"], check=True)

    while retry_count < max_retries:
        if subprocess.run(['qvm-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
        else:
            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)
    subprocess.run(['qvm-shutdown', '--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):
    new_num = current_num + 1
    max_retries = 3
    retry_count = 0
    success = False

    if clone:
        message(f"Cloning {template} to {new_template_name}...")
        subprocess.run(['qvm-clone', 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...")
    subprocess.run(['qvm-start', '--skip-if-running', template], check=True)

    while retry_count < max_retries:
        if subprocess.run(['qvm-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
        else:
            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)
    subprocess.run(['qvm-shutdown', '--wait', template], check=True)

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

def main():
    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 == "unknown":
        print(f"Error: Could not determine template type for {template}.", file=sys.stderr)
        sys.exit(1)

    if not template_version or template_version == "unknown":
        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 == "unknown":
            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()