#!/usr/bin/python
# -*- coding: utf-8 -*-
#
# LXC container setup script for Debian VMs
#
# Author: Jan Dittberner <jandd@cacert.org>
#
# Minimum required Python version: 2.6
#
# inspired by lxc-debian by Daniel Lezcano <daniel.lezcano@free.fr>
#
# 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 3 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, see
# <http://www.gnu.org/licenses/>.
"""
Configure a LXC container.
"""

from ConfigParser import RawConfigParser
from hashlib import sha1
from subprocess import Popen, PIPE, STDOUT
import StringIO
import apt
import argparse
import ipcalc
import logging
import logging.config
import os
import os.path
import select
import shutil
import sys
import time


config = RawConfigParser()
config.read('lxc-setup.ini')

# global configuration
SUBNET = config.getint('network', 'subnet')
GATEWAY = config.get('network', 'gateway')
IP6_NETSIZE = config.getint('network', 'ip6_netsize')
IP6_GATEWAY = config.get('network', 'ip6_gateway')
NETDEVICE = config.get('network', 'device')
DOMAIN = config.get('network', 'domain')
SMTPRELAY = config.get('network', 'smtprelay')

DEBRELEASE = config.get('debian', 'release')
DEBMIRROR = config.get('debian', 'mirror')
PACKAGES = [
        name.strip() for name in
        config.get('debian', 'packages').split(',')]
EXCLUDED_PACKAGES = [
        name.strip() for name in
        config.get('debian', 'excluded').split(',')]

CACHE = config.get('lxc', 'cachedir')

VOLGROUP = config.get('host', 'volumegroup')

DEBCONF_SEED = open('templates/debconf_seed.txt').read()
TMPL_INITTAB = open('templates/inittab.txt').read()
TMPL_NETWORK_INTERFACES = open('templates/network_interfaces.txt').read()
TMPL_LXC_CONFIG = open('templates/lxc_container.txt').read()
TMPL_FERM_CONF = open('templates/ferm_conf.txt').read()

logging.config.fileConfig('cacerttools.ini')
logger = logging.getLogger('lxcsetup')


class ProcessError(Exception):
    def __init__(self, exitcode, cmdline, message):
        self.exitcode = exitcode
        self.cmdline = cmdline
        self.message = message

    def __repr__(self):
        return "ProcessError<exitcode=%d, cmdline='%s', message='%s'>" % (
            self.exitcode, self.cmdline, self.message)


def run_process(args, stdin=None, message=None):
    inputbuf = None
    if (stdin is not None):
        inputbuf = StringIO.StringIO(stdin)
    logger.debug("starting process %s", " ".join(args))
    try:
        proc = Popen(args, stdout=PIPE, stderr=PIPE, stdin=PIPE)
        readables = [proc.stdout, proc.stderr]
        writeables = [proc.stdin]
        running = True
        while running:
            time.sleep(0.1)
            if proc.poll() is not None:
                running = False
            try:
                reads, writes, excs = select.select(readables,
                                                    writeables, [])
            except select.error, e:
                proc.wait()
                break
            for readable in reads:
                output = readable.readline().strip().decode('UTF-8')
                if output:
                    logger.info(output)
            for writeable in writes:
                if inputbuf and not inputbuf.closed:
                    lines = inputbuf.read()
                    logger.debug("send %s to process", lines)
                    writeable.write(lines)
                    writeable.close()
                    writeables.remove(writeable)
                    inputbuf.close()

        if proc.returncode != 0:
            raise ProcessError(proc.returncode, " ".join(args), message)
        return proc.returncode
    except KeyboardInterrupt, e:
        logger.warning("Interrupted")
        return 1


def download_debian(cache, arch):
    cachedir = os.path.join(cache, 'partial-%s' % arch)
    # check the mini debian was not already downloaded
    if not os.path.isdir(cachedir):
        os.makedirs(cachedir)

    # download a mini debian into a cache
    logger.info("Downloading debian minimal ...")
    run_process(['debootstrap', '--foreign', '--verbose', '--variant=minbase',
                 '--arch=%s' % arch, '--include=%s' % ','.join(PACKAGES),
                 '--exclude=%s' % ','.join(EXCLUDED_PACKAGES),
                 DEBRELEASE, cachedir, DEBMIRROR],
                message="Failed to download the rootfs")

    shutil.move(cachedir, os.path.join(cache, 'rootfs-%s' % arch))


class AlreadySetup(Exception):
    def __init__(self, rootfs):
        self.message = u"Container already setup with root fs %s" % rootfs

    def __str__(self):
        return self.message


class HostSetup(object):
    """
    Represents the setup of a Debian LXC container.
    """
    arch = None

    def __init__(
            self, name, ipaddress, ip6address, rootpassword, lvsize,
            adminaddress, extip, norootfs):
        self.name = name
        self.adminaddress = adminaddress
        self.ipaddress = ipcalc.Network(ipaddress, SUBNET)
        self.ip6address = ipcalc.Network(ip6address, IP6_NETSIZE)
        self.extip = extip
        self.rootpassword = rootpassword

        self.hostname = "%s.%s" % (self.name, DOMAIN)
        self.path = '/var/lib/lxc/vm-%s' % self.name
        self.rootfs = os.path.join(self.path, 'rootfs')
        self.norootfs = norootfs
        if not os.path.exists(self.path):
            os.makedirs(self.path)
        self.volumedev = '/dev/mapper/%s-vm--%s' % (VOLGROUP, self.name.replace('-', '--'))
        if (lvsize != '0') and not os.path.exists(self.volumedev):
            run_process(['lvcreate', '-L', lvsize, '-v',
                         '-n', 'vm-%s' % self.name, VOLGROUP],
                        message="Failed to create logical volume")
            run_process(['mkfs.ext4', self.volumedev],
                        message="Failed to create ext4 file system")
            if not [line for line in open('/etc/fstab') if line.startswith(
                    "%s " % self.volumedev)]:
                logger.debug("adding %s to container /etc/fstab",
                             self.volumedev)
                with open('/etc/fstab', 'a') as fstab:
                    print >> fstab, self.volumedev, self.path, \
                        'ext4', 'noatime', 0, 0
            else:
                logger.debug("%s already in container /etc/fstab",
                             self.volumedev)
            run_process(['mount', self.volumedev],
                        message="Failed to mount logical volume")
        self.determine_architecture()

        self.install()
        if not self.norootfs:
            self.configure_container()
        self.create_container_configuration()
        self.create_ferm_stub()

    def determine_architecture(self):
        self.arch = Popen(['arch'], stdout=PIPE).communicate()[0].strip()
        if self.arch == 'x86_64':
            self.arch = 'amd64'
        elif self.arch == 'i686':
            self.arch = 'i386'
        logger.debug("determined architecture as %s", self.arch)

    def copy_debian(self):
        """Make a local copy of the minidebian."""
        logger.info("Copying rootfs to %s...", self.rootfs)
        run_process(['cp', '-a', os.path.join(CACHE, 'rootfs-%s' % self.arch),
                     self.rootfs], message="failed to copy rootfs")

    def install(self):
        if os.path.exists(self.rootfs):
            raise AlreadySetup(self.rootfs)
        if not os.path.isdir('/var/lock/subsys'):
            os.makedirs('/var/lock/subsys')
        lxclock = os.open('/var/lock/subsys/lxc',
                          os.O_WRONLY | os.O_EXCL | os.O_CREAT)
        try:
            logger.info("Checking cache download in %s/rootfs-%s ... ",
                        CACHE, self.arch)
            if self.norootfs:
                logger.info("Skipping root filesystem setup")
                os.makedirs(self.rootfs)
                return
            if not os.path.exists(
                os.path.join(CACHE, 'rootfs-%s' % self.arch)):
                download_debian(CACHE, self.arch)
            self.copy_debian()
            logger.debug("Finished Debian base installation.")
        finally:
            os.unlink('/var/lock/subsys/lxc')

    def configure_container(self):
        """Configure the Debian LXC container."""
        pid = os.fork()
        if pid != 0:
            (wpid, status) = os.waitpid(pid, 0)
            return status
        os.chroot(self.rootfs)
        os.chdir("/")
        # configure inittab
        with open('/etc/inittab', 'w') as f:
            f.write(TMPL_INITTAB)
        # make /dev/pts
        if not os.path.isdir('/dev/pts'):
            os.makedirs('/dev/pts')
        # disable selinux
        if not os.path.isdir('/selinux'):
            os.makedirs('/selinux')
        with open('/selinux/enforce', 'w') as f:
            print >> f, "0"
        # set the hostname
        with open('/etc/hostname', 'w') as f:
            f.write(self.hostname)
        try:
            if ('LANG' in os.environ):
                del os.environ['LANG']
            run_process(['debconf-set-selections'], DEBCONF_SEED % {
                    'name': self.name,
                    'hostname': self.hostname,
                    'adminaddress': self.adminaddress,
                    'smtprelay': SMTPRELAY})
            run_process(['/debootstrap/debootstrap', '--second-stage'])
            # remove pointless services in a container
            for service in ['checkroot.sh', 'umountfs', 'hwclock.sh', 'hwclockfirst.sh', 'umountroot']:
                run_process(['/usr/sbin/update-rc.d', '-f',
                             service, 'remove'])
            if self.rootpassword:
                run_process(['chpasswd'], 'root:%s\n' % self.rootpassword)
        except ProcessError, e:
            logger.error("%r", e)
            sys.exit(1)
        with open('/etc/apt/sources.list', 'w') as f:
            print >> f, "deb", DEBMIRROR, DEBRELEASE, 'main'
            print >> f, "deb", 'http://security.debian.org/debian-security', \
                "%s/updates" % DEBRELEASE, 'main'
        logger.debug("Finished container configuration")
        sys.exit(0)

    def generate_mac_address(self):
        sha = sha1()
        sha.update(DOMAIN)
        sha.update(str(self.ipaddress))
        ipencoded = sha.hexdigest()[-8:].upper()
        machine_part = ipencoded[:2]
        for i in range(2, len(ipencoded), 2):
            machine_part += ":" + ipencoded[i:i + 2]
        return "00:FF:" + machine_part

    def create_container_configuration(self):
        if os.path.exists(os.path.join(self.path, 'config')):
            raise AlreadyConfigured(self.path)
        with open(os.path.join(self.path, 'config'), 'w') as f:
            f.write(TMPL_LXC_CONFIG % {
                    'hostname': self.name,
                    'networkdev': NETDEVICE,
                    'hwaddr': self.generate_mac_address(),
                    'ipv4network': '%s/%d' % self.ipaddress.to_tuple(),
                    'ipv4gateway': GATEWAY,
                    'ipv6network': '%s/%d' % self.ip6address.to_tuple(),
                    'ipv6gateway': IP6_GATEWAY,
                    'rootfs': self.rootfs})
        os.symlink(os.path.join(self.path, 'config'),
                   '/etc/lxc/auto/vm-%s' % self.name)
        logger.debug("Created container configuration file")

    def create_ferm_stub(self):
        if not self.extip:
            logger.info(u"no external ip defined, skipping ferm setup.")
            return
        with open(os.path.join('/etc/ferm/ferm.d',
                               '%s.conf' % self.name), 'w') as f:
            f.write(TMPL_FERM_CONF % {
                    'name': self.name,
                    'extip': self.extip,
                    'intip': self.ipaddress})
        logger.debug(u"Create container ferm stub.")


def clean_cache():
    if not os.path.exists(CACHE):
        logger.info("cache %s does not exist", CACHE)
        sys.exit(0)

    lxclock = os.open('/var/lock/subsys/lxc',
                      os.O_WRONLY | os.O_EXCL | os.O_CREAT)
    try:
        logger.info("Purging the download cache ... ")
        run_process(['rm', '--preserve-root', '--one-file-system', '-rf',
                     CACHE], message="Error purging the download cache")
    finally:
        os.unlink('/var/lock/subsys/lxc')


class NameArgumentAction(argparse.Action):
    def __call__(self, parser, namespace, values, option_string=None):
        setattr(namespace, self.dest, values)
        section = u'host-%s' % values
        if config.has_section(section):
            for opt in ('adminaddress', 'ip', 'ip6', 'extip', 'rootpw', 'lvsize'):
                if config.has_option(section, opt):
                    setattr(namespace, opt, config.get(section, opt))
                else:
                    logger.debug(
                        u'option %s not found in [%s] in lxc-setup.ini',
                        opt, section)
        else:
            logger.info(u'section [%s] not found in lxc-setup.ini', section)


if __name__ == '__main__':
    aptcache = apt.Cache()
    for required in 'lxc', 'lvm2', 'debootstrap':
        if not aptcache[required].is_installed:
            logger.error('Required Debian package %s is not installed',
                         required)
            sys.exit(1)

    if not [line for line in open('/etc/mtab') if line.startswith('cgroup')]:
        logger.error("no cgroup support detected in /etc/mtab")
        sys.exit(1)

    if os.getuid() != 0:
        logger.error("This script should be run as 'root'")
        sys.exit(1)

    # command line argument parser
    parser = argparse.ArgumentParser(
        description=u'''Build configuration and (optionally chroot)
for new LXC.''',
        epilog=u'''If a section [host-<name>] exists in lxc-setup.ini
it is used to provide initial parameters for the container.'''
        u' containers')
    parser.add_argument(
        '-n', '--name', metavar='name', required=True,
        action=NameArgumentAction,
        help=u"name of the LXC container")
    parser.add_argument(
        '-a', '--adminaddress', metavar='adminaddress',
        help=u"email address of the container's administrator, defaults to"
        u" <name>-admin@%s" % DOMAIN)
    parser.add_argument(
        '-i', '--ip', metavar='ip',
        help=u"internal IPv4 address to be assigned to the container")
    parser.add_argument(
        '-6', '--ip6', metavar='ip6',
        help=u"internal IPv6 address to be assigned to the container")
    parser.add_argument(
        '-e', '--extip', metavar='extip',
        help=u"externally visible IP address that should be forwarded from the"
        u"host to the container")
    parser.add_argument(
        '-r', '--rootpw', metavar='rootpw',
        help=u"password for the container's root account")
    parser.add_argument(
        '-c', '--clean', action='store_true',
        help=u"clean the debootstrap cache directory")
    parser.add_argument(
        '-l', '--lvsize', metavar='lvsize',
        help=u"size of the logical volume for the container")
    parser.add_argument(
        '--norootfs', action='store_true',
        help=u"don't create a root filesystem")
    args = parser.parse_args()

    for argname in ('ip', 'ip6', 'lvsize'):
        if not getattr(args, argname):
            logger.error(u"parameter '%s' is required as command line option"
                         u" or entry in the [host-%s] section in"
                         u" lxc-setup.ini.", argname, args.name)
            parser.print_usage(sys.stderr)
            sys.exit(1)
    if not getattr(args, 'adminaddress'):
        setattr(args, 'adminaddress', '%s-admin@%s' % (args.name, DOMAIN))

    try:
        setuphost = HostSetup(
                args.name, args.ip, args.ip6, args.rootpw, args.lvsize,
                args.adminaddress, args.extip, args.norootfs)
    except AlreadySetup, e:
        logger.error("%s", e)

    if args.clean:
        clean_cache()