""" 知识库路径配置管理器 提供灵活的知识库路径配置，支持自定义目录结构。 """ import os from typing import Dict, Any, Optional from pathlib import Path from app.models.database.models import CategoryConfig class KnowledgePathConfig: """知识库路径配置管理器""" # 默认路径配置（统一为data目录） DEFAULT_PATHS = { 'data': 'data', # 统一的数据目录 'prompts': 'prompts', # Prompt文件目录 'output': 'output' # 输出目录 } # 旧版本兼容路径（向后兼容） LEGACY_PATHS = { 'writeups': ['docs/writeups', 'raw'], # 旧版本：docs/writeups或raw目录 'knowledge_db': ['json/knowledge_db.json'], # 旧版本：json目录 'choice_script': ['tools/scripts/choice.py', 'choice.py'] # 旧版本：tools/scripts或根目录 } def __init__(self, category_id: str, base_path: str = None): """ 初始化路径配置管理器 Args: category_id: 方向ID base_path: 知识库基础路径（默认从项目根目录查找ge10） """ self.category_id = category_id self.category = CategoryConfig.query.get(category_id) if base_path: self.base_path = base_path else: # 从 ctf 目录开始，找到 ge10 目录 current_file = os.path.abspath(__file__) # 从 app/services/category/knowledge_path_config.py 向上找到 ctf 目录 # 路径: ctf/app/services/category/knowledge_path_config.py # parent(1): category, parent(2): services, parent(3): app, parent(4): ctf ctf_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(current_file)))) self.base_path = os.path.join(ctf_root, 'ge10') self.category_dir = os.path.join(self.base_path, category_id) # 加载路径配置 self.paths_config = self._load_paths_config() def _load_paths_config(self) -> Dict[str, Any]: """加载路径配置""" if not self.category: return self.DEFAULT_PATHS.copy() # 从数据库读取路径配置（如果存在） # 可以存储在 advanced_config 或新增字段中 # 这里先使用默认配置，后续可以扩展 # 检查是否有自定义配置 advanced_config = self.category.get_advanced_config() if advanced_config and 'knowledge_paths' in advanced_config: # 合并自定义配置和默认配置 custom_paths = advanced_config['knowledge_paths'] paths = self.DEFAULT_PATHS.copy() self._deep_update(paths, custom_paths) return paths return self.DEFAULT_PATHS.copy() def _deep_update(self, base: Dict, update: Dict): """深度更新字典""" for key, value in update.items(): if key in base and isinstance(base[key], dict) and isinstance(value, dict): self._deep_update(base[key], value) else: base[key] = value def get_writeups_path(self, subcategory: str = None) -> str: """ 获取 Writeup 文件路径（从data目录查找） Args: subcategory: 子分类（如：sql_injection, xss等），可选 Returns: Writeup 文件目录路径 """ # 统一从data目录查找writeups data_dir = os.path.join(self.category_dir, self.paths_config['data']) # 优先查找writeups子目录 writeups_dir = os.path.join(data_dir, 'writeups') if os.path.exists(writeups_dir): if subcategory: return os.path.join(writeups_dir, subcategory) return writeups_dir # 如果不存在writeups子目录，返回data目录 if subcategory: return os.path.join(data_dir, subcategory) return data_dir def get_writeup_file_path(self, filename: str, subcategory: str = None) -> str: """ 获取单个 Writeup 文件路径 Args: filename: 文件名 subcategory: 子分类，可选 Returns: 文件完整路径 """ writeups_dir = self.get_writeups_path(subcategory) return os.path.join(writeups_dir, filename) def find_writeup_file(self, filename: str) -> Optional[str]: """ 查找 Writeup 文件（从data目录递归搜索） Args: filename: 文件名（支持完整文件名或部分匹配） Returns: 文件路径，如果不存在返回None """ data_dir = os.path.join(self.category_dir, self.paths_config['data']) if not os.path.exists(data_dir): return None # 递归搜索data目录（支持子目录） # 精确匹配 for root, dirs, files in os.walk(data_dir): if filename in files: return os.path.join(root, filename) # 模糊匹配（文件名包含关键字） filename_lower = filename.lower() for root, dirs, files in os.walk(data_dir): for file in files: if file.endswith('.md') and (file.lower().startswith(filename_lower) or filename_lower in file.lower()): return os.path.join(root, file) return None def get_data_file_path(self, data_type: str, auto_detect: bool = True) -> Optional[str]: """ 获取数据文件路径（从data目录查找） Args: data_type: 数据类型（knowledge_db, categories, items, taxonomy, metadata） auto_detect: 是否自动检测多种可能的文件名（默认True） Returns: 数据文件路径，如果不存在且auto_detect=False则返回None """ data_dir = os.path.join(self.category_dir, self.paths_config['data']) # 自动检测多种可能的文件名 if auto_detect and os.path.exists(data_dir): possible_names = self._get_possible_data_filenames(data_type) for filename in possible_names: test_path = os.path.join(data_dir, filename) if os.path.exists(test_path): return test_path # 返回默认路径（即使不存在） return os.path.join(data_dir, f'{data_type}.json') def _get_possible_data_filenames(self, data_type: str) -> list: """ 获取数据类型可能的文件名列表（支持多种命名约定） Args: data_type: 数据类型 Returns: 可能的文件名列表 """ # 根据方向类型和数据类型，返回可能的文件名 category_id = self.category_id.lower() # 通用映射 name_mapping = { 'knowledge_db': [ f'{category_id}_knowledge_db.json', # web_knowledge_db.json f'{category_id}_db.json', # web_db.json 'knowledge_db.json', # 通用名称 'vulnerability_db.json', # Web方向常用 'algorithm_db.json', # Crypto方向常用 'technique_db.json', # Reverse/Pwn方向常用 'skill_db.json', # Misc方向常用 ], 'categories': [ f'{category_id}_categories.json', 'categories.json', 'vulnerability_taxonomy.json', 'taxonomy.json', ], 'items': [ f'{category_id}_items.json', 'items.json', ], 'taxonomy': [ f'{category_id}_taxonomy.json', 'taxonomy.json', 'vulnerability_taxonomy_v3.json', ], 'metadata': [ f'{category_id}_metadata.json', 'metadata.json', ] } return name_mapping.get(data_type, [f'{data_type}.json']) def get_tool_script_path(self, script_name: str = None) -> str: """ 获取工具脚本路径（从data目录查找） Args: script_name: 脚本名称（如：choice.py），如果为None返回脚本目录 Returns: 脚本路径 """ data_dir = os.path.join(self.category_dir, self.paths_config['data']) # 优先查找scripts子目录 scripts_dir = os.path.join(data_dir, 'scripts') if not os.path.exists(scripts_dir): scripts_dir = data_dir if script_name: return os.path.join(scripts_dir, script_name) return scripts_dir def find_tool_script(self, script_name: str) -> Optional[str]: """ 查找工具脚本（从data目录递归搜索） Args: script_name: 脚本名称 Returns: 脚本路径，如果不存在返回None """ data_dir = os.path.join(self.category_dir, self.paths_config['data']) if not os.path.exists(data_dir): return None # 递归搜索data目录 for root, dirs, files in os.walk(data_dir): if script_name in files: return os.path.join(root, script_name) return None def get_prompt_template_path(self) -> str: """ 获取 Prompt 模板路径 注意：CLAUDE.md 文件会在运行时由 guideline_manager 根据难度自动生成， 不需要手动创建或编辑。此方法仅保留用于向后兼容。 """ if self.category and self.category.prompt_template_path: return self.category.prompt_template_path return os.path.join(self.category_dir, self.paths_config['prompt_template']) def list_all_writeups(self, recursive: bool = True) -> list: """ 列出所有 Writeup 文件（从data目录查找） Args: recursive: 是否递归搜索子目录 Returns: Writeup 文件列表 """ writeups = [] data_dir = os.path.join(self.category_dir, self.paths_config['data']) if not os.path.exists(data_dir): return writeups if recursive: for root, dirs, files in os.walk(data_dir): for filename in files: if filename.endswith('.md'): rel_path = os.path.relpath(os.path.join(root, filename), data_dir) writeups.append({ 'filename': filename, 'relative_path': rel_path, 'full_path': os.path.join(root, filename), 'subcategory': os.path.relpath(root, data_dir) if root != data_dir else None }) else: for filename in os.listdir(data_dir): file_path = os.path.join(data_dir, filename) if filename.endswith('.md') and os.path.isfile(file_path): writeups.append({ 'filename': filename, 'relative_path': filename, 'full_path': file_path, 'subcategory': None }) return sorted(writeups, key=lambda x: x['filename']) def ensure_directories(self): """确保所有必要的目录存在""" directories = [ os.path.join(self.category_dir, self.paths_config['data']), # data目录 os.path.join(self.category_dir, 'output') # 输出目录 # 注意：不再自动创建 prompts 目录，由用户按需创建 ] for directory in directories: os.makedirs(directory, exist_ok=True) def get_paths_summary(self) -> Dict[str, Any]: """获取路径配置摘要（用于前端显示）""" return { 'base_path': self.base_path, 'category_dir': self.category_dir, 'data': os.path.join(self.category_dir, self.paths_config['data']), 'prompts': os.path.join(self.category_dir, 'prompts'), 'output': os.path.join(self.category_dir, 'output') }