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)