mirror of
https://github.com/SunZhimin2021/AIPentest.git
synced 2025-11-05 10:53:16 +00:00
259 lines
8.9 KiB
Python
259 lines
8.9 KiB
Python
import sys
|
||
import os
|
||
from typing import Dict, List, Optional, Tuple
|
||
from bs4 import BeautifulSoup
|
||
import json
|
||
import requests
|
||
import re
|
||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
||
from common.llmcaller import LLMCaller
|
||
from common.sshconnection import SSHConnection
|
||
|
||
import time
|
||
import logging
|
||
from dotenv import load_dotenv
|
||
load_dotenv()
|
||
deepseek_key=os.getenv('DEEPSEEK_API_KEY')
|
||
|
||
logger = logging.getLogger(__name__)
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
|
||
handlers=[
|
||
logging.FileHandler('bandit.log'),
|
||
logging.StreamHandler() # 同时输出到控制台
|
||
]
|
||
)
|
||
class RecordFileManager:
|
||
"""记录文件管理类，用于追加和读取最后一条记录"""
|
||
|
||
def __init__(self, filepath: str):
|
||
"""
|
||
初始化记录文件管理器
|
||
|
||
Args:
|
||
filepath: 记录文件路径
|
||
"""
|
||
self.filepath = filepath
|
||
# 确保目录存在
|
||
directory = os.path.dirname(filepath)
|
||
if directory and not os.path.exists(directory):
|
||
os.makedirs(directory)
|
||
|
||
def append(self, n: int, content: str) -> bool:
|
||
"""
|
||
添加一条记录到文件末尾
|
||
|
||
Args:
|
||
n: 数字部分
|
||
content: 内容部分
|
||
|
||
Returns:
|
||
bool: 是否成功添加
|
||
"""
|
||
try:
|
||
with open(self.filepath, 'a', encoding='utf-8') as f:
|
||
f.write(f"{n} {content}\n")
|
||
return True
|
||
except Exception as e:
|
||
print(f"添加记录失败: {e}")
|
||
return False
|
||
|
||
def load_last_record(self) -> Optional[Tuple[int, str]]:
|
||
"""
|
||
读取最后一条记录
|
||
|
||
Returns:
|
||
Optional[Tuple[int, str]]: (数字, 内容) 或 None（如果文件不存在或为空）
|
||
"""
|
||
try:
|
||
# 如果文件不存在，返回None
|
||
if not os.path.exists(self.filepath):
|
||
return None
|
||
|
||
with open(self.filepath, 'r', encoding='utf-8') as f:
|
||
lines = f.readlines()
|
||
|
||
# 如果文件为空，返回None
|
||
if not lines:
|
||
return None
|
||
|
||
# 获取最后一行，去除换行符
|
||
last_line = lines[-1].strip()
|
||
|
||
# 如果最后一行为空，向前查找非空行
|
||
for line in reversed(lines):
|
||
line = line.strip()
|
||
if line:
|
||
last_line = line
|
||
break
|
||
else:
|
||
# 所有行都为空
|
||
return None
|
||
|
||
# 解析最后一行
|
||
parts = last_line.split(' ', 1) # 只分割第一个空格
|
||
if len(parts) >= 2:
|
||
try:
|
||
n = int(parts[0])
|
||
content = parts[1]
|
||
return (n, content)
|
||
except ValueError:
|
||
print(f"解析数字失败: {parts[0]}")
|
||
return None
|
||
else:
|
||
print(f"记录格式错误: {last_line}")
|
||
return None
|
||
|
||
except Exception as e:
|
||
print(f"读取记录失败: {e}")
|
||
return None
|
||
|
||
|
||
class BanditLevelFetcher:
|
||
"""获取Bandit级别信息"""
|
||
|
||
def __init__(self):
|
||
self.base_url = "https://overthewire.org/wargames/bandit"
|
||
|
||
def get_level_info(self, level: int) -> Tuple[str, str]:
|
||
"""获取指定级别的目标和帮助信息"""
|
||
try:
|
||
if level == 0:
|
||
url = f"{self.base_url}/bandit0.html"
|
||
else:
|
||
url = f"{self.base_url}/bandit{level}.html"
|
||
|
||
response = requests.get(url, timeout=10)
|
||
response.raise_for_status()
|
||
|
||
soup = BeautifulSoup(response.content, 'html.parser')
|
||
|
||
# 查找Level Goal
|
||
goal = ""
|
||
goal_header = soup.find('h2', string='Level Goal')
|
||
if goal_header:
|
||
goal_content = goal_header.find_next_sibling('p')
|
||
if goal_content:
|
||
goal = goal_content.get_text().strip()
|
||
|
||
# 查找Commands you may need to solve this level
|
||
help_info = ""
|
||
help_header = soup.find('h2', string='Commands you may need to solve this level')
|
||
if help_header:
|
||
help_content = help_header.find_next_sibling('p')
|
||
if help_content:
|
||
help_info = help_content.get_text().strip()
|
||
|
||
# 如果没有找到标准格式，尝试其他方式
|
||
if not goal:
|
||
# 尝试查找包含关键词的段落
|
||
paragraphs = soup.find_all('p')
|
||
for p in paragraphs:
|
||
text = p.get_text().lower()
|
||
if 'password' in text or 'level' in text or 'goal' in text:
|
||
goal = p.get_text().strip()
|
||
break
|
||
|
||
return goal, help_info
|
||
|
||
except Exception as e:
|
||
logger.error(f"获取级别 {level} 信息失败: {e}")
|
||
return "", ""
|
||
|
||
|
||
def getcommand(level:int,goal:str,help_info:str,userprompt:str):
|
||
systemprompt=f"""
|
||
你是一个Linux专家，熟悉使用Linux,包括命令行的各种技巧，正在解决OverTheWire Bandit第{level}级挑战。
|
||
游戏规则:
|
||
1. 每个级别都有一个目标，你需要通过执行一系列命令来完成目标。最终目标是在系统中的某个文件中找到一个password.
|
||
1. 每个级别都有一个帮助信息，帮助信息中会包含一些有用的信息，比如可能用到的系统命令。
|
||
|
||
本关目标: {goal}
|
||
|
||
帮助信息,可能用到的系统命令: {help_info}
|
||
游戏的难点是如何找到文件并读取文件内容。有时候文件是特殊字符，需要使用技巧来读取，有时候文件藏在很多目录里，需要技巧才能找到。你要发挥聪明才智才能完成任务。
|
||
游戏的方法是：
|
||
我会登录系统，并执行初始命令 ls -l,然后把所有的结果返回给你。
|
||
你来输出下一个命令，我执行后，把结果返回给你。
|
||
只输出可以直接执行的命令，不要输出任何解释，以便程序执行。
|
||
有的文件里密码可能有说明，有些文件则只包括一个字符串，那这个字符串就是password
|
||
直到你告诉我目标完成。目标完成的标志是：<goal_done>password</goal_done>
|
||
必须取得真正的密码才能完成，真正的密码肯定不是'password'
|
||
如果部分命令返回异常，可以要求重新执行。
|
||
|
||
|
||
"""
|
||
|
||
openai_llm = LLMCaller('deepseek', api_key=deepseek_key)
|
||
response = openai_llm.call(
|
||
system_prompt=systemprompt,
|
||
user_prompt=userprompt,
|
||
temperature=0.7
|
||
)
|
||
print(response)
|
||
return response
|
||
|
||
def main():
|
||
recorder = RecordFileManager("records.txt")
|
||
bandit_level_fetcher = BanditLevelFetcher()
|
||
sshclient=SSHConnection()
|
||
lastpassword=""
|
||
hostname=f"bandit.labs.overthewire.org"
|
||
last_record = recorder.load_last_record()
|
||
if last_record:
|
||
n, lastpassword = last_record
|
||
print(f"最后一条记录: {n} {lastpassword}")
|
||
else:
|
||
n=0
|
||
for i in range(n,33):
|
||
if i==0:
|
||
lastpassword="bandit0"
|
||
|
||
|
||
username=f"bandit{i}"
|
||
|
||
print(f"正在获取级别 {i} 的信息...")
|
||
logger.info(f"正在获取级别 {i} 的信息...")
|
||
goal,info = bandit_level_fetcher.get_level_info(i+1)
|
||
#print(goal,info)
|
||
logger.info(f"level {i} goal: {goal}")
|
||
if not sshclient.connect(username,lastpassword,hostname=hostname):
|
||
logger.error(f"level {i} connect failed")
|
||
break
|
||
result=sshclient.execute_command("ls -l")
|
||
userprompt=f"""
|
||
command: ls -l
|
||
return value"{result}
|
||
|
||
"""
|
||
while True:
|
||
command=getcommand(i,goal,info,userprompt)
|
||
print(command)
|
||
logger.info(f"level {i} command: {command}")
|
||
match = re.search(r"<goal_done>(.*?)</goal_done>", command)
|
||
if match:
|
||
lastpassword = match.group(1).strip()
|
||
print(f"get password: {lastpassword}")
|
||
recorder.append(i+1,lastpassword)
|
||
sshclient.close()
|
||
logger.info(f"level {i} get password: {lastpassword}")
|
||
break
|
||
|
||
result=sshclient.execute_command(command)
|
||
print(result)
|
||
userprompt+=f"""
|
||
command: {command}
|
||
return value: {result}
|
||
"""
|
||
print(userprompt)
|
||
logger.info(f"level {i} userprompt: {userprompt}")
|
||
|
||
|
||
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|
||
# 使用示例
|