#!/usr/bin/python

# Copyright (C) 2013:
#    Gabes Jean, j.gabes@shinken-solutions.com
#
# This file is part of Shinken Enterprise, all rights reserved.

import sys
import optparse
import subprocess
import socket
import re
import os
import shutil
import httplib

try:
    import shinken
except ImportError:
    print 'Cannot import shinken lib, please install it before launching this tool'
    sys.exit(2)

from shinkensolutions.localinstall import get_local_daemons, get_local_addons, set_local_daemon_instance, set_local_addon, VERSION, WARNING_COLOR, ERROR_COLOR, OK_COLOR, check_root, POSSIBLE_DAEMONS, POSSIBLE_ADDONS, get_local_instances_for_type, \
    get_local_daemon_configuration_file_path, \
    get_instance_name, stop_daemon

YELLOW_COLOR = 33
DEFAULT_UMASK = 022

# Set umask to avoid problems when creating files
os.umask(DEFAULT_UMASK)

def get_major_linux_version():
    cmd = """cat /etc/redhat-release"""
    p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
    out = p.stdout.read()
    err = p.stderr.read()
    if err:
        return False
    
    match = re.search(r'release (\d)\.\d', out)
    if match:
        return int(match.group(1))
    else:
        return None


def get_local_daemon_cfg_path(daemon_type):
    # try to search also in local if we have such a daemon on our address
    cmd = """hostname -I"""
    p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
    out = p.stdout.read()
    err = p.stderr.read()
    if err:
        print ' \033[%dm ERROR: %s\033[0m' % (ERROR_COLOR, err)
        return None
    local_addresses = out.split(' ')
    local_addresses.append('localhost')
    local_addresses.append('127.0.0.1')
    
    # now look at local daemons
    cmd = """grep -H address /etc/shinken/%ss/*cfg | grep -v '#' """ % daemon_type
    p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
    out = p.stdout.read()
    err = p.stderr.read()
    if err:
        out = ''
    for line in out.splitlines():
        elts = [s.strip() for s in line.split(' ') if s.strip()]
        pth = elts[0].split(':')[0]
        try:
            addr = elts[2]
        except IndexError:
            return None
        
        # try to resolve addr
        try:
            addr = socket.gethostbyname(addr)
        except Exception, exp:
            pass
        
        if addr in local_addresses:
            return pth
    
    return None


def is_module_enabled(cfg_file, module_name):
    cmd = """grep -e '^\s*modules.*%s' %s > /dev/null""" % (module_name, cfg_file)
    p = subprocess.call(cmd, shell=True)
    return not bool(p)


def enable_module_in_file(cfg_file, module_name):
    if is_module_enabled(cfg_file, module_name):
        return True
    
    cmd = """sed -i "/^\s*modules/ s/$/, %s/" %s""" % (module_name, cfg_file)
    p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
    err = p.stderr.read()
    if err:
        return False

    # sed does not return an error exit code when no substitution has been done, so we check that the wanted module is present in the file
    cmd = """grep "^\s*modules.*%s" %s""" % (module_name, cfg_file)
    p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
    p.wait()
    if p.returncode != 0:
        return False
    
    return True


def disable_module_in_file(cfg_file, module_name):
    if not is_module_enabled(cfg_file, module_name):
        return True
    
    cmd = """sed -n \"s/^\s*modules\s*\(.*\)\s*/\\1/p\" %s""" % (cfg_file)
    p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
    module_list = p.stdout.read()
    err = p.stderr.read()
    if err or not module_list:
        return False
    
    module_list_without_module_to_remove = [m.strip() for m in module_list.split(',') if m.strip() != module_name]
    cmd = """sed -i "s/\(^\s*modules\s*\).*/\\1%s/" %s""" % (','.join(module_list_without_module_to_remove), cfg_file)
    p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
    err = p.stderr.read()
    if err:
        return False

    # sed does not return an error exit code when no substitution has been done, so we check that the wanted module is not present in the file
    cmd = """grep -E "^\s*modules.*%s" %s""" % (module_name, cfg_file)
    p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
    p.wait()
    if p.returncode == 0:
        return False
    
    return True


# enable_module(daemon_type string, module_name string)
# Returns (module_enabled bool, return_message string, daemons_to_restart set)
def enable_module(daemon_type, module_name):
    daemons_to_restart = set()
    daemon_cfg_path = get_local_daemon_cfg_path(daemon_type)
    if daemon_cfg_path is None:
        return True, "    - %(status)s Cannot find local %(daemon)s daemon configuration in /etc/shinken/%(daemon)ss. The '%(module)s' module must be enabled manually in your %(daemon)s configuration file for this addon to work properly.\n" % {
            "status": status("warning"),
            "daemon": daemon_type,
            "module": module_name
        }, set()
    
    if is_module_enabled(daemon_cfg_path, module_name):
        return True, "", set()
    
    module_enabled = enable_module_in_file(daemon_cfg_path, module_name)
    if not module_enabled:
        return False, "    - %s Cannot enable %s module in %s configuration file\n" % (status("error"), module_name, daemon_type), set()
    
    daemons_to_restart.add("Arbiter")
    return True, "    - %s %s module enabled in %s.\n" % (status("success"), module_name, daemon_cfg_path), daemons_to_restart


# disable_module(daemon_type string, module_name string)
# Returns (module_enabled bool, return_message string , daemons_to_restart set)
def disable_module(daemon_type, module_name):
    daemons_to_restart = set()
    daemon_cfg_path = get_local_daemon_cfg_path(daemon_type)
    if daemon_cfg_path is None:
        return True, "    - %(status)s Cannot find local %(daemon)s daemon configuration in /etc/shinken/%(daemon)ss. The '%(module)s' module must be disabled manually in your %(daemon)s configuration file for this addon to be complety disabled.\n" % {
            "status": status("warning"),
            "daemon": daemon_type,
            "module": module_name
        }, set()
    
    if not is_module_enabled(daemon_cfg_path, module_name):
        return True, "", set()
    
    module_enabled = disable_module_in_file(daemon_cfg_path, module_name)
    if not module_enabled:
        return False, "   - %s Cannot disable %s module in %s configuration file\n" % (status("error"), module_name, daemon_type), set()

    daemons_to_restart.add(daemon_type.capitalize())
    return True, "    - %s %s module disabled in %s.\n" % (status("success"), module_name, daemon_cfg_path), daemons_to_restart


def restart_apache():
    if get_major_linux_version() == 7:
        cmd = """systemctl restart httpd"""
    else:
        cmd = """/etc/init.d/httpd restart"""
    
    subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
    return True


def get_architecture_export_configuration():
    conf = {}
    file = open("/etc/shinken/modules/architecture-export.cfg")
    for line in file:
        if not line.startswith("#"):
            m = re.match(r"\s*(?P<key>\w+)\s*(?P<value>.*)", line)
            if m is not None:
                match = m.groupdict()
                conf[match['key']] = match['value']
                
    return conf


def notify_architecture_export_disabled():
    architecture_export_cfg = get_architecture_export_configuration()
    
    port = int(architecture_export_cfg.get("port", 7780))
    use_ssl = False if int(architecture_export_cfg.get('use_ssl', 0)) == 0 else True
    ssl_cert = architecture_export_cfg.get('ssl_cert', '')
    ssl_key = architecture_export_cfg.get('ssl_key', '')
    
    try:
        if use_ssl:
            conn = httplib.HTTPSConnection("localhost", port, key_file=ssl_key, cert_file=ssl_cert, timeout=3)
        else:
            conn = httplib.HTTPConnection("localhost", port, timeout=3)
            
        conn.request("GET", "/v1/architecture/disabled")
        response = conn.getresponse()
        response.read()
    
        if response.status not in (httplib.CREATED, httplib.OK):
            return False, "    - %s Could not remove hosts in all other installations listed in architecture-export module configuration:  %s \n" % (status("warning"), response.status)
        else:
            return True, ""
    except (httplib.HTTPException, socket.error) as exp:
        return False, "    - %s Could not remove hosts in all other installations listed in architecture-export module configuration:  %s\n" % (status("warning"), exp)


def enable_apache_conf(filename=""):
    src="/etc/httpd/conf.d/%s.disabled" % filename
    dest="/etc/httpd/conf.d/%s" % filename
    try:
        shutil.move(src, dest)
    except:
        pass
        
    apache_restarted = restart_apache()
    return apache_restarted


def disable_apache_conf(filename=""):
    src="/etc/httpd/conf.d/%s" % filename
    dest="/etc/httpd/conf.d/%s.disabled" % filename
    try:
        shutil.move(src, dest)
    except:
        pass
        
    apache_restarted = restart_apache()
    return apache_restarted


def recover_addons_configuration_files():
    if not os.path.isfile('/etc/shinken/modules/architecture-export.cfg'):
        # Copying
        shutil.copyfile('/etc/shinken-skeletons/modules/architecture-export.cfg', '/etc/shinken/modules/architecture-export.cfg')
        
    if not os.path.islink('/etc/shinken/external/nagvis'):
        if not os.path.isdir('/etc/shinken/external'):
            os.mkdir('/etc/shinken/external')
        os.symlink('/var/lib/shinken-nagvis', '/etc/shinken/external/nagvis')
    
    
def display_summary_message(addon, enabled, status_message=""):
    if enabled:
        print ' %s:\033[%dm ENABLED\033[0m' % (addon.ljust(30), OK_COLOR)
    else:
        print ' %s:\033[%dm DISABLED\033[0m' % (addon.ljust(30), WARNING_COLOR)
        
    if status_message != "":
        print status_message


def display_instructions(daemons_to_restart=None):
    if daemons_to_restart:
        print "\n\033[%dm Some changes have been made in Shinken daemons configuration files. To apply these changes, the following daemons must be restarted:\033[0m" % YELLOW_COLOR
        for daemon in daemons_to_restart:
            print "    - \033[%dm%s\033[0m" % (YELLOW_COLOR, daemon)
            

def status(type="success"):
    msg = ""
    color = 0
    if type == "success":
        msg="OK"
        color = OK_COLOR
    elif type == "warning":
        msg="WARNING"
        color=WARNING_COLOR
    elif type == "error":
        msg="ERROR"
        color = ERROR_COLOR

    return '[ \033[%dm%-8s\033[0m ]' % (color, msg)


if __name__ == '__main__':
    parser = optparse.OptionParser("%prog [nagvis nagvis-shinken-architecture]", version="%prog: " + VERSION, description='This tool is used to enable addons on the local server')
    parser.add_option('--enable', dest='mode_enable', action='store_true',
                      help="Enable a local addon.")
    parser.add_option('--disable', dest='mode_disable', action='store_true',
                      help="Disable a local addon.")
    parser.add_option('--force', dest='mode_force', action='store_true',
                      help="Do not look at previous state before doing the change.")
    
    opts, args = parser.parse_args()
    
    # If not root, exit
    check_root()
    
    if len(args) == 0:
        print "Missing addon to enable/disable"
        sys.exit(2)
    
    if not opts.mode_enable and not opts.mode_disable:
        print "Please select at least a --enable or --disable option"
        sys.exit(2)
    
    if opts.mode_force is None:
        opts.mode_force = False
    
    instance_number, is_arbiter = get_local_instances_for_type('arbiter')[0]
    
    addons = [a.strip() for a in args if a.strip()]
    
    for a in addons:
        if a not in POSSIBLE_ADDONS:
            print "The addon %s is unknown, exiting" % a
            sys.exit(2)

    daemons_to_restart = set()
    addons_status = get_local_addons()
    for (a, enabled) in addons_status.iteritems():
        status_message = ""
        
        if not a in addons:
            continue

        # Check if configuration files and folder structure for addons exists. If not create them to have a working addon once enabled
        recover_addons_configuration_files()

        # Look at previous state, but only if we are not in a force mode
        if not opts.mode_force:
            # Maybe already enabled
            if opts.mode_enable:
                if enabled:
                    print ' %s:\033[%dm ALREADY ENABLED, skipping\033[0m' % (a.ljust(30), WARNING_COLOR)
                    continue
            
            # Maybe already disable
            if opts.mode_disable:
                if not enabled:
                    print ' %s:\033[%dm ALREADY DISABLED, skipping\033[0m' % (a.ljust(30), WARNING_COLOR)
                    continue
        
        if opts.mode_enable:
            if a == "nagvis":
                livestatus_enabled = True
                apache_conf_enabled = enable_apache_conf("nagvis_opt.conf")
                
                if is_arbiter:
                    livestatus_enabled, return_message, to_restart = enable_module("broker", "Livestatus")
                    status_message += return_message
                    if livestatus_enabled:
                        daemons_to_restart = daemons_to_restart.union(to_restart)
                else:
                    status_message += "    - %s 'nagvis' addon has been enabled locally. To be able to see hosts and checks state in NagVis, the 'Livestatus' module must be enabled on the corresponding broker (via configuration file on the Arbiter server).\n" % status("warning")
                
                if livestatus_enabled and apache_conf_enabled:
                    set_local_addon(a, True)
                    display_summary_message(a, True, status_message=status_message)
                else:
                    display_summary_message(a, enabled, status_message=status_message)
                    
            elif a == "nagvis-shinken-architecture":
                livestatus_enabled = False
                arbiter_module_enabled = False
                apache_conf_enabled = False
                
                if is_arbiter:
                    apache_conf_enabled = enable_apache_conf("nagvis_etc.conf")
                    livestatus_enabled, return_message, to_restart = enable_module("broker", "Livestatus")
                    status_message += return_message
                    if livestatus_enabled:
                        daemons_to_restart = daemons_to_restart.union(to_restart)
                    
                    arbiter_module_enabled, return_message, to_restart = enable_module("arbiter", "architecture-export")
                    status_message += return_message
                    if arbiter_module_enabled:
                        daemons_to_restart = daemons_to_restart.union(to_restart)
                else:
                    status_message += "    - %s 'nagvis-shinken-architecture' addon needs to be enabled on a server hosting an enabled Arbiter daemon. Please enable this addon on the server hosting the Arbiter.\n" % status("error")
                
                if livestatus_enabled and apache_conf_enabled and arbiter_module_enabled:
                    set_local_addon(a, True)
                    display_summary_message(a, True, status_message=status_message)
                else:
                    display_summary_message(a, enabled, status_message=status_message)
        else:  # disable
            if a == "nagvis":
                # Disable Livestatus only if the other addon using Livestatus is not enabled and we didn't ask for him to be disabled
                apache_conf_disabled = disable_apache_conf("nagvis_opt.conf")
                
                if is_arbiter:
                    status_message += "    - %s Livestatus module may still be enabled on the broker. If this module is not needed by any other addon or external tool, you may want to disable it in the Broker configuration file.\n" % status("warning")
                else:
                    status_message += "    - %s 'nagvis' addon has been disabled locally.  'Livestatus' module may still be enabled on the broker. If this module is not needed by any other addon or external tool, you may want to disable it in the Broker configuration file.\n" % status("warning")
                
                if apache_conf_disabled:
                    set_local_addon(a, False)
                    display_summary_message(a, False, status_message=status_message)
                else:
                    display_summary_message(a, enabled, status_message=status_message)
                    
            elif a == "nagvis-shinken-architecture":
                # Disable Livestatus only if the other addon using Livestatus is not enabled and we didn't ask for him to be disabled
                arbiter_module_disabled = True
                apache_conf_disabled = disable_apache_conf("nagvis_etc.conf")
                
                if is_arbiter:
                    status_message += "    - %s Livestatus module may still be enabled on the broker. If this module is not needed by any other addon or external tool, you may want to disable it in the Broker configuration file.\n" % status("warning")
                    arbiter_module_disabled, return_message, to_restart = disable_module("arbiter", "architecture-export")
                    status_message += return_message
                    if arbiter_module_disabled:
                        daemons_to_restart = daemons_to_restart.union(to_restart)
                        
                    architecture_modules_notified, return_message = notify_architecture_export_disabled()
                    status_message += return_message
                
                if apache_conf_disabled and arbiter_module_disabled:
                    set_local_addon(a, False)
                    display_summary_message(a, False, status_message=status_message)
                else:
                    display_summary_message(a, enabled, status_message=status_message)

    display_instructions(daemons_to_restart)