#!/usr/bin/python3
import configparser
import json
import logging
import os
import platform
import re
import shlex
import shutil
import stat
import subprocess
import sys
import uuid
import zipfile
from functools import wraps

import psutil

try:
    import win32api
    import win32con
except ImportError:
    pass

conf = configparser.ConfigParser()
SYSTEM_LINUX = 'Linux'
SYSTEM_WINDOWS = 'Windows'

PERMISSION_700 = stat.S_IRWXU
PERMISSION_500 = stat.S_IRUSR | stat.S_IXUSR
PERMISSION_600 = stat.S_IRUSR | stat.S_IWUSR


def rollback(func):
    @wraps(func)
    def wrapper(self, *args, **kwargs):
        func_return = func(self, *args, **kwargs)
        if not func_return:
            rm_result = self.remove_path(self._agent_assist_path,
                                         self._agent_assist_python_path)
            if rm_result:
                logging.info('Success to rollback.')
        return func_return

    return wrapper


class CommonMethod(object):
    def __init__(self, install_path):
        self._cur_system = platform.system()
        self._install_path = os.path.realpath(install_path)
        self._parameter_pattern = re.compile(r'^[\w+-.]+$')
        self._agent_assist_path = os.path.join(self._install_path,
                                               'AgentAssist')
        self._agent_assist_python_path = os.path.join(self._install_path,
                                                      'AgentAssistPython')
        if self._cur_system == SYSTEM_WINDOWS:
            self._python_path = os.path.join(self._install_path,
                                             'AgentAssistPython/python.exe')
        else:
            self._python_path = os.path.join(self._install_path,
                                             'AgentAssistPython/bin/python3')

    @staticmethod
    def remove_path(*paths):
        if not paths:
            logging.info(f"The remove path is None.")
            return False
        for path in paths:
            try:
                if os.path.exists(path):
                    shutil.rmtree(path)
                    logging.info(f"Success to remove the {path}.")
            except Exception as err:
                logging.error(f"Failed to remove the {path}. {err}.")
                return False
        return True

    def _mod_chmod(self, path, mode):
        if self._cur_system != SYSTEM_WINDOWS:
            os.chmod(path, mode)

    def _set_permissions(self, folder):
        for root, dirs, files in os.walk(folder):
            # folder
            for dir_name in dirs:
                self._set_dirs_permissions(root, dir_name)
            # files
            for file in files:
                self._set_files_permissions(root, file)

    def _set_files_permissions(self, path, file_name):
        file_type = os.path.splitext(file_name)[1]
        file_path = os.path.join(path, file_name)
        if os.path.isfile(file_path):
            list_permission_500 = ['.py', '.sh', '.pyc']
            # 500
            if file_type in list_permission_500:
                self._mod_chmod(file_path, PERMISSION_500)
            # 600
            else:
                self._mod_chmod(file_path, PERMISSION_600)

    def _set_dirs_permissions(self, dir_path, dir_name):
        dir_path = os.path.join(dir_path, dir_name)
        list_permission_700 = ['conf', 'log', 'log_bak_log', 'logbak', 'cert']
        # 700
        if dir_name in list_permission_700:
            self._mod_chmod(dir_path, stat.S_IMODE(PERMISSION_700))
        # 500
        else:
            self._mod_chmod(dir_path, stat.S_IMODE(PERMISSION_500))

    def _check_conf_parameter(self, check_parameter):
        if self._parameter_pattern.findall(check_parameter):
            return True
        logging.error('Failed to check parameter.')
        return False

    def _set_python_permissions(self, folder_path):
        for root, dirs, files in os.walk(folder_path):
            for file in files:
                file_path = os.path.join(root, file)
                file_type = os.path.splitext(file_path)[1]
                if file_type == ".py" or file_type == ".so":
                    self._mod_chmod(file_path, PERMISSION_500)

    @rollback
    def _create_config(self, is_upgrade=False):
        log_bak_path = os.path.join(self._install_path, 'AgentAssist',
                                    'logbak/log_bak_log')
        try:
            os.makedirs(log_bak_path, mode=0o700, exist_ok=True)
            if self._cur_system == SYSTEM_WINDOWS:
                return True
            # software Environment Variables
            assist_folder = os.path.join(self._install_path, 'AgentAssist')
            python_folder = os.path.join(self._install_path,
                                         'AgentAssistPython')
            if is_upgrade:
                assist_folder += '_Image'
                python_folder += '_Image'
                self._mod_chmod(assist_folder, PERMISSION_500)
            self._set_permissions(assist_folder)
            self._set_python_permissions(python_folder)
            logging.info('The file permission is set successfully.')
            return True
        except Exception as err:
            logging.error(f'Failed to set file permission {err}')
            return False

    @staticmethod
    def _find_pkg(path):
        pkg_name = ""
        for file in os.listdir(path):
            pkg_file = re.match('BCManager.*AgentAssist', file)
            if pkg_file:
                logging.info("Succeeded in searching for the package.")
                return file
        logging.error("The package does not exist.")
        return pkg_name

    @staticmethod
    def _read_conf_without_section(file_name, key):
        if not os.path.isfile(file_name):
            return ""
        try:
            with open(file_name) as f:
                file_content = '[dummy_section]\n' + f.read()
        except Exception as error:
            logging.error(f"Read conf failed, err:{error}.")
            return ''
        config_parser = configparser.ConfigParser(allow_no_value=True)
        config_parser.read_string(file_content)
        value = config_parser.get("dummy_section", key)
        return value


class AgentAssistInstallHandle(CommonMethod):
    def __init__(self, install_path):
        super(AgentAssistInstallHandle, self).__init__(install_path)
        self._agent_assist_install_log = "AgentAssist_install.log"
        self._log_path = os.path.join(os.getcwd(),
                                      self._agent_assist_install_log)

        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s %(filename)s[line:%(lineno)d] %(levelname)s %(message)s',
            datefmt='%a, %d %b %Y %H:%M:%S',
            filename=self._log_path,
            filemode='a')

    # Decompress the installation package.
    @rollback
    def _unzip_pkg(self):
        cur_path = os.path.join(os.getcwd())
        pkg_name = AgentAssistInstallHandle._find_pkg(cur_path)
        if not os.path.isfile(pkg_name):
            print("The installation package does not exist.")
            logging.error('The installation package does not exist.')
            return False

        install_path = os.path.join(self._install_path, 'AgentAssist')
        logging.info(f'The path of the running package is: '
                     f'{install_path}. pkg_name: {pkg_name}.')
        # 限制最大文件数2000
        max_file_num = 2000
        # 设置解压内容最大值(一般平均最大的压缩率20，再高就很可能是异常文件了！)
        max_file_size = 1024 * 1024 * 100 * 20
        total_size = 0
        self.remove_path(install_path)

        try:
            install_pkg_path = os.path.join(cur_path, pkg_name)
            arc_file = zipfile.ZipFile(install_pkg_path)
            if len(arc_file.namelist()) > max_file_num:
                raise ValueError("The compressed file nums exceed maximum.")
            for file_info in arc_file.filelist:
                total_size += file_info.file_size
                if total_size > max_file_size:
                    raise IOError("The compressed package has an exception.")
                if self._cur_system == SYSTEM_WINDOWS:
                    free_disk_usage = psutil.disk_usage(cur_path).free
                else:
                    free_disk_usage = psutil.disk_usage(install_pkg_path).free
                if total_size >= free_disk_usage:
                    raise IOError(f'The zipfile size({total_size}) exceed remain '
                                  f'target disk space({free_disk_usage}).')
            arc_file.extractall(install_path)
            arc_file.close()
            return True
        except Exception as err:
            print("Failed to decompress the running package.")
            logging.error(f'Failed to decompress the running package. {err}.')
            return False

    @staticmethod
    def _mod_registry(name, path):
        key_name = 'SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Run'  # 注册表项名
        try:
            key = win32api.RegOpenKey(win32con.HKEY_CURRENT_USER, key_name, 0,
                                      win32con.KEY_ALL_ACCESS)
            win32api.RegSetValueEx(key, name, 0, win32con.REG_SZ, path)
            win32api.RegCloseKey(key)
        except Exception as err:
            logging.error(f'Failed to add registry. {err}')
            return False

        logging.info('Registry added successfully.')
        return True

    def _set_auto_start_windows(self):
        name = 'AgentAssist'  # 要添加的项目名
        file_path = os.path.join(self._install_path, 'AgentAssist\\start.bat')
        agent_assist_path = ['call',
                             os.path.realpath(file_path)]
        path = ' '.join(agent_assist_path)
        # 计算机\HKEY_USERS\***\Software\Microsoft\Windows\CurrentVersion\Run
        if not AgentAssistInstallHandle._mod_registry(name, path):
            return False
        return True

    @staticmethod
    def execute_cmd_linux(str_command):
        try:
            process = subprocess.Popen(str_command, shell=False,
                                       stdout=subprocess.PIPE,
                                       stderr=subprocess.PIPE,
                                       encoding="utf-8")
            process.wait()
            ret_code = process.returncode
            res = process.stdout.read()
            process.stdout.close()
            return ret_code, res
        except Exception as err:
            logging.error(f"Exec command failed, error:{err}.")
            return -1, ""

    def _is_least_ubuntu18(self):
        """
        判断当前系统是否是Ubuntu18以上版本（包含Ubuntu18）
        """
        config_file = r'/etc/os-release'
        if os.path.exists(config_file):
            os_type = self._read_conf_without_section(config_file,
                                                      'NAME').replace('"', '')
            os_ver = self._read_conf_without_section(config_file,
                                                     'VERSION_ID').replace('"',
                                                                           '')
            if os_type == "Ubuntu" and float(os_ver) >= 18:
                logging.info('The os type is Ubuntu, os version is least 18.')
                return True
            else:
                logging.info('The os type or os version are not met.')
                return False
        else:
            logging.error('The os_info config_file not exist.')
        return False

    def _get_auto_start_file(self):
        system_file_rcd_local = '/etc/rc.d/rc.local'
        system_file_boot_local = '/etc/init.d/boot.local'
        system_file_rc_local = '/etc/rc.local'
        link_cmd = 'ln -fs /lib/systemd/system/rc-local.service /etc/systemd/system/rc-local.service'
        if os.path.exists(system_file_rcd_local):
            system_file = system_file_rcd_local
            new_path = '/etc/rc.d/rc.local.bak'
        elif os.path.exists(system_file_boot_local):
            system_file = system_file_boot_local
            new_path = '/etc/init.d/boot.local.bak'
        elif os.path.exists(system_file_rc_local):
            system_file = system_file_rc_local
            new_path = '/etc/rc.local.bak'
            if self._is_least_ubuntu18():
                self.execute_cmd_linux(shlex.split(link_cmd))
        else:
            # 针对ubuntu 18.0 以上版本，开机自启动方式不一样，需要增加软连接和rc.loacal配置文件（系统默认没有）
            if self._is_least_ubuntu18():
                self.execute_cmd_linux(shlex.split(link_cmd))
                with open(system_file_rc_local, 'w') as f:
                    f.write('#!/bin/bash\n\n')
                system_file = system_file_rc_local
                new_path = '/etc/rc.local.bak'
                logging.info('Success to make link and rc.local autostart file.')
            else:
                logging.error('Failed to get auto start file.')
                return False, '', ''

        return True, system_file, new_path

    def _set_auto_start_linux(self):
        find_key = 'AgentAssist'
        # monitor.py的路径
        file_path = os.path.join(self._install_path,
                                 'AgentAssist/start.sh')
        msg = f'bash {file_path} & \n'  #
        # python的路径
        ret, system_file, new_path = self._get_auto_start_file()
        if not ret:
            return False
        if not self.write_file(new_path, system_file, msg, find_key):
            return False
        return True

    def write_file(self, src_file, des_file, msg, key):
        auto_start_exists = False
        with open(des_file, 'r') as fr, open(src_file, 'w') as g:
            for line in fr.readlines():
                if 'exit 0' == line.strip():
                    g.write(msg)
                    g.write(line)
                    auto_start_exists = True
                    break
                elif key not in line:
                    g.write(line)
                else:
                    g.write(msg)
                    auto_start_exists = True
            if not auto_start_exists:
                g.write(msg)
            try:
                shutil.move(src_file, des_file)
                logging.info('Succeeded in adding automatic startup.')
                self._mod_chmod(des_file, PERMISSION_700)
            except Exception as err:
                logging.error(f'Failed to add automatic startup, err: {err}.')
                return False
            return True

    @rollback
    def _set_auto_start(self):
        if self._cur_system == SYSTEM_WINDOWS:
            self._set_auto_start_windows()
        elif self._cur_system == SYSTEM_LINUX:
            self._set_auto_start_linux()
        else:
            print(f"Not supported by {self._cur_system} systems.")
            logging.error(f'The {self._cur_system} system is not supported.')
            return False
        return True

    def _add_agent_id(self):
        if self._cur_system == SYSTEM_WINDOWS:
            file_name = 'C:\\.hostID'
        else:
            file_name = '/etc/.hostID'

        try:
            if os.path.exists(file_name):
                with open(file_name, 'r') as f:
                    r = f.readlines()
                    return True, r[0]
            else:
                agent_uuid = str(uuid.uuid4())
                with open(file_name, 'w') as f:
                    f.write(agent_uuid)
                    return True, agent_uuid
        except Exception as err:
            logging.error(f'Failed to add agent ID. {err}.')
        return False, ""

    # Write the common.conf file
    @rollback
    def _write_common_conf(self, language, ip, activation_code):
        file_path = os.path.join(self._install_path,
                                 'AgentAssist/conf/common.conf')
        try:
            if os.path.isfile(file_path):
                os.remove(file_path)
            conf.read(file_path)
            conf.add_section("input_info")
            conf.set("input_info", "language", f"{language}")
            conf.set("input_info", "ha_address", f"{ip}")
            conf.set("input_info", "activation_code", f"{activation_code}")
            conf.set("input_info", "install_path", f"{self._install_path}")
            conf.set("input_info", "python_path", f"{self._python_path}")
            with open(file_path, 'w') as f:
                conf.write(f)
            logging.info('The common.conf is updated successfully.')
        except Exception as err:
            logging.error(f'Failed to update common.conf. {err}.')
            return False
        return True

    def _save_path_con(self):
        path = os.path.realpath(__file__).split('scripts')[0]
        conf_path = os.path.realpath(os.path.join(path, 'conf'))
        file_path = os.path.join(self._install_path,
                                 'AgentAssist/conf/common.conf')
        self.remove_path(conf_path)
        os.makedirs(conf_path, mode=0o700, exist_ok=True)
        shutil.copy(file_path, conf_path)

    # write agent_list.json file
    # agent id, assist: uuid，sub agent machine code
    @rollback
    def _write_agent_list_json(self):
        file_path = os.path.join(self._install_path,
                                 'AgentAssist/conf/agent_list.json')
        ret, agent_uuid = self._add_agent_id()
        if not ret:
            return False
        version_file = os.path.join(self._install_path,
                                    'AgentAssist/conf/VersionDetails')
        if not os.path.exists(version_file):
            return False
        assist_name = self._read_conf_without_section(version_file,
                                                      'Assist_Name')
        if not self._check_conf_parameter(assist_name):
            logging.error('The assist name parameter illegal.')
            return False
        assist_version = self._read_conf_without_section(version_file,
                                                         'Package_Version')
        if not self._check_conf_parameter(assist_version):
            logging.error('The assist version parameter illegal.')
            return False
        info_dict = {
            "agent_list": [{
                "file_name": {
                    "agent_name": assist_name,  # agent install status
                    "agent_version": assist_version,  # agent version
                    "agent_id": agent_uuid
                }
            }]
        }

        try:
            json_str = json.dumps(info_dict)
            with open(file_path, "a", encoding='utf8') as json_file:
                json_file.write(json_str)
            logging.info('The agent_list.json is updated successfully.')
        except Exception as err:
            logging.error(f'Failed to update agent_list.json. {err}.')
            return False
        return True

    @staticmethod
    def _check_process(process_name):
        """
        stop a certain process by name
        :param process_name: process name to be killed
        :return: none
        """
        for process in psutil.process_iter():
            try:
                if len(process.cmdline()) > 1 \
                        and re.search(process_name.replace('.', '\\.') + '$',
                                      process.cmdline()[-1]) \
                        and not re.search('stop\\.py', process.cmdline()[-2]):
                    logging.info(f'{process_name} : {process.cmdline()}, '
                                 f'pid : {process.pid}')
                    return True
            except Exception as err:
                logging.error(f'Failed to check {process_name}. {err}.')
        return False

    def _copy_log_file(self):
        new_path = os.path.join(self._install_path, 'AgentAssist/log')

        try:
            shutil.copy(self._log_path, new_path)
            # AgentAssist_install.log
            file_path = os.path.join(new_path, self._agent_assist_install_log)
            # 640
            self._mod_chmod(file_path, PERMISSION_600)
            os.remove(self._log_path)
            logging.info('Succeeded in copying AgentAssist_install.log.')
        except Exception as err:
            logging.error(f'Failed to copy AgentAssist_install.log. {err}.')
            return False
        return True

    def install_handle(self, address, activation_code, language):
        if any((AgentAssistInstallHandle._check_process('rdagent.py'),
                AgentAssistInstallHandle._check_process('monitor.py'),
                AgentAssistInstallHandle._check_process('log_package.py'))):
            print("The client assistant has been installed on the computer.")
            logging.error(
                'The client assistant has been installed on the computer.')
            return False

        if not self._unzip_pkg():
            print("Failed to install the client assistant.")
            return False

        if not self._set_auto_start():
            return False

        if not self._write_common_conf(language, address, activation_code):
            return False
        self._save_path_con()
        # 保存用户输入参数
        if not self._write_agent_list_json():
            return False

        if not self._create_config():
            return False

        self._copy_log_file()
        return True


if __name__ == "__main__":
    if len(sys.argv) <= 4:
        logging.error('The input parameters are incorrect.')
        sys.exit(1)
    install = AgentAssistInstallHandle(sys.argv[4])
    if install.install_handle(sys.argv[1], sys.argv[2], sys.argv[3]):
        sys.exit(0)
    sys.exit(1)
