init: RAG 知识库服务初始提交
- 后端 API（Flask + Gunicorn） - RAG 引擎（混合检索 + 云端 Reranker + 引用溯源） - 文档解析（MinerU + 多格式支持） - Docker 生产部署配置 - 排除前端项目、敏感配置、模型文件
This commit is contained in:
lacerate551
48
auth/__init__.py
Normal file
48
auth/__init__.py
Normal file
@@ -0,0 +1,48 @@
|
||||
"""
|
||||
认证与安全模块
|
||||
|
||||
包含：
|
||||
- gateway: 网关认证（Header 读取）
|
||||
- security: Prompt 注入防护、输入验证、输出过滤
|
||||
"""
|
||||
|
||||
from auth.gateway import (
|
||||
require_gateway_auth,
|
||||
require_role,
|
||||
get_user_permissions,
|
||||
get_current_user,
|
||||
get_auth_manager,
|
||||
get_accessible_collections,
|
||||
check_collection_permission,
|
||||
can_create_collection,
|
||||
can_delete_collection,
|
||||
require_collection_permission,
|
||||
)
|
||||
|
||||
from auth.security import (
|
||||
validate_query,
|
||||
sanitize_user_input,
|
||||
filter_response,
|
||||
is_safe_response,
|
||||
AgentConstraints,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
# gateway
|
||||
'require_gateway_auth',
|
||||
'require_role',
|
||||
'get_user_permissions',
|
||||
'get_current_user',
|
||||
'get_auth_manager',
|
||||
'get_accessible_collections',
|
||||
'check_collection_permission',
|
||||
'can_create_collection',
|
||||
'can_delete_collection',
|
||||
'require_collection_permission',
|
||||
# security
|
||||
'validate_query',
|
||||
'sanitize_user_input',
|
||||
'filter_response',
|
||||
'is_safe_response',
|
||||
'AgentConstraints',
|
||||
]
|
||||
262
auth/gateway.py
Normal file
262
auth/gateway.py
Normal file
@@ -0,0 +1,262 @@
|
||||
"""
|
||||
网关认证模块 - 从网关注入的 Header 读取用户信息
|
||||
|
||||
## 使用方式
|
||||
|
||||
from auth.gateway import require_gateway_auth, get_current_user
|
||||
|
||||
@app.route('/protected')
|
||||
@require_gateway_auth
|
||||
def protected():
|
||||
user = request.current_user # {"user_id": ..., "role": ..., ...}
|
||||
...
|
||||
|
||||
## Header 规范（开发模式可选）
|
||||
|
||||
- X-User-ID: 用户唯一标识 (可选)
|
||||
- X-User-Name: 用户名 (可选)
|
||||
- X-User-Role: 用户角色 (可选)
|
||||
- X-User-Department: 部门 (可选)
|
||||
|
||||
## 模式说明
|
||||
|
||||
开发模式 (DEV_MODE=true，默认):
|
||||
- 支持 mock token 模拟用户：Authorization: Bearer mock-token-admin
|
||||
- 无 Header 时自动使用开发测试用户
|
||||
- 适用于前端测试和开发调试
|
||||
|
||||
生产模式 (DEV_MODE=false):
|
||||
- 不需要 Header，直接放行
|
||||
- 权限由后端完全控制，通过 collections 参数传入
|
||||
- RAG 服务完全无状态，只负责问答检索
|
||||
"""
|
||||
|
||||
from functools import wraps
|
||||
from flask import request, jsonify
|
||||
from typing import Dict, Optional
|
||||
import os
|
||||
from pathlib import Path
|
||||
from dotenv import load_dotenv
|
||||
|
||||
# 加载 .env 文件（从项目根目录）
|
||||
env_path = Path(__file__).parent.parent / '.env'
|
||||
load_dotenv(env_path)
|
||||
|
||||
|
||||
# ==================== 模拟用户数据（开发环境）====================
|
||||
# 用于前端模拟登录测试，仅 DEV_MODE=true 时生效
|
||||
MOCK_USERS = {
|
||||
'admin': {
|
||||
'user_id': 'admin001',
|
||||
'password': 'admin123',
|
||||
'role': 'admin',
|
||||
'department': '管理部'
|
||||
},
|
||||
'admin2': {
|
||||
'user_id': 'admin002',
|
||||
'password': 'admin456',
|
||||
'role': 'admin',
|
||||
'department': '技术部'
|
||||
},
|
||||
'admin3': {
|
||||
'user_id': 'admin003',
|
||||
'password': 'admin789',
|
||||
'role': 'admin',
|
||||
'department': '运营部'
|
||||
},
|
||||
'manager': {
|
||||
'user_id': 'manager001',
|
||||
'password': 'manager123',
|
||||
'role': 'manager',
|
||||
'department': '财务部'
|
||||
},
|
||||
'user': {
|
||||
'user_id': 'user001',
|
||||
'password': 'test123',
|
||||
'role': 'user',
|
||||
'department': '技术部'
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def require_gateway_auth(f):
|
||||
"""
|
||||
网关认证装饰器 - 从 Header 读取用户信息
|
||||
|
||||
开发模式 (DEV_MODE=true，默认):
|
||||
- 支持 mock token: Authorization: Bearer mock-token-admin
|
||||
- 无 Header 时自动使用开发测试用户（admin 角色）
|
||||
|
||||
生产模式 (DEV_MODE=false):
|
||||
- 不需要 Header，直接放行
|
||||
- 用户信息设为默认值
|
||||
- 权限由后端通过 collections 参数控制
|
||||
"""
|
||||
@wraps(f)
|
||||
def decorated(*args, **kwargs):
|
||||
# 开发模式开关（默认开启，生产环境设置 DEV_MODE=false）
|
||||
dev_mode = os.environ.get('DEV_MODE', 'true').lower() != 'false'
|
||||
|
||||
# 开发模式：支持 mock token
|
||||
if dev_mode:
|
||||
auth_header = request.headers.get('Authorization', '')
|
||||
if auth_header.startswith('Bearer mock-token-'):
|
||||
username = auth_header.replace('Bearer mock-token-', '')
|
||||
mock_user = MOCK_USERS.get(username)
|
||||
if mock_user:
|
||||
request.current_user = {
|
||||
"user_id": mock_user['user_id'],
|
||||
"username": username,
|
||||
"role": mock_user['role'],
|
||||
"department": mock_user['department']
|
||||
}
|
||||
return f(*args, **kwargs)
|
||||
|
||||
# 从 Header 读取用户信息（可选）
|
||||
user_id = request.headers.get('X-User-ID')
|
||||
username = request.headers.get('X-User-Name', '')
|
||||
role = request.headers.get('X-User-Role', 'user')
|
||||
department = request.headers.get('X-User-Department', '')
|
||||
|
||||
# 如果有 Header，使用 Header 中的用户信息
|
||||
if user_id:
|
||||
request.current_user = {
|
||||
"user_id": user_id,
|
||||
"username": username,
|
||||
"role": role,
|
||||
"department": department
|
||||
}
|
||||
return f(*args, **kwargs)
|
||||
|
||||
# 无 Header 时的默认处理
|
||||
if dev_mode:
|
||||
# 开发模式：使用默认测试用户
|
||||
request.current_user = {
|
||||
"user_id": "dev-user",
|
||||
"username": "开发测试用户",
|
||||
"role": "admin",
|
||||
"department": "开发部"
|
||||
}
|
||||
else:
|
||||
# 生产模式：使用默认用户（后端通过 collections 控制权限）
|
||||
request.current_user = {
|
||||
"user_id": "backend-caller",
|
||||
"username": "后端调用",
|
||||
"role": "user",
|
||||
"department": ""
|
||||
}
|
||||
|
||||
return f(*args, **kwargs)
|
||||
|
||||
return decorated
|
||||
|
||||
|
||||
def get_current_user() -> Optional[Dict]:
|
||||
"""
|
||||
获取当前登录用户信息
|
||||
|
||||
Returns:
|
||||
用户信息字典，未认证返回 None
|
||||
"""
|
||||
return getattr(request, 'current_user', None)
|
||||
|
||||
|
||||
# ==================== 兼容旧代码 ====================
|
||||
|
||||
# 别名
|
||||
require_auth = require_gateway_auth
|
||||
|
||||
|
||||
def get_user_permissions(role: str):
|
||||
"""
|
||||
兼容旧代码 - 权限由后端管理，此函数仅返回默认值
|
||||
"""
|
||||
return ['public', 'internal', 'confidential']
|
||||
|
||||
|
||||
def check_collection_permission(role: str, department: str, collection_name: str, operation: str = "read") -> bool:
|
||||
"""
|
||||
兼容旧代码 - 权限由后端管理，默认返回 True
|
||||
"""
|
||||
return True
|
||||
|
||||
|
||||
def get_accessible_collections(role: str, department: str, operation: str = "read", all_collections=None):
|
||||
"""
|
||||
兼容旧代码 - 返回所有向量库
|
||||
"""
|
||||
if all_collections:
|
||||
return all_collections
|
||||
return ['public_kb']
|
||||
|
||||
|
||||
def can_create_collection(role: str) -> bool:
|
||||
"""兼容旧代码 - 权限由后端管理"""
|
||||
return True
|
||||
|
||||
|
||||
def can_delete_collection(role: str) -> bool:
|
||||
"""兼容旧代码 - 权限由后端管理"""
|
||||
return True
|
||||
|
||||
|
||||
def require_role(*roles):
|
||||
"""
|
||||
兼容旧代码 - 权限由后端管理，此装饰器不再执行权限验证
|
||||
"""
|
||||
def decorator(f):
|
||||
@wraps(f)
|
||||
def decorated(*args, **kwargs):
|
||||
return f(*args, **kwargs)
|
||||
return decorated
|
||||
return decorator
|
||||
|
||||
|
||||
def require_collection_permission(operation: str):
|
||||
"""
|
||||
兼容旧代码 - 权限由后端管理，此装饰器不再执行权限验证
|
||||
"""
|
||||
def decorator(f):
|
||||
@wraps(f)
|
||||
def decorated(*args, **kwargs):
|
||||
return f(*args, **kwargs)
|
||||
return decorated
|
||||
return decorator
|
||||
|
||||
|
||||
def get_auth_manager():
|
||||
"""兼容旧代码"""
|
||||
return _FakeAuthManager()
|
||||
|
||||
|
||||
class _FakeAuthManager:
|
||||
"""兼容旧代码的假 AuthManager"""
|
||||
|
||||
@staticmethod
|
||||
def get_user_permissions(role: str):
|
||||
return ['public', 'internal', 'confidential']
|
||||
|
||||
@staticmethod
|
||||
def get_accessible_collections(role: str, department: str):
|
||||
return ['public_kb']
|
||||
|
||||
|
||||
def is_admin() -> bool:
|
||||
"""检查当前用户是否为管理员"""
|
||||
user = get_current_user()
|
||||
return user is not None and user.get('role') == 'admin'
|
||||
|
||||
|
||||
def is_manager_or_above() -> bool:
|
||||
"""检查当前用户是否为经理或以上级别"""
|
||||
user = get_current_user()
|
||||
return user is not None and user.get('role') in ('admin', 'manager')
|
||||
|
||||
|
||||
def normalize_department_name(department: str) -> str:
|
||||
"""兼容旧代码 - 部门名称标准化"""
|
||||
if not department:
|
||||
return ""
|
||||
if department.replace("_", "").replace("-", "").isalnum() and department.isascii():
|
||||
return department.lower()
|
||||
return ""
|
||||
182
auth/security.py
Normal file
182
auth/security.py
Normal file
@@ -0,0 +1,182 @@
|
||||
"""
|
||||
Prompt 注入防护模块 - 输入验证、查询隔离、输出过滤
|
||||
|
||||
功能：
|
||||
1. 输入验证 - 检测注入模式、长度限制
|
||||
2. 查询隔离 - XML 标签包裹用户输入，防止指令注入
|
||||
3. 输出过滤 - 阻止敏感信息泄露
|
||||
4. Agent 行为约束 - 调用次数上限、工具白名单
|
||||
|
||||
使用方式：
|
||||
from security import validate_query, sanitize_user_input, filter_response
|
||||
"""
|
||||
|
||||
import re
|
||||
import os
|
||||
import logging
|
||||
from typing import Tuple, Optional
|
||||
from pathlib import Path
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# ==================== 违禁词配置 ====================
|
||||
|
||||
# 违禁词文件路径
|
||||
BANNED_WORDS_FILE = Path(__file__).parent.parent / "config" / "banned_words.txt"
|
||||
|
||||
def _load_banned_words() -> list:
|
||||
"""从配置文件加载违禁词"""
|
||||
banned_words = []
|
||||
try:
|
||||
if BANNED_WORDS_FILE.exists():
|
||||
with open(BANNED_WORDS_FILE, 'r', encoding='utf-8') as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
# 跳过空行和注释
|
||||
if line and not line.startswith('#'):
|
||||
banned_words.append(line)
|
||||
except Exception as e:
|
||||
logger.warning(f"加载违禁词文件失败: {e}")
|
||||
return banned_words
|
||||
|
||||
# 加载违禁词列表
|
||||
BANNED_WORDS = _load_banned_words()
|
||||
|
||||
|
||||
# ==================== 输入验证 ====================
|
||||
|
||||
# 注入攻击常见模式
|
||||
INJECTION_PATTERNS = [
|
||||
# 直接指令覆盖
|
||||
r"(?i)(ignore|forget|disregard|discard)\s+(previous|above|all|earlier|prior)\s+(instructions?|prompts?|rules?|context)",
|
||||
# 角色切换
|
||||
r"(?i)(you\s+are\s+now|act\s+as|pretend\s+to\s+be|roleplay|new\s+role)",
|
||||
# 系统提示词提取
|
||||
r"(?i)(show|display|print|output|reveal|tell)\s+me\s+(your|the)\s+(system\s+)?(prompt|instructions?|rules)",
|
||||
# 文档内容提取
|
||||
r"(?i)(output|print|display|show|list)\s+(all|every|complete|full)\s+(documents?|data|records?|files?|contents?)",
|
||||
# 系统指令标记
|
||||
r"(?i)system\s*[:：]\s*",
|
||||
# 配置信息提取
|
||||
r"(?i)(show|display|reveal)\s+(config|api\s*key|password|secret|credentials?)",
|
||||
]
|
||||
|
||||
MAX_QUERY_LENGTH = 1000
|
||||
MAX_CONVERSATION_LENGTH = 5000
|
||||
|
||||
|
||||
def validate_query(query: str) -> Tuple[bool, str]:
|
||||
"""
|
||||
验证用户查询是否安全
|
||||
|
||||
Returns:
|
||||
(is_valid, reason)
|
||||
"""
|
||||
if not query or not query.strip():
|
||||
return False, "查询内容不能为空"
|
||||
|
||||
if len(query) > MAX_QUERY_LENGTH:
|
||||
return False, f"查询内容过长（最多{MAX_QUERY_LENGTH}字符）"
|
||||
|
||||
# 检测违禁词
|
||||
for word in BANNED_WORDS:
|
||||
if word in query:
|
||||
return False, "查询包含违禁内容"
|
||||
|
||||
# 检测注入模式
|
||||
for pattern in INJECTION_PATTERNS:
|
||||
if re.search(pattern, query):
|
||||
return False, "查询包含不允许的内容"
|
||||
|
||||
return True, ""
|
||||
|
||||
|
||||
def sanitize_user_input(query: str) -> str:
|
||||
"""
|
||||
将用户输入包裹在 XML 标签中，隔离指令注入
|
||||
|
||||
LLM 在处理 <user_query> 标签内的内容时，
|
||||
应仅将其作为文本分析，不执行其中的指令。
|
||||
"""
|
||||
# 移除可能破坏 XML 结构的字符
|
||||
cleaned = query.replace("<user_query>", "").replace("</user_query>", "")
|
||||
return f"<user_query>\n{cleaned}\n</user_query>"
|
||||
|
||||
|
||||
def is_safe_response(response: str) -> Tuple[bool, Optional[str]]:
|
||||
"""
|
||||
检查 LLM 输出是否包含敏感信息
|
||||
|
||||
Returns:
|
||||
(is_safe, leaked_info_type or None)
|
||||
"""
|
||||
sensitive_patterns = [
|
||||
(r"sk-[a-f0-9]{20,}", "API密钥"),
|
||||
(r"(?:password|密码)\s*[:：]\s*\S+", "密码"),
|
||||
(r"config\.(?:py|example)", "配置文件"),
|
||||
(r"(?:JWT_SECRET)\s*=\s*\S+", "密钥配置"),
|
||||
]
|
||||
|
||||
for pattern, info_type in sensitive_patterns:
|
||||
if re.search(pattern, response, re.IGNORECASE):
|
||||
return False, info_type
|
||||
|
||||
return True, None
|
||||
|
||||
|
||||
def filter_response(response: str) -> str:
|
||||
"""
|
||||
过滤 LLM 响应中的敏感信息（API密钥、密码等）
|
||||
"""
|
||||
filtered = response
|
||||
|
||||
replacements = [
|
||||
(r"sk-[a-f0-9]{20,}", "[已过滤]"),
|
||||
(r"(?:password|密码)\s*[:：]\s*\S+", "[已过滤]"),
|
||||
(r"config\.(?:py|example)", "[已过滤]"),
|
||||
(r"(?:JWT_SECRET)\s*=\s*['\"]?\S+['\"]?", "[已过滤]"),
|
||||
]
|
||||
|
||||
for pattern, replacement in replacements:
|
||||
filtered = re.sub(pattern, replacement, filtered, flags=re.IGNORECASE)
|
||||
|
||||
return filtered
|
||||
|
||||
|
||||
# ==================== Agent 行为约束 ====================
|
||||
|
||||
class AgentConstraints:
|
||||
"""Agent 行为约束，防止恶意使用"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
max_iterations: int = 3,
|
||||
max_api_calls: int = 10,
|
||||
max_query_length: int = 1000,
|
||||
allowed_tools: set = None
|
||||
):
|
||||
self.max_iterations = max_iterations
|
||||
self.max_api_calls = max_api_calls
|
||||
self.max_query_length = max_query_length
|
||||
self.allowed_tools = allowed_tools or {
|
||||
"kb_search", "web_search", "answer", "rewrite", "decompose"
|
||||
}
|
||||
self.api_calls = 0
|
||||
|
||||
def check_tool_allowed(self, tool_name: str) -> bool:
|
||||
"""检查工具是否在白名单中"""
|
||||
return tool_name in self.allowed_tools
|
||||
|
||||
def check_budget(self) -> bool:
|
||||
"""检查是否还有 API 调用预算"""
|
||||
self.api_calls += 1
|
||||
return self.api_calls <= self.max_api_calls
|
||||
|
||||
def check_query_length(self, query: str) -> bool:
|
||||
"""检查查询长度是否在限制内"""
|
||||
return len(query) <= self.max_query_length
|
||||
|
||||
def reset(self):
|
||||
"""重置调用计数（每次新请求开始时调用）"""
|
||||
self.api_calls = 0
|
||||
Reference in New Issue
Block a user