#!/usr/pkg/bin/python3.11
#
# autopkgtest-virt-qemu is part of autopkgtest
# autopkgtest is a tool for testing Debian binary packages
#
# autopkgtest is Copyright (C) 2006-2014 Canonical Ltd.
#
# autopkgtest-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 time
import socket
import errno
import fcntl
import re
import argparse

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(
    os.path.abspath(__file__))), 'lib'))

from reprotest.lib import VirtSubproc
from reprotest.lib import adtlog


args = None
workdir = None
p_qemu = None
ssh_port = None
normal_user = None
qemu_cmd_default = None


def parse_args():
    global args, qemu_cmd_default

    uname_to_qemu_suffix = {'i[3456]86$': 'i386'}
    arch = os.uname()[4]
    for pattern, suffix in uname_to_qemu_suffix.items():
        if re.match(pattern, arch):
            qemu_cmd_default = 'qemu-system-' + suffix
            break
    else:
        qemu_cmd_default = 'qemu-system-' + arch

    parser = argparse.ArgumentParser()

    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('--timeout-reboot', type=int, metavar='SECONDS', default=60,
                        help='timeout for waiting for reboot (default: %(default)ss)')
    parser.add_argument('--show-boot', action='store_true',
                        help='Show boot messages from serial console')
    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: ', args.timeout_reboot, 'login prompt on ttyS0',
                       echo=args.show_boot)
    # 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') + b'\n')
    VirtSubproc.expect(term, b'assword:', 10, 'password prompt')
    # send password
    passwd_b = args.password.encode('UTF-8')
    term.send(passwd_b + b'\n')
    VirtSubproc.expect(term, None, 10, 'acked password')
    term.send(b'echo "LOG""IN""_"OK\n')
    adtlog.debug('login_tty: logged in')
    VirtSubproc.expect(term, b'LOGIN_OK', 120, 'logged in')

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

    # if we are a non-root user, run through sudo
    if args.user != 'root':
        cmd = b"echo '%s' | sudo --background --stdin sh -c '" % passwd_b + cmd + b"'"
    else:
        cmd = b'setsid ' + cmd

    term.send(cmd + b'\n')
    VirtSubproc.expect(term, None, 10, 'accepted ttyS1 shell command')
    term.send(b'exit\n')
    VirtSubproc.expect(term, b'\nlogout', 10)


def setup_baseimage():
    '''setup /dev/baseimage in VM'''

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

    # Setup udev rules for /dev/baseimage; set link_priority to -1024 so
    # that the duplicate UUIDs of the partitions will have no effect.
    term.send(b'''mkdir -p -m 0755 /run/udev/rules.d ; printf '# Created by autopkgtest-virt-qemu\\n%s\\n%s\\n%s\\n' 'KERNEL=="vd*[!0-9]", ENV{ID_SERIAL}=="BASEIMAGE", OPTIONS+="link_priority=-1024", SYMLINK+="baseimage", MODE="0664"' 'KERNEL=="vd*[0-9]",  ENV{ID_SERIAL}=="BASEIMAGE", OPTIONS+="link_priority=-1024"' 'KERNEL=="vd*", ENV{ID_SERIAL}=="BASEIMAGE", ENV{ID_FS_TYPE}:="", ENV{ID_FS_USAGE}:="", ENV{ID_FS_UUID}:=""' > /run/udev/rules.d/61-baseimage.rules\n''')
    VirtSubproc.expect(term, b'#', 10)
    # Reload udev to make sure the rules take effect (udev only auto-
    # rereads rules every 3 seconds)
    term.send(b'udevadm control --reload\n')
    VirtSubproc.expect(term, b'#', 10)

    # Add the base image as an additional drive
    monitor = VirtSubproc.get_unix_socket(os.path.join(workdir, 'monitor'))
    monitor.send(('drive_add 0 file=%s,if=none,readonly=on,serial=BASEIMAGE,id=drive-baseimage\n' % args.image[0]).encode())
    VirtSubproc.expect(monitor, b'(qemu)', 10)
    monitor.send(b'device_add virtio-blk-pci,drive=drive-baseimage,id=virtio-baseimage\n')
    VirtSubproc.expect(monitor, b'(qemu)', 10)

    term.send(b'udevadm settle --exit-if-exists=/dev/baseimage\n')
    VirtSubproc.expect(term, b'#', 10)


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

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

    term.send(b'''mkdir -p -m 1777 /run/autopkgtest/shared
mount -t 9p -o trans=virtio,access=any autopkgtest /run/autopkgtest/shared
chmod 1777 /run/autopkgtest/shared
touch /run/autopkgtest/shared/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)

    # 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 /run/autopkgtest/shared 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 > /tmp/eofcat; chmod 755 /tmp/eofcat
#!$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 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 line in f:
                if line.startswith(b'#'):
                    continue
                line = line.strip()
                if line:
                    tz = line
                    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'#', 30)
        else:
            adtlog.debug('Could not determine host timezone')

    # ensure that we have Python for our the auxverb helpers
    term.send(b'type python3 2>/dev/null || type python 2>/dev/null\n')
    try:
        out = VirtSubproc.expect(term, b'/python', 5)
    except VirtSubproc.Timeout:
        VirtSubproc.bomb('Neither python3 nor python is installed in the VM, '
                         'one of them is required by autopkgtest')
    if b'\n# ' not in out:
        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 = '/run/autopkgtest/shared/' + 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 = 'PYTHONHASHSEED=0 /tmp/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()
# code 255 means that the auxverb itself failed, so translate
sys.exit(rc == 255 and 253 or 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 = VirtSubproc.execute_timeout(None, 5, VirtSubproc.auxverb + ['true'])[0]
    if status == 0:
        adtlog.debug('can connect to autopkgtest sh in VM')
    else:
        VirtSubproc.bomb('failed to connect to VM')


def get_cpuflag():
    '''Return QEMU cpu option list suitable for host CPU'''

    try:
        with open('/proc/cpuinfo', 'r') as f:
            for line in f:
                if line.startswith('flags'):
                    words = line.split()
                    if 'vmx' in words:
                        adtlog.debug('Detected KVM capable Intel host CPU, enabling nested KVM')
                        return ['-cpu', 'kvm64,+vmx,+lahf_lm']
                    elif 'svm' in words:  # AMD kvm
                        adtlog.debug('Detected KVM capable AMD host CPU, enabling nested KVM')
                        # FIXME: this should really be the one below for more
                        # reproducible testbeds, but nothing except -cpu host works
                        # return ['-cpu', 'kvm64,+svm,+lahf_lm']
                        return ['-cpu', 'host']
    except IOError as e:
        adtlog.warning('Cannot read /proc/cpuinfo to detect CPU flags: %s' % e)
        # fetching CPU flags isn't critical (only used to enable nested KVM),
        # so don't fail here
        pass

    return []


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

    for p in range(start, start + 50):
        adtlog.debug('find_free_port: trying %i' % p)
        try:
            lockfile = '/tmp/autopkgtest-virt-qemu.port.%i' % p
            f = None
            try:
                f = open(lockfile, 'x')
                os.unlink(lockfile)
                fcntl.flock(f, fcntl.LOCK_EX | fcntl.LOCK_NB)
            except (IOError, OSError):
                adtlog.debug('find_free_port: %i is locked' % p)
                continue
            finally:
                if f:
                    f.close()

            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:
                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"> /run/autopkgtest/shared/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='autopkgtest-virt-qemu.')
    os.chmod(workdir, 0o755)

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

    overlay = prepare_overlay()

    # 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)
        nic_opt = ',hostfwd=tcp::%i-:22' % ssh_port
    else:
        nic_opt = ''

    # start QEMU
    argv = [args.qemu_command,
            '-m', str(args.ram_size),
            '-smp', str(args.cpus),
            '-nographic',
            '-net', 'nic,model=virtio',
            '-net', 'user' + nic_opt,
            '-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))

    if os.path.exists('/dev/kvm'):
        argv.append('-enable-kvm')
        # Enable nested KVM by default on x86_64
        if os.uname()[4] == 'x86_64' and args.qemu_command == qemu_cmd_default and \
                (not args.qemu_options or '-cpu' not in args.qemu_options.split()):
            argv += get_cpuflag()

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

    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_baseimage()
        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 = '/run/autopkgtest/shared/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_prepare_reboot():
    # Remove baseimage drive again, so that it does not break the subsequent
    # boot due to the duplicate UUID
    monitor = VirtSubproc.get_unix_socket(os.path.join(workdir, 'monitor'))
    monitor.send(b'device_del virtio-baseimage\n')
    VirtSubproc.expect(monitor, b'(qemu)', 10)


def hook_wait_reboot():
    global workdir
    shareddir = os.path.join(workdir, 'shared')
    os.unlink(os.path.join(shareddir, 'done_shared'))
    wait_boot()
    setup_shell()
    setup_shared(shareddir)
    setup_baseimage()


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 running tests.
''' % {'tty0': os.path.join(workdir, 'ttyS0'), 'dir': dir, 'ssh': ssh})
    with open('/dev/tty', 'r') as f:
        f.readline()


parse_args()
VirtSubproc.main()
