16 changed files with 589 additions and 275 deletions
@ -1,86 +1,72 @@ |
|||||
import openai |
import re |
||||
import subprocess |
import subprocess |
||||
|
import tempfile |
||||
|
import os |
||||
|
import pexpect |
||||
|
|
||||
# 设置OpenAI API密钥 |
|
||||
openai.api_key = "fGBYaQLHykBOQsFwVrQdIFTsYr8YDtDVDQWFU41mFsmvfNPc" |
|
||||
|
|
||||
# 初始化阶段和已知信息 |
def do_worker(str_instruction): |
||||
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: |
try: |
||||
result = subprocess.run(instruction, shell=True, capture_output=True, text=True) |
# 使用 subprocess 执行 shell 命令 |
||||
return result.stdout if result.stdout else result.stderr |
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: |
except Exception as e: |
||||
return f"执行失败：{str(e)}" |
return {"error": str(e)} |
||||
elif "dirb" in instruction: |
|
||||
try: |
def do_worker_ftp_pexpect(str_instruction): |
||||
result = subprocess.run(instruction, shell=True, capture_output=True, text=True) |
# 解析指令 |
||||
return result.stdout if result.stdout else result.stderr |
lines = str_instruction.strip().split('\n') |
||||
except Exception as e: |
cmd_line = lines[0].split('<<')[0].strip() # 提取 "ftp -n 192.168.204.137" |
||||
return f"执行失败：{str(e)}" |
inputs = [line.strip() for line in lines[1:] if line.strip() != 'EOF'] |
||||
else: |
|
||||
return "未知指令，请重新生成。" |
# 使用 pexpect 执行命令 |
||||
|
child = pexpect.spawn(cmd_line) |
||||
|
for input_line in inputs: |
||||
# 主循环 |
child.expect('.*') # 等待任意提示 |
||||
while current_stage != "报告生成": |
child.sendline(input_line) # 发送输入 |
||||
# 构建提示并获取指令 |
child.expect(pexpect.EOF) # 等待命令结束 |
||||
if not results: # 第一次执行 |
output = child.before.decode() # 获取输出 |
||||
prompt = build_initial_prompt(current_stage, known_info) |
child.close() |
||||
else: # 反馈结果 |
return output |
||||
prompt = build_feedback_prompt(last_instruction, last_result) |
|
||||
|
def do_worker_ftp_script(str_instruction): |
||||
instruction = get_llm_instruction(prompt) |
# 创建临时文件保存输出 |
||||
print(f"生成的指令：{instruction}") |
with tempfile.NamedTemporaryFile(delete=False) as tmpfile: |
||||
|
output_file = tmpfile.name |
||||
# 执行指令 |
|
||||
task_result = execute_instruction(instruction) |
# 构建并执行 script 命令 |
||||
print(f"任务结果：{task_result}") |
script_cmd = f"script -c '{str_instruction}' {output_file}" |
||||
results.append({"instruction": instruction, "result": task_result}) |
result = subprocess.run(script_cmd, shell=True, text=True) |
||||
|
|
||||
# 更新变量 |
# 读取输出文件内容 |
||||
last_instruction = instruction |
with open(output_file, 'r') as f: |
||||
last_result = task_result |
output = f.read() |
||||
|
|
||||
# 示例阶段更新逻辑（可根据实际结果调整） |
# 删除临时文件 |
||||
if current_stage == "信息收集" and "开放端口" in task_result: |
os.remove(output_file) |
||||
current_stage = "漏洞扫描" |
return output |
||||
known_info["ports"] = "80, 443" # 示例更新已知信息 |
|
||||
elif current_stage == "漏洞扫描" and "扫描完成" in task_result: |
if __name__ == "__main__": |
||||
current_stage = "漏洞利用" |
# 示例使用 |
||||
# 添加更多阶段切换逻辑 |
str_instruction = """ |
||||
|
ftp -n 192.168.204.137 << EOF |
||||
# 生成测试报告 |
user anonymous anonymous@example.com |
||||
report = "渗透测试报告\n" |
ls |
||||
report += f"目标网站：{known_info['url']}\n" |
bye |
||||
report += "测试结果：\n" |
EOF |
||||
for res in results: |
""" |
||||
report += f"指令：{res['instruction']}\n结果：{res['result']}\n\n" |
output = do_worker(str_instruction) |
||||
print(report) |
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