#!/usr/bin/env python3


import os
import shlex
import shutil
import signal
import subprocess
import tempfile
import time


CLOUD_INIT_META_TEMPLATE = """\
# Amazon EC2 style metadata
local-hostname: {hostname}
"""

CLOUD_INIT_USER_TEMPLATE = """\
#cloud-config
ssh_authorized_keys:
- {pubkey}
"""


class Qemu:
    def __init__(self, name, base_image, disk_size, pubkey):
        self._name = name
        self._dirname = tempfile.mkdtemp()
        self._image = self._copy_image(base_image, disk_size)
        self._pubkey = pubkey
        self._memory = 512
        self._vcpus = 1
        self._username = None
        self._port = None
        self._pid = None

    def _join(self, *names):
        return os.path.join(self._dirname, *names)

    def set_disk_size(self, size):
        self._disk_size = str(size * 1024 * 1024)

    def set_memory(self, memory):
        self._memory = memory

    def set_vcpus(self, vcpus):
        self._vcpus = vcpus

    def set_username(self, username):
        self._username = username

    def _copy_image(self, base_image, size):
        image = self._join("qemu.img")
        shutil.copy(base_image, image)
        if size is not None:
            subprocess.check_call(["qemu-img", "resize", "-q", image, str(size)])
        return image

    def _cloud_init_iso(self):
        iso = self._join("cloud-init.iso")

        config = self._join("cloud-init")
        os.mkdir(config)

        meta = os.path.join(config, "meta-data")
        write_file(meta, CLOUD_INIT_META_TEMPLATE.format(hostname=self._name))

        user = os.path.join(config, "user-data")
        write_file(user, CLOUD_INIT_USER_TEMPLATE.format(pubkey=self._pubkey))

        if os.path.exists(iso):
            os.remove(iso)
        subprocess.check_call(
            [
                "genisoimage",
                "-quiet",
                "-joliet",
                "-rock",
                "-output",
                iso,
                "-volid",
                "cidata",
                config,
            ]
        )

        return iso

    def start_qemu(self):
        iso = self._cloud_init_iso()

        self._port = 2222  # FIXME
        pid_file = self._join("qemu.pid")
        out_file = self._join("qemu.out")
        err_file = self._join("qemu.err")

        argv = [
            "/usr/sbin/daemonize",
            "-c",
            self._dirname,
            "-p",
            pid_file,
            "-e",
            err_file,
            "-o",
            out_file,
            "/usr/bin/qemu-system-x86_64",
            "-name",
            self._name,
            "-m",
            str(self._memory),
            "-cpu",
            "host",
            "-smp",
            f"cpus={self._vcpus}",
            "-drive",
            f"file={self._image},format=qcow2,cache=none,if=virtio",
            "-drive",
            f"file={iso},if=ide,media=cdrom",
            "-device",
            "virtio-net,netdev=user.0",
            "-netdev",
            f"user,id=user.0,hostfwd=tcp::{self._port}-:22",
            "-nographic",
            "-enable-kvm",
        ]

        subprocess.check_call(argv, stdout=None, stderr=None)
        self._pid = int(wait_for(lambda: got_pid(pid_file), 1))

    def stop_qemu(self):
        if self._pid is not None:
            os.kill(self._pid, signal.SIGTERM)

    def wait_for_ssh(self):
        return wait_for(lambda: self.ssh(["hostname"]), 60)

    def ssh(self, argv):
        assert self._username is not None
        known = self._join("known.hosts")
        p = subprocess.Popen(
            [
                "ssh",
                "-p",
                str(self._port),
                "-l",
                self._username,
                f"-ouserknownhostsfile={known}",
                f"-ostricthostkeychecking=accept-new",
                "localhost",
                "--",
            ]
            + [shlex.quote(x) for x in argv],
            stdout=subprocess.PIPE,
            stderr=subprocess.DEVNULL,
        )
        output, _ = p.communicate()
        if output:
            return output, p.returncode
        time.sleep(5)


def read_file(filename):
    with open(filename, "r") as f:
        return f.read()


def write_file(filename, content):
    with open(filename, "w") as f:
        f.write(content)


def wait_for(func, timeout):
    start = time.time()
    while time.time() < start + timeout:
        val = func()
        if val is not None:
            return val
        time.sleep(0.1)


def got_pid(pid_file):
    if os.path.exists(pid_file):
        pid = read_file(pid_file)
        if pid:
            return pid.strip()


vmname = "scap-dev"
pubkey = read_file("/home/liw/.ssh/liw-openpgp.pub")
base_image = "/home/liw/tmp/debian-10-openstack-amd64.qcow2"
username = "debian"
memory = 1024
vcpus = 1

qemu = Qemu(vmname, base_image, 100 * 1024 ** 3, pubkey)
qemu.set_memory(memory)
qemu.set_vcpus(vcpus)
qemu.set_username(username)
print("starting qemu")
qemu.start_qemu()
print("waiting for ssh")
if qemu.wait_for_ssh():
    print("got ssh")

    out, exit = qemu.ssh(["id"])
    print(out.decode("UTF8"))

    out, exit = qemu.ssh(["find", "-ls"])
    print(out.decode("UTF8"))

    qemu.stop_qemu()
else:
    print("did not get ssh, boo")