path: root/bin/sbuild-qemu-update

#!/usr/bin/python3
# SPDX-License-Identifier: GPL-2.0-or-later
#
# Copyright © 2020-2024 Christian Kastner <ckk@debian.org>
#             2021      Simon McVittie <smcv@debian.org>
#             2024      Johannes Schauer Marin Rodrigues <josch@debian.org>
#
# 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, see
# <http://www.gnu.org/licenses/>.
#
#######################################################################


# Note that there is significant overlap between this program and
# sbuild-qemu-boot. Both are in their developmental stages and I'd prefer to
# wait and see where this goes before refactoring them. --ckk


import argparse
import datetime
import os
import subprocess
import sys

import pexpect


SUPPORTED_ARCHS = [
    'amd64',
    'arm64',
    'armhf',
    'i386',
    'ppc64el',
]

IMAGEDIR = os.environ.get(
    'IMAGEDIR',
    os.path.join(os.path.expanduser('~'), '.cache', 'sbuild'),
)


def make_snapshot(image):
    iso_stamp = datetime.datetime.now().strftime('%Y-%m-%d_%H%M%S')
    run = subprocess.run(
        ['qemu-img', 'snapshot', '-l', image],
        capture_output=True
    )
    tags = [t.split()[1].decode('utf-8') for t in run.stdout.splitlines()[2:]]

    if iso_stamp in tags:
        print(
            f"Error: snapshot for {iso_stamp} already exists.",
            file=sys.stderr
        )
        return False

    run = subprocess.run(['qemu-img', 'snapshot', '-c', iso_stamp, image])
    return True if run.returncode == 0 else False


def get_qemu_base_args(image, guest_arch=None, boot="auto"):
    host_arch = subprocess.check_output(
        ['dpkg', '--print-architecture'],
        text=True,
    ).strip()

    if not guest_arch:
        # This assumes that images are named foo-bar-ARCH.img
        root, _ = os.path.splitext(os.path.basename(image))
        components = root.split('-')
        for c in reversed(components):
            if c in SUPPORTED_ARCHS:
                guest_arch = c
                break
        if not guest_arch:
            print(
                f"Could not guess guest architecture, please use --arch",
                file=sys.stderr,
            )
            return
    else:
        if not guest_arch in SUPPORTED_ARCHS:
            print(f"Unsupported architecture: {guest_arch}", file=sys.stderr)
            print("Supported architectures are: ", file=sys.stderr, end="")
            print(f"{', '.join(SUPPORTED_ARCHS)}", file=sys.stderr)
            return

    if guest_arch == 'amd64' :
        argv = ['qemu-system-x86_64']
        if host_arch == 'amd64':
            argv.append('-enable-kvm')
    elif guest_arch == 'i386':
        argv = ['qemu-system-i386', '-machine', 'q35']
        if host_arch in ['amd64', 'i386']:
            argv.append('-enable-kvm')
    elif guest_arch == 'ppc64el':
        argv = ['qemu-system-ppc64le']
        if host_arch == 'ppc64el':
            argv.append('-enable-kvm')
    elif guest_arch == 'arm64':
        argv = [
            'qemu-system-aarch64',
            '-machine', 'virt',
        ]
        if host_arch == 'arm64':
            argv.extend(['-cpu', 'host', '-enable-kvm'])
        else:
            argv.extend(['-cpu', 'cortex-a53'])
    elif guest_arch == 'armhf':
            if host_arch == 'arm64':
                argv = [
                    'qemu-system-aarch64',
                    '-cpu', 'host,aarch64=off',
                    '-enable-kvm'
                ]
            else:
                argv = ['qemu-system-arm']
            argv.extend([
                '-machine', 'virt',
            ])

    if boot == "auto":
        match guest_arch:
            case 'amd64'|'i386':
                boot = "bios"
            case 'arm64'|'armhf':
                boot = "efi"
            case 'ppc64el':
                boot = "ieee1275"

    eficode = None
    match boot:
        case "bios"|"none":
            pass
        case "efi":
            match guest_arch:
                case 'amd64':
                    eficode = "/usr/share/OVMF/OVMF_CODE_4M.fd"
                    if not os.path.exists(eficode):
                        eficode = "/usr/share/OVMF/OVMF_CODE.fd"
                case 'i386':
                    eficode = "/usr/share/OVMF/OVMF32_CODE_4M.secboot.fd"
                case 'arm64':
                    eficode = "/usr/share/AAVMF/AAVMF_CODE.fd"
                case 'armhf':
                    eficode = "/usr/share/AAVMF/AAVMF32_CODE.fd"
                case 'ppc64el':
                    print("efi not supported on ppc64el")
    if eficode:
        argv.extend(["-drive", f"if=pflash,format=raw,unit=0,read-only=on,file={eficode}"])


    return argv


def update_interaction(child, quiet):
    child.expect('host login: ')
    child.sendline('root')
    if not quiet:
        child.logfile = sys.stdout
    child.expect('root@host:~# ')
    child.sendline('DEBIAN_FRONTEND=noninteractive apt-get --quiet update')
    child.expect('root@host:~# ')
    child.sendline('DEBIAN_FRONTEND=noninteractive apt-get --quiet --assume-yes dist-upgrade')
    child.expect('root@host:~# ')
    child.sendline('DEBIAN_FRONTEND=noninteractive apt-get --quiet --assume-yes clean')
    child.expect('root@host:~# ')
    child.sendline('DEBIAN_FRONTEND=noninteractive apt-get --quiet --assume-yes autoremove')
    child.expect('root@host:~# ')
    child.sendline('sync')
    child.expect('root@host:~# ')
    child.sendline('mount -o remount,discard /')
    child.expect('root@host:~# ')
    child.sendline('fstrim /')
    child.expect('root@host:~# ')
    child.sendline('sync')
    child.expect('root@host:~# ')
    # Don't recall what issue this solves, but it solves it
    child.sendline('sleep 1')
    child.expect('root@host:~# ')
    child.sendline('shutdown -h now')


def main():
    parser = argparse.ArgumentParser(
        description='sbuild-update analog for QEMU images.',
    )
    parser.add_argument(
        '--snapshot',
        action='store_true',
        help="Create a snapshot of the image before changing it. Useful for "
             "reproducibility purposes.",
    )
    parser.add_argument(
        '--arch',
        help="Architecture to use (instead of attempting to auto-guess based "
             "on the image name).",
    )
    parser.add_argument(
        '--timeout',
        type=int,
        default=600,
        metavar='SECONDS',
        action='store',
        help="Maximum time to wait for command to finish with expected "
             "result. Mostly relevant for foreign architectures, where "
             "`apt-get update` can take quite a while. Default: 600s."
    )
    parser.add_argument(
        '--noexec',
        action='store_true',
        help="Don't actually do anything. Just print the command string that "
             "would be executed, and then exit.",
    )
    parser.add_argument(
        '--boot',
        choices=['auto', 'bios', 'efi', 'ieee1275', 'none'],
        default='auto',
        help="How to boot the image. Default is BIOS on amd64 and i386, EFI "
             "on arm64 and armhf, and IEEE1275 on ppc64el.",
        )
    parser.add_argument(
        'image',
        help="Image. Will first be interpreted as a path. If no suitable "
        "image exists at that location, then $IMAGEDIR\<image> is tried.",
    )
    dbglvlgroup = parser.add_mutually_exclusive_group()
    dbglvlgroup.add_argument("-v", "--verbose", action="store_true")
    dbglvlgroup.add_argument("-q", "--quiet", action="store_true")

    parser.add_argument("extra_args", nargs='*')

    parsed_args = parser.parse_args()

    if os.path.exists(parsed_args.image):
        image = parsed_args.image
    elif os.path.exists(os.path.join(IMAGEDIR, parsed_args.image)):
        image = os.path.join(IMAGEDIR, parsed_args.image)
    else:
        print("Image does not exist", file=sys.stderr)
        sys.exit(1)

    args = get_qemu_base_args(parsed_args.image, parsed_args.arch, parsed_args.boot)
    if not args:
        sys.exit(1)

    args.extend([
            '-object', 'rng-random,filename=/dev/urandom,id=rng0',
            '-device', 'virtio-rng-pci,rng=rng0,id=rng-device0',
            '-device', 'virtio-serial',
            '-nic',    'user,model=virtio',
            '-m',      '1024',
            '-smp',    '1',
            '-nographic',
            '-drive',  f'file={image},discard=unmap,if=virtio',
    ])
    if parsed_args.extra_args:
        args.extend(parsed_args.extra_args)

    print(' '.join(str(a) for a in args))
    if parsed_args.noexec:
        return
    elif parsed_args.snapshot and not make_snapshot(image):
        return

    logfile = None
    if parsed_args.verbose:
        logfile = sys.stdout
    child = pexpect.spawn(args[0], args[1:], logfile=logfile, encoding="utf-8")
    child.timeout = parsed_args.timeout
    try:
        update_interaction(child, parsed_args.quiet)
    except pexpect.TIMEOUT:
        print("Update timed out. Consider using --timeout.", file=sys.stderr)
        child.terminate()
    child.wait()

    try:
       subprocess.check_call(["qemu-img", "convert", "-O", "qcow2", image, f"{image}.tmp"])
    except subprocess.CalledProcessError as err:
        print(f"Error converting image: {err}", file=sys.stderr)
        return
    os.replace(f"{image}.tmp", image)


if __name__ == '__main__':
    main()