mirror of
https://github.com/SunZhimin2021/AIPentest.git
synced 2025-06-21 10:21:03 +00:00
312 lines
10 KiB
Python
312 lines
10 KiB
Python
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()
|