import pexpect
import shlex
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_pexpect(self,str_instruction,timeout,ext_params):
        try:
            safe_command = shlex.quote(str_instruction)
            cmd = f"bash -c {safe_command}"
            exc_do = pexpect.spawn(cmd,timeout=timeout,encoding='utf-8')
            index = exc_do.expect([
                pexpect.TIMEOUT,
                pexpect.EOF,
                'Password for user postgres:',
                '用户 postgres 的口令：'
            ])
            #strout = exc_do.before.decode('utf-8', errors='replace')
            strout = exc_do.before
            if index == 0 or index ==1:
                return strout
            elif index ==2 or index==3:
                strout += '用户 postgres 的口令：'
                exc_do.sendline('') #输入空密码后不知道会有多少种情况，密码不对，密码对
                index = exc_do.expect([pexpect.TIMEOUT,pexpect.EOF])
                print(index)
                strout += exc_do.before
                return strout
        except Exception as e:
            return f"执行错误: {str(e)}"

    def execute_instruction(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)
        output = self.do_worker_pexpect(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