#!/usr/bin/python
#
# adt-virt-lxc is part of autopkgtest
# autopkgtest is a tool for testing Debian binary packages
#
# autopkgtest is Copyright (C) 2006-2014 Canonical Ltd.
#
# adt-virt-lxc was derived from adt-virt-schroot and modified to suit LXC by
# Robie Basak <robie.basak@canonical.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 string
from optparse import OptionParser
import random
import subprocess
import time
import tempfile
import shutil

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

import VirtSubproc

capabilities = ['revert', 'revert-full-system', 'root-on-testbed',
                'downtmp-shared-fifo', 'isolation-container']

ephemeral = False
eatmydata = False
use_sudo = False
lxc_template = None
lxc_container_name = None
normal_user = None
shared_dir = None


def parse_args():
    global lxc_template, eatmydata, ephemeral, use_sudo

    usage = '%prog [<options>] <template>'
    parser = OptionParser(usage=usage)

    parser.add_option('-d', '--debug', action='store_true')
    parser.add_option('--eatmydata', action='store_true')
    parser.add_option('-e', '--ephemeral', action='store_true',
                      help='Use ephemeral overlays instead of cloning (much '
                      'faster, but might cause errors in some corner cases)')
    parser.add_option('-s', '--sudo', action='store_true',
                      help='Run lxc-* commands with sudo; use if you run '
                      'adt-run as normal user')

    (opts, args) = parser.parse_args()
    if len(args) != 1:
        parser.error('need exactly one arg, template lxc container name')

    lxc_template = args[0]
    eatmydata = opts.eatmydata
    ephemeral = opts.ephemeral
    use_sudo = opts.sudo

    VirtSubproc.debuglevel = opts.debug


def sudoify(command):
    '''Prepend sudo to command with the --sudo option'''

    if use_sudo:
        if isinstance(command, list):
            return ['sudo'] + command
        else:
            return 'sudo ' + command
    else:
        return command


def check_for_eatmydata(container, template_name):
    with open('/dev/null', 'wb') as null:
        result = subprocess.call(
            sudoify(['lxc-attach', '-n', container, '--', 'sh', '-c',
                     'command -v eatmydata']),
            close_fds=True, stdout=null, stderr=null)
    if result:
        raise RuntimeError(
            ('Container %s does not contain eatmydata, which is required '
             'for the --eatmydata option.' % repr(template_name))
        )


def get_available_lxc_container_name():
    '''Return an LXC container name that isn't already taken.

    There is a race condition between returning this name and creation of
    the container. Ideally lxc-start-ephemeral would generate a name in a
    race free way and return the name used in machine-readable way, so that
    this function would not be necessary. See LP: #1197754.
    '''
    while True:
        # generate random container name
        rnd = [random.choice(string.ascii_lowercase) for i in range(6)]
        candidate = 'adt-virt-lxc-' + ''.join(rnd)

        containers = subprocess.check_output(sudoify(['lxc-ls']),
                                             universal_newlines=True)
        if candidate not in containers:
            return candidate


def wait_booted(lxc_name):
    '''Wait until the container has sufficiently booted to interact with it

    Do this by checking that the runlevel is someting numeric, i. e. not
    "unknown" or "S".
    '''
    timeout = 60
    while timeout > 0:
        try:
            out = subprocess.check_output(
                sudoify(['lxc-attach', '-n', lxc_name, 'runlevel'])).strip()
            if out.split()[-1].isdigit():
                return

            VirtSubproc.debug('wait_booted: runlevel "%s", retrying...' % out)
        except subprocess.CalledProcessError:
            VirtSubproc.debug('wait_booted: lxc-attach failed, retrying...')
            pass

        timeout -= 1
        time.sleep(1)

    VirtSubproc.bomb('timed out waiting for container %s to start; '
                     'last runlevel "%s"' % (lxc_name, out))


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

    global capabilities, normal_user

    # get the first UID >= 500
    out = subprocess.check_output(
        sudoify(['lxc-attach', '-n', lxc_name,
                 '--', 'sh', '-c', "getent passwd | sort -t: -nk3 | "
                 "awk -F: '{if ($3 >= 500) { print $1; exit } }'"]),
        universal_newlines=True).strip()
    if out:
        normal_user = out
        capabilities.append('suggested-normal-user=' + normal_user)
        VirtSubproc.debug('determine_normal_user: got user "%s"' % normal_user)
    else:
        VirtSubproc.debug('determine_normal_user: no uid >= 500 available')


def hook_open():
    global lxc_container_name, shared_dir

    lxc_container_name = get_available_lxc_container_name()
    VirtSubproc.debug('using container name %s' % lxc_container_name)
    shared_dir = tempfile.mkdtemp(prefix='adt-virt-lxc.shared.')
    os.chmod(shared_dir, 0o755)
    if ephemeral:
        VirtSubproc.execute(sudoify('lxc-start-ephemeral -n %s -k -d --bdir %s -o')
                            % (lxc_container_name, shared_dir), [lxc_template],
                            downp=False, outp=True)
    else:
        VirtSubproc.execute(sudoify('lxc-clone -n %s -o' % lxc_container_name),
                            [lxc_template], downp=False, outp=True)
        subprocess.check_call(sudoify([
            'lxc-start', '-n', lxc_container_name, '-d',
            '-s', 'lxc.mount.entry=%s %s none bind,create=dir 0 0' % (shared_dir, shared_dir[1:])]))
    try:
        VirtSubproc.debug('waiting for lxc guest start')
        wait_booted(lxc_container_name)
        VirtSubproc.debug('lxc guest started')
        determine_normal_user(lxc_container_name)
        VirtSubproc.down = ['lxc-attach', '-n', lxc_container_name, '--']
        if use_sudo:
            VirtSubproc.down = ['sudo', '-E'] + VirtSubproc.down
        if eatmydata:
            check_for_eatmydata(lxc_container_name, lxc_template)
            VirtSubproc.down.extend(['eatmydata', '--'])
        VirtSubproc.downkind = 'auxverb'
    except:
        # Clean up on failure
        VirtSubproc.execute(sudoify('lxc-stop -n'), [lxc_container_name],
                            downp=False, outp=True)
        VirtSubproc.execute(sudoify('lxc-destroy -n'), [lxc_container_name],
                            downp=False, outp=True)
        raise


def hook_downtmp():
    global capabilities, shared_dir

    if shared_dir:
        d = os.path.join(shared_dir, 'downtmp')
        # these permissions are ugly, but otherwise we can't clean up files
        # written by the testbed when running as user
        VirtSubproc.execute('mkdir -m 777 %s' % d, downp=True)
        capabilities.append('downtmp-host=' + d)
    else:
        d = VirtSubproc.downtmp_mktemp()
    return d


def hook_revert():
    hook_cleanup()
    hook_open()


def hook_cleanup():
    global capabilities, shared_dir

    VirtSubproc.downtmp_remove()
    capabilities = [c for c in capabilities if not c.startswith('downtmp-host')]
    if shared_dir:
        shutil.rmtree(shared_dir)
        shared_dir = None

    VirtSubproc.execute(sudoify('lxc-stop -n'), [lxc_container_name],
                        downp=False, outp=True)
    VirtSubproc.execute(sudoify('lxc-destroy -n'), [lxc_container_name],
                        downp=False, outp=True)


def hook_forked_inchild():
    pass


def hook_capabilities():
    return capabilities


def hook_shell(stdin, stdout, stderr, dir):
    argv = ['lxc-attach', '-n', lxc_container_name, '--', 'sh', '-c', 'cd %s; exec bash -i' % dir]
    if use_sudo:
        argv = ['sudo', '-E'] + argv
    with open(stdin, 'rb') as sin:
        with open(stdout, 'wb') as sout:
            with open(stderr, 'wb') as serr:
                subprocess.call(argv, stdin=sin, stdout=sout, stderr=serr)


parse_args()
VirtSubproc.main()
