#!/usr/bin/python
#
# 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 atexit
import socket
import errno

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

import VirtSubproc

args = None
workdir = None
p_qemu = None
ssh_port = 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('--ram-size', type=int, default=1024,
                        help='VM RAM size in MiB (default: 1024)')
    parser.add_argument('-d', '--debug', action='store_true',
                        help='Enable debugging output')
    parser.add_argument('image', nargs='+',
                        help='disk image to add to the VM (in order)')

    args = parser.parse_args()

    VirtSubproc.debuglevel = args.debug


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())
        atexit.register(os.remove, overlay)
    else:
        overlay = os.path.join(workdir, 'overlay.img')
    VirtSubproc.debug('Creating temporary overlay image in %s' % overlay)
    subprocess.check_call(['qemu-img', 'create', '-q', '-f', 'qcow2', '-b',
                           os.path.abspath(args.image[0]), overlay])
    return overlay


def wait_boot():
    term = VirtSubproc.get_unix_socket(os.path.join(workdir, 'ttyS0'))
    VirtSubproc.expect(term, b' login: ', 120, 'login prompt on ttyS0')


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():
        VirtSubproc.debug('setup_shell(): there already is a shell on ttyS1')
        return
    else:
        VirtSubproc.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
        VirtSubproc.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')
    VirtSubproc.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 /autopkgtest
mount -t 9p -o trans=virtio,access=any autopkgtest /autopkgtest
chmod 777 /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)


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):
    fcntl.fcntl(fin, fcntl.F_SETFL,
                fcntl.fcntl(fin, fcntl.F_GETFL) | os.O_NONBLOCK)
    while running:
        try:
            block = fin.read()
        except IOError as e:
            if e.errno != errno.EAGAIN:
                raise
            block = None
        if not block:
            fout.flush()
            time.sleep(0.01)
            continue
        while True:
            try:
                fout.write(block)
                break
            except IOError 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')
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, open(fin, 'w')))
t_stdin.start()
t_stdout = threading.Thread(None, shovel, 'copyout', (open(fout), sys.stdout))
t_stdout.start()
t_stderr = threading.Thread(None, shovel, 'copyerr', (open(ferr), sys.stderr))
t_stderr.start()

# give the stdin reading some head start
time.sleep(0.1)

# run command through QEMU shell
s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
s.connect('%(tty)s')
cmd = '%%(c)s < %%(d)s/stdin >> %%(d)s/stdout 2>> %%(d)s/stderr; echo $?  > %%(d)s/exit\\n' %% \\
       {'d': job_guest, 'c': ' '.join(map(pipes.quote, sys.argv[1:]))}
s.send(cmd.encode())
s.close()

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

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.down = [auxverb]
    VirtSubproc.downkind = 'auxverb'

    # verify that we can connect
    status, out = VirtSubproc.execute_raw('down_check', None, 5,
                                          VirtSubproc.down + ['runlevel'],
                                          stdout=subprocess.PIPE)
    VirtSubproc.debug('runlevel: exit %i, out "%s"' % (status, out))
    if status == 0 and out.endswith('2\n'):
        VirtSubproc.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)'''

    for p in range(start, start + 50):
        VirtSubproc.debug('find_free_port: trying %i' % p)
        try:
            s = socket.create_connection(('127.0.0.1', p))
            # if that works, the port is taken
            s.close()
            continue
        except socket.error as e:
            if e.errno == errno.ECONNREFUSED:
                VirtSubproc.debug('find_free_port: %i is free' % p)
                return p
            else:
                pass

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


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),
            '-localtime',
            '-no-reboot',
            '-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,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))

    # find free port to forward VM port 22 (for SSH access)
    ssh_port = find_free_port(10022)
    if ssh_port:
        VirtSubproc.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:
        wait_boot()
        setup_shell()
        setup_shared(shareddir)
        make_auxverb(shareddir)
    except:
        # Clean up on failure
        hook_cleanup()
        raise


def hook_downtmp():
    downtmp = '/autopkgtest/tmp'
    VirtSubproc.execute('mkdir %s' % downtmp, downp=True)
    return downtmp


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

    shutil.rmtree(workdir)
    workdir = None


def hook_forked_inchild():
    pass


def hook_capabilities():
    caps = ['revert', 'revert-full-system', 'root-on-testbed',
            'isolation-machine',
            'downtmp-host=%s' % os.path.join(workdir, 'shared', 'tmp')]
    if args.user and args.user != 'root':
        caps.append('suggested-normal-user=' + args.user)
    return caps


def hook_shell(stdin, stdout, stderr, dir):
    global ssh_port

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

    with open(stdout, '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(stdin, 'r') as f:
        f.readline()


parse_args()
VirtSubproc.main()
