#!/usr/bin/python3
# -*- coding: utf-8 -*-

# Copyright (C) 2014-2017 Canonical Ltd.
# Author: Christopher Townsend <christopher.townsend@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; version 3 of the License.
#
# 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/>.

import argparse
import getpass
import json
import libertine.utils
import os
import sys
import re

import gettext
gettext.textdomain('libertine')
_ = gettext.gettext

from libertine import ContainerRunning, LibertineContainer
from libertine.ContainersConfig import ContainersConfig
from libertine.HostInfo import HostInfo


class LibertineContainerManager(object):

    def __init__(self):
        self.containers_config = ContainersConfig()
        self.host_info = HostInfo()


    def _container(self, container_id):
        try:
            return LibertineContainer(container_id, self.containers_config)
        except ImportError as e:
            container_type = self.containers_config.get_container_type(container_id)
            libertine.utils.get_logger().error(_("Backend for container '{id}' not installed. Install "
                                               "'python3-libertine-{type}' and try again.").format(id=container_id, type=container_type))
            sys.exit(1)

    def _get_updated_locale(self, container_id):
        host_locale = self.host_info.get_host_locale()

        if host_locale == self.containers_config.get_container_locale(container_id):
            return None
        else:
            return host_locale

    def create(self, args):
        password = None

        if args.distro and not self.host_info.is_distro_valid(args.distro, args.force):
            libertine.utils.get_logger().error("Invalid distro %s" % args.distro)
            sys.exit(1)

        if self.containers_config.container_exists(args.id):
            libertine.utils.get_logger().error("Container id '%s' is already used." % args.id)
            sys.exit(1)
        elif re.match("^[a-z0-9][a-z0-9+.-]+$", args.id) is None:
            libertine.utils.get_logger().error("Container id '%s' invalid. ID must be of form ([a-z0-9][a-z0-9+.-]+)." % args.id)
            sys.exit(1)

        if not args.type:
            container_type = self.host_info.select_container_type_by_kernel()
        else:
            if (args.type == 'lxc' and not self.host_info.has_lxc_support()) or \
               (args.type == 'lxd' and not self.host_info.has_lxd_support()):
                libertine.utils.get_logger().error("System kernel does not support %s type containers. "
                                                   "Please either use chroot or omit the -t option." % args.type)
                sys.exit(1)
            container_type = args.type

        if not args.distro:
            args.distro = self.host_info.get_host_distro_release()
        elif container_type == "chroot":
            host_distro = self.host_info.get_host_distro_release()

            if args.distro != host_distro:
                libertine.utils.get_logger().error("The container distribution needs to match the host ditribution for chroot"
                                                   " based containers. Please either use \'%s\' or omit the -d/--distro option."
                      % host_distro)
                sys.exit(1)

        if not args.name:
            args.name = "Ubuntu \'" + (self.host_info.get_distro_codename(args.distro) or args.distro) + "\'"

        if container_type == "lxc" or container_type == "lxd":
            if args.password:
                password = args.password
            elif sys.stdin.isatty():
                print("Enter password for your user in the Libertine container or leave blank for no password:")
                password = getpass.getpass()
            else:
                password = sys.stdin.readline().rstrip()

        self.containers_config.add_new_container(args.id, args.name, container_type, args.distro)

        multiarch = 'disabled'
        if args.multiarch == 'enable':
            multiarch = 'enabled'
        self.containers_config.update_container_multiarch_support(args.id, multiarch)

        try:
            self.containers_config.update_container_locale(args.id, self.host_info.get_host_locale())
            container = LibertineContainer(args.id, self.containers_config)
            try:
                self.containers_config.update_container_install_status(args.id, "installing")
                if not container.create_libertine_container(password, args.multiarch):
                    libertine.utils.get_logger().error("Failed to create container")
                    self.containers_config.delete_container(args.id)
                    sys.exit(1)
            except Exception as e:
                container.destroy_libertine_container(force=True)
                raise
        except Exception as e:
            libertine.utils.get_logger().error("Failed to create container: '{}'".format(str(e)))

            self.containers_config.delete_container(args.id)
            sys.exit(1)

        self.containers_config.update_container_install_status(args.id, "ready")

        libertine.utils.refresh_libertine_scope()

    def destroy_container(self, container, force):
        fallback = self.containers_config.get_container_install_status(container.container_id)

        self.containers_config.update_container_install_status(container.container_id, "removing")
        if not container.destroy_libertine_container(force):
            self.containers_config.update_container_install_status(container.container_id, fallback)
            return

        self.containers_config.update_container_install_status(container.container_id, "removed")
        self.containers_config.delete_container(container.container_id)

    def destroy(self, args):
        container_id = self.containers_config.check_container_id(args.id)
        container = self._container(container_id)

        self.destroy_container(container, args.force)

        libertine.utils.refresh_libertine_scope()

    def install_package(self, args):
        container_id = self.containers_config.check_container_id(args.id)
        container = self._container(container_id)
        failure = False

        with ContainerRunning(container.container):
            for i, pkg in enumerate(args.package):
                if not pkg:
                    continue

                is_debian_package = pkg.endswith('.deb')

                if is_debian_package:
                    if os.path.exists(pkg):
                        package = libertine.utils.get_deb_package_name(pkg)
                    else:
                        libertine.utils.get_logger().error("%s does not exist." % pkg)
                        failure = True
                        continue
                else:
                    package = pkg

                if self.containers_config.package_exists(container_id, package):
                    if not is_debian_package:
                        libertine.utils.get_logger().error("Package '%s' is already installed." % package)
                        failure = True
                        continue
                else:
                    self.containers_config.add_new_package(container_id, package)

                self.containers_config.update_package_install_status(container_id, package, "installing")
                if not container.install_package(pkg, args.no_dialog, update_cache=i==0):
                    libertine.utils.get_logger().error("Package '{}' failed to install in container '{}'"
                                                       .format(package, container_id))
                    self.containers_config.delete_package(container_id, package)
                    failure = True
                    continue

                self.containers_config.update_package_install_status(container_id, package, "installed")

        libertine.utils.refresh_libertine_scope()

        if failure:
            sys.exit(1)

    def remove_package_by_name(self, container, package_name, no_dialog=False):
        fallback_status = self.containers_config.get_package_install_status(container.container_id, package_name)
        self.containers_config.update_package_install_status(container.container_id, package_name, "removing")

        if not container.remove_package(package_name, no_dialog) and fallback_status == 'installed':
            self.containers_config.update_package_install_status(container.container_id, package_name, fallback_status)
            return False

        self.containers_config.update_package_install_status(container.container_id, package_name, "removed")
        self.containers_config.delete_package(container.container_id, package_name)

        return True

    def remove_package(self, args):
        container_id = self.containers_config.check_container_id(args.id)
        container = self._container(container_id)
        failure = False

        with ContainerRunning(container.container):
            for pkg in args.package:
                if not pkg:
                    continue

                if self.containers_config.get_package_install_status(container_id, pkg) != 'installed':
                    libertine.utils.get_logger().error("Package \'%s\' is not installed." % pkg)
                    failure = True
                    continue

                if not self.remove_package_by_name(container, pkg, args.no_dialog):
                    libertine.utils.get_logger().error("Package '{}' failed to be removed from container '{}'"
                                                       .format(pkg, container_id))
                    failure = True
                    continue

        libertine.utils.refresh_libertine_scope()

        if failure:
            sys.exit(1)

    def search_cache(self, args):
        container_id = self.containers_config.check_container_id(args.id)
        container = self._container(container_id)

        if container.search_package_cache(args.search_string) is not 0:
            libertine.utils.get_logger().error("Search for '{}' in container '{}' exited with non-zero status"
                                               .format(args.id, args.search_string))
            sys.exit(1)

    def update(self, args):
        container_id = self.containers_config.check_container_id(args.id)
        container = self._container(container_id)

        new_locale = self._get_updated_locale(container_id)

        if not container.update_libertine_container(new_locale):
            sys.exit(1)

        if new_locale:
            self.containers_config.update_container_locale(container_id, new_locale)

    def list(self, args):
        for container in ContainersConfig().get_containers():
            print("%s" % container)

    def list_apps(self, args):
        container_id = self.containers_config.check_container_id(args.id)

        app_ids = self._container(container_id).list_app_ids()
        if args.json:
            print(json.dumps(app_ids))
        else:
            for app in app_ids:
                print(app)

    def exec(self, args):
        container_id = self.containers_config.check_container_id(args.id)

        container = self._container(container_id)

        if not container.exec_command(args.command):
            sys.exit(1)

    def delete_archive_by_name(self, container, archive_name):
        if self.containers_config.get_archive_install_status(container.container_id, archive_name) == 'installed':
            self.containers_config.update_archive_install_status(container.container_id, archive_name, 'removing')
            if container.configure_remove_archive("\"" + archive_name + "\"") is not 0:
                self.containers_config.update_archive_install_status(container.container_id, archive_name, 'installed')
                return False

        self.containers_config.delete_container_archive(container.container_id, archive_name)
        return True

    def configure(self, args):
        container_id = self.containers_config.check_container_id(args.id)
        container = self._container(container_id)

        if args.multiarch and self.host_info.get_host_architecture() == 'amd64':
            multiarch = 'disabled'
            if args.multiarch == 'enable':
                multiarch = 'enabled'

            current_multiarch = self.containers_config.get_container_multiarch_support(container_id)
            if current_multiarch == multiarch:
                libertine.utils.get_logger().error("i386 multiarch support is already %s" % multiarch)
                sys.exit(1)

            if container.configure_multiarch(args.multiarch) is not 0:
                sys.exit(1)

            self.containers_config.update_container_multiarch_support(container_id, multiarch)

        elif args.archive is not None:
            if args.archive_name is None:
                libertine.utils.get_logger().error("Configure archive called with no archive name. See configure --help for usage.")
                sys.exit(1)

            archive_name = args.archive_name.strip("\'\"")
            archive_name_esc = "\"" + archive_name + "\""

            if args.archive == 'add':
                if self.containers_config.archive_exists(container_id, archive_name):
                    libertine.utils.get_logger().error("%s already added in container." % archive_name)
                    sys.exit(1)

                self.containers_config.add_container_archive(container_id, archive_name)
                self.containers_config.update_archive_install_status(container_id, archive_name, 'installing')
                if container.configure_add_archive(archive_name_esc, args.public_key_file) is not 0:
                    self.containers_config.delete_container_archive(container_id, archive_name)
                    sys.exit(1)

                self.containers_config.update_archive_install_status(container_id, archive_name, 'installed')

            elif args.archive == 'remove':
                if not self.containers_config.archive_exists(container_id, archive_name):
                    libertine.utils.get_logger().error("%s is not added in container." % archive_name)
                    sys.exit(1)

                if not self.delete_archive_by_name(container, archive_name):

                    libertine.utils.get_logger().error("%s was not properly deleted." % archive_name)
                    sys.exit(1)

        elif args.bind_mount is not None:
            if args.mount_path is None:
                libertine.utils.get_logger().error("Configure bind-mounts called without mount path. See configure --help for usage")
                sys.exit(1)

            mount_path = args.mount_path.rstrip('/').strip('"')

            # validate bind-mount
            if not mount_path.startswith(os.environ['HOME']) and not mount_path.startswith('/media/%s' % os.environ['USER']):
                libertine.utils.get_logger().error("Cannot mount {}, mount path must be in {} or /media/{}.".format(mount_path, os.environ['HOME'], os.environ['USER']))
                sys.exit(1)
            if mount_path.startswith('/media/%s' % os.environ['USER']) and \
                   self.containers_config.get_container_type(container_id) == 'lxc':
                libertine.utils.get_logger().error("/media mounts not currently supported in lxc.")
                sys.exit(1)
            if not os.path.isdir(mount_path):
                libertine.utils.get_logger().error("Cannot mount '%s', mount path must be an existing directory." % mount_path)
                sys.exit(1)

            # update database with new bind-mount
            container_bind_mounts = self.containers_config.get_container_bind_mounts(container_id)
            if args.bind_mount == 'add':
                if mount_path in container_bind_mounts:
                    libertine.utils.get_logger().error("Cannot add mount '%s', bind-mount already exists." % mount_path)
                    sys.exit(1)
                self.containers_config.add_new_bind_mount(container_id, mount_path)
            elif args.bind_mount == 'remove':
                if mount_path not in container_bind_mounts:
                    libertine.utils.get_logger().error("Cannot remove mount '%s', bind-mount does not exist." % mount_path)
                    sys.exit(1)
                self.containers_config.delete_bind_mount(container_id, mount_path)

            container_type = self.containers_config.get_container_type(container_id)

            if (container_type == 'lxc' or container_type == 'lxd' and
                self.containers_config.get_freeze_on_stop(container_id)):
                if not container.restart_libertine_container():
                    libertine.utils.get_logger().warning("Container cannot be restarted at this time.  You will need to "
                                                         "restart the container at a later time using the \'restart\' subcommand.")

        elif args.freeze is not None:
            container_type = self.containers_config.get_container_type(container_id)

            if container_type != 'lxc' and container_type != 'lxd':
                libertine.utils.get_logger().error("Configuring freeze is only valid on LXC and LXD container types.")
                sys.exit(1)

            self.containers_config.update_freeze_on_stop(container_id, args.freeze == 'enable')

        else:
            libertine.utils.get_logger().error("Configure called with no subcommand. See configure --help for usage.")
            sys.exit(1)


    def merge(self, args):
        self.containers_config.merge_container_config_files(args.file)

    def fix_integrity(self, args):
        if 'containerList' in self.containers_config.container_list:
            for container in self.containers_config.container_list['containerList']:
                libertine_container = self._container(container['id'])

                if 'installStatus' not in container or container['installStatus'] == 'removing':
                    self.destroy_container(libertine_container)
                    continue
                libertine_container.exec_command('dpkg --configure -a')

                for package in container['installedApps']:
                    if package['appStatus'] != 'installed':
                        self.remove_package_by_name(libertine_container, package['packageName'])

                if 'extraArchives' in container:
                    for archive in container['extraArchives']:
                        if archive['archiveStatus'] != 'installed':
                            self.delete_archive_by_name(libertine_container, archive['archiveName'])

    def set_default(self, args):
        if args.clear:
            self.containers_config.clear_default_container_id(True)
            sys.exit(0)

        container_id = self.containers_config.check_container_id(args.id)

        self.containers_config.set_default_container_id(container_id, True)

    def restart(self, args):
        container_id = self.containers_config.check_container_id(args.id)

        container_type = self.containers_config.get_container_type(container_id)

        if container_type != 'lxc' and container_type != 'lxd':
            libertine.utils.get_logger().error("The restart subcommand is only valid for LXC and LXD type containers.")
            sys.exit(1)

        container = self._container(container_id)

        container.restart_libertine_container()


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="Legacy X application support for Unity 8")

    if not os.geteuid():
        libertine.utils.get_logger().error("Please do not run %s using sudo" % parser.prog)
        sys.exit(1)

    container_manager = LibertineContainerManager()

    parser.add_argument('-q', '--quiet',
                        action='store_const', dest='verbosity', const=0,
                        help=('disables all non-vital output'))
    parser.add_argument('-v', '--verbosity',
                        action='store_const', dest='verbosity', const=2,
                        help=('enables debug output'))
    subparsers = parser.add_subparsers(dest="subparser_name",
                                       title="subcommands",
                                       metavar='create, destroy, install-package, remove-package, search-cache, update, list, list-apps, configure')

    # Handle the create command and its options
    parser_create = subparsers.add_parser(
        'create',
        help=("Create a new Libertine container."))
    parser_create.add_argument(
        '-i', '--id',
        required=True,
        help=("Container identifier of form ([a-z0-9][a-z0-9+.-]+). Required."))
    parser_create.add_argument(
        '-t', '--type',
        help=("Type of Libertine container to create. Either 'lxd', 'lxc' or 'chroot'."))
    parser_create.add_argument(
        '-d', '--distro',
        help=("Ubuntu distro series to create."))
    parser_create.add_argument(
        '-n', '--name',
        help=("User friendly container name."))
    parser_create.add_argument(
        '--force', action='store_true',
        help=("Force the installation of the given valid Ubuntu distro even if "
              "it is no longer supported."))
    parser_create.add_argument(
        '-m', '--multiarch', action='store_true',
        help=("Add i386 support to amd64 Libertine containers.  This option has "
              "no effect when the Libertine container is i386."))
    parser_create.add_argument(
        '--password',
        help=("Pass in the user's password when creating an LXC container.  This "
              "is intended for testing only and is very insecure."))
    parser_create.set_defaults(func=container_manager.create)

    # Handle the destroy command and its options
    parser_destroy = subparsers.add_parser(
        'destroy',
        help=("Destroy any existing environment entirely."))
    parser_destroy.add_argument(
        '-i', '--id',
        help=("Container identifier.  Default container is used if omitted."))
    parser_destroy.add_argument(
        '-f', '--force', action='store_true', required=False,
        help=("Force destroy.  Forces running containers to stop before destruction."))
    parser_destroy.set_defaults(func=container_manager.destroy)

    # Handle the install-package command and its options
    parser_install = subparsers.add_parser(
        'install-package',
        help=("Install a package or packages in the specified Libertine container."))
    parser_install.add_argument(
        '-p', '--package',
        required=True,
        nargs='+',
        help=("Name of package or full path to a Debian package. Multiple packages "
              "can be entered, separated by a space. Required."))
    parser_install.add_argument(
        '-i', '--id',
        help=("Container identifier.  Default container is used if omitted."))
    parser_install.add_argument(
        '-n', '--no-dialog', action='store_true',
        help=("No dialog mode. Use text-based frontend during debconf interactions."))
    parser_install.set_defaults(func=container_manager.install_package)

    # Handle the remove-package command and its options
    parser_remove = subparsers.add_parser(
        'remove-package',
        help=("Remove a package in the specified Libertine container."))
    parser_remove.add_argument(
        '-p', '--package',
        required=True,
        nargs='+',
        help=("Name of package to remove. Multiple packages can be entered, separated "
              "by a space. Required."))
    parser_remove.add_argument(
        '-i', '--id',
        help=("Container identifier.  Default container is used if omitted."))
    parser_remove.add_argument(
        '-n', '--no-dialog', action='store_true',
        help=("No dialog mode. Use text-based frontend during debconf interactions."))
    parser_remove.set_defaults(func=container_manager.remove_package)

    # Handle the search-cache command and its options
    parser_search = subparsers.add_parser(
        'search-cache',
        help=("Search for packages based on the search string in the specified Libertine container."))
    parser_search.add_argument(
        '-s', '--search-string',
        required=True,
        help=("String to search for in the package cache. Required."))
    parser_search.add_argument(
        '-i', '--id',
        help=("Container identifier.  Default container is used if omitted."))
    parser_search.set_defaults(func=container_manager.search_cache)

    # Handle the update command and its options
    parser_update = subparsers.add_parser(
        'update',
        help=("Update the packages in the Libertine container.  Also updates the container's "
              "locale and installs necessary language packs if the host's locale has changed."))
    parser_update.add_argument(
        '-i', '--id',
        help=("Container identifier.  Default container is used if omitted."))
    parser_update.set_defaults(func=container_manager.update)

    # Handle the list command
    parser_list = subparsers.add_parser(
        "list",
        help=("List all Libertine containers."))
    parser_list.set_defaults(func=container_manager.list)

    # Handle the list-apps command and its options
    parser_list_apps = subparsers.add_parser(
        'list-apps',
        help=("List available app launchers in a container."))
    parser_list_apps.add_argument(
        '-i', '--id',
        help=("Container identifier.  Default container is used if omitted."))
    parser_list_apps.add_argument(
        '-j', '--json',
        action='store_true',
        help=("use JSON output format."))
    parser_list_apps.set_defaults(func=container_manager.list_apps)

    # Handle the execute command and it's options
    parser_exec = subparsers.add_parser(
        'exec',
        add_help=False)
        #help=("Run an arbitrary command in the specified Libertine container."))
    parser_exec.add_argument(
        '-i', '--id',
        help=("Container identifier.  Default container is used if omitted."))
    parser_exec.add_argument(
        '-c', '--command',
        help=("The command to run in the specified container."))
    parser_exec.set_defaults(func=container_manager.exec)

    # Handle the configure command and it's options
    parser_configure = subparsers.add_parser(
        'configure',
        help=("Configure various options in the specified Libertine container."))
    parser_configure.add_argument(
        '-i', '--id',
        help=("Container identifier.  Default container is used if omitted."))
    multiarch_group = parser_configure.add_argument_group("Multiarch support",
                      "Enable or disable multiarch support for a container.")
    multiarch_group.add_argument(
         '-m', '--multiarch',
         choices=['enable', 'disable'],
         help=("Enables or disables i386 multiarch support for amd64 Libertine "
               "containers. This option has no effect when the Libertine "
               "container is i386."))

    archive_group = parser_configure.add_argument_group("Additional archive support",
                    "Add or delete an additional archive (PPA).")
    archive_group.add_argument(
        '-a', '--archive',
        choices=['add', 'remove'],
        help=("Adds or removes an archive (PPA) in the specified Libertine container."))
    archive_group.add_argument(
      '-n', '--archive-name',
      metavar='Archive name',
      help=("Archive name to be added or removed."))
    archive_group.add_argument(
        '-k', '--public-key-file',
        metavar='Public key file',
        help=("File containing the key used to sign the given archive. "
              "Useful for third-party or private archives."))

    mount_group = parser_configure.add_argument_group("Additional bind-mounts",
                    "Add or delete an additional bind-mount.")
    mount_group.add_argument(
        '-b', '--bind-mount',
        choices=['add', 'remove'],
        help="Adds or removes a bind-mount in the specified Libertine container.")
    mount_group.add_argument(
      '-p', '--mount-path',
      metavar='Mount path',
      help=("The absolute host path to bind-mount."))

    freeze_group = parser_configure.add_argument_group("Freeze container support",
                   "Enable or disable freezing LXC/LXD containers when not in use.")
    freeze_group.add_argument(
        '-f', '--freeze',
        choices=['enable', 'disable'],
        help=("Enables or disables freezing of LXC/LXD containers when not in use."
              " When disabled, the container will stop."))

    parser_configure.set_defaults(func=container_manager.configure)

    # Handle merging another ContainersConfig.json file into the main ContainersConfig.json file
    parser_merge = subparsers.add_parser(
        'merge-configs',
        add_help=False)
    parser_merge.add_argument(
        '-f', '--file',
        required=True)
    parser_merge.set_defaults(func=container_manager.merge)

    # Indiscriminately destroy containers, packages, and archives which are not fully installed
    parser_integrity = subparsers.add_parser(
        'fix-integrity',
        add_help=False)
    parser_integrity.set_defaults(func=container_manager.fix_integrity)

    # Set the default container in ContainersConfig
    parser_default = subparsers.add_parser(
        'set-default',
        help=("Set the default container."))
    parser_default.add_argument(
        '-i', '--id',
        metavar='Container id',
        help=("Container identifier.  Default container is used if omitted."))
    parser_default.add_argument(
        '-c', '--clear', action='store_true',
        help=("Clear the default container."))
    parser_default.set_defaults(func=container_manager.set_default)

    # Handle the restart command and its options
    parser_update = subparsers.add_parser(
        'restart',
        help=("Restart a frozen Libertine container.  This only works on LXC "
              "and LXD type containers."))
    parser_update.add_argument(
        '-i', '--id',
        help=("Container identifier.  Default container is used if omitted."))
    parser_update.set_defaults(func=container_manager.restart)

    # Actually parse the args
    args = parser.parse_args()

    libertine.utils.set_environmental_verbosity(args.verbosity)

    if args.subparser_name == None:
        parser.print_help()
    else:
        args.func(args)
