zfsafe/InstructionManager.py

'''
对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,TM):
        self.TM  = TM   #任务类对象穿透
        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(self.TM)    #类对象穿透
                except ImportError as e:
                    print(f"加载工具 {module_name} 失败：{str(e)}")


    # 执行指令
    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"分析结果：{result}")
            print(f"*****指令：{instr}，执行完毕")
        else:
            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

    #过来指令：合规、判重、待执行等
    def _instruction_filter(self,instruction):
        pass

    #根据指令执行结果，重新构建提示词
    def _fetch_prompt(self):
        pass

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)