#!/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()

#!/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


class TransparentDict(dict):

    def __missing__(self, key):
        return key


ARCHITECTURE_MAP = TransparentDict({
    'x86_64': 'amd64',
    'x86': 'i386',
    'armv7l': 'armhf',
    'armv8l': 'arm64'
})


class User:

    def __init__(self, name, password, salt, admin):
        self.name = name
        self.passord = password
        self.salt = salt
        self.admin = admin

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

        return crypt.crypt(self.password, ('$6${}$'.format(self.salt)))


class Container:
    create_cmd = "/usr/bin/lxc-create -t download -n {0.name} -- -d {0.distribution} -r {0.release} -a {0.architecture}"
    start_cmd = 'lxc-start -n {0.name}'
    attach_cmd = "/usr/bin/lxc-attach -n {0.name} -- "
    del_user_cmd = attach_cmd + "deluser --remove-all-files ubuntu"
    del_group_cmd = attach_cmd + "delgroup --only-if-empty ubuntu"
    create_group_cmd = attach_cmd + "groupadd -g {gid} {user.name}"
    create_user_cmd = attach_cmd +
        "useradd --create-home -u {uid} -g {gid} -p {user.shadow_password} --shell=/bin/bash {user.name}"
    mod_user_cmd = attach_cmd + "usermod -a -G sudo {user.name}"
    apt_get_install_cmd = attach_cmd + "apt-get install -y {}"
    apt_get_remove_cmd = attach_cmd + "apt-get remove -y --purge {}"
    apt_get_autoremove_cmd = attach_cmd + "apt-get autoremove -y --purge"

    def __init__(self,
                 name=None,
                 architecture=None,
                 distribution=None,
                 release=None):
        self.name = name
        self.architecture = architecture
        self.distribution = distribution
        self.release = release

    def create(self):
        cmd = self.create_cmd.format(self).split()
        try:
            sp.check_call(cmd, stdout=sys.stdout, stderr=sys.stderr)
        except sp.CalledProcessError as e:
            print "Something went wrong when creating the container."
            print e
            exit()

    def start(self):
        cmd = self.start_cmd.format(self).split()
        try:
            sp.check_call(cmd, stdout=sys.stdout, stderr=sys.stderr)
        except sp.CalledProcessError as e:
            print "Could not start the container, cannot continue with bootstrap."
            print e
            exit()

    def bootstrap_users(self, users):
        # Default comes with 'ubuntu' user and group; let's nuke them
        try:
            sp.check_call(self.del_user_cmd.format(self).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:
            sp.check_call(self.del_group.format(self).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:
                # return code of 3 means that the group doesn't exist, so we can
                # move on.
                pass

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

        for user in users:

            groupcreatecmd = self.create_group_cmd.format(
                self, user=user, gid=gid, uid=uid).split()

            usrcreatecmd = self.create_user_cmd.format(
                self, user=user, gid=gid, uid=uid).split()

            try:
                sp.check_call(groupcreatecmd, stdout=sys.stdout,
                              stderr=sys.stderr)
                sp.check_call(usrcreatecmd, stdout=sys.stdout,
                              stderr=sys.stderr)
                if user.admin:
                    sp.check_call(self.mod_user_cmd.format(self, user=user).split(),
                                  stdout=sys.stdout, stderr=sys.stderr)
            except sp.CalledProcessError:
                print "Something went wrong when bootstrapping user '{0.name}'...".format(user)

            uid += 1
            gid += 1

    def bootstrap_packages(self, packages, autoremove=False):
        try:
            sp.call(self.apt_get_install_cmd.format(self, " ".join(packages.to_add)).split(),
                    stdout=sys.stdout, stderr=sys.stderr)
        except sp.CalledProcessError as e:
            print "Something went wrong installing additional packages."
            print e
            exit()

        try:
            sp.call(self.apt_get_remove_cmd.format(self, " ".join(data.to_remove))
                    .split(),
                    stdout=sys.stdout, stderr=sys.stderr)
        except sp.CalledProcessError as e:
            print "Something went wrong removing specified packages."
            print e
            exit()

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


class Packages:


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

    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'")

    args = argparser.parse_args()
    args.container_name = args.container_name.strip(''')

    if not args.container_bootstrap_existing:
        if not args.container_arch:
            args.container_arch = ARCHITECTURE_MAP[
                platform.machine()]

        if not args.container_dist:
            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:
                args.container_dist = lsb_dist[0].lower().strip('rn')

        if not args.container_release:
            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")
            args.container_release = lsb_codename[0].lower().strip('rn')

    args.to_add = {'openssh-server', 'software-properties-common', 'haveged', 'python', 'python-dev',
                   'python3', 'python3-dev', 'perl-modules', 'ubuntu-server', 'iptables', 'libnetfilter-conntrack3'}
    args.to_remove = {'lxd', 'lxd-client', 'lxd-tools', 'lxc'}
    args.autoremove_after = False
    if args.packages_add:
        args.to_add |= set(map(str.lower, args.packages_add.split(',')))

    if args.packages_exclude:
        delpackages = set(map(str.lower, args.packages_exclude.split(',')))
        args.to_add -= delpackages
        args.to_remove |= delpackages
    return args


def get_users(user_datafile):
    users = []
    if user_datafile:
        with open(user_datafile) as datafile:
            for row in datafile:
                data = row.split(':')
                if len(data) != 4:
                    raise IOError("Invalid data provided from datafile.")
                users.append(User(data[0], data[1], data[2], int(data[3])))
    return users


def run():
    args = _parse_arguments()
    container = Container(args.container_name,
                          args.container_arch,
                          args.container_dist,
                          args.container_release)

    users = [User('teward', '47AlphaTango@2012', None, True)]
    users += get_users(args.user_datafile)

    if not args.bootstrap_existing:
        container.create()

    container.start()
    container.bootstrap_users(users)
    container.bootstrap_packages(args.to_add, args.to_exclude)


if __name__ == "__main__":
    run()