You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
102 lines
3.9 KiB
102 lines
3.9 KiB
'''
|
|
对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)
|
|
|
|
|
|
|
|
|
|
|