v0.1.1 node_tree_0.4 修改全局提示词前备份

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

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  # 漏洞类型

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}")  #丢弃该节点

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为补充输出信息；

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

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")
+        #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):
         #返回结果

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',

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)

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:

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