#!/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-update. 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


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',
            '-drive',   'if=pflash,format=raw,unit=0,read-only=on,'
                        'file=/usr/share/AAVMF/AAVMF_CODE.fd',
        ]
        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',
                '-drive',   'if=pflash,format=raw,unit=0,read-only=on,'
                            'file=/usr/share/AAVMF/AAVMF32_CODE.fd',
            ])

    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 main():
    parser = argparse.ArgumentParser(
        description='Boot a VM using a QEMU image.',
    )
    parser.add_argument('--read-write',
        action='store_true',
        help="Write changes back to the image, instead of using the image "
             "read-only.",
    )
    parser.add_argument(
        '--snapshot',
        action='store_true',
        help="Create a snapshot of the image before changing it. Useful for "
             "reproducibility purposes. Ignored if the image is not booted in "
             "read-write mode, which is the default.",
    )
    parser.add_argument(
        '--shared-dir',
        help="Share this directory on the host with the guest. This will only "
             "work when the image was created with sbuild-qemu-create(1).",
    )
    parser.add_argument(
        '--arch',
        help="Architecture to use (instead of attempting to auto-guess based "
             "on the image name).",
    )
    parser.add_argument(
        '--ram-size',
        metavar='MiB',
        action='store',
        default=2048,
        help=f"VM memory size in MB. Default: 2048",
    )
    parser.add_argument(
        '--cpus',
        metavar='CPUs',
        action='store',
        default=2,
        help="VM CPU count. Default: 2",
    )
    parser.add_argument(
        '--ssh-port',
        metavar='PORT',
        action='store',
        help="Forward local port PORT to port 22 within the guest. Package "
             "'openssh-server' must be installed within the guest for this "
             "to be useful.",
    )
    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.",
    )
    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)

    nic = 'user,model=virtio'
    if parsed_args.ssh_port:
        nic += f',hostfwd=tcp:127.0.0.1:{parsed_args.ssh_port}-:22'

    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',    nic,
            '-m',      str(parsed_args.ram_size),
            '-smp',    str(parsed_args.cpus),
            '-nographic',
    ])

    if parsed_args.shared_dir:
        args.extend([
            '-virtfs', f'local,path={parsed_args.shared_dir},id=sbuild-qemu,'
                        'mount_tag=sbuild-qemu,security_model=none',
        ])

    # Pass on host terminal rows/columns to guest
    # FIXME: qemu-system-pp64le doesn't support fw_cfg?
    if args[0] not in ['qemu-system-ppc64le']:
        termsize = os.get_terminal_size()
        args.extend([
            '-fw_cfg', f'name=opt/sbuild-qemu/tty-rows,string={termsize.lines}',
            '-fw_cfg', f'name=opt/sbuild-qemu/tty-cols,string={termsize.columns}',
            ])

    args.extend(['-drive', f'file={image},discard=unmap,if=virtio'])

    print(' '.join(str(a) for a in args))
    if parsed_args.noexec:
        return

    if parsed_args.read_write:
        if parsed_args.snapshot and not make_snapshot(image):
            return
    else:
        args.append('-snapshot')

    os.execvp(args[0], args)


if __name__ == '__main__':
    main()
