""" 基于RAG的安全知识检索系统 利用现有的RAG模块实现安全知识的向量检索 """ import logging from typing import List, Dict, Any, Optional from .base import KnowledgeDocument, KnowledgeCategory logger = logging.getLogger(__name__) class SecurityKnowledgeRAG: """ 安全知识RAG检索系统 使用现有的RAG模块进行向量检索 """ COLLECTION_NAME = "security_knowledge" def __init__( self, persist_directory: Optional[str] = None, ): self.persist_directory = persist_directory self._indexer = None self._retriever = None self._initialized = False # 内置知识库 - 从模块化文件加载 self._builtin_knowledge = self._load_builtin_knowledge() async def initialize(self): """初始化RAG组件""" if self._initialized: return try: from ...rag import CodeIndexer, CodeRetriever, EmbeddingService embedding_service = EmbeddingService() self._indexer = CodeIndexer( collection_name=self.COLLECTION_NAME, embedding_service=embedding_service, persist_directory=self.persist_directory, ) self._retriever = CodeRetriever( collection_name=self.COLLECTION_NAME, embedding_service=embedding_service, persist_directory=self.persist_directory, ) await self._indexer.initialize() await self._retriever.initialize() # 检查是否需要索引内置知识 count = await self._indexer.get_chunk_count() if count == 0: await self._index_builtin_knowledge() self._initialized = True logger.info("SecurityKnowledgeRAG initialized") except Exception as e: logger.warning(f"Failed to initialize RAG: {e}, using fallback") self._initialized = True # 标记为已初始化，使用fallback def _load_builtin_knowledge(self) -> List[KnowledgeDocument]: """从模块化文件加载内置安全知识""" all_docs = [] # 加载漏洞知识 try: from .vulnerabilities import ALL_VULNERABILITY_DOCS all_docs.extend(ALL_VULNERABILITY_DOCS) logger.debug(f"Loaded {len(ALL_VULNERABILITY_DOCS)} vulnerability docs") except ImportError as e: logger.warning(f"Failed to load vulnerability docs: {e}") # 加载框架知识 try: from .frameworks import ALL_FRAMEWORK_DOCS all_docs.extend(ALL_FRAMEWORK_DOCS) logger.debug(f"Loaded {len(ALL_FRAMEWORK_DOCS)} framework docs") except ImportError as e: logger.warning(f"Failed to load framework docs: {e}") logger.info(f"Total knowledge documents loaded: {len(all_docs)}") return all_docs async def _index_builtin_knowledge(self): """索引内置知识到向量数据库""" if not self._indexer: return logger.info("Indexing builtin security knowledge...") # 转换为RAG可索引的格式 files = [] for doc in self._builtin_knowledge: files.append({ "path": f"knowledge/{doc.category.value}/{doc.id}.md", "content": doc.to_embedding_text(), }) async for progress in self._indexer.index_files(files, base_path="knowledge"): pass logger.info(f"Indexed {len(files)} knowledge documents") async def search( self, query: str, category: Optional[KnowledgeCategory] = None, top_k: int = 5, ) -> List[Dict[str, Any]]: """ 搜索安全知识 Args: query: 搜索查询 category: 知识类别过滤 top_k: 返回数量 Returns: 匹配的知识文档列表 """ await self.initialize() # 如果RAG可用，使用向量检索 if self._retriever: try: results = await self._retriever.retrieve( query=query, top_k=top_k, ) return [ { "id": r.chunk_id, "content": r.content, "score": r.score, "file_path": r.file_path, } for r in results ] except Exception as e: logger.warning(f"RAG search failed: {e}, using fallback") # Fallback: 简单关键词匹配 return self._fallback_search(query, category, top_k) def _fallback_search( self, query: str, category: Optional[KnowledgeCategory], top_k: int, ) -> List[Dict[str, Any]]: """简单的关键词匹配搜索（fallback）""" query_lower = query.lower() query_terms = query_lower.split() results = [] for doc in self._builtin_knowledge: if category and doc.category != category: continue # 计算匹配分数 score = 0 content_lower = doc.content.lower() title_lower = doc.title.lower() # 标题匹配权重更高 for term in query_terms: if term in title_lower: score += 0.3 if term in content_lower: score += 0.1 # 完整查询匹配 if query_lower in title_lower: score += 0.5 if query_lower in content_lower: score += 0.2 # 标签匹配 for tag in doc.tags: if query_lower in tag.lower() or any(t in tag.lower() for t in query_terms): score += 0.15 # CWE/OWASP匹配 for cwe in doc.cwe_ids: if query_lower in cwe.lower(): score += 0.25 for owasp in doc.owasp_ids: if query_lower in owasp.lower(): score += 0.25 if score > 0: results.append({ "id": doc.id, "title": doc.title, "content": doc.content, "category": doc.category.value, "score": min(score, 1.0), "tags": doc.tags, "cwe_ids": doc.cwe_ids, "severity": doc.severity, }) # 按分数排序 results.sort(key=lambda x: x["score"], reverse=True) return results[:top_k] async def get_vulnerability_knowledge( self, vuln_type: str, ) -> Optional[Dict[str, Any]]: """ 获取特定漏洞类型的知识 Args: vuln_type: 漏洞类型（如sql_injection, xss等） Returns: 漏洞知识文档 """ # 标准化漏洞类型名称 vuln_type_normalized = vuln_type.lower().replace("-", "_").replace(" ", "_") # 先尝试精确匹配 for doc in self._builtin_knowledge: if doc.id == f"vuln_{vuln_type_normalized}" or doc.id == vuln_type_normalized: return doc.to_dict() # 尝试部分匹配 for doc in self._builtin_knowledge: if vuln_type_normalized in doc.id: return doc.to_dict() # 使用搜索 results = await self.search(vuln_type, top_k=1) return results[0] if results else None async def get_framework_knowledge( self, framework: str, ) -> Optional[Dict[str, Any]]: """ 获取特定框架的安全知识 Args: framework: 框架名称（如fastapi, django等） Returns: 框架安全知识文档 """ framework_normalized = framework.lower().replace("-", "_").replace(" ", "_") for doc in self._builtin_knowledge: if doc.category == KnowledgeCategory.FRAMEWORK: if doc.id == f"framework_{framework_normalized}" or framework_normalized in doc.id: return doc.to_dict() # 使用搜索 results = await self.search(framework, category=KnowledgeCategory.FRAMEWORK, top_k=1) return results[0] if results else None def get_all_vulnerability_types(self) -> List[str]: """获取所有支持的漏洞类型""" return [ doc.id.replace("vuln_", "") for doc in self._builtin_knowledge if doc.category == KnowledgeCategory.VULNERABILITY ] def get_all_frameworks(self) -> List[str]: """获取所有支持的框架""" return [ doc.id.replace("framework_", "") for doc in self._builtin_knowledge if doc.category == KnowledgeCategory.FRAMEWORK ] def get_knowledge_by_tags(self, tags: List[str]) -> List[Dict[str, Any]]: """根据标签获取知识""" results = [] tags_lower = [t.lower() for t in tags] for doc in self._builtin_knowledge: doc_tags_lower = [t.lower() for t in doc.tags] if any(tag in doc_tags_lower for tag in tags_lower): results.append(doc.to_dict()) return results def get_knowledge_stats(self) -> Dict[str, Any]: """获取知识库统计信息""" stats = { "total": len(self._builtin_knowledge), "by_category": {}, "by_severity": {}, } for doc in self._builtin_knowledge: cat = doc.category.value stats["by_category"][cat] = stats["by_category"].get(cat, 0) + 1 if doc.severity: sev = doc.severity stats["by_severity"][sev] = stats["by_severity"].get(sev, 0) + 1 return stats # 全局实例 security_knowledge_rag = SecurityKnowledgeRAG()