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.
86 lines
2.9 KiB
86 lines
2.9 KiB
import openai
|
|
import subprocess
|
|
|
|
# 设置OpenAI API密钥
|
|
openai.api_key = "fGBYaQLHykBOQsFwVrQdIFTsYr8YDtDVDQWFU41mFsmvfNPc"
|
|
|
|
# 初始化阶段和已知信息
|
|
current_stage = "信息收集"
|
|
known_info = {"url": "www.test.com"}
|
|
results = [] # 存储所有任务结果，用于生成报告
|
|
|
|
|
|
# 构建初始提示
|
|
def build_initial_prompt(stage, known_info):
|
|
return f"你是一个渗透测试专家，正在对网站{known_info['url']}进行渗透测试。当前阶段是{stage}，已知信息是{known_info}。请生成下一步的指令。"
|
|
|
|
|
|
# 构建反馈提示
|
|
def build_feedback_prompt(instruction, result):
|
|
return f"执行指令“{instruction}”的结果是“{result}”。请根据这个结果生成下一步的指令。"
|
|
|
|
|
|
# 调用LLM生成指令
|
|
def get_llm_instruction(prompt):
|
|
response = openai.Completion.create(
|
|
engine="deepseek-r1", # 替换为你的模型
|
|
prompt=prompt,
|
|
max_tokens=100
|
|
)
|
|
return response.choices[0].text.strip()
|
|
|
|
|
|
# 执行指令
|
|
def execute_instruction(instruction):
|
|
# 示例：支持Nmap和dirb指令
|
|
if "nmap" in instruction:
|
|
try:
|
|
result = subprocess.run(instruction, shell=True, capture_output=True, text=True)
|
|
return result.stdout if result.stdout else result.stderr
|
|
except Exception as e:
|
|
return f"执行失败：{str(e)}"
|
|
elif "dirb" in instruction:
|
|
try:
|
|
result = subprocess.run(instruction, shell=True, capture_output=True, text=True)
|
|
return result.stdout if result.stdout else result.stderr
|
|
except Exception as e:
|
|
return f"执行失败：{str(e)}"
|
|
else:
|
|
return "未知指令，请重新生成。"
|
|
|
|
|
|
# 主循环
|
|
while current_stage != "报告生成":
|
|
# 构建提示并获取指令
|
|
if not results: # 第一次执行
|
|
prompt = build_initial_prompt(current_stage, known_info)
|
|
else: # 反馈结果
|
|
prompt = build_feedback_prompt(last_instruction, last_result)
|
|
|
|
instruction = get_llm_instruction(prompt)
|
|
print(f"生成的指令：{instruction}")
|
|
|
|
# 执行指令
|
|
task_result = execute_instruction(instruction)
|
|
print(f"任务结果：{task_result}")
|
|
results.append({"instruction": instruction, "result": task_result})
|
|
|
|
# 更新变量
|
|
last_instruction = instruction
|
|
last_result = task_result
|
|
|
|
# 示例阶段更新逻辑（可根据实际结果调整）
|
|
if current_stage == "信息收集" and "开放端口" in task_result:
|
|
current_stage = "漏洞扫描"
|
|
known_info["ports"] = "80, 443" # 示例更新已知信息
|
|
elif current_stage == "漏洞扫描" and "扫描完成" in task_result:
|
|
current_stage = "漏洞利用"
|
|
# 添加更多阶段切换逻辑
|
|
|
|
# 生成测试报告
|
|
report = "渗透测试报告\n"
|
|
report += f"目标网站：{known_info['url']}\n"
|
|
report += "测试结果：\n"
|
|
for res in results:
|
|
report += f"指令：{res['instruction']}\n结果：{res['result']}\n\n"
|
|
print(report)
|