zfsafe/tools/PythonTool.py

#python代码动态执行

from tools.ToolBase import ToolBase

class PythonTool(ToolBase):
    def validate_instruction(self, instruction):
        #指令过滤
        timeout = 0
        return "",timeout

    def execute_instruction(self, instruction_old):
        '''
        执行指令：验证合法性 -> 执行 -> 分析结果
        :param instruction_old:
        :return:
            bool:true-正常返回给大模型，false-结果不返回给大模型
            str:执行的指令
            str:执行指令的结果
        '''
        ext_params = self.create_extparams()
        # 定义允许的内置函数集合
        allowed_builtins = {
            "abs": abs,
            "all": all,
            "any": any,
            "bool": bool,
            "chr": chr,
            "dict": dict,
            "float": float,
            "int": int,
            "len": len,
            "list": list,
            "max": max,
            "min": min,
            "print": print,
            "range": range,
            "set": set,
            "str": str,
            "sum": sum,
            "type": type,
            # 根据需要可以添加其他安全的内置函数
        }
        # 第一步：验证指令合法性
        instruction,time_out = self.validate_instruction(instruction_old)
        if not instruction:
            return False, instruction_old, "该指令暂不执行！","",ext_params
        # 过滤修改后的指令是否需要判重？同样指令再执行结果一致？待定---#？

        # 第二步：执行指令
        output = ""
        # 构造安全的全局命名空间，只包含我们允许的 __builtins__
        safe_globals = {
            "__builtins__": allowed_builtins,
        }
        try:
            # 编译代码
            code_obj = compile(instruction, filename="<dynamic>", mode="exec")
            # 在限制环境中执行代码
            exec(code_obj, safe_globals)
        except Exception as e:
            print(f"执行动态代码时出错: {e}")


        # 第三步：分析执行结果
        analysis = self.analyze_result(output, instruction,"","")
        # 指令和结果入数据库
        # ？
        if not analysis:  # analysis为“” 不提交LLM
            return False, instruction, analysis,"",ext_params
        return True, instruction, analysis,"",ext_params

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

if __name__ == "__main__":
    llm_code = """
    def run_test():
        return 'Penetration test executed successfully!'
    """