''' 对LLM回复的指令进行管理(过滤判重等) 1.curl--保存到待用户确认类 2.nmap 执行若是增加端口要考虑是否有效? 负责指令的执行和针对执行子进程的管理 1.suprocess 执行指令,对返回结果进行提炼 2.实时显示suprocess的状态,可以人工干预停止 ''' import subprocess import importlib import os from tools.ToolBase import ToolBase from myutils.ReturnParams import ReturnParams class InstructionManager: def __init__(self): self.tool_registry = {} # 安全工具list self.load_tools() # 加载工具类 def init_data(self): pass # 动态加载工具类 def load_tools(self): tools_dir = "tools" if not os.path.exists(tools_dir): os.makedirs(tools_dir) return for filename in os.listdir(tools_dir): if filename.endswith(".py") and not filename.startswith("__"): module_name = filename[:-3] try: module = importlib.import_module(f"{tools_dir}.{module_name}") cls = getattr(module,module_name) if(issubclass(cls, ToolBase) and cls != ToolBase): tool_name = module_name.lower()[:-4] self.tool_registry[tool_name] = cls() except ImportError as e: print(f"加载工具 {module_name} 失败:{str(e)}") def get_tool_out_time(self,tool_name): tool = self.tool_registry[tool_name] out_time = tool.get_time_out() return out_time # 执行指令 def execute_instruction(self,instruction): bres = False instr = None result = None source_result = None ext_params = None tool_name_tmp = instruction.split()[0] # 提取工具名称 tool_name = tool_name_tmp.replace("-","") tool_name = tool_name.replace("_", "") # 检查是否存在对应工具 if tool_name in self.tool_registry: tool = self.tool_registry[tool_name] #print(f"*****开始执行指令:{instruction}") bres,instr, result,source_result,ext_params = tool.execute_instruction(instruction) #print(f"*****指令:{instr},执行完毕") else: other_tool = self.tool_registry["other"] #尝试直接执行(处理一些组合指令) bres, instr, result, source_result, ext_params = other_tool.execute_instruction(instruction) # bres = False # instr = instruction #保障后续代码的一致性 # source_result = result = f"{tool_name}-该工具暂不支持" # ext_params = ReturnParams() # ext_params["is_user"] = True # 是否要提交用户确认 -- 默认False # ext_params["is_vulnerability"] = False # 是否是脆弱点 print(f"执行指令:{instr}") print(f"未知工具:{tool_name}") #return bres,instr,result,source_result,ext_params return instr, result, source_result, ext_params #取消bres的返回,所有指令执行结果多需要返回到Llm,用于控制节点的状态 #过来指令:合规、判重、待执行等 def _instruction_filter(self,instruction): pass #根据指令执行结果,重新构建提示词 def _fetch_prompt(self): pass g_instrM = InstructionManager() #全局唯一 if __name__ == "__main__": instrM = InstructionManager() instrS = ['whois haitutech.cn', 'dig haitutech.cn ANY +noall +answer', 'nslookup -query=MX haitutech.cn', 'sublist3r -d haitutech.cn -o subdomains.txt', 'amass enum -d haitutech.cn', 'nmap -sV -p- -T4 --min-rate 1000 haitutech.cn -oN nmap_scan.txt' ] for instr in instrS: bres,instr,reslut =instrM.execute_instruction(instr)