diff --git a/TaskManager.py b/TaskManager.py index a620a68..48be73a 100644 --- a/TaskManager.py +++ b/TaskManager.py @@ -134,7 +134,8 @@ if __name__ == "__main__": current_path = os.path.dirname(os.path.realpath(__file__)) strMsg = FM.read_file("test",1) - test_type = 2 + test_type = 5 + iput_index = 6 # 0是根节点 if test_type == 0: #新目标测试 # 启动--初始化指令 node_list = TM.CCM.start_do("192.168.204.137", 0) @@ -164,7 +165,6 @@ if __name__ == "__main__": with open("attack_tree", "rb") as f: TM.CCM.attack_tree = pickle.load(f) #遍历node，查看有res的数据 - iput_index = 2 #0是根节点 iput_max_num = 1 iput_num = 0 nodes = TM.CCM.attack_tree.traverse_bfs() @@ -191,13 +191,9 @@ if __name__ == "__main__": for t in TM.CCM.llmth_list: t.join() elif test_type ==3: #执行指定指令 - instrlist=[''' -ftp -n 192.168.204.137 << EOF -user anonymous anonymous@example.com -ls -bye -EOF - ''',"wget -m --no-passive ftp://anonymous:anonymous@192.168.204.137/"] + instrlist=[ + "msfconsole -q -x \"use auxiliary/scanner/smb/smb_version; set RHOSTS 192.168.204.137; run; exit\""] + for instr in instrlist: start_time = get_local_timestr() # 指令执行开始时间 bres, instr, reslut, source_result, ext_params = TM.InstrM.execute_instruction(instr) @@ -208,12 +204,27 @@ EOF ext_params, "独立命令执行") else: TM.logger.error("数据库连接失败！！") - elif test_type == 4: #读取messages + elif test_type == 4: #修改Message with open("attack_tree", "rb") as f: TM.CCM.attack_tree = pickle.load(f) + #创建一个新的节点 + from mycode.AttackMap import TreeNode + testnode = TreeNode("test",0) + TM.CCM.LLM.build_initial_prompt(testnode)#新的Message + systems = testnode.messages[0]["content"] + #print(systems) # 遍历node，查看有instr的ndoe nodes = TM.CCM.attack_tree.traverse_bfs() - print(nodes[0].messages) + for node in nodes: + node.messages[0]["content"] = systems + with open("attack_tree", 'wb') as f: + pickle.dump(TM.CCM.attack_tree, f) + elif test_type ==5: #显示指令和结果list + with open("attack_tree", "rb") as f: + TM.CCM.attack_tree = pickle.load(f) + nodes = TM.CCM.attack_tree.traverse_bfs() + print(f"********\n{','.join(nodes[iput_index].instr_queue)}\n********") + print(f"&&&&&&&&\n{','.join(nodes[iput_index].res_quere)}\n&&&&&&&&") else: #完整过程测试---要设定终止条件 pass diff --git a/mycode/AttackMap.py b/mycode/AttackMap.py index 8c0d893..bd38ef3 100644 --- a/mycode/AttackMap.py +++ b/mycode/AttackMap.py @@ -63,7 +63,7 @@ class AttackTree: if child_node.name == node_name: return child_node #走到这说明没有匹配到-则新建一个节点 - newNode = TreeNode(node_name) + newNode = TreeNode(node_name,node.task_id) node.add_child(newNode) return newNode else: @@ -129,7 +129,7 @@ class AttackTree: class TreeNode: - def __init__(self, name, status="未完成", vul_type="未发现"): + def __init__(self, name,task_id,status="未完成", vul_type="未发现"): self.name = name # 节点名称 self.status = status # 节点状态 self.vul_type = vul_type # 漏洞类型 diff --git a/mycode/ControlCenter.py b/mycode/ControlCenter.py index 720ad8d..179a8b3 100644 --- a/mycode/ControlCenter.py +++ b/mycode/ControlCenter.py @@ -51,7 +51,7 @@ class ControlCenter: #创建/初始化测试树 if self.attack_tree: self.attack_tree = None #释放 - root_node = TreeNode(target) + root_node = TreeNode(target,task_id) self.attack_tree = AttackTree(root_node)#创建测试树，同时更新根节点相关内容 #初始化启动提示信息 know_info = self.get_user_init_info() @@ -137,7 +137,8 @@ class ControlCenter: # 构造本次提交的prompt ext_Prompt = f''' 上一步结果：{res_str} -任务：生成下一步渗透测试指令或判断是否完成该节点测试。 +任务：生成下一步渗透测试指令或判断是否完成该节点测试 +请确保生成的指令满足以下 ''' elif llm_type ==2: #llm返回的指令存在问题，需要再次请求返回 ext_Prompt = f''' @@ -256,7 +257,7 @@ class ControlCenter: break if not bfind: #添加节点 - new_node = TreeNode(node_name,status) + new_node = TreeNode(node_name,node.task_id,status) node.add_child(new_node) #message的传递待验证 else: self.logger.error(f"添加子节点时，遇到父节点名称不一致的，需要介入！！{node_json}") #丢弃该节点 diff --git a/mycode/LLMManager.py b/mycode/LLMManager.py index b3818f7..a5fc2f6 100644 --- a/mycode/LLMManager.py +++ b/mycode/LLMManager.py @@ -48,7 +48,7 @@ class LLMManager: # 初始化messages def build_initial_prompt(self,node): if not node: - return None + return '''暂时未添加到提示词 **核心要求**： - 每次新增节点时，必须在同一响应中为该节点提供测试指令。 @@ -71,12 +71,15 @@ class LLMManager: 4. 若服务扫描发现服务版本或漏洞，则新增漏洞测试节点并提供测试指令； 5. 若漏洞验证成功，则根据结果决定是否需要进一步测试，若需要进一步测试，则新增子节点并提供测试指令； 6. 若节点测试无新信息和测试指令，则更新状态为“已完成”。 +**测试指令生成要求** +1.明确每个测试指令的测试目标，并优先尝试最简单、最直接的办法； +2.对于复杂的测试点，使用递进逻辑组织指令：先尝试基础测试方法，根据执行结果决定是否进行更深入的测试。 **节点指令格式** -- 新增节点：{\"action\":\"add_node\", \"parent\": \"80端口\", \"nodes\": \"21端口,80端口,445端口,3306端口,1000端口,http://192.168.1.100/index.php?id=1\", \"status\": \"未完成\"}； +- 新增节点：{\"action\":\"add_node\", \"parent\": \"父节点\", \"nodes\": \"子节点1，子节点2\", \"status\": \"未完成\"}； - 已新增但未生成指令的节点列表：{\"action\": \"no_instruction\", \"nodes\": \"3306端口,1000端口\"}； - 未新增的节点列表：{\"action\": \"no_create\", \"nodes\": \"8080端口,8081端口,9000端口\"}； -- 更新节点未发现漏洞：{\"action\": \"update_status\", \"node\": \"21端口\", \"status\": \"已完成\"}； -- 更新节点发现漏洞：{\"action\": \"update_status\", \"node\": \"21端口\", \"status\": \"已完成\"，\"vulnerability\": {\"name\":\"ftp匿名登录\",\"risk\":\"高\"}}； +- 节点完成测试未发现漏洞：{\"action\": \"update_status\", \"node\": \"21端口\", \"status\": \"已完成\"}； +- 节点完成测试发现漏洞：{\"action\": \"update_status\", \"node\": \"21端口\", \"status\": \"已完成\"，\"vulnerability\": {\"name\":\"ftp匿名登录\",\"risk\":\"高\"}}； **测试指令格式** - shell指令：```bash-[节点路径](.*?)```包裹，需要避免用户交互； - python指令：```python-[节点路径](.*?)```包裹，主函数名为dynamic_fun，需包含错误处理，执行结束后必须返回一个tuple (status, output)，其中status为'success'或'failure'，output为补充输出信息； diff --git a/pipfile b/pipfile index e0be570..1b5b6d3 100644 --- a/pipfile +++ b/pipfile @@ -11,5 +11,8 @@ git clone https://github.com/defparam/smuggler.git pip install beautifulsoup4 pip install cryptography pip install loguru -i https://pypi.tuna.tsinghua.edu.cn/simple/ +pip install paramiko -i https://pypi.tuna.tsinghua.edu.cn/simple/ +pip install impacket -i https://pypi.tuna.tsinghua.edu.cn/simple/ - +cd /usr/share/wordlists/ +gzip -d rockyou.txt.gz \ No newline at end of file diff --git a/tools/HydraTool.py b/tools/HydraTool.py index 930bf6b..3f56a9e 100644 --- a/tools/HydraTool.py +++ b/tools/HydraTool.py @@ -1,37 +1,33 @@ import os import shlex +import re from tools.ToolBase import ToolBase class HydraTool(ToolBase): def validate_instruction(self, instruction): timeout = 0 - #hydra过滤 - #hydra -L emails.txt -P passwords.txt pop3://haitutech.cn 像这样针对邮箱爆破，邮箱名不是用户名，需要特殊处理 - # 分割指令为参数列表 - cmd_parts = shlex.split(instruction) - new_cmd = [] - # 获取当前程序所在目录 current_path = os.path.dirname(os.path.realpath(__file__)) - #new_pass_path = os.path.join(current_path, "payload", "passwords") - new_pass_path = os.path.join(current_path, "../payload", "passwords") - new_user_path = os.path.join(current_path, "../payload", "users") + #hydra过滤 需要判断指令中添加字典文件存不存在 + match_p = re.search(r'-P\s+([^\s]+)', instruction) + match_l = re.search(r'-L\s+([^\s]+)', instruction) + if match_p: + str_p = match_p.group(1) + #判断文件是否存在 + if not os.path.exists(str_p): #文件不存在要替换 + new_pass_path = os.path.join(current_path, "../payload", "passwords") + instruction = instruction.replace(str_p,new_pass_path) + if match_l: + str_l = match_l.group(1) + #判断文件是否存在 + if not os.path.exists(str_l): + new_user_path = os.path.join(current_path, "../payload", "users") + instruction = instruction.replace(str_l, new_user_path) - i = 0 - while i < len(cmd_parts): - part = cmd_parts[i] - new_cmd.append(part) - # 检测到-P参数 - if part == "-P" and i + 1 < len(cmd_parts): #密码 - # 替换下一参数为指定路径 - new_cmd.append(new_pass_path) - i += 1 # 跳过原路径参数 - elif part == "-L" and i + 1 < len(cmd_parts): #用户名 - # 替换下一参数为指定路径 - new_cmd.append(new_user_path) - i += 1 # 跳过原路径参数 - i += 1 + if "-l" in instruction or "-p" in instruction: + if "-f" not in instruction: + instruction = instruction + " -f" #当是单密码，或单用户名时，使用成功即停止模式 - return " ".join(shlex.quote(p) for p in new_cmd),timeout + return instruction,timeout def analyze_result(self, result,instruction,stderr,stdout): #返回结果 diff --git a/tools/MsfconsoleTool.py b/tools/MsfconsoleTool.py index ce20d39..da8ccbe 100644 --- a/tools/MsfconsoleTool.py +++ b/tools/MsfconsoleTool.py @@ -105,6 +105,8 @@ class MsfconsoleTool(ToolBase): cmds = [cmd.strip() for cmd in instruction.splitlines() if cmd.strip()] results = {} for cmd in cmds: + if cmd == "exit": #exit这里不输入 + continue msf.sendline(cmd) # 等待指令执行完成（匹配提示符或超时） #这里要不要用try待验证 index = msf.expect([pexpect.TIMEOUT,'msf6','Command shell session','Exploit completed', diff --git a/tools/PythoncodeTool.py b/tools/PythoncodeTool.py index c93618f..a6384f1 100644 --- a/tools/PythoncodeTool.py +++ b/tools/PythoncodeTool.py @@ -4,6 +4,8 @@ import subprocess import json import builtins import re +import paramiko +import impacket from tools.ToolBase import ToolBase class PythoncodeTool(ToolBase): @@ -29,6 +31,7 @@ class PythoncodeTool(ToolBase): #指令过滤 timeout = 0 instr = instruction.replace("python_code ","") + instr = instruction.replace("python-code ", "") # Safety check if not self.is_safe_code(instr): return "", timeout @@ -91,7 +94,9 @@ class PythoncodeTool(ToolBase): safe_globals = {"__builtins__": allowed_builtins, 'subprocess':subprocess, 'json':json, - 're':re,} + 're':re, + 'paramiko':paramiko, + 'impacket':impacket} safe_locals = {} #不需要预设局部参数 # 在限制环境中执行代码 exec(instruction, safe_globals,safe_locals) diff --git a/tools/SearchsploitTool.py b/tools/SearchsploitTool.py index ddcdda8..7fa25a4 100644 --- a/tools/SearchsploitTool.py +++ b/tools/SearchsploitTool.py @@ -121,7 +121,6 @@ class SearchsploitTool(ToolBase): else: #暂时只利用python脚本 return False - def do_ext(self,filepath,payload_dir) -> str: bdownload = self.do_download(filepath,payload_dir) if bdownload: diff --git a/tools/SmbclientTool.py b/tools/SmbclientTool.py new file mode 100644 index 0000000..8139c27 --- /dev/null +++ b/tools/SmbclientTool.py @@ -0,0 +1,11 @@ +from tools.ToolBase import ToolBase + +class SmbclientTool(ToolBase): + def validate_instruction(self, instruction): + #指令过滤 + timeout = 0 + return instruction,timeout + + def analyze_result(self, result,instruction,stderr,stdout): + #指令结果分析 + return result \ No newline at end of file