16 changed files with 589 additions and 275 deletions
@ -1,86 +1,72 @@ |
|||
import openai |
|||
import re |
|||
import subprocess |
|||
import tempfile |
|||
import os |
|||
import pexpect |
|||
|
|||
# 设置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: |
|||
def do_worker(str_instruction): |
|||
try: |
|||
result = subprocess.run(instruction, shell=True, capture_output=True, text=True) |
|||
return result.stdout if result.stdout else result.stderr |
|||
# 使用 subprocess 执行 shell 命令 |
|||
result = subprocess.run(str_instruction, shell=True, text=True,capture_output=True) |
|||
|
|||
return { |
|||
"returncode": result.returncode, |
|||
"stdout": result.stdout, |
|||
"stderr": 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) |
|||
return {"error": str(e)} |
|||
|
|||
def do_worker_ftp_pexpect(str_instruction): |
|||
# 解析指令 |
|||
lines = str_instruction.strip().split('\n') |
|||
cmd_line = lines[0].split('<<')[0].strip() # 提取 "ftp -n 192.168.204.137" |
|||
inputs = [line.strip() for line in lines[1:] if line.strip() != 'EOF'] |
|||
|
|||
# 使用 pexpect 执行命令 |
|||
child = pexpect.spawn(cmd_line) |
|||
for input_line in inputs: |
|||
child.expect('.*') # 等待任意提示 |
|||
child.sendline(input_line) # 发送输入 |
|||
child.expect(pexpect.EOF) # 等待命令结束 |
|||
output = child.before.decode() # 获取输出 |
|||
child.close() |
|||
return output |
|||
|
|||
def do_worker_ftp_script(str_instruction): |
|||
# 创建临时文件保存输出 |
|||
with tempfile.NamedTemporaryFile(delete=False) as tmpfile: |
|||
output_file = tmpfile.name |
|||
|
|||
# 构建并执行 script 命令 |
|||
script_cmd = f"script -c '{str_instruction}' {output_file}" |
|||
result = subprocess.run(script_cmd, shell=True, text=True) |
|||
|
|||
# 读取输出文件内容 |
|||
with open(output_file, 'r') as f: |
|||
output = f.read() |
|||
|
|||
# 删除临时文件 |
|||
os.remove(output_file) |
|||
return output |
|||
|
|||
if __name__ == "__main__": |
|||
# 示例使用 |
|||
str_instruction = """ |
|||
ftp -n 192.168.204.137 << EOF |
|||
user anonymous anonymous@example.com |
|||
ls |
|||
bye |
|||
EOF |
|||
""" |
|||
output = do_worker(str_instruction) |
|||
print(f"*****\n{output}\n*****") |
|||
|
|||
output = do_worker_ftp_script(str_instruction) |
|||
lines = output.splitlines() |
|||
# 跳过第一行（Script started）和最后一行（Script done） |
|||
ftp_output = lines[1:-1] |
|||
strout = '\n'.join(ftp_output) |
|||
print("111111111") |
|||
print(strout) |
|||
@ -0,0 +1,11 @@ |
|||
from tools.ToolBase import ToolBase |
|||
|
|||
class WgetTool(ToolBase): |
|||
def validate_instruction(self, instruction): |
|||
#指令过滤 |
|||
timeout = 0 |
|||
return instruction,timeout |
|||
|
|||
def analyze_result(self, result,instruction,stderr,stdout): |
|||
#指令结果分析 |
|||
return result |
|||
Loading…
Reference in new issue