#!/usr/bin/python3.3 -S
import argparse
import os
import shutil
import sys
import urllib.parse

def path_resolve(cmd):
    cmd0 = cmd

    if '/' not in cmd:
        cmd = shutil.which(cmd)
        assert cmd, cmd0

    cmd = os.path.abspath(cmd)
    assert os.path.exists(cmd), cmd0
    return cmd

script = path_resolve(sys.argv[0])

assert script.endswith('/vido')
runner = os.path.dirname(script) + '/virt-stub'
assert os.path.exists(runner)

def kopt_safe(st):
    # The kernel cmdline doesn't support any kind of quoting
    # XXX Check tabs and newlines as well?
    return ' ' not in st

assert kopt_safe(runner)

quote = urllib.parse.quote

parser = argparse.ArgumentParser()

# HOME and TERM seem to be set by the kernel (as well as root and the kernel
# cmdline), the rest would be set by a shell
preserve_default = 'HOME TERM PWD PATH SHELL'.split()

parser.add_argument(
    '--preserve-envvars', dest='preserve_envvars', nargs='+', default=[], metavar='ENVVAR')
parser.add_argument('--rw-dirs', dest='rw_dirs', nargs='+', metavar='DIR')
parser.add_argument('--clear-dirs', dest='clear_dirs', default=[], nargs='+', metavar='DIR')
parser.add_argument('--kopts', dest='kopts', nargs='+', metavar='KOPT')
parser.add_argument('--drives', dest='drives', nargs='+', metavar='DRIVE')
parser.add_argument('--kernel', metavar='KERNEL')
parser.add_argument('--mem', default='128M')
parser.add_argument('--gdb', action='store_true')
parser.add_argument('--kvm', action='store_true')
parser.add_argument('--net', action='store_true')
parser.add_argument('--qemu-runner')
parser.add_argument('cmd', nargs='*')

args = parser.parse_args()

if 'PWD' not in os.environ:
    os.environ['PWD'] = os.getcwd()

# Require that all vars in preserve_default be defined
preserve_envvars = preserve_default + [
    var for var in args.preserve_envvars if var in os.environ]

cmd = args.cmd or [os.environ['SHELL']]
cmd[0] = path_resolve(cmd[0])

kcmd = (
    [
        'rw', 'quiet',
        'init=' + runner,
        'UML_ENV=' + ';'.join(
            quote(var) + '=' + quote(os.environ[var])
            for var in preserve_envvars),
        'UML_CMD=' + '+'.join(map(quote, cmd))])

if args.rw_dirs is not None:
    for dn in args.rw_dirs:
        assert os.path.exists(dn), dn
    kcmd.append('UML_RWDIRS=' + ':'.join(
        quote(dn) for dn in args.rw_dirs))

if args.net:
    args.clear_dirs.append('/var/lib/dhcp')

if args.clear_dirs:
    for dn in args.clear_dirs:
        assert os.path.exists(dn), dn
    kcmd.append('UML_CLEARDIRS=' + ':'.join(
        quote(dn) for dn in args.clear_dirs))

if args.kopts is not None:
    assert all(map(kopt_safe, args.kopts))
    kcmd.extend(args.kopts)

if args.kvm:
    if args.qemu_runner is None:
        qemu = 'qemu-system-' + os.uname().machine
    else:
        qemu = args.qemu_runner

    if args.kernel is None:
        # Those may not work out of the box, need 9p+virtio built-in
        # building a custom initramfs would work around that
        args.kernel = '/boot/vmlinuz-' + os.uname().release

    qcmd = [
        '-m', args.mem, '-enable-kvm',
        '-fsdev', 'local,id=root,path=/,security_model=none',
        '-device', 'virtio-9p-pci,fsdev=root,mount_tag=/dev/root',
        '-kernel', args.kernel]

    if args.net:
        qcmd += ['-net', 'user']

    if args.gdb:
        qcmd += ['-s']
        # We can't use the -S flag to wait for gdb, any breakpoints set
        # at early startup would be unusable:
        # http://thread.gmane.org/gmane.comp.emulators.qemu/80327
        print('Please run: gdb -ex "target remote :1234" ./vmlinux')

    if args.drives:
        for dri in args.drives:
            qcmd += ['-drive', 'file={},if=virtio'.format(dri.replace(',', ',,'))]

    kcmd += ['rootfstype=9p', 'rootflags=trans=virtio']
    os.execvp(qemu, [qemu] + qcmd + ['-append', ' '.join(kcmd)])
else:
    if args.net:
        print('UML net support is not implemented', file=sys.stderr)
        sys.exit(2)
    gdb_cmd = [
        'gdb',
        '-ex', 'handle SIGSEGV pass nostop noprint',
        '-ex', 'handle SIGUSR1 pass nopass stop print',
        '--args']

    kcmd += ['rootfstype=hostfs', 'mem=' + args.mem]
    if args.drives:
        kcmd += ['ubd{}={}'.format(*el) for el in enumerate(args.drives)]
    if args.kernel is None:
        args.kernel = '/usr/bin/linux.uml'

    if args.gdb:
        os.execvp('gdb', gdb_cmd + [args.kernel] + kcmd)
    else:
        os.execvp(args.kernel, [args.kernel] + kcmd)

