#!/usr/bin/python3

import sys
import os
import errno
import pwd
import grp
import tempfile
import stat
import subprocess
import struct
import socket
import fcntl
import json
import signal
import resource


config_file = "/etc/guests.json"
lock_directory = "/run/guest-locks"
sockets_parent_directory = "/run/guest-socks"


def load_and_validate_config():
    global guests

    with open(config_file) as f:
        guests = json.load(f)

    users = { g["user"] for g in guests.values() }
    if len(users) != len(guests):
        raise ValueError("duplicate user")

    ifnames = { g["ifname"] for g in guests.values() }
    if len(ifnames) != len(guests):
        raise ValueError("duplicate ifname")

    for ifname in ifnames:
        if len(ifname.encode("ascii")) >= 16:
            raise ValueError("tap device name too long", ifname)

    all_disk_names = set()
    for guest in guests.values():
        guest_disks = frozenset(guest["disks"])
        inter = guest_disks & all_disk_names
        if inter:
            raise ValueError("disk appears on two vms", next(iter(inter)))
        all_disk_names.update(guest_disks)

    for guest in guests.values():
        ok = ("kernel" in guest) == ("initrd" in guest) == ("append" in guest)
        if not ok:
            raise ValueError("supply all three of kernel, initrd, append or none")

def open_lockfile_without_locking(guest_name):
    os.makedirs(lock_directory, exist_ok=True)
    return os.open(os.path.join(lock_directory, guest_name), os.O_WRONLY | os.O_CREAT)

def acquire_lock(lock_fd, block=False):
    flags = fcntl.LOCK_EX
    if not block: flags |= fcntl.LOCK_NB

    try:
        fcntl.lockf(lock_fd, flags)
    except OSError as e:
        if e.errno == errno.EAGAIN or e.errno == errno.EACCES:
            raise Exception("Failed to acquire lockfile. Is the guest running?") from e
        else:
            raise

def ip(*args, output=False):
    if output:
        f = subprocess.check_output
    else:
        f = subprocess.check_call

    return f(("/sbin/ip",) + args)

def open_tap(guest_name):
    guest = guests[guest_name]
    ifname = guest["ifname"]
    ips = guest["ips"]

    routes = ip("route", output=True)
    if 'dev {}'.format(ifname).encode('ascii') in routes:
        raise Exception("There are already routes for {}?".format(ifname))
    for addr in ips:
        if addr.encode("ascii") in routes:
            raise Exception("There are already routes for {}?".format(addr))

    IFF_TAP      = 0x0002
    IFF_VNET_HDR = 0x4000
    IFF_NO_PI    = 0x1000
    TUNSETIFF    = 0x400454ca

    fd = os.open("/dev/net/tun", os.O_RDWR)
    os.set_inheritable(fd, True)

    buf = bytearray(1024)
    flags = IFF_TAP | IFF_VNET_HDR | IFF_NO_PI
    ifreq = struct.pack_into('16sh', buf, 0, ifname.encode('ascii'), flags)
    res = fcntl.ioctl(fd, TUNSETIFF, buf)
    assert res == 0

    ip("link", "set", ifname, "up")
    ip("addr", "add", "fe80::1", "dev", ifname)
    for addr in ips:
        ip("route", "add", addr, "dev", ifname)
        if ":" in addr:
            ip("neigh", "add", "proxy", addr, "dev", "eth0")

    return fd

def open_disks(guest_name):
    guest = guests[guest_name]
    disk_fds = []

    for disk_name in guest["disks"]:
        fd = os.open("/dev/vg0/{}".format(disk_name), os.O_RDWR)
        os.set_inheritable(fd, True)
        disk_fds.append(fd)

    return disk_fds

def user_pwd(guest_name):
    return pwd.getpwnam(guests[guest_name]["user"])

def hugepages_folder(guest_name):
    user = guests[guest_name]["user"]
    return "/dev/hugepages/{}".format(user)

def make_hugepages_folder(guest_name):
    name = hugepages_folder(guest_name)
    user = user_pwd(guest_name)

    os.makedirs(name, exist_ok=True)
    os.chmod(name, 0o700)
    os.chown(name, user.pw_uid, user.pw_gid)

def allow_unlimited_mlock():
    resource.setrlimit(resource.RLIMIT_MEMLOCK, (resource.RLIM_INFINITY, resource.RLIM_INFINITY))

def drop_privileges(guest_name, kvm_group=False):
    user = user_pwd(guest_name)
    kvm = grp.getgrnam("kvm")
    os.setgid(user.pw_gid)
    if kvm_group:
        os.setgroups([kvm.gr_gid])
    else:
        os.setgroups([])
    os.setuid(user.pw_uid)  
    os.chdir("/")

def socket_dir(guest_name):
    return os.path.join(sockets_parent_directory, guest_name)

def make_socket_dir(guest_name):
    d = socket_dir(guest_name)
    os.makedirs(d, exist_ok=True)
    os.chmod(d, 0o770)
    user = user_pwd(guest_name)
    os.chown(d, os.getuid(), user.pw_gid)

def socket_path(guest_name, socket_name):
    return os.path.join(socket_dir(guest_name), socket_name)

def fd_path(fd):
    return "/proc/self/fd/{}".format(fd)

def build_composite_initrd(initrd_spec, chown_to):
    with tempfile.TemporaryFile(prefix='guest-manager-initrd-') as temp:
        for entry in initrd_spec:
            mode = os.stat(entry).st_mode

            if stat.S_ISDIR(mode):
                # directories are packed up with cpio and included
                with subprocess.Popen(["/usr/bin/find", "-L", "."], cwd=entry, stdout=subprocess.PIPE) as find:
                    cmd = ["/bin/cpio", "--quiet", "-R", "root", "-L", "-o", "-H", "newc"]
                    subprocess.check_call(cmd, cwd=entry, stdin=find.stdout, stdout=temp)
                    if find.wait() != 0:
                        raise Exception("find failed while building initrd")

            elif stat.S_ISREG(mode):
                # files (which must be cpio archives) are appended
                subprocess.check_call(["/bin/cat", "--", entry], stdout=temp, close_fds=True)

            else:
                raise Exception("can't include non-file/directory in initrd", entry)

            # uncompressed CPIO archives must be 4-byte aligned
            # (if we write 4 useless bytes, it will be ignored; that's fine)
            temp.seek(0, os.SEEK_END)
            padding = 4 - temp.tell() % 4
            temp.write(b'\0' * padding)
            temp.flush()

        os.fchown(temp.fileno(), chown_to, -1)

        # return a read only FD instead (and moreover one that will outlive the 'with')
        return os.open(fd_path(temp.fileno()), os.O_RDONLY)

def setup_boot_spec(guest_name):
    guest = guests[guest_name]

    if "kernel" not in guest:
        return None

    kernel_fd = os.open(guest["kernel"], os.O_RDONLY)

    initrd = guest["initrd"]
    if isinstance(initrd, list):
        user = user_pwd(guest_name)
        initrd_fd = build_composite_initrd(initrd, chown_to=user.pw_uid)
    else:
        initrd_fd = os.open(initrd, os.O_RDONLY)

    os.set_inheritable(kernel_fd, True)
    os.set_inheritable(initrd_fd, True)

    return {"kernel_fd": kernel_fd, "initrd_fd": initrd_fd, "append": guest["append"]}

def exec_kvm(guest_name, disk_fds, tap_fd, boot_spec):
    guest = guests[guest_name]

    prog = "/usr/bin/qemu-system-x86_64"
    args = \
        [ "qemu-KVM:{}".format(guest_name)
        , "-cpu", "host"
        , "-M", "pc-1.1" # emulated machine type
        , "-enable-kvm"
        , "-m", str(guest["ram_mbs"])
        , "-mem-path", hugepages_folder(guest_name)
        , "-mem-prealloc"
        , "-overcommit", "mem-lock=on"
        , "-smp", "{0},sockets={0},cores=1,threads=1".format(guest["cpus"])
        , "-name", guest_name
        , "-nographic"        # don't try to open a gui
        , "-nodefaults"       # don't create default emulated devices
        , "-chardev", "socket,id=charmonitor,path={},server,nowait".format(socket_path(guest_name, "monitor"))
        , "-mon", "chardev=charmonitor,id=monitor,mode=readline"
        , "-chardev", "socket,id=charcontrol,path={},server,nowait".format(socket_path(guest_name, "control"))
        , "-mon", "chardev=charcontrol,id=control,mode=control"
        , "-rtc", "base=utc"
        ]

    for idx, disk_fd in enumerate(disk_fds):
        fdset_id = 100 + idx
        args += \
            [ "-add-fd", "fd={},set={}".format(disk_fd, fdset_id)
            , "-blockdev", "node-name=disk-{},driver=raw,file.driver=host_device,file.filename=/dev/fdset/{}".format(idx, fdset_id)
            , "-device", "virtio-blk,drive=disk-{}".format(idx)
            ]

    args += \
        [ "-netdev", "tap,fd={},id=netdev1".format(tap_fd)
        , "-device", "virtio-net-pci,netdev=netdev1,id=net1,mac={}".format(guest["mac"])
        , "-chardev", "socket,id=charconsole,path={},server,nowait".format(socket_path(guest_name, "console"))
        , "-serial", "chardev:charconsole"
        ]

    if boot_spec is not None:
        args += \
            [ "-kernel", fd_path(boot_spec["kernel_fd"])
            , "-initrd", fd_path(boot_spec["initrd_fd"])
            , "-append", boot_spec["append"]
            ]

    os.execve(prog, args, {})

def exec_socat(socket_path):
    saved, = subprocess.check_output(["stty", "-g"]).splitlines()
    subprocess.check_call(["/bin/stty", "raw", "-echo", "isig", "intr", "^_", "kill", "^-", "susp", "^-", "quit", "^-"])
    print("Use Ctrl-_ to exit.", file=sys.stderr)
    try:
        subprocess.call(["/usr/bin/socat", "stdio", "unix:{}".format(socket_path)])
    except:
        pass
    subprocess.check_call(["/bin/stty", saved])
    print()

def acpi_shutdown(guest_name, lock_fd):
    def timeout(_x, _y): raise Exception("timeout")
    signal.signal(signal.SIGALRM, timeout)
    signal.alarm(300)

    with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as s:
        s.connect(socket_path(guest_name, "control"))
        f = s.makefile("rw", 1)

        def tx(**kwargs):
            f.write(json.dumps(kwargs) + "\n")
        def lines():
            for s in f:
                r = json.loads(s)
                if "event" in r:
                    continue
                elif "error" in r:
                    raise Exception(r)
                else:
                    yield r

        lines = iter(lines())

        greeting = next(lines)
        if "QMP" not in greeting:
            raise Exception("bad greeting", greeting)

        tx(execute="qmp_capabilities")
        resp = next(lines)
        if resp != {"return": {}}:
            raise Exception("hello failed", resp)

        tx(execute="system_powerdown")
        resp = next(lines)
        if resp != {"return": {}}:
            raise Exception("shutdown failed", resp)

    # we'll only be able to get the lockfile once the other process has actually
    # exited.
    acquire_lock(lock_fd, block=True)

def main():
    usage = "Usage: {} guest-name (boot|monitor|console|acpi-shutdown)".format(sys.argv[0])

    if len(sys.argv) != 3:
        raise Exception(usage)

    _, guest_name, action = sys.argv

    if action == "boot":
        lock_fd = open_lockfile_without_locking(guest_name)
        os.set_inheritable(lock_fd, True) # we want qemu to hold the lock
        acquire_lock(lock_fd)
        tap_fd = open_tap(guest_name)
        disk_fds = open_disks(guest_name)
        boot_spec = setup_boot_spec(guest_name)
        make_socket_dir(guest_name)
        make_hugepages_folder(guest_name)
        allow_unlimited_mlock()
        drop_privileges(guest_name, kvm_group=True)
        exec_kvm(guest_name, disk_fds, tap_fd, boot_spec)
    elif action in ("monitor", "console"):
        drop_privileges(guest_name)
        exec_socat(socket_path(guest_name, action))
    elif action == "acpi-shutdown":
        lock_fd = open_lockfile_without_locking(guest_name)
        drop_privileges(guest_name)
        acpi_shutdown(guest_name, lock_fd)
    else:
        raise Exception(usage)

if __name__ == "__main__":
    load_and_validate_config()
    main()
