1 files changed, 216 insertions, 0 deletions
diff --git a/subplot/qemumgr.py b/subplot/qemumgr.py
new file mode 100755
index 0000000..c4fd385
--- /dev/null
+++ b/subplot/qemumgr.py
@@ -0,0 +1,216 @@
+#!/usr/bin/env python3
+
+
+import logging
+import os
+import random
+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 QemuSystem:
+    def __init__(self, name, base_image, disk_size, pubkey):
+        self._name = name
+        self._dirname = "."
+        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
+
+        fd, self._knownhosts = tempfile.mkstemp()
+        os.close(fd)
+
+    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 get_port(self):
+        return self._port
+
+    def _copy_image(self, base_image, size):
+        image = self._join("qemu.img")
+        logging.debug(f"QemuSystem: actual disk image: {image}")
+        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(self):
+        iso = self._cloud_init_iso()
+
+        self._port = random.randint(2000, 30000)
+        pid_file = self._join("qemu.pid")
+        out_file = self._join("qemu.out")
+        err_file = self._join("qemu.err")
+
+        MiB = 1024 ** 2
+        memory = int(self._memory / MiB)
+
+        argv = [
+            "/usr/bin/daemonize",
+            "-c",
+            self._dirname,
+            "-p",
+            pid_file,
+            "-e",
+            err_file,
+            "-o",
+            out_file,
+            "/usr/bin/qemu-system-x86_64",
+            "-name",
+            self._name,
+            "-m",
+            str(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))
+
+        logging.debug("started qemu-system")
+        logging.debug(f"  argv: {argv}")
+        logging.debug(f"  pid: {self._pid}")
+
+    def stop(self):
+        if self._pid is not None:
+            logging.debug(f"killing qemu-system process {self._pid}")
+            os.kill(self._pid, signal.SIGTERM)
+
+    def wait_for_ssh(self):
+        def ssh_ok():
+            output, exit = self.ssh(["true"])
+            logging.debug(f"tried ssh, exit={exit}")
+            if exit == 0:
+                return True
+            return None
+
+        return wait_for(ssh_ok, 300, sleep=5)
+
+    def ssh(self, argv):
+        assert self._username is not None
+
+        srcdir = globals()["srcdir"]
+
+        ssh_opts = [
+            "-i",
+            os.path.join(srcdir, "ssh", "id"),
+            "-p",
+            str(self._port),
+            "-l",
+            self._username,
+            f"-ouserknownhostsfile={self._knownhosts}",
+            "-ostricthostkeychecking=accept-new",
+            "-opasswordauthentication=no",
+        ]
+
+        real_argv = (
+            ["ssh"] + ssh_opts + ["localhost", "--"] + [shlex.quote(x) for x in argv]
+        )
+
+        logging.debug(f"running on VM: {real_argv}")
+        p = subprocess.Popen(
+            real_argv, stdout=subprocess.PIPE, stderr=subprocess.STDOUT
+        )
+        output, _ = p.communicate()
+        logging.debug(f"exit code: {p.returncode}")
+        logging.debug(f"output: {output}")
+        return output, p.returncode
+
+
+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, sleep=0.1):
+    start = time.time()
+    while time.time() < start + timeout:
+        val = func()
+        if val is not None:
+            return val
+        time.sleep(sleep)
+
+
+def got_pid(pid_file):
+    if os.path.exists(pid_file):
+        pid = read_file(pid_file)
+        if pid:
+            return pid.strip()