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.
87 lines
3.0 KiB
87 lines
3.0 KiB
2 months ago
|
'''
|
||
|
|
对LLM回复的指令进行管理（过滤判重等）
|
||
|
|
1.curl--保存到待用户确认类
|
||
|
|
2.nmap 执行若是增加端口要考虑是否有效？
|
||
|
|
负责指令的执行和针对执行子进程的管理
|
||
|
|
1.suprocess 执行指令，对返回结果进行提炼
|
||
|
|
2.实时显示suprocess的状态，可以人工干预停止
|
||
|
|
'''
|
||
|
|
import subprocess
|
||
|
|
import importlib
|
||
|
|
import os
|
||
|
|
from tools.ToolBase import ToolBase
|
||
|
|
|
||
|
|
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 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("-","")
|
||
|
|
# 检查是否存在对应工具
|
||
|
|
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:
|
||
|
|
instr = instruction #保障后续代码的一致性
|
||
|
|
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)
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
|