#! /usr/bin/python3
#
# Copyright (c) 2018 Ultimum Technologies s.r.o.
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License. You may obtain
# a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations
# under the License.
#
# Author: Michal Arbet <michal.arbet@ultimum.io>

from stevedore.named import ExtensionManager
from oslo_config import cfg
import argparse
import os
import fileinput
import re
import logging
import sys
import subprocess
#
# Logging
LOG = logging.getLogger()
LOG.setLevel(logging.INFO)
ch = logging.StreamHandler(sys.stdout)
ch.setLevel(logging.DEBUG)
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
ch.setFormatter(formatter)
LOG.addHandler(ch)

class NeutronArgsParser(object):

    def __init__(self, neutron_plugin_enabler):
        self.configs = neutron_plugin_enabler.configs
        self.available_service_plugins = list(neutron_plugin_enabler.available_service_plugins.keys())
        self.available_l3_extensions = list(neutron_plugin_enabler.available_l3_extensions.keys())
        self.neutron_plugin_manager = neutron_plugin_enabler
        self._init_parser()


    def _enable(self, args):
        if args.service_plugin is not None:
            self.neutron_plugin_manager.enable(args.service_plugin, 'service_plugin')
            self.neutron_plugin_manager.link_configs(args.service_plugin, 'service_plugin')
        if args.l3_extension is not None:
            self.neutron_plugin_manager.enable(args.l3_extension, 'l3_extension')
            if args.service_plugin is None:
                self.neutron_plugin_manager.link_configs(args.l3_extension, 'l3_extension')
        self.neutron_plugin_manager.restart_services()

    def _disable(self, args):
        if args.service_plugin is not None:
            self.neutron_plugin_manager.disable(args.service_plugin, 'service_plugin')
            self.neutron_plugin_manager.unlink_configs(args.service_plugin, 'service_plugin')
        if args.l3_extension is not None:
            self.neutron_plugin_manager.disable(args.l3_extension, 'l3_extension')
            if args.service_plugin is None:
                self.neutron_plugin_manager.unlink_configs(args.l3_extension, 'l3_extension')
        self.neutron_plugin_manager.restart_services()

    def _init_parser(self):
        self.parser = argparse.ArgumentParser(prog='neutron-plugin-manage', add_help=False)
        self.parser.add_argument('--service-plugin', action='store',
                              dest='service_plugin', help='Define service_plugin.', choices=self.available_service_plugins)
        self.parser.add_argument('--l3-extension', action='store',
                              dest='l3_extension', help='Define l3_extension.', choices=self.available_l3_extensions)

        sp = self.parser.add_subparsers()
        sp_enable = sp.add_parser('enable', help='Enable.', parents=[self.parser])
        sp_disable = sp.add_parser('disable', parents=[self.parser],
                                help='Disable.')
        sp_enable.set_defaults(func=self._enable)
        sp_disable.set_defaults(func=self._disable)

    def parse(self):
        args = self.parser.parse_args()
        args.func(args)

class NeutronPluginManager(object):

    def __init__(self, configs):
        self.configs = configs
        self.available_service_plugins = self._stevedore_discover('neutron.service_plugins')
        self.available_l3_extensions = self._stevedore_discover('neutron.agent.l3.extensions')
        self.CONF = self._load_neutron_config()
        self.restart_trigger = True
        # Some neutron paths
        self.neutron_config_dir = "/etc/neutron/"
        self.neutron_server_config_dir = self.neutron_config_dir + "server.conf.d/"
        self.neutron_agent_config_dir = self.neutron_config_dir + "agent.conf.d/"

    def restart_services(self):
        if self.restart_trigger:
            LOG.info("Restarting neutron services.")
            os.system("sudo systemctl restart neutron-*.service")
            return
        LOG.info("Restart is not needed.")

    def _restart(self, restart=True):
        if self.restart_trigger and restart:
            return
        self.restart_trigger = restart

    def _stevedore_discover(self, namespace):
        stevedore_elements = ExtensionManager(namespace=namespace, invoke_on_load=False)
        elements = {}
        for name, entry_point in stevedore_elements.items():
            try:
                element = {}
                element['name'] = name
                element['module'] = entry_point.entry_point_target.split(".")[0]
                element['target'] = entry_point.entry_point_target
                elements.update({name: element})
            except Exception:
                pass
        return elements

    def _stevedore_elements(self, type, module):
        elements = {}
        if type == 'service_plugin':
            for k,v in self.available_service_plugins.items():
                if v.get('module') == module:
                    elements.update({k:self.available_service_plugins.get(k)})
        if type == 'l3_extension':
            for k, v in self.available_l3_extensions.items():
                if v.get('module') == module:
                    elements.update({k:self.available_l3_extensions.get(k)})
        return elements

    def _load_neutron_config(self):
        oslo_config_opts = ExtensionManager(namespace='oslo.config.opts', invoke_on_load=True)
        for k, v in oslo_config_opts.items():
            for group, items in v.obj:
                if not k.startswith('neutron'):
                    continue
                cfg.CONF.register_opts(items, group)
        args_oslo = ["--config-file=" + i for i in self.configs]
        cfg.CONF(args=args_oslo)
        return cfg.CONF

    def _service_plugin(self, name):
        return self.available_service_plugins.get(name)

    def _l3_extension(self, name):
        return self.available_l3_extensions.get(name)

    def _service_plugins(self, module):
        return self._stevedore_elements('service_plugin', module)

    def _l3_extensions(self, module):
        return self._stevedore_elements('l3_extension', module)

    # Replace line inplace
    def _replace(self, file, searchExp, replaceExp):
        uid = os.stat(file).st_uid
        gid = os.stat(file).st_gid
        config = fileinput.FileInput(file, inplace=True)
        for line in config:
            line = re.sub(searchExp, replaceExp, line.rstrip())
            print(line)
        os.chown(file, uid, gid)

    def _check(self, name, type):
        if type == 'service_plugin':
            if name in self.available_service_plugins:
                return True
            return False
        if type == 'l3_extension':
            if name in self.available_l3_extensions:
                return True
            return False

    def enable(self, element, type):
        if type == 'service_plugin':
            if self._check(element, 'service_plugin'):
                neutron_module = self._service_plugin(element).get('module')
                elements_enabled = self.CONF.service_plugins
            else:
                LOG.error("Enabling service_plugin failed. {} is not neutron's service_plugin.".format(element))
                self._restart(False)
                return
        if type == 'l3_extension':
            if self._check(element, 'l3_extension'):
                neutron_module = self._l3_extension(element).get('module')
                elements_enabled = self.CONF.agent.extensions
            else:
                LOG.error("Enabling l3 extension failed. {} is not neutron's l3 extension.".format(element))
                self._restart(False)
                return
        if neutron_module != 'neutron':
            for i in elements_enabled:
                if type == 'service_plugin':
                    element_dict = self.available_service_plugins.get(i)
                if type == 'l3_extension':
                    element_dict = self.available_l3_extensions.get(i)
                if element_dict.get('module') == neutron_module:
                    if element_dict.get('name') == element:
                        if type == 'service_plugin':
                            LOG.info("Enabling service_plugin {} failed. This service_plugin is already enabled.".format(
                                element))
                            self._restart(False)
                            return
                        if type == 'l3_extension':
                            LOG.info("Enabling l3 extension {} failed. This l3 extension is already enabled.".format(
                                element))
                            self._restart(False)
                            return
                    if type == 'service_plugin':
                        # Another service_plugin enabled from neutron_module.
                        LOG.error("Enabling service_plugin {} failed. Another service_plugin "
                                  "{} from {} module is currently "
                                  "enabled".format(element, element_dict.get('name'), neutron_module))
                        self._restart(False)
                        return
                    if type == 'l3_extension':
                        # Another service_plugin enabled from neutron_module.
                        LOG.error("Enabling l3 extension {} failed. Another l3 extension "
                                  "{} from {} module is currently "
                                  "enabled".format(element, element_dict.get('name'), neutron_module))
                        self._restart(False)
                        return
            if type == 'service_plugin':
                LOG.info("Enabling service_plugin {}".format(element))
                elements_enabled.append(element)
                self._replace(self.configs[0], '^#?service_plugins.*',
                              'service_plugins = {}'.format(",".join(elements_enabled)))
                self._restart(True)
            if type == 'l3_extension':
                LOG.info("Enabling l3 extension {}".format(element))
                elements_enabled.append(element)
                self._replace(self.configs[1], '^#?extensions.*',
                              'extensions = {}'.format(",".join(elements_enabled)))
                self._restart(True)
        else:
            if type == 'service_plugin':
                if element in elements_enabled:
                    LOG.info("Enabling service_plugin {} failed. This service_plugin is already enabled.".format(element))
                    self._restart(False)
                    return
            else:
                LOG.info("Enabling service_plugin {}".format(element))
                elements_enabled.append(element)
                self._replace(self.configs[0], '^#?service_plugins.*',
                              'service_plugins = {}'.format(",".join(elements_enabled)))
                self._restart(True)
                return
            if type == 'l3_extension':
                if element in elements_enabled:
                    LOG.info("Enabling l3 extension {} failed. This l3 extension is already enabled.".format(element))
                    self._restart(False)
                    return
                else:
                    LOG.info("Enabling l3 extension {}".format(element))
                    elements_enabled.append(element)
                    self._replace(self.configs[1], '^#?extensions.*',
                                  'extensions = {}'.format(",".join(elements_enabled)))
                    self._restart(True)
                    return

    def disable(self, element, type):
        if type == 'service_plugin':
            if not self._check(element, 'service_plugin'):
                LOG.error("Disabling service_plugin failed. {} is not neutron's service_plugin".format(element))
                self._restart(False)
            else:
                elements_enabled = self.CONF.service_plugins
                if element in elements_enabled:
                    LOG.info("Disabling service_plugin {}".format(element))
                    elements_enabled.remove(element)
                    self._replace(self.configs[0], '^#?service_plugins.*',
                                  'service_plugins = {}'.format(",".join(elements_enabled)))
                    self._restart()
                else:
                    if len(elements_enabled) == 0:
                        self._replace(self.configs[0], '^service_plugins.*', '#service_plugins = ')
                    LOG.info("Disabling service_plugin {} failed. {} is "
                              "not currently enabled.".format(element, element))
                    self._restart(False)
        if type == 'l3_extension':
            if not self._check(element, 'service_plugin'):
                LOG.error("Disabling l3 extension failed. {} is not neutron's l3 extension.".format(element))
                self._restart(False)
            else:
                elements_enabled = self.CONF.agent.extensions
                if element in elements_enabled:
                    LOG.info("Disabling l3 extension {}".format(element))
                    elements_enabled.remove(element)
                    self._replace(self.configs[1], '^#?extensions.*',
                                  'extensions = {}'.format(",".join(elements_enabled)))
                    self._restart()
                else:
                    if len(elements_enabled) == 0:
                        self._replace(self.configs[1], '^extensions.*', '#extensions = ')
                    LOG.info("Disabling l3 extension {} failed. {} is "
                              "not currently enabled. agent.extensions directive is commented now.".format(element, element))
                    self._restart(False)

    def link_configs(self, element, type):
        if type == 'service_plugin':
            if not self._check(element, 'service_plugin'):
                LOG.error("Linking configs failed. {} is not neutron's service_plugin.".format(element))
                self._restart(False)
                return
            neutron_module = self._service_plugin(element).get('module')
        if type == 'l3_extension':
            if not self._check(element, 'l3_extension'):
                LOG.error("Linking configs failed. {} is not neutron's l3_extension.".format(element))
                self._restart(False)
                return
            neutron_module = self._l3_extension(element).get('module')

        if neutron_module != "neutron":
            if type == 'service_plugin':
                neutron_module = self.available_service_plugins.get(element).get('module').split("_")[1]
            if type == 'l3_extension':
                neutron_module = self.available_l3_extensions.get(element).get('module').split("_")[1]
            config_files = subprocess.check_output(['bash', '-c',
                                                    "dpkg -L neutron-{}-common|"
                                                    "egrep '{}.*\.conf|{}.*\.ini'|"
                                                    "sed -e 's/.*\///g'"
                                                   .format(neutron_module, self.neutron_config_dir,
                                                           self.neutron_config_dir)]).decode('utf-8').strip().split('\n')
            for cfg in config_files:
                if "conf" in cfg:
                    if not os.path.islink(self.neutron_server_config_dir + cfg):
                        os.symlink(self.neutron_config_dir + cfg, self.neutron_server_config_dir + cfg)
                        LOG.info("Creating symlink: {} -> {}".format(self.neutron_config_dir + cfg,
                                                                     self.neutron_server_config_dir + cfg))
                        self._restart()
                    else:
                        LOG.info("Symlink already exist: {} -> {}".format(self.neutron_config_dir + cfg,
                                                                          self.neutron_server_config_dir + cfg))
                        self._restart(False)
                if "ini" in cfg:
                    if not os.path.islink(self.neutron_agent_config_dir + cfg + ".conf"):
                        os.symlink(self.neutron_config_dir + cfg, self.neutron_agent_config_dir + cfg + ".conf")
                        LOG.info("Creating symlink: {} -> {}".format(self.neutron_config_dir + cfg,
                                                                     self.neutron_agent_config_dir + cfg + ".conf"))
                        self._restart()
                    else:
                        LOG.info("Symlink already exist: {} -> {}".format(self.neutron_config_dir + cfg,
                                                                          self.neutron_agent_config_dir + cfg + ".conf"))
                        self._restart(False)

    def unlink_configs(self, element, type):
        if type == 'service_plugin':
            if not self._check(element, 'service_plugin'):
                LOG.error("Unlinking configs failed. {} is not neutron's l3 extension.".format(element))
                self._restart(False)
                return
            neutron_module = self._service_plugin(element).get('module')
        if type == 'l3_extension':
            if not self._check(element, 'l3_extension'):
                LOG.error("Unlinking configs failed. {} is not neutron's l3 extension.".format(element))
                self._restart(False)
                return
            neutron_module = self._l3_extension(element).get('module')
        if neutron_module != "neutron":
            if type == 'service_plugin':
                neutron_module = self.available_service_plugins.get(element).get('module').split("_")[1]
            if type == 'l3_extension':
                neutron_module = self.available_l3_extensions.get(element).get('module').split("_")[1]
            config_files = subprocess.check_output(['bash', '-c',
                                                    "dpkg -L neutron-{}-common|"
                                                    "egrep '{}.*\.conf|{}.*\.ini'|"
                                                    "sed -e 's/.*\///g'"
                                                   .format(neutron_module, self.neutron_config_dir,
                                                           self.neutron_config_dir)]).decode('utf-8').strip().split(
                '\n')

            for cfg in config_files:
                    if "conf" in cfg:
                        try:
                            os.remove(self.neutron_server_config_dir + cfg)
                            LOG.info(
                                "Removing symlink: {} -> {}".format(self.neutron_config_dir + cfg,
                                                                    self.neutron_server_config_dir + cfg))
                            self._restart()
                        except OSError:
                            LOG.info("Not removing symlinks. They are already removed.")
                            pass
                    if "ini" in cfg:
                        try:
                            os.remove(self.neutron_agent_config_dir + cfg + ".conf")
                            LOG.info("Removing symlink: {} -> {}".format(self.neutron_config_dir + cfg,
                                                                         self.neutron_agent_config_dir + cfg + ".conf"))
                            self._restart()
                        except OSError:
                            pass

def main():
    manager = NeutronPluginManager(['/etc/neutron/neutron.conf', '/etc/neutron/l3_agent.ini'])
    parser = NeutronArgsParser(manager)
    parser.parse()

if __name__ == "__main__":
    main()
