import re
import os
import sys
from enum import Enum
import yaml


# 路径信息
class FilePath(Enum):
 # 服务启动脚本
 base_script_path = f"/usr/onesec/java-onesec-base/scripts/start-base-service.sh"
 log_script_path = f"/usr/onesec/java-onesec-server/scripts/start-log-service.sh"
 le_script_path = f"/usr/onesec/java-onesec-server/scripts/start-le-service.sh"
 # 引擎配置文件
 engine_path = f"/data/onesec/store/java-onesec-server/le-service/ti_engine/engine.cfg"
 # 病毒库更新配置文件
 bd_conf_path = f"/usr/onesec/bd-update/mirror.cnf"
 # socat命令定期执行脚本
 socat_proxy_scripts = f"/data/onesec/scripts/file_proxy_socat.sh"
 # nginx配置文件
 nginx_conf = f"/etc/onesec/system/nginx/nginx.conf"
 # 测试代理认证的脚本路径
 test_proxy_path = f"/data/onesec/scripts/test_proxy.sh"


# 获取源文件数据列表
def get_file_list(file_path: str):
 """
 :param file_path: 传入源文件路径包含文件名
 :return: 获取成功返回列表,失败返回False,异常
 """
 try:
 if os.path.exists(file_path):
 file_size = os.stat(file_path).st_size
 if file_size <= 1024:
 with open(file_path, "r", encoding='utf-8') as file:
 read_files = file.readlines()
 return True, read_files
 return False, f"The uploaded file is too large, exceeding 1024 bytes"
 return False, f"{file_path} The file or path does not exist"
 except PermissionError:
 return False, f"{str(PermissionError)} Insufficient permissions"
 except RuntimeError:
 return False, f"{str(RuntimeError)} Failed to read"


# 文件中写入内容
def write_data_tofile(file_path, data_list):
 """
 :param file_path: 要写入的文件全路径
 :param data_list: 数据列表
 :return: 返回执行True,结果|False,结果
 """
 try:
 if os.path.exists(file_path):
 with open(file_path, 'w', encoding='utf-8') as file:
 file.writelines(data_list)
 return True, f"写入完成"
 return False, f"{file_path} The file or path does not exist"
 except PermissionError:
 return False, f"{str(PermissionError)} Insufficient permissions"
 except RuntimeError:
 return False, f"{str(RuntimeError)} Failed to read"


# 检测路径是否存在不存在就创建
def mkdir_path(path):
 if not os.path.isdir(path):
 print(f"{path}不存在进行创建...")
 os.makedirs(path, exist_ok=True)


# 去除java参数中代理并重新写入
def remove_java_proxy_argument(file_path):
 """
 :param file_path: 文件全路径
 :return: 清理代理后写入结果,完成返回True,异常返回False和失败原因
 """
 status, get_file_list_result = get_file_list(file_path)
 if status:
 for index in range(len(get_file_list_result)):
 result = re.search(r'^java', get_file_list_result[index])
 if result:
 java_config = index
 break
 java_config = get_file_list_result[index:len(get_file_list_result)]
 del get_file_list_result[index:len(get_file_list_result)]
 for index in range(len(java_config)):
 java_config[index] = java_config[index].rstrip(' \\\n')
 java_config = ' '.join(java_config)
 java_config = java_config.split()
 for index in range(len(java_config)):
 filter_proxy_result = re.search(r'^(-Dhttp|-Dhttps).proxy|-Dhttp.nonProxyHosts', java_config[index])
 if filter_proxy_result is not None:
 java_config[index] = None
 if java_config.count(None) != 5:
 return False, f"未检测到足够的代理条件不做修改"
 java_config = [x for x in java_config if x is not None]
 for index in range(len(java_config)):
 if java_config[index] == '2>&1':
 get_file_list_result[get_file_list_result.index(java_config[index - 1] + ' \\\n')] = java_config[
 index - 1].rstrip('\\\n')
 get_file_list_result.append(' ' + java_config[index])
 break
 if index != 0 and index % 2 == 0:
 if java_config[index] == "-jar":
 get_file_list_result.append(java_config[index] + ' ')
 continue
 get_file_list_result.append(java_config[index] + ' \\\n')
 continue
 if len(java_config[index]) > 38:
 get_file_list_result.append(java_config[index] + ' \\\n')
 continue

 get_file_list_result.append(java_config[index] + ' ')
 result = write_data_tofile(file_path, get_file_list_result)
 return result, f"执行完成"
 return status, get_file_list_result


# java参数中添加代理并重新写入
def add_java_proxy_argument(proxy_host, proxy_port, file_path):
 """
 :param proxy_host: 代理地址
 :param proxy_port: 代理端口
 :param file_path: 文件全路径
 :return: 写入代理配置结果完成返回True和执行结果,异常返回False和异常原因
 """
 status, get_file_list_result = get_file_list(file_path)
 if status:
 for index in range(len(get_file_list_result)):
 result = re.search(r'^java', get_file_list_result[index])
 if result:
 java_config = index
 break
 java_config = get_file_list_result[index:len(get_file_list_result)]
 del get_file_list_result[index:len(get_file_list_result)]
 for index in range(len(java_config)):
 java_config[index] = java_config[index].rstrip(' \\\n')
 java_config = ' '.join(java_config)
 java_config = java_config.split()
 proxy_data = []
 proxy_data.append(f"-Dhttp.proxyHost={proxy_host}")
 proxy_data.append(f"-Dhttp.proxyPort={proxy_port}")
 proxy_data.append(f"-Dhttps.proxyHost={proxy_host}")
 proxy_data.append(f"-Dhttps.proxyPort={proxy_port}")
 proxy_data.append(
 f"-Dhttp.nonProxyHosts=\"127.0.0.1|onesec-redis-service|onesec-kafka-service|onesec-es-service|onesec-log-service|onesec-le-service|onesec-backend-service|onesec-base-service|onesec-monitor-service|onesec-data-service\"")
 java_config[3:3] = proxy_data
 for index in range(len(java_config)):
 if java_config[index] == '2>&1':
 get_file_list_result[get_file_list_result.index(java_config[index - 1] + ' \\\n')] = java_config[
 index - 1].rstrip('\\\n')
 get_file_list_result.append(' ' + java_config[index])
 break
 if index != 0 and index % 2 == 0:
 if java_config[index] == "-jar":
 get_file_list_result.append(java_config[index] + ' ')
 continue
 get_file_list_result.append(java_config[index] + ' \\\n')
 continue
 if len(java_config[index]) > 38:
 get_file_list_result.append(java_config[index] + ' \\\n')
 continue

 get_file_list_result.append(java_config[index] + ' ')
 result, result_data = write_data_tofile(file_path, get_file_list_result)
 return result, result_data
 return status, get_file_list_result


# 获取java参数中的代理配置
def get_java_proxy_host_port(file_path):
 """
 :param file_path: 文件全路径
 :return: 获取到的代理地址和端口,未获取到返回为空列表
 """
 proxy_info = []
 status, get_file_list_result = get_file_list(file_path)
 if status:
 for index in range(len(get_file_list_result)):
 result = re.search(r'^java', get_file_list_result[index])
 if result:
 java_config = index
 break
 java_config = get_file_list_result[index:len(get_file_list_result)]
 del get_file_list_result[index:len(get_file_list_result)]
 for index in range(len(java_config)):
 java_config[index] = java_config[index].rstrip(' \\\n')
 java_config = ' '.join(java_config)
 java_config = java_config.split()
 for index in range(len(java_config)):
 pattern = re.search('-Dhttp.proxyHost=(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})|-Dhttp.proxyPort=(\d{2,5})',
 java_config[index])
 if pattern is not None:
 if pattern.group(1) is not None:
 proxy_info.append(pattern.group(1))
 if pattern.group(2) is not None:
 proxy_info.append(pattern.group(2))
 return proxy_info


# 获取引擎engine配置,检测代理是否存在不存在则写入
def add_engine_proxy(proxy_host, proxy_port, file_path):
 """
 :param proxy_host: 代理地址
 :param proxy_port: 代理端口
 :param file_path: 文件全路径
 :return: 写入代理配置结果完成返回True和执行结果,异常返回False和异常原因
 """
 status, get_file_list_result = get_file_list(file_path)
 temp_num = 0
 if status:
 for index in range(len(get_file_list_result)):
 engine_proxy = [f"proxy_type http", f"proxy_addr {proxy_host}", f"proxy_port {proxy_port}"]
 result = re.search(r'^proxy_type|^proxy_addr|^proxy_port', get_file_list_result[index])
 if result is not None:
 temp_num = temp_num + 1
 if temp_num == 0:
 get_file_list_result.append(engine_proxy)
 result_status, result_data = write_data_tofile(file_path, get_file_list_result)
 return result_status, result_data
 return True, f"数据已存在不写入"
 return status, get_file_list_result


# 获取病毒库代理配置,检测代理不存在则写入
def add_bd_proxy(proxy_host, proxy_port, file_path):
 """
 :param proxy_host: 代理地址
 :param proxy_port: 代理端口
 :param file_path: 文件全路径
 :return: 写入代理配置结果完成返回True和执行结果,异常返回False和异常原因
 """
 status, get_file_list_result = get_file_list(file_path)
 temp_num = 0
 if status:
 for index in range(len(get_file_list_result)):
 result = re.search(r'^USE_PROXY', get_file_list_result[index])
 if result is not None:
 temp_num = temp_num + 1
 if temp_num == 0:
 result_status, result_data = write_data_tofile(file_path, get_file_list_result)
 return result_status, result_data
 return True, f"数据已存在不写入"
 return status, get_file_list_result


# 获取nginx配置, 未配置代理替换成本地代理地址
def add_nginx_proxy(proxy_host, proxy_port, file_path):
 status, get_file_list_result = get_file_list(file_path)
 if status:
 for index in range(len(get_file_list_result)):
 result = re.search(r'server\safsa.onesec.net:443;', get_file_list_result[index])
 if result is not None:
 get_file_list_result[index] = f" server 127.0.0.1:443;"
 break
 result_status, result_data = write_data_tofile(file_path, get_file_list_result)
 return result_status, result_data
 return status, get_file_list_result


# 写入定时任务-重复操作会直接覆盖原有内容
def add_crontab_file(crontab_list):
 data = []
 cron_path = '/var/spool/cron/root'
 for i in crontab_list:
 data.append(i)
 mkdir_path(os.path.dirname(cron_path))
 result_status, result_data = write_data_tofile(cron_path, data)
 return result_status, result_data


# 配置代理时定时任务用的脚本-必须
def write_data_toscript(proxy_user=None, proxy_password=None, proxy_host=None, proxy_port=None, listen_port=None):
 """
 # 不传入proxy_user和proxy_password /data/onesec/scripts/test_proxy.sh 不写入内容
 :param proxy_user: 代理使用账号名
 :param proxy_password: 代理使用的密码
 :param proxy_host: 代理地址
 :param proxy_port: 代理地址端口
 :param listen_port: socat监听端口
 :return: 返回写入结果,和具体写入内容日志
 """
 socat_script_path = FilePath.socat_proxy_scripts
 mkdir_path(os.path.dirname(socat_script_path))
 data = []
 script_content_1 = f'''#!/bin/bash
 COUNT=$(ps -ef | grep socat | grep tcp4-listen:21443 | wc -l)
 [[ $COUNT -eq 0 ]] && {{
 echo will start proxy...
 setsid nohup socat tcp4-listen:{listen_port},reuseaddr,fork proxy:{proxy_host}:{proxy_port},proxyport={proxy_port} > /dev/null &
 }}
 '''
 data.append(script_content_1)
 if proxy_user is not None and proxy_password is not None:
 test_proxy_script_path = FilePath.test_proxy_path
 mkdir_path(os.path.dirname(test_proxy_script_path))
 script_content_2 = f'''#!/bin/bash
 https_proxy=http://{proxy_user}:{proxy_password}@{proxy_host}:{proxy_port} curl https://static-js.threatbook.cn
 '''
 data.append(script_content_2)
 result_status, result_data = write_data_tofile(socat_script_path, data)
 return result_status, result_data


# 读取 yaml 中的文件
def get_yaml_data(yaml_file_path):
 if os.path.exists(yaml_file_path):
 try:
 with open(yaml_file_path, 'r', encoding='utf-8') as f:
 data = yaml.safe_load(f)
 except yaml.error.YAMLError as e:
 return False, f"配置文件格式错误{e}"
 return True, data
 return False, f"engine_cfg_proxy_host"


# 获取代理中的配置
def get_proxy_config(yaml_file_path, query_data):
 result_status, result_data = get_yaml_data(yaml_file_path)
 if result_status:
 if result_data["proxy"]["enable"]:
 try:
 return True, result_data["proxy"][query_data]
 except KeyError as e:
 return False, f"索引不存在 {KeyError}"
 return False, f"未检测到代理配置"
 return False, result_data


# 验证端口是否被占用
def is_port_in_use(port):
 shell_command = f"""#!/bin/bash
 num = ss -tln | grep -c ':{port}'
 if [ $num == 0 ];then
 exit 1
 fi
 exit 0
 """
 result = run_command(shell_command)
 return result


# 验证进程是否存在
def is_process_running_by_name(process_name):
 shell_command = ("""#!/bin/bash
 num = ps -ef | grep {process_name}| grep -v grep | wc -l
 if [ $num != 1 ];then
 exit 1
 fi
 exit 0
 """)
 result = run_command(shell_command)
 return result


# 执行命令
def run_command(command):
 """
 :param command: 命令参数
 :return: 执行状态码,成功 0 不成功 -1
 """
 result = os.system(command)
 return result


if __name__ == "__main__":
 if sys.platform == "linux":
 print(f"[ERROR] Only allowed to run on Linux systems")
 exit(1)

temp = './config.yaml'
status, data = get_proxy_config(temp, "enable")
if status:
 if data:
 print("start setting proxy...")
 # 修改java配置文件的代理
 java_proxy_host = get_proxy_config(temp, 'java_proxy_host')
 java_proxy_port = get_proxy_config(temp, 'java_proxy_port')
 if java_proxy_host and java_proxy_port:
 result = add_java_proxy_argument(java_proxy_host[1], java_proxy_port[1],
 FilePath.base_script_path.value)
 print(
 f"setting server proxy {FilePath.base_script_path.value}{(lambda result: 'Success' if result[0] else 'Failure')(result)}: {result[1]}")
 result = add_java_proxy_argument(java_proxy_host[1], java_proxy_port[1],
 FilePath.le_script_path.value)
 print(
 f"setting server proxy {FilePath.le_script_path.value}{(lambda result: 'Success' if result[0] else 'Failure')(result)}: {result[1]}")
 result = add_java_proxy_argument(java_proxy_host[1], java_proxy_port[1],
 FilePath.log_script_path.value)
 print(
 f"setting server proxy {FilePath.log_script_path.value}{(lambda result: 'Success' if result[0] else 'Failure')(result)}: {result[1]}")
 # 设置引擎engine的代理
 engine_cfg_proxy_host = get_proxy_config(temp, 'engine_cfg_proxy_host')
 engine_cfg_proxy_port = get_proxy_config(temp, 'engine_cfg_proxy_port')
 if engine_cfg_proxy_host and engine_cfg_proxy_port:
 result = add_engine_proxy(engine_cfg_proxy_host[1], engine_cfg_proxy_port[1],
 FilePath.engine_path.value)
 print(
 f"setting engine proxy {FilePath.engine_path.value}{(lambda result: 'Success' if result[0] else 'Failure')(result)}: {result[1]}")

 bd_mirror_proxy_host = get_proxy_config(temp, 'bd_mirror_proxy_host')
 bd_mirror_proxy_port = get_proxy_config(temp, 'bd_mirror_proxy_port')
 # 设置BD病毒库配置代理
 if bd_mirror_proxy_host and bd_mirror_proxy_port:
 result = add_bd_proxy(bd_mirror_proxy_host[1], bd_mirror_proxy_port[1],
 FilePath.bd_conf_path.value)
 print(
 f"setting bd proxy {FilePath.bd_conf_path.value}{(lambda result: 'Success' if result[0] else 'Failure')(result)}: {result[1]}")
 print("start setting crontab...")
 crontab_list = get_proxy_config(temp, 'crontab_list')
 # 添加crontab定时任务
 result = add_crontab_file(crontab_list[1])
 print(f"setting crontab {result[0]}:{result[1]}")
 # 写入定时任务使用的脚本