import subprocess
import tempfile
import os
from tools.ToolBase import ToolBase

class PsqlTool(ToolBase):
    def validate_instruction(self, instruction):
        #指令过滤
        timeout = 60
        return instruction,timeout

    def analyze_result(self, result,instruction,stderr,stdout):
        #指令结果分析
        return result

    def do_worker_script(self,str_instruction,timeout,ext_params):
        # 创建临时文件保存输出
        with tempfile.NamedTemporaryFile(delete=False) as tmpfile:
            output_file = tmpfile.name

        # 构建并执行 script 命令
        script_cmd = f"script -c '{str_instruction}' {output_file}"
        try:
            result = subprocess.run(script_cmd, shell=True, text=True,timeout=timeout)
            # 读取输出文件内容
            with open(output_file, 'r') as f:
                output = f.read()
                lines = output.splitlines()
                # 跳过第一行（Script started）和最后一行（Script done）
                ftp_output = lines[1:-1]
                output = '\n'.join(ftp_output)
        except subprocess.TimeoutExpired:
            output = "命令超时返回"
            try:
                with open(output_file, 'r') as f:
                    partial_output = f.read()
                if partial_output:
                    output += f"\n部分输出:\n{partial_output}"
            except FileNotFoundError:
                pass  # 文件可能未创建
        except subprocess.CalledProcessError as e:
            output = f"错误: {e}"
        finally:
            # 删除临时文件
            try:
                os.remove(output_file)
            except FileNotFoundError:
                pass  # 文件可能未创建
        return output

    def execute_instruction1(self, instruction_old):
        '''
        执行指令：验证合法性 -> 执行 -> 分析结果
        *****如果指令要做验证，只做白名单，所有逻辑不是全放开就是白名单*****
        :param instruction_old:
        :return:
            bool:true-正常返回给大模型处理下一步，false-结果不返回给大模型,2--需要人工确认的指令
            str:执行的指令
            str:执行指令的结果-解析过滤后的结果--也是提交给LLM的结果
            str:执行指令的结果-原结果
            object:补充参数-封装一个对象： 0-不知是否攻击成功，1-明确存在漏洞，2-明确不存在漏洞
        '''

        ext_params = self.create_extparams()
        # 第一步：验证指令合法性
        instruction,timeout = self.validate_instruction(instruction_old)
        if not instruction:
            ext_params.is_user= True
            return False,instruction_old,"该指令暂不执行！由用户确认是否要兼容支持","",ext_params      #未
        #过滤修改后的指令是否需要判重？同样指令再执行结果一致？待定---#？

        # 第二步：执行指令
        output = self.do_worker_script(instruction,timeout,ext_params)

        # 第三步：分析执行结果
        if isinstance(output,bytes):#若是bytes则转成str
            output = output.decode('utf-8', errors='ignore')

        analysis = self.analyze_result(output,instruction,"","")

        if not analysis:    #analysis为“” 不提交LLM
            ext_params.is_user = True
            return False,instruction,analysis,output,ext_params
        return True,instruction, analysis,output,ext_params