#!/usr/bin/python3
#
# adt-virt-qemu is part of autopkgtest
# autopkgtest is a tool for testing Debian binary packages
#
# autopkgtest is Copyright (C) 2006-2014 Canonical Ltd.
#
# adt-virt-qemu was developed by
# Martin Pitt <martin.pitt@ubuntu.com>
#
# 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, write to the Free Software
# Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
#
# See the file CREDITS for a full list of credits information (often
# installed as /usr/share/doc/autopkgtest/CREDITS).

import sys
import os
import subprocess
import tempfile
import shutil
import argparse
import time
import socket
import errno
import fcntl

try:
    our_base = os.environ['AUTOPKGTEST_BASE'] + '/lib'
except KeyError:
    our_base = '/usr/share/autopkgtest/python'
sys.path.insert(1, our_base)

import VirtSubproc
import adtlog


args = None
workdir = None
p_qemu = None
ssh_port = None
ssh_port_lock = None
normal_user = None


def parse_args():
    global args

    parser = argparse.ArgumentParser(fromfile_prefix_chars='@')

    qemu_cmd_default = 'qemu-system-' + os.uname()[4]

    parser.add_argument('-q', '--qemu-command', default=qemu_cmd_default,
                        help='QEMU command (default: %s)' % qemu_cmd_default)
    parser.add_argument('-o', '--overlay-dir',
                        help='Temporary overlay directory (default: in /tmp)')
    parser.add_argument('-u', '--user',
                        help='user to log into the VM on ttyS0 (must be able '
                        'to sudo if not "root")')
    parser.add_argument('-p', '--password',
                        help='password for user to log into the VM on ttyS0')
    parser.add_argument('-c', '--cpus', type=int, default=1,
                        help='Number of (virtual) CPUs in the VM (default: %(default)s)')
    parser.add_argument('--ram-size', type=int, default=1024,
                        help='VM RAM size in MiB (default: %(default)s)')
    parser.add_argument('-d', '--debug', action='store_true',
                        help='Enable debugging output')
    parser.add_argument('--qemu-options',
                        help='Pass through arguments to QEMU command.')
    parser.add_argument('image', nargs='+',
                        help='disk image to add to the VM (in order)')

    args = parser.parse_args()

    if args.debug:
        adtlog.verbosity = 2


def prepare_overlay():
    '''Generate a temporary overlay image'''

    # generate a temporary overlay
    if args.overlay_dir:
        overlay = os.path.join(args.overlay_dir, os.path.basename(
            args.image[0]) + '.overlay-%s' % time.time())
    else:
        overlay = os.path.join(workdir, 'overlay.img')
    adtlog.debug('Creating temporary overlay image in %s' % overlay)
    VirtSubproc.check_exec(['qemu-img', 'create', '-f', 'qcow2', '-b',
                            os.path.abspath(args.image[0]), overlay],
                           outp=True, timeout=300)
    return overlay


def wait_boot():
    term = VirtSubproc.get_unix_socket(os.path.join(workdir, 'ttyS0'))
    VirtSubproc.expect(term, b' login: ', 300, 'login prompt on ttyS0')
    # this is really ugly, but runlevel, "service status hwclock" etc. all
    # don't help to determine if the system is *really* booted; running
    # commands too early causes the system time to be all wrong
    time.sleep(3)


def check_ttyS1_shell():
    '''Check if there is a shell running on ttyS1'''

    term = VirtSubproc.get_unix_socket(os.path.join(workdir, 'ttyS1'))
    term.send(b'echo -n o; echo k\n')
    try:
        VirtSubproc.expect(term, b'ok', 1)
        return True
    except VirtSubproc.Timeout:
        return False


def setup_shell():
    '''Log into the VM and set up root shell on ttyS1'''

    # if the VM is already prepared to start a root shell on ttyS1, just use it
    if check_ttyS1_shell():
        adtlog.debug('setup_shell(): there already is a shell on ttyS1')
        return
    else:
        adtlog.debug('setup_shell(): no default shell on ttyS1')

    if args.user and args.password:
        # login on ttyS0 and start a root shell on ttyS1 from there
        adtlog.debug('Shell setup: have user and password, logging in..')
        login_tty_and_setup_shell()
    else:
        VirtSubproc.bomb('The VM does not start a root shell on ttyS1 already.'
                         ' The only other supported login mechanism is '
                         'through --user and --password on the guest ttyS0')

    if not check_ttyS1_shell():
        VirtSubproc.bomb('setup_shell(): failed to setup shell on ttyS1')


def login_tty_and_setup_shell():
    '''login on ttyS0 and start a root shell on ttyS1 from there'''

    term = VirtSubproc.get_unix_socket(os.path.join(workdir, 'ttyS0'))

    # send user name
    term.send(args.user.encode('UTF-8'))
    term.send(b'\n')
    # wait until we get some more data for the password prompt
    VirtSubproc.expect(term, None, 10, 'password prompt')
    # send password
    term.send(args.password.encode('UTF-8'))
    term.send(b'\n')
    adtlog.debug('login_tty: sent password')

    cmd = b'setsid sh </dev/ttyS1 >/dev/ttyS1 2>&1 &'

    # if we are a non-root user, run through sudo
    if args.user != 'root':
        cmd = b"sudo sh -c '" + cmd + "'"

    term.send(cmd)
    term.send(b'\nexit\n')
    VirtSubproc.expect(term, b'\nlogout', 10)


def setup_shared(shared_dir):
    '''Set up shared dir'''

    term = VirtSubproc.get_unix_socket(os.path.join(workdir, 'ttyS1'))

    term.send(b'''mkdir -m 1777 /autopkgtest
mount -t 9p -o trans=virtio,access=any autopkgtest /autopkgtest
chmod 1777 /autopkgtest
touch /autopkgtest/done_shared
''')

    with VirtSubproc.timeout(10, 'timed out on client shared directory setup'):
        flag = os.path.join(shared_dir, 'done_shared')
        while not os.path.exists(flag):
            time.sleep(0.2)
    VirtSubproc.expect(term, b'#', 30)

    # ensure that root has $HOME set
    term.send(b'[ -n "$HOME" ] || export HOME=`getent passwd root|cut -f6 -d:`\n')
    VirtSubproc.expect(term, b'#', 5)


def setup_config(shared_dir):
    '''Set up configuration files'''

    term = VirtSubproc.get_unix_socket(os.path.join(workdir, 'ttyS1'))

    # copy our timezone, to avoid time skews with the host
    if os.path.exists('/etc/timezone'):
        tz = None
        with open('/etc/timezone', 'rb') as f:
            for l in f:
                if l.startswith(b'#'):
                    continue
                l = l.strip()
                if l:
                    tz = l
                    break

        if tz:
            adtlog.debug('Copying host timezone %s to VM' % tz.decode())
            term.send(b'echo ' + tz + b' > /etc/timezone; DEBIAN_FRONTEND=noninteractive dpkg-reconfigure tzdata\n')
            VirtSubproc.expect(term, b'#', 5)
        else:
            adtlog.debug('Could not determine host timezone')

    # create helper for runcmd: cat data from its stdin (from a file) to stdout
    # eternally (like tail -f), but stop once either an "EOF" file exists and
    # we copied at least as many bytes as given in that EOF file (the first
    # arg), or an "exit flag" file exists.
    # We don't run that from /autopkgtest/ as 9p from older QEMU versions is
    # buggy and causes "invalid numeric result" errors on that.
    term.send(b'PYTHON=$(which python3) || PYTHON=$(which python); cat <<EOF > /bin/eofcat; chmod 755 /bin/eofcat\n')
    term.send(b'''#!$PYTHON
import sys, os, fcntl, time, errno
(feof, fexit) = sys.argv[1:]
count = 0
limit = None
fcntl.fcntl(0, fcntl.F_SETFL, fcntl.fcntl(0, fcntl.F_GETFL) | os.O_NONBLOCK)
while not os.path.exists(fexit):
    try:
        block = os.read(0, 1000000)
        if block:
            os.write(1, block)
            count += len(block)
            continue
    except OSError as e:
        if e.errno != errno.EAGAIN:
            raise

    time.sleep(0.05)
    if limit is None:
        try:
            with open(feof, 'r') as f:
                limit = int(f.read())
        except (IOError, ValueError):
            pass

    if limit is not None and count >= limit:
        break
EOF
''')
    VirtSubproc.expect(term, b'#', 5)


def make_auxverb(shared_dir):
    '''Create auxverb script'''

    auxverb = os.path.join(workdir, 'runcmd')
    with open(auxverb, 'w') as f:
        f.write('''#!%(py)s
import sys, os, tempfile, threading, time, atexit, shutil, fcntl, errno, pipes
import socket

dir_host = '%(dir)s'
job_host = tempfile.mkdtemp(prefix='job.', dir=dir_host)
atexit.register(shutil.rmtree, job_host)
os.chmod(job_host, 0o755)
job_guest = '/autopkgtest/' + os.path.basename(job_host)
running = True

def shovel(fin, fout, flagfile_on_eof=None):
    fcntl.fcntl(fin, fcntl.F_SETFL,
                fcntl.fcntl(fin, fcntl.F_GETFL) | os.O_NONBLOCK)
    count = 0
    while running:
        try:
            block = os.read(fin, 1000000)
            if flagfile_on_eof and not block:
                os.fsync(fout)
                os.close(fout)
                with open(flagfile_on_eof, 'w') as f:
                    f.write('%%i' %% count)
                return
            count += len(block)
        except OSError as e:
            if e.errno != errno.EAGAIN:
                raise
            block = None
        if not block:
            time.sleep(0.01)
            continue
        while True:
            try:
                os.write(fout, block)
                break
            except OSError as e:
                if e.errno != errno.EAGAIN:
                    raise
                continue


# redirect the guest process stdin/out/err files to our stdin/out/err
fin = os.path.join(job_host, 'stdin')
stdin_eof = os.path.join(job_host, 'stdin_eof')
fout = os.path.join(job_host, 'stdout')
ferr = os.path.join(job_host, 'stderr')
with open(fout, 'w'):
    pass
with open(ferr, 'w'):
    pass
t_stdin = threading.Thread(None, shovel, 'copyin', (sys.stdin.fileno(), os.open(fin, os.O_CREAT|os.O_WRONLY), stdin_eof))
t_stdin.start()
t_stdout = threading.Thread(None, shovel, 'copyout', (os.open(fout, os.O_RDONLY), sys.stdout.fileno()))
t_stdout.start()
t_stderr = threading.Thread(None, shovel, 'copyerr', (os.open(ferr, os.O_RDONLY), sys.stderr.fileno()))
t_stderr.start()

# Run command through QEMU shell. We can't directly feed the stdin file into
# the process as we'd hit EOF too soon; so funnel it through eofcat to get a
# "real" stdin behaviour.
s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
s.connect('%(tty)s')
cmd = '/bin/eofcat %%(d)s/stdin_eof %%(d)s/exit.tmp < %%(d)s/stdin | ' \\
      '(%%(c)s >> %%(d)s/stdout 2>> %%(d)s/stderr; echo $? > %%(d)s/exit.tmp);' \\
      'mv %%(d)s/exit.tmp %%(d)s/exit\\n' %% \\
       {'d': job_guest, 'c': ' '.join(map(pipes.quote, sys.argv[1:]))}
s.send(cmd.encode())

# wait until command has exited
path_exit = os.path.join(job_host, 'exit')
while not os.path.exists(path_exit) or os.path.getsize(path_exit) == 0:
    time.sleep(0.2)
running = False

# mop up terminal response
while True:
    try:
        block = s.recv(4096, socket.MSG_DONTWAIT)
        if not block:
            break
    except IOError:
        break
    time.sleep(0.05)
s.close()

with open(path_exit) as f:
    rc = int(f.read().strip())

t_stdin.join()
t_stdout.join()
t_stderr.join()
sys.exit(rc)
''' % {'py': sys.executable, 'tty': os.path.join(workdir, 'ttyS1'), 'dir': shared_dir})

    os.chmod(auxverb, 0o755)

    VirtSubproc.auxverb = [auxverb]

    # verify that we can connect
    status, out, err = VirtSubproc.execute_timeout(
        None, 5, VirtSubproc.auxverb + ['runlevel'],
        stdout=subprocess.PIPE)
    adtlog.debug('runlevel: exit %i, out "%s", err "%s"' % (status, out, err))
    if status == 0 and out.split()[-1] in ['2', '3', '4', '5']:
        adtlog.debug('can connect to autopkgtest sh in VM')
    else:
        VirtSubproc.bomb('failed to connect to VM')


def find_free_port(start):
    '''Find an unused port in the range [start, start+50)'''

    global ssh_port_lock

    for p in range(start, start + 50):
        adtlog.debug('find_free_port: trying %i' % p)
        try:
            try:
                ssh_port_lock = open('/run/lock/adt-virt-qemu.port.%i' % p, 'w')
                fcntl.flock(ssh_port_lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
            except (IOError, OSError):
                adtlog.debug('find_free_port: %i is locked' % p)
                if ssh_port_lock:
                    ssh_port_lock.close()
                ssh_port_lock = None
                continue

            s = socket.create_connection(('127.0.0.1', p))
            # if that works, the port is taken
            s.close()
            ssh_port_lock.close()
            ssh_port_lock = None
            continue
        except socket.error as e:
            if e.errno == errno.ECONNREFUSED:
                adtlog.debug('find_free_port: %i is free' % p)
                return p
            else:
                pass

    adtlog.debug('find_free_port: all ports are taken')
    return None


def determine_normal_user(shared_dir):
    '''Check for a normal user to run tests as.'''

    global normal_user

    if args.user and args.user != 'root':
        normal_user = args.user
        return

    # get the first UID >= 500
    term = VirtSubproc.get_unix_socket(os.path.join(workdir, 'ttyS1'))
    term.send(b"getent passwd | sort -t: -nk3 | "
              b"awk -F: '{if ($3 >= 500) { print $1; exit } }'"
              b"> /autopkgtest/normal_user\n")
    with VirtSubproc.timeout(5, 'timed out on determining normal user'):
        outfile = os.path.join(shared_dir, 'normal_user')
        while not os.path.exists(outfile):
            time.sleep(0.2)
    with open(outfile) as f:
        out = f.read()
        if out:
            normal_user = out.strip()
            adtlog.debug('determine_normal_user: got user "%s"' % normal_user)
        else:
            adtlog.debug('determine_normal_user: no uid >= 500 available')


def hook_open():
    global workdir, p_qemu, ssh_port

    workdir = tempfile.mkdtemp(prefix='adt-virt-qemu.')
    os.chmod(workdir, 0o755)

    shareddir = os.path.join(workdir, 'shared')
    os.mkdir(shareddir)

    overlay = prepare_overlay()

    # start QEMU
    argv = [args.qemu_command,
            '-enable-kvm',
            '-m', str(args.ram_size),
            '-smp', str(args.cpus),
            '-localtime',
            '-nographic',
            '-net', 'user',
            '-net', 'nic,model=virtio',
            '-monitor', 'unix:%s/monitor,server,nowait' % workdir,
            '-serial', 'unix:%s/ttyS0,server,nowait' % workdir,
            '-serial', 'unix:%s/ttyS1,server,nowait' % workdir,
            '-virtfs',
            'local,id=autopkgtest,path=%s,security_model=none,mount_tag=autopkgtest' % shareddir,
            '-drive', 'file=%s,cache=unsafe,if=virtio,index=0' % overlay]
    for i, image in enumerate(args.image[1:]):
        argv.append('-drive')
        argv.append('file=%s,if=virtio,index=%i,readonly' % (image, i + 1))

    # pass through option to qemu
    if args.qemu_options:
        argv.extend(args.qemu_options.split())

    # find free port to forward VM port 22 (for SSH access)
    ssh_port = find_free_port(10022)
    if ssh_port:
        adtlog.debug('Forwarding local port %i to VM ssh port 22' % ssh_port)
        argv.append('-redir')
        argv.append('tcp:%i::22' % ssh_port)

    p_qemu = subprocess.Popen(argv)

    try:
        try:
            wait_boot()
        finally:
            # remove overlay as early as possible, to avoid leaking large
            # files; let QEMU run with the deleted inode
            os.unlink(overlay)
        setup_shell()
        setup_shared(shareddir)
        setup_config(shareddir)
        make_auxverb(shareddir)
        determine_normal_user(shareddir)
    except:
        # Clean up on failure
        hook_cleanup()
        raise


def hook_downtmp(path):
    # we would like to do this, but 9p is currently way too slow for big source
    # trees
    # downtmp = '/autopkgtest/tmp'
    # VirtSubproc.check_exec(['mkdir', '-m', '1777', downtmp], downp=True)
    return VirtSubproc.downtmp_mktemp(path)


def hook_revert():
    VirtSubproc.downtmp_remove()
    hook_cleanup()
    hook_open()


def hook_cleanup():
    global p_qemu, workdir

    if p_qemu:
        p_qemu.terminate()
        p_qemu.wait()
        p_qemu = None

    if workdir:
        shutil.rmtree(workdir)
        workdir = None


def hook_reboot():
    global workdir
    VirtSubproc.check_exec(['reboot'], downp=True, timeout=5)
    wait_boot()
    shareddir = os.path.join(workdir, 'shared')
    setup_shared(shareddir)


def hook_forked_inchild():
    pass


def hook_capabilities():
    global normal_user
    caps = ['revert', 'revert-full-system', 'root-on-testbed',
            'isolation-machine', 'reboot']
    # disabled, see hook_downtmp()
    # caps.append('downtmp-host=%s' % os.path.join(workdir, 'shared', 'tmp'))
    if normal_user:
        caps.append('suggested-normal-user=' + normal_user)
    return caps


def hook_shell(dir, *extra_env):
    global ssh_port, normal_user

    if ssh_port:
        user = normal_user or '<user>'
        ssh = '    ssh -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no -p %i %s@localhost\n' % (
            ssh_port, user)
    else:
        ssh = ''

    with open('/dev/tty', 'w') as f:
        f.write('''You can now log into the VM through the serial terminal.
Depending on which terminal program you have installed, you can use one of

%(ssh)s    minicom -D unix#%(tty0)s
    nc -U %(tty0)s
    socat - UNIX-CONNECT:%(tty0)s

The tested source package is in %(dir)s

Press Enter to resume adt-run.
''' % {'tty0': os.path.join(workdir, 'ttyS0'), 'dir': dir, 'ssh': ssh})
    with open('/dev/tty', 'r') as f:
        f.readline()


parse_args()
VirtSubproc.main()
