#!/usr/bin/python

# lxc_bootstrap
# Copyright (C) 2017  Thomas Ward <teward@ubuntu.com>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero 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 Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

# LXC Bootstrapper, around the lxc-create 'Download' template for userspace
# containers; creates then modifies the container based on specifications.
#
# Designed for Ubuntu / Debian systems.

# import os
import sys
import argparse
import crypt
import subprocess as sp
import random
import platform
from string import ascii_letters, digits

SHADOW_SALT_CHARSET = ascii_letters + digits

ARCHITECTURE_MAP = {
    'amd64': 'amd64',  # Explicit definition
    'i386': 'i386',    # Explicit definition
    'armhf': 'armhf',  # Explicit definition
    'arm64': 'arm64',  # Explicit definition
    'x86_64': 'amd64',
    'x86': 'i386',
    'armv7l': 'armhf',
    'armv8l': 'arm64'
}


# noinspection PyClassHasNoInit
class Bootstrap:
    # noinspection PyClassHasNoInit
    class Container:
        bootstrap_existing = False
        name = None
        architecture = None
        distribution = None
        release = None

    # noinspection PyClassHasNoInit
    class Customization:
        # noinspection PyClassHasNoInit
        class Packages:
            to_add = ['openssh-server', 'software-properties-common', 'haveged', 'python', 'python-dev',
                      'python3', 'python3-dev', 'perl-modules', 'ubuntu-server', 'iptables', 'libnetfilter-conntrack3']
            to_remove = ['lxd', 'lxd-client', 'lxd-tools', 'lxc']
            autoremove_after = False

        # noinspection PyClassHasNoInit
        class Users:
            userdata = [('teward', '47AlphaTango@2012', None, True)]


def _parse_arguments():
    argparser = argparse.ArgumentParser(description="LXC Container Bootstrapper Assistant", add_help=True)
    argparser.add_argument('-e', '--existing', '--bootstrap-existing', dest="container_bootstrap_existing",
                           default=False, required=False, action='store_true',
                           help="Don't create a container, run bootstrapping on "
                                "an already-existing container.")
    argparser.add_argument('-n', '--name', type=str, dest="container_name", required=True,
                           help="The name to assign to the LXC container.")
    argparser.add_argument('-a', '--arch', type=str, dest="container_arch", default=None, required=False,
                           help="The architecture for the container")
    argparser.add_argument('-d', '--dist', '--distro', type=str, dest="container_dist", default=None, required=False,
                           help="The distribution for the container")
    argparser.add_argument('-r', '--release', '--codename', type=str, dest="container_release", default=None,
                           required=False, help="The specific release of the container")
    argparser.add_argument('--add-packages', type=str, dest="packages_add", required=False,
                           default='openssh-server,software-properties-common,haveged,python,python-dev,'
                                   'python3,python3-dev,perl-modules,ubuntu-server,iptables',
                           help="Comma-separated list of packages to add to the container.")
    argparser.add_argument('--exclude-packages', type=str, dest="packages_exclude", required=False,
                           default='lxd,lxd-client,lxd-tools,lxc',
                           help="Comma-separated list of packages to exclude from the container.")
    argparser.add_argument('--users', '--userdata', '--userfile', type=str, dest="user_datafile", required=False,
                           default=None,
                           help="Path to a file containing user data, one user per line in USERNAME:PASSWORD:SALT:ADMIN"
                                " format, where SALT is an optional 8-character alpha numeric string, and ADMIN is "
                                "'True' or 'False'")

    return argparser.parse_args()


# noinspection PyShadowingNames
def _get_bootstrap(args):
    if type(args) is not argparse.Namespace:
        raise TypeError("Invalid arguments provided.")
    tmpbootstrap = Bootstrap()

    tmpbootstrap.Container.name = str(args.container_name).strip(''')

    tmpbootstrap.Container.bootstrap_existing = args.container_bootstrap_existing

    if not tmpbootstrap.Container.bootstrap_existing:
        if 'Windows' in platform.platform():
            raise OSError("LXC doesn't work on Windows, so we can't use this script. Sorry!")
        elif 'Linux' not in platform.platform():
            raise OSError("This script only works for Linux OSes, sorry!")

        if args.container_arch:
            tmpbootstrap.Container.architecture = args.container_arch
        else:
            tmpbootstrap.Container.architecture = ARCHITECTURE_MAP[platform.machine()]

        if args.container_dist:
            tmpbootstrap.Container.distribution = args.container_dist
        else:
            lsb_dist = sp.Popen('/usr/bin/lsb_release -s -i'.split(), stdout=sp.PIPE, stderr=sp.PIPE).communicate()
            if lsb_dist[1]:
                raise SystemError("Error getting release distributor ID.")
            else:
                tmpbootstrap.Container.distribution = str(lsb_dist[0]).lower().strip('rn')

        if args.container_release:
            tmpbootstrap.Container.release = args.container_release
        else:
            try:
                lsb_codename = sp.Popen('/usr/bin/lsb_release -s -c'.split(),
                                        stdout=sp.PIPE, stderr=sp.PIPE).communicate()
                if lsb_codename[1]:
                    raise SystemError("Error getting release codename from lsb_release")
                tmpbootstrap.Container.release = str(lsb_codename[0]).lower().strip('rn')
            except Exception as e:
                raise e

    if args.packages_add:
        addpackages = args.packages_add.split(',')
        for pkg in addpackages:
            if pkg.lower() not in tmpbootstrap.Customization.Packages.to_add:
                tmpbootstrap.Customization.Packages.to_add.append(pkg.lower())

    if args.packages_exclude:
        delpackages = args.packages_exclude.split(',')
        for pkg in delpackages:
            if pkg.lower() in tmpbootstrap.Customization.Packages.to_add:
                tmpbootstrap.Customization.Packages.to_add.remove(pkg.lower())
            else:
                if pkg.lower() not in tmpbootstrap.Customization.Packages.to_remove:
                    tmpbootstrap.Customization.Packages.to_remove.append(pkg.lower())

    if args.user_datafile:
        datafile = open(args.user_datafile, 'r')
        for data in datafile:
            data = data.split(':')
            if len(data) != 4:
                raise IOError("Invalid data provided from datafile.")
            else:
                datatuple = (data[0], data[1], data[2], bool(data[3]))
                # noinspection PyTypeChecker
                tmpbootstrap.Customization.Users.userdata.append(datatuple)

    return tmpbootstrap


def _create_container(data):
    lxc_create_cmd = str("/usr/bin/lxc-create -t download -n %s -- -d %s -r %s -a %s" % (
        data.name, data.distribution, data.release, data.architecture)).split()

    try:
        # noinspection PyUnusedLocal
        ecode = sp.check_call(lxc_create_cmd, stdout=sys.stdout, stderr=sys.stderr)
    except sp.CalledProcessError:
        print "Something went wrong when creating the container."
        exit()


def _start_container(container_name):
    execstr = str('lxc-start -n %s' % container_name).split()
    try:
        # noinspection PyUnusedLocal
        ecode = sp.check_call(execstr, stdout=sys.stdout, stderr=sys.stderr)
    except sp.CalledProcessError:
        print "Could not start the container, cannot continue with bootstrap."
        exit()


def _shadow_pw_str(password, salt=None):
    if not salt or len(salt) != 8:
        # Create a random salt
        for _ in range(8):
            try:
                salt += random.SystemRandom().choice(SHADOW_SALT_CHARSET)
            except (TypeError, UnboundLocalError):
                salt = random.SystemRandom().choice(SHADOW_SALT_CHARSET)

    hashed = crypt.crypt(password, ('$6$%s$' % salt))
    return hashed


# noinspection PyUnreachableCode,PyUnusedLocal
def _container_bootstrap_users(container_name, data):
    # Default comes with 'ubuntu' user and group; let's nuke them
    del_default_usr = "/usr/bin/lxc-attach -n %s -- deluser --remove-all-files ubuntu" % container_name
    del_default_grp = "/usr/bin/lxc-attach -n %s -- delgroup --only-if-empty ubuntu" % container_name

    try:
        deldefusr = sp.check_call(del_default_usr.split(), stdout=sys.stdout, stderr=sys.stderr)
    except sp.CalledProcessError:
        print "Could not delete default user and/or group, please refer to error logs."

    try:
        deldefgrp = sp.check_call(del_default_grp.split(), stdout=sys.stdout, stderr=sys.stderr)
    except sp.CalledProcessError as e:
        if e.returncode == 5:
            print "Could not delete default group, it's not empty!"
        elif e.returncode == 3:
            pass  # return code of 3 means that the group doesn't exist, so we can move on.

    # Since we just nuked the '1000' group and user (default Ubuntu), let's start there now.
    uid = 1000
    gid = 1000

    for (username, password, salt, admin) in data.userdata:
        hashpw = _shadow_pw_str(password, salt)

        groupcreatecmd = "lxc-attach -n %s -- groupadd -g %s %s" % (container_name, gid, username)

        usrcreatecmd = ("lxc-attach -n %s -- useradd --create-home -u %s -g %s -p '%s' --shell=/bin/bash %s"
                        % (container_name, uid, gid, hashpw, username))

        try:
            groupcmd = sp.check_call(groupcreatecmd.split(), stdout=sys.stdout, stderr=sys.stderr)
            usercmd = sp.check_call(usrcreatecmd.split(), stdout=sys.stdout, stderr=sys.stderr)
            if admin:
                usermodcmd = "lxc-attach -n %s -- usermod -a -G sudo %s" % (container_name, username)
                usermod = sp.check_call(usermodcmd.split(), stdout=sys.stdout, stderr=sys.stderr)

        except sp.CalledProcessError:
            print "Something went wrong when bootstrapping user '%s'..." % username

        uid += 1
        gid += 1


# noinspection PyUnusedLocal
def _container_bootstrap_packages(container_name, data, autoremove=False):
    # Construct string
    basecommand = '/usr/bin/lxc-attach -n %s -- apt-get ' % container_name
    installstr = basecommand + "install -y %s" % " ".join(data.to_add)
    removestr = basecommand + "remove -y --purge %s" % " ".join(data.to_remove)
    autoremove = basecommand + "autoremove -y --purge"

    try:
        ecode = sp.call(installstr.split(), stdout=sys.stdout, stderr=sys.stderr)
    except sp.CalledProcessError:
        print "Something went wrong installing additional packages."
        exit()

    try:
        ecode = sp.call(removestr.split(), stdout=sys.stdout, stderr=sys.stderr)
    except sp.CalledProcessError:
        print "Something went wrong removing specified packages."
        exit()

    if autoremove:
        try:
            ecode = sp.call(autoremove.split(), stdout=sys.stdout, stderr=sys.stderr)
        except sp.CalledProcessError:
            print "Something went wrong cleaning up after removal with 'autoremove'."


def run():
    # Get args first, we'll pass this to the Bootstrapper.
    # noinspection PyUnreachableCode
    args = _parse_arguments()

    # Take the arguments, throw it into _get_bootstrap, and get a Bootstrap object.
    bootstrap = _get_bootstrap(args)

    if not bootstrap.Container.bootstrap_existing:
        # Create the container
        _create_container(bootstrap.Container)
    else:
        # Do nothing, continue on with bootstrap process.
        pass

    _start_container(bootstrap.Container.name)

    _container_bootstrap_packages(bootstrap.Container.name, bootstrap.Customization.Packages)

    _container_bootstrap_users(bootstrap.Container.name, bootstrap.Customization.Users)


if __name__ == "__main__":
    run()