#Metasploit 工具类
import pexpect
import re
import threading
import json
from tools.ToolBase import ToolBase

class MsfconsoleTool(ToolBase):
    '''
    Metasploit 这样的交互工具，保持一个会话，析构时结束会话
    '''
    def __init__(self,TM):
        super().__init__(TM)
        self.bOK = False
        self.msf_lock = threading.Lock()  # 线程锁

    def __del__(self):
        if self.bOK:
            self.bOK = False
            print("Metasploit Exit!")

    def validate_instruction(self, instruction):
        timeout = 0
        modified_code = ""
        #针对有分号的指令情况，一般是一行
        if ";" in instruction: #举例：msfconsole -q -x "use exploit/unix/ftp/vsftpd_234_backdoor; set RHOST 192.168.204.137; exploit"
            # 正则表达式匹配双引号内的内容
            pattern = r'"([^"]*)"'
            match = re.search(pattern, instruction)
            if match:
                modified_code = match.group(1)
                #统一把单行，换成多行内容
                modified_code = modified_code.replace(";", "\n")
                print(modified_code)
        else: #针对多行命令格式，
            '''
            msfconsole
            use auxiliary/scanner/smtp/smtp_vuln_cve2011_1764
            set RHOSTS 58.216.217.67
            run
            '''
            lines = instruction.splitlines()
            if(len(lines) > 1):
                modified_code = "\n".join(lines[1:])    #去除第一行的msfconsole命令字符
            else:   #其他情况，暂不兼容
                modified_code = ""
        #print(modified_code)
        return modified_code,timeout

    def parse_msf_output(self,output):
        """
        解析MSF运行结果
        :param output: run命令的输出文本
        :return: 结构化结果字典
        """
        result = {
            "vulnerable": False,
            "details": [],
            "raw_output": output
        }

        # 匹配漏洞状态（示例正则，需根据实际模块输出调整）
        vuln_pattern = re.compile(r'$\+$\s+(Vulnerable|Exploit\s+succeeded)')
        if vuln_pattern.search(output):
            result["vulnerable"] = True

        # 提取关键详情（如Banner信息）
        detail_pattern = re.compile(r'$\*$\s+(.*?)\n')
        result["details"] = detail_pattern.findall(output)
        #转换成字符串
        result = json.dumps(result)
        return result

    def parse_exploit_output(self, output):
        """
        解析 exploit 命令的输出，判断 exploit 是否成功
        判断依据：
         - 如果输出中包含 "failed" 则认为 exploit 失败
         - 如果输出中出现 "opened session" 或 "shell"（大小写不敏感），则认为 exploit 成功
        """
        low_output = output.lower()
        if "failed" in low_output:
            return False, output
        if re.search(r'opened\s+session', low_output) or re.search(r'command\s+shell\s+session', low_output) \
                or re.search(r'found\s+shell',low_output):
            return True, output
        # 如无明显标志，可视为失败
        return False, output

    def remove_ansi_escape_sequences(self,text):
        # 这个正则表达式匹配 ANSI 转义序列
        ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')
        return ansi_escape.sub('', text)

    def execute_msf(self,msf,instruction_old):
        #创建扩展参数-局部变量
        ext_params = self.create_extparams()
        #with self.msf_lock:  # 指令执行解释才能交另一个线程处理--如果有的话
        #print("开始执行-----")
        instruction,time_out = self.validate_instruction(instruction_old)
        if not instruction:
            ext_params.is_user = True
            return False,instruction_old,"该指令暂不执行！","",ext_params

        cmds = [cmd.strip() for cmd in instruction.splitlines() if cmd.strip()]
        results = {}
        for cmd in cmds:
            msf.sendline(cmd)
            # 等待指令执行完成（匹配提示符或超时） #这里要不要用try待验证
            index = msf.expect([pexpect.TIMEOUT,'msf6','Command shell session','Exploit completed',
                                     'exploit failed'],timeout=5*60)
            if index == 0:#超时
                print(f"{cmd}执行超时！")
            # 捕获输出
            output = self.remove_ansi_escape_sequences(msf.before)
            #提取输出信息
            if cmd.startswith("use") and "Failed to load module" in output:
                ext_params.is_user = True
                return True, instruction, "Failed to load module", "",ext_params

            if cmd == "run":
                results = self.parse_msf_output(output)
            elif cmd == "exploit":
                success = False
                if index != 2: #执行成功，建立了会话
                    success,exploit_output = self.parse_exploit_output(output)
                    if not success:
                        return True, instruction, f"Exploit 执行失败--{exploit_output}", exploit_output,ext_params
                if index ==2 or success:#暂时认定commadn shell session 就认定是是利用成功
                    # 如果 exploit 成功，发送 background 命令以返回 msfconsole 提示符
                    ext_params["is_vulnerability"] = True      #识别出弱点
                    # msf.sendline("background") #判断成功是因为需要提交background
                    # try:
                    #     msf.expect("msf6", timeout=60)
                    # except pexpect.TIMEOUT:
                    #     print("****msf退出子会话失败****")
                    #     #return False, instruction, "发送 background 命令超时", ""
                    results = "exploit 利用成功！"
                else:
                    results = f"exploit_output:{output}"

        # 第三步：分析执行结果
        analysis = self.analyze_result(results,instruction,"","")
        return True,instruction, analysis,results,ext_params


    #对于非sh命令调用的工具，自己实现命令执行的内容
    def execute_instruction(self, instruction_old):
        try:
            # 启动msfconsole进程
            msf = pexpect.spawn('msfconsole', encoding='utf-8', timeout=120)
            msf.expect('msf6')
            #执行msf指令
            res,instr, result,source_result,ext_params = self.execute_msf(msf,instruction_old)
            #结束msf
            msf.sendline('exit')
            msf.close()
            #返回结果
            return res,instr, result,source_result,ext_params
        except pexpect.TIMEOUT:
            print("[错误] msfconsole启动超时，可能原因：环境配置错误或资源不足")
            msf.close()
        except pexpect.EOF:
            print("[错误] msfconsole进程异常终止")
            if msf:
                msf.close()

        ext_params = self.create_extparams()
        ext_params.is_user = True  # 提交人工确认异常
        return False, instruction_old, "Metasploit 未成功运行！", "", ext_params

    def analyze_result(self, result,instruction,stderr,stdout):
        #
        return result