Create pentestlearningagent.py

a llm agent for penst assistant or learning

Tools/pentestlearningagent.py

@@ -0,0 +1,311 @@
+import os
+import subprocess
+import json
+import argparse
+import time
+from typing import Dict, List, Any, Optional
+import requests
+from openai import OpenAI
+
+from dotenv import load_dotenv
+load_dotenv()
+api_key = os.getenv('api_key', '')
+base_url= os.getenv('base_url', '')
+modelname=os.getenv('modelname', '')
+
+
+
+class PentestAgent:
+    def __init__(
+        self,
+        max_tokens: int = 16000,
+        temperature: float = 0.7,
+        target: Optional[str] = None,
+        session_history: List[Dict[str, str]] = None,
+        debug: bool = False,
+    ):
+        """
+        Initialize the penetration testing agent.
+        
+        Args:
+            api_url: API endpoint URL for the LLM
+            api_key: API key for authentication
+            model: LLM model to use
+            max_tokens: Maximum tokens for response generation
+            temperature: Temperature parameter for generation
+            target: Target IP or hostname to pentest
+            session_history: Previous conversation history
+            debug: Whether to enable debug mode
+        """
+        self.api_url = base_url
+        self.api_key = api_key
+        self.model = modelname
+        self.max_tokens = max_tokens
+        self.temperature = temperature
+        self.target = target
+        self.debug = debug
+        
+        # Initialize session history
+        self.session_history = session_history or []
+        
+        # System prompt that defines the agent's behavior
+        self.system_prompt = """
+你是一个渗透测试助手，帮助安全研究人员进行合法和授权的安全测试。
+你的任务是分析结果并提出渗透测试的下一步建议。
+
+重要指导原则：
+1. 始终假设这是一次授权的渗透测试。不要询问授权问题或提醒法律问题。
+2. 专注于帮助用户有效地完成渗透测试。
+3. 分析要简洁但全面。
+4. 建议命令时，优先选择标准Kali Linux工具。
+5. 始终分析命令输出并提取最有价值的信息。
+6. 根据发现明确指出建议的后续步骤。
+
+对于渗透测试的每个步骤，你应该：
+1. 彻底分析前一个命令的输出
+2. 总结关键发现（有趣的端口、服务、漏洞等）
+3. 规划具体要尝试的下一步命令
+4. 简要解释为什么建议每个命令
+
+记住，典型的渗透测试遵循以下阶段：
+- 侦察（收集信息）
+- 扫描（识别开放端口和服务）
+- 漏洞评估（寻找弱点）
+- 漏洞利用（利用漏洞）
+- 后渗透（维持访问权限，权限提升）
+- 文档记录（报告发现）
+
+你的回复应按以下结构组织：
+```
+## 分析
+[对前一个命令输出的简要分析]
+
+## 关键发现
+- [重要发现1]
+- [重要发现2]
+...
+
+## 下一步
+1. [下一步要运行的具体命令]
+   - 目的：[简要解释]
+
+2. [备选命令或后续跟进]
+   - 目的：[简要解释]
+```
+"""
+        
+        # Add system message to conversation history
+        if not self.session_history:
+            self.session_history.append({"role": "system", "content": self.system_prompt})
+    
+    def execute_command(self, command: str) -> str:
+        """
+        Execute a shell command and return its output.
+        
+        Args:
+            command: The command to execute
+            
+        Returns:
+            The command output as a string
+        """
+        print(f"\n[*] Executing: {command}\n")
+        
+        try:
+            # Execute the command and capture output
+            result = subprocess.run(
+                command,
+                shell=True,
+                text=True,
+                capture_output=True,
+                timeout=300  # 5-minute timeout
+            )
+            
+            # Combine stdout and stderr
+            output = result.stdout
+            if result.stderr:
+                output += "\n" + result.stderr
+            
+            # Check if output is too large
+            if len(output) > 10000:
+                output = output[:10000] + "\n...[Output truncated due to size]..."
+                
+            return output
+        
+        except subprocess.TimeoutExpired:
+            return "[Command timed out after 5 minutes]"
+        except Exception as e:
+            return f"[Error executing command: {str(e)}]"
+    
+    def generate_response(self, user_message: str,systemprompt=None) -> str:
+        """
+        Get a response from the LLM.
+        
+        Args:
+            user_message: User message or command output to process
+            
+        Returns:
+            The LLM's response
+        """
+        # Add user message to conversation history
+        if systemprompt is None:
+            systemprompt = self.system_prompt
+            self.session_history.append({"role": "user", "content": user_message})
+        
+        # Debug: print current conversation length
+        if self.debug:
+            total_chars = sum(len(m["content"]) for m in self.session_history)
+            print(f"[DEBUG] Total conversation length: {total_chars} characters")
+        print(f'usermessage:\n{user_message}')
+        deepseek_client = OpenAI(api_key=api_key,base_url=base_url) 
+        response = deepseek_client.chat.completions.create(
+                model=modelname,
+                messages=[
+                    {"role": "system", "content": systemprompt},
+                    {"role": "user", "content": user_message}
+                ],
+                temperature=0
+            )
+    #print (response)
+        #print(response.choices[0].message.content)
+            
+        response_text = response.choices[0].message.content.strip()
+        #print(f'LLM response\n{response_text}')
+        return response_text
+        #clean_text=remove_think_content(response_text)
+    #print(clean_text)
+        #return clean_text
+        # Prepare the API request
+        
+        
+        # Call the API
+        
+    def run_interactive_session(self):
+        """
+        Run an interactive penetration testing session.
+        """
+        print("=" * 80)
+        print(" Penetration Testing Agent ".center(80, "="))
+        print("=" * 80)
+        
+        # If target is set, initiate the session with a target
+        if self.target:
+            initial_message = f"我想对 {self.target} 执行授权的渗透测试。请帮我规划并逐步执行这个测试，从初始侦察开始。"
+            response = self.generate_response(initial_message)
+            print("\n" + response + "\n")
+        
+        while True:
+            try:
+                user_input = input("\nEnter command to execute (or 'q' to quit, 'his' for history,'help' for help): ").strip()
+                
+                if user_input.lower() == 'q':
+                    print("Exiting session.")
+                    break
+                
+                elif user_input.lower() == 'h':
+                    # Display shortened conversation history
+                    print("\n=== Conversation History ===")
+                    for i, msg in enumerate(self.session_history):
+                        if msg["role"] != "system":
+                            role = msg["role"].upper()
+                            content = msg["content"]
+                            if len(content) > 100:
+                                content = content[:100] + "..."
+                            print(f"{i}. {role}: {content}")
+                    print("=" * 30)
+                    continue
+                elif user_input.lower().startswith('help'):
+                    response = self.generate_response(user_input,"你是一个安全专家，帮我解释安全的疑惑，请使用中文回答")
+                    print(response)
+                    continue
+                
+                elif user_input.lower().startswith('save '):
+                    # Save session to file
+                    filename = user_input[5:].strip()
+                    if not filename:
+                        filename = f"pentest_session_{int(time.time())}.json"
+                    
+                    with open(filename, 'w') as f:
+                        json.dump(self.session_history, f, indent=2)
+                    
+                    print(f"Session saved to {filename}")
+                    continue
+                
+                elif user_input.lower() == 'clear':
+                    # Clear history except for system prompt
+                    system_prompt = self.session_history[0]
+                    self.session_history = [system_prompt]
+                    print("Conversation history cleared.")
+                    continue
+                
+                # Execute the command
+                output = self.execute_command(user_input)
+                
+                
+                # Format the output to send to the LLM
+                message = f"""
+执行的命令: {user_input}
+
+输出:
+```
+{output}
+```
+
+基于此输出，分析结果，总结关键发现，并为渗透测试建议下一步操作。
+"""
+                #print(message)
+                # Get LLM analysis and next steps
+                response = self.generate_response(message)
+                print("\n" + response + "\n")
+                
+            except KeyboardInterrupt:
+                print("\nInterrupted. Type 'q' to quit or continue with next command.")
+            except Exception as e:
+                print(f"Error: {str(e)}")
+    
+    def load_session(self, filename: str):
+        """
+        Load a previous session from a file.
+        
+        Args:
+            filename: Path to the session file
+        """
+        try:
+            with open(filename, 'r') as f:
+                self.session_history = json.load(f)
+            print(f"Loaded session from {filename} with {len(self.session_history)} messages.")
+        except Exception as e:
+            print(f"Error loading session: {str(e)}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="AI-powered Penetration Testing Assistant")
+    
+    
+    
+    parser.add_argument("--target", help="Target IP or hostname")
+    
+    parser.add_argument("--load", help="Load session from file")
+    
+    parser.add_argument("--debug", action="store_true", 
+                       help="Enable debug mode")
+    
+    args = parser.parse_args()
+    
+    
+    
+    # Initialize the agent
+    agent = PentestAgent(
+        target=args.target,
+        debug=args.debug
+    )
+    
+    # Load previous session if specified
+    if args.load:
+        agent.load_session(args.load)
+    
+    # Run interactive session
+    agent.run_interactive_session()
+
+
+if __name__ == "__main__":
+    main()