#!/usr/bin/python3

description = """
Bootstrap a minimal Debian rootfs to be used as base for LXC containers
with the 'local' template.

Requires cdebootstrap.

Produces a rootfs directory with the specified name and a _meta directory with
LXC metadata.
"""

import argparse
from fnmatch import fnmatch
import os
import os.path
import subprocess
import sys

import os

def open_dirfd(path, mode = 'r', perms = 0o644, dir_fd = None, **kwargs):
    """
    Same as the built-in open(), but with support for file permissions
    and operation with respect to a directory FD.
    """
    flags = 0

    if '+' in mode:
        flags |= os.O_RDWR
    else:
        if 'r' in mode:
            flags |= os.O_RDONLY
        elif 'w' in mode or 'a' in mode:
            flags |= os.O_WRONLY

    if 'w' in mode:
        flags |= os.O_CREAT | os.O_TRUNC
    elif 'a' in mode:
        flags |= os.O_CREAT

    if 'x' in mode:
        flags |= os.O_EXCL

    fd = os.open(path, flags, mode = perms, dir_fd = dir_fd)
    try:
        return os.fdopen(fd, mode, **kwargs)
    except:
        os.close(fd)
        raise

# make sure these packages are present
include_pkgs = (
    # init
    'init',
    'sysvinit-core',

    # packaging
    'apt-utils',
    'aptitude',
    'whiptail',

    # networking
    'netbase',
    'iproute2',
    'dhcpcd',
    'ifupdown',
    'openssh-client',
    'openssh-server',

    # system infrastructure
    'cron',
    'procps',
    'ncurses-term',
    'locales',
    'tzdata',

    # misc utils
    'vim',
    'less',
    'htop',
    'bash-completion',
)

# packages we do NOT want, that would be installed by default
exclude_pkgs = (
    # systemd crap
    'systemd-sysv',
    'systemd',

    # networking: we use dhcpcd and do not need filtering
    'isc-dhcp-client',
    'isc-dhcp-common',
    'iptables',

    # stuff not useful inside a container
    'udev',
    'dmsetup',
    'dmidecode',
    'kmod',

    # misc utils
    'nano',
    'tasksel',
)

# written as /etc/apt/sources.list
apt_sources_template = """
deb     http://deb.debian.org/debian            {release}               main contrib non-free
deb-src http://deb.debian.org/debian            {release}               main contrib non-free

deb     http://deb.debian.org/debian            {release}-backports     main contrib non-free
deb-src http://deb.debian.org/debian            {release}-backports     main contrib non-free

deb     http://deb.debian.org/debian-security   {release}-security      main contrib non-free
deb-src http://deb.debian.org/debian-security   {release}-security      main contrib non-free
"""

# written as /etc/hosts
hosts_template = """
::1     LXC_NAME LXC_NAME.{domain}
::1     localhost ip6-localhost ip6-loopback

fe00::0 ip6-localnet
ff00::0 ip6-mcastprefix
ff02::1 ip6-allnodes
ff02::2 ip6-allrouters
"""

# disable initscripts that do not make sense for a container
initscripts_disable = [
    'dhcpcd',
]
    #'mountnfs-bootclean.sh', # XXX
    #'mountnfs.sh',
    #'mountall-bootclean.sh',
    #'mountall.sh',
    #'checkroot-bootclean.sh',
    #'checkfs.sh',
    #'checkroot.sh',
    #'hwclock.sh',
    #'mountdevsubfs.sh',

lxc_templates = """
etc/hosts
etc/hostname
"""

class Bootstrapper:

    _args      = None
    _verbose   = None

    _meta_path = None
    _root_path = None
    _root_fd   = None
    _meta_fd   = None
    _etc_fd    = None

    def __init__(self, args):
        self._args = args

        self._verbose = args.verbose

        # resolve symlinks and normalize the path
        out_path = os.path.normpath(os.path.realpath(args.rootfs_dir))

        # strip possible trailing slash
        if len(out_path) > 1 and out_path[-1] == '/':
            out_path = out_path[:-1]

        # sanitize output directory
        if len(out_path) == 0 or out_path == '/':
            raise ValueError('Invalid output directory: %s' % out_path)

        self._root_path = out_path
        self._meta_path = out_path + '_meta'

        # raise errors if the output directory exists
        os.makedirs(self._root_path)
        os.makedirs(self._meta_path)

        self._root_fd = os.open(self._root_path, os.O_DIRECTORY)
        self._meta_fd = os.open(self._meta_path, os.O_DIRECTORY)

    def close(self):
        for fd_attr in ('root', 'meta', 'etc'):
            name =  '_%s_fd' % fd_attr
            val = getattr(self, name)
            if val is not None:
                os.close(val)
                setattr(self, name, None)

    def _debootstrap(self, release, dst_dir):
        cmdline = ['cdebootstrap', '-f', 'minimal',
                   '--include=%s' % ','.join(include_pkgs), '--exclude=%s' % ','.join(exclude_pkgs)]
        if self._verbose > 1:
            cmdline += ['--verbose']
        elif self._verbose == 0:
            cmdline += ['--quiet']
        cmdline += [release, dst_dir]

        if self._verbose > 1:
            sys.stderr.write('Executing: ' + ' '.join(cmdline) + '\n')
        subprocess.run(cmdline, check = True)

    def _clean_dev(self):
        dev_dir = os.open('dev', os.O_DIRECTORY, dir_fd = self._root_fd)
        try:
            for node in os.listdir(dev_dir):
                if self._verbose > 1:
                    sys.stderr.write('Removing /dev node: %s\n' % node)

                try:
                    os.remove(node, dir_fd = dev_dir)
                except IsADirectoryError:
                    os.rmdir(node, dir_fd = dev_dir)
        finally:
            os.close(dev_dir)

    def _config_locales(self, locales):
        locales = locales.replace(',', '\n')

        with open_dirfd('locale.gen', 'w', dir_fd = self._etc_fd) as f:
            f.write(locales)

        self._exec_chroot(['dpkg-reconfigure', '-fnoninteractive', 'locales'])

    def _config_tz(self, tz):
        tz_path = '/usr/share/zoneinfo/' + tz
        if not os.path.exists(self._root_path + tz_path):
            raise ValueError('Timezone not known: %s' % tz)

        os.remove('localtime', dir_fd = self._etc_fd)
        os.symlink(tz_path, 'etc/localtime', dir_fd = self._root_fd)

        self._exec_chroot(['dpkg-reconfigure', '-fnoninteractive', 'tzdata'])

    def _pkg_mark_auto(self):
        # mark all packages except explicitly included ones as auto-installed
        res = self._exec_chroot(['aptitude', 'search', '-F', '%p', '~i'],
                                capture_output = True, text = True)
        packages_all  = set(res.stdout.split())
        packages_auto = packages_all - set(include_pkgs)

        cmdline = ['apt-mark', 'auto'] + (self._verbose == 0) * ['-qq'] + list(packages_auto)

        self._exec_chroot(cmdline)

    def _clean_ssh_server(self):
        etc_ssh_fd = os.open('ssh', os.O_DIRECTORY, dir_fd = self._etc_fd)

        # clear SSH host keys so they are not shared by all containers created
        # from this image
        try:
            for item in os.listdir(etc_ssh_fd):
                if fnmatch(item, 'ssh_host_*_key*'):
                    if self._verbose > 1:
                        sys.stderr.write('Removing SSH host key: %s\n' % item)
                    os.remove(item, dir_fd = etc_ssh_fd)
        finally:
            os.close(etc_ssh_fd)

    def run(self):
        self._debootstrap(self._args.release, self._root_path)

        self._etc_fd = os.open('etc', os.O_DIRECTORY, dir_fd = self._root_fd)

        self._clean_dev()

        # remove the cdebootstrap helper that blocks rc operations
        self._exec_chroot(['dpkg', '--remove', 'cdebootstrap-helper-rc.d'])

        # set up initscripts
        self._exec_chroot(['insserv', '-r'] + initscripts_disable)

        if self._args.locales:
            self._config_locales(self._args.locales)
        if self._args.timezone:
            self._config_tz(self._args.timezone)

        # write /etc/apt/sources.list
        with open_dirfd('apt/sources.list', 'w', dir_fd = self._etc_fd) as f:
            f.write(apt_sources_template.format(release = self._args.release))

        # write /etc/network/interfaces
        if self._args.interfaces:
            for iface in self._args.interfaces.split(','):
                with open_dirfd('network/interfaces.d/' + iface, 'w', dir_fd = self._etc_fd) as f:
                    f.write('auto %s\niface %s inet dhcp\n' % (iface, iface))

        # write /etc/hostname
        with open_dirfd('hostname', 'w', dir_fd = self._etc_fd) as f:
            f.write('LXC_NAME\n')

        # write /etc/hosts
        # cdebootstrap writes a minimal default, but we overwrite it
        with open_dirfd('hosts', 'w', dir_fd = self._etc_fd) as f:
            f.write(hosts_template.format(domain = self._args.domain))

        # cdebootstrap copies host's resolv.conf, get rid of it
        os.remove('resolv.conf', dir_fd = self._etc_fd)

        self._pkg_mark_auto()
        self._clean_ssh_server()

        # record extra files for which lxc-local performs template processing
        with open_dirfd('templates', 'w', dir_fd = self._meta_fd) as f:
            f.write(lxc_templates)

        with open_dirfd('config', 'w', dir_fd = self._meta_fd) as f:
            f.write("")

    def _exec_chroot(self, cmdline, **kwargs):
        if self._verbose > 1:
            sys.stderr.write('Executing in chroot: %s\n' % cmdline)

        # discard command's stdout/err on zero verbosity, as there's no general
        # way to make all of them quiet
        # failures should still be detected via return codes
        if (self._verbose == 0 and
            not any((it in kwargs for it in ('stdout', 'stderr', 'capture_output')))):
            kwargs['stderr'] = subprocess.DEVNULL
            kwargs['stdout'] = subprocess.DEVNULL

        return subprocess.run(['chroot', self._root_path] + cmdline, check = True, **kwargs)

parser = argparse.ArgumentParser(description = description)

parser.add_argument('release')
parser.add_argument('rootfs_dir')

parser.add_argument('-v', '--verbose', action = 'count', default = 0)

parser.add_argument('-l', '--locales',
    default = 'en_US.UTF-8 UTF-8,en_DK.UTF-8 UTF-8',
    help = 'comma-separated list of locales to generate (as in /etc/locale.gen)')
parser.add_argument('-t', '--timezone', default = 'Europe/Prague')
parser.add_argument('-d', '--domain', default = 'khirnov.net')
parser.add_argument('-i', '--interfaces', default = 'eth0',
    help = 'comma-separated list of network interfaces to set up')

args = parser.parse_args()

bootstrapper = Bootstrapper(args)

try:
    bootstrapper.run()
finally:
    bootstrapper.close()