323 lines
10 KiB
Python
323 lines
10 KiB
Python
|
"""
|
|||
|
|
基于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()
|