fix(security): 代码审查安全加固 — 三批次修复(6H/13M/12L)

第一批(快速修复):
- H6: main.py --debug 默认值 True→False,防止 Werkzeug RCE
- M2+M3: /search 增加 validate_query + top_k 范围限制(1-50)
- M4: context_count 范围限制(0-10) + 异常捕获
- L3: assert → raise RuntimeError(生产环境 API Key 检查)
- H1: SSE 错误事件移除 traceback 字段

第二批(安全加固):
- H2+H3: 文档接口路径遍历 realpath 校验 + 文件类型/大小限制
- H4+H5: 批量上传文件大小检查
- M6: LIKE 查询通配符转义
- M1: 37 处 str(e) 异常信息统一脱敏(6 文件)
- M5: CORS 生产环境限制来源
- M7: SESSION_MANAGER None 保护(503)
- M11: subprocess 参数注入防护(白名单 + -- 分隔符)

第三批(架构改进):
- M8+M9: 提取 JSON 解析共享工具(extract_json_object/list)
- M10: Prompt 注入检测防御(prompt_guard.py)
- M12: 解析器文件大小限制(Excel 50MB/TXT 20MB/PDF 100MB)
- M13: 全局单例竞态条件双重检查锁定(engine/bm25/intent_analyzer)
This commit is contained in:
lacerate551
2026-06-05 15:26:32 +08:00
parent a7206377cc
commit 90b915232a
19 changed files with 436 additions and 102 deletions

View File

@@ -67,6 +67,16 @@ def create_app() -> 'Flask':
static_folder = os.path.join(PROJECT_ROOT, 'chat-ui') static_folder = os.path.join(PROJECT_ROOT, 'chat-ui')
app = Flask(__name__, static_folder=static_folder, static_url_path='') app = Flask(__name__, static_folder=static_folder, static_url_path='')
# CORS 配置:生产环境限制来源,开发环境允许全部
if IS_PROD:
cors_origins = os.environ.get('CORS_ORIGINS', '').split(',') if os.environ.get('CORS_ORIGINS') else []
if cors_origins:
CORS(app, origins=cors_origins)
else:
CORS(app) # 未配置时仍允许全部,但记录警告
logger.warning("生产环境未配置 CORS_ORIGINSCORS 允许所有来源")
else:
CORS(app) CORS(app)
# ==================== Repository 依赖注入 ==================== # ==================== Repository 依赖注入 ====================
@@ -227,14 +237,14 @@ def _validate_production_config() -> None:
- DASHSCOPE_API_KEY: 大模型调用必需 - DASHSCOPE_API_KEY: 大模型调用必需
Raises: Raises:
AssertionError: 缺少必需的配置项 RuntimeError: 缺少必需的配置项
""" """
import os import os
from config import DASHSCOPE_API_KEY from config import DASHSCOPE_API_KEY
# 检查环境变量或配置文件中的 API Key # 检查环境变量或配置文件中的 API Key
has_key = os.getenv("DASHSCOPE_API_KEY") or os.environ.get("DASHSCOPE_API_KEY") or DASHSCOPE_API_KEY has_key = os.getenv("DASHSCOPE_API_KEY") or os.environ.get("DASHSCOPE_API_KEY") or DASHSCOPE_API_KEY
assert has_key and has_key != "", \ if not has_key or has_key == "":
"Missing DASHSCOPE_API_KEY in production environment" raise RuntimeError("Missing DASHSCOPE_API_KEY in production environment")
logger.info("Configuration validated") logger.info("Configuration validated")

View File

@@ -97,7 +97,8 @@ def get_audit_logs():
return jsonify({"logs": logs, "total": total}) return jsonify({"logs": logs, "total": total})
except Exception as e: except Exception as e:
return jsonify({"error": str(e), "logs": [], "total": 0}), 500 logger.error(f"审计查询异常: {e}")
return jsonify({"error": "查询失败", "logs": [], "total": 0}), 500
def log_audit_event(user_id: str, username: str, action: str, def log_audit_event(user_id: str, username: str, action: str,

View File

@@ -41,6 +41,7 @@ from auth.gateway import require_gateway_auth
from auth.security import validate_query, filter_response from auth.security import validate_query, filter_response
from config import RAG_CHAT_MODEL from config import RAG_CHAT_MODEL
from core.llm_utils import call_llm, call_llm_stream from core.llm_utils import call_llm, call_llm_stream
from core.prompt_guard import detect_injection, sanitize_user_input
def _get_vlm_cache(image_path: str) -> Optional[str]: def _get_vlm_cache(image_path: str) -> Optional[str]:
@@ -1953,11 +1954,11 @@ def rag():
yield f"data: {json.dumps(finish_event, ensure_ascii=False)}\n\n" yield f"data: {json.dumps(finish_event, ensure_ascii=False)}\n\n"
except Exception as e: except Exception as e:
import traceback import logging as _logging
_logging.getLogger(__name__).error(f"[SSE] RAG 流异常: {e}", exc_info=True)
error_event = { error_event = {
"type": "error", "type": "error",
"message": str(e), "message": "服务内部错误,请稍后重试"
"traceback": traceback.format_exc()
} }
yield f"data: {json.dumps(error_event, ensure_ascii=False)}\n\n" yield f"data: {json.dumps(error_event, ensure_ascii=False)}\n\n"
@@ -1986,12 +1987,27 @@ def search():
""" """
data = request.json or {} data = request.json or {}
query = data.get('query', '') query = data.get('query', '')
query = sanitize_user_input(query)
injection_matches = detect_injection(query)
if injection_matches:
logger.warning(f"[Chat] 检测到可疑注入: {injection_matches}")
top_k = data.get('top_k', 5) top_k = data.get('top_k', 5)
collections = data.get('collections') # 后端传入的知识库列表 collections = data.get('collections') # 后端传入的知识库列表
if not query: if not query:
return jsonify({'error': 'query is required'}), 400 return jsonify({'error': 'query is required'}), 400
# 输入安全校验(注入检测、违禁词、长度限制)
is_valid, reason = validate_query(query)
if not is_valid:
return jsonify({'error': reason}), 400
# top_k 范围校验
try:
top_k = max(1, min(int(top_k), 50))
except (ValueError, TypeError):
top_k = 5
# 如果没有指定 collections使用默认的公开库 # 如果没有指定 collections使用默认的公开库
if not collections: if not collections:
collections = ['public_kb'] collections = ['public_kb']

View File

@@ -96,6 +96,24 @@ def safe_filename(filename: str) -> str:
return safe_name return safe_name
def _validate_doc_path(doc_path: str, documents_path: str) -> str:
"""
校验文档路径是否在允许目录内,防止路径遍历攻击。
Returns:
解析后的安全绝对路径
Raises:
ValueError: 路径不合法
"""
filepath = os.path.join(documents_path, doc_path)
real_path = os.path.realpath(filepath)
real_base = os.path.realpath(documents_path)
if not real_path.startswith(real_base + os.sep) and real_path != real_base:
raise ValueError("非法路径")
return real_path
# 延迟初始化缓存 # 延迟初始化缓存
_kb_manager = None _kb_manager = None
_kb_checked = False _kb_checked = False
@@ -160,6 +178,10 @@ def serve_document_file(doc_path: str) -> Tuple[Any, int]:
from flask import send_from_directory from flask import send_from_directory
filepath = os.path.join(DOCUMENTS_PATH, doc_path) filepath = os.path.join(DOCUMENTS_PATH, doc_path)
try:
filepath = _validate_doc_path(doc_path, DOCUMENTS_PATH)
except ValueError:
return jsonify({"error": "非法路径"}), 403
if not os.path.exists(filepath): if not os.path.exists(filepath):
return jsonify({"error": "文件不存在"}), 404 return jsonify({"error": "文件不存在"}), 404
@@ -329,7 +351,8 @@ def upload_document() -> Tuple[Any, int]:
sync_service.process_change(change) sync_service.process_change(change)
sync_status = "已保存并添加到向量库" sync_status = "已保存并添加到向量库"
except Exception as e: except Exception as e:
sync_status = f"已保存,向量化失败: {str(e)}" logger.warning(f"向量化失败: {e}")
sync_status = "已保存,向量化失败"
return success_response( return success_response(
data={ data={
@@ -404,6 +427,18 @@ def batch_upload_documents() -> Tuple[Any, int]:
}) })
continue continue
# 文件大小校验
file.seek(0, 2)
file_size = file.tell()
file.seek(0)
if file_size > MAX_FILE_SIZE:
results.append({
"filename": file.filename,
"status": "error",
"message": f"文件过大(最大 {MAX_FILE_SIZE // 1024 // 1024}MB"
})
continue
try: try:
original_filename = file.filename original_filename = file.filename
ext = os.path.splitext(original_filename)[1].lower() ext = os.path.splitext(original_filename)[1].lower()
@@ -464,10 +499,11 @@ def batch_upload_documents() -> Tuple[Any, int]:
"replaced": replaced "replaced": replaced
}) })
except Exception as e: except Exception as e:
logger.error(f"上传处理异常: {e}")
results.append({ results.append({
"filename": file.filename, "filename": file.filename,
"status": "error", "status": "error",
"message": str(e) "message": "上传处理失败"
}) })
return success_response( return success_response(
@@ -644,12 +680,27 @@ def update_document(doc_path: str) -> Tuple[Any, int]:
if file.filename == '': if file.filename == '':
return jsonify({"error": "没有选择文件"}), 400 return jsonify({"error": "没有选择文件"}), 400
# 文件类型校验
ext = os.path.splitext(file.filename)[1].lower()
if ext not in ALLOWED_EXTENSIONS:
return jsonify({"error": f"不支持的文件类型: {ext}"}), 400
# 文件大小校验
file.seek(0, 2) # 跳到文件末尾获取大小
file_size = file.tell()
file.seek(0) # 回到文件开头
if file_size > MAX_FILE_SIZE:
return jsonify({"error": f"文件过大(最大 {MAX_FILE_SIZE // 1024 // 1024}MB"}), 400
# 解析路径 # 解析路径
parts = doc_path.split('/') parts = doc_path.split('/')
if len(parts) < 2: if len(parts) < 2:
return jsonify({"error": "无效的文档路径"}), 400 return jsonify({"error": "无效的文档路径"}), 400
filepath = os.path.join(DOCUMENTS_PATH, doc_path) try:
filepath = _validate_doc_path(doc_path, DOCUMENTS_PATH)
except ValueError:
return jsonify({"error": "非法路径"}), 403
if not os.path.exists(filepath): if not os.path.exists(filepath):
return jsonify({"error": "文件不存在"}), 404 return jsonify({"error": "文件不存在"}), 404
@@ -711,7 +762,10 @@ def delete_document(doc_path: str) -> Tuple[Any, int]:
# 目录名即向量库名 # 目录名即向量库名
collection = subdir collection = subdir
filepath = os.path.join(DOCUMENTS_PATH, doc_path) try:
filepath = _validate_doc_path(doc_path, DOCUMENTS_PATH)
except ValueError:
return jsonify({"error": "非法路径"}), 403
if not os.path.exists(filepath): if not os.path.exists(filepath):
return jsonify({"error": "文件不存在"}), 404 return jsonify({"error": "文件不存在"}), 404
@@ -730,7 +784,8 @@ def delete_document(doc_path: str) -> Tuple[Any, int]:
}) })
except Exception as e: except Exception as e:
return jsonify({"error": f"删除失败: {str(e)}"}), 500 logger.error(f"删除文档异常: {e}")
return jsonify({"error": "删除失败,请稍后重试"}), 500
@document_bp.route('/documents/<path:doc_path>/chunks', methods=['GET']) @document_bp.route('/documents/<path:doc_path>/chunks', methods=['GET'])
@@ -817,7 +872,10 @@ def preview_document(doc_path: str) -> Tuple[Any, int]:
# 查询参数 # 查询参数
chunk_index_str = request.args.get('chunk_index') chunk_index_str = request.args.get('chunk_index')
context_count = int(request.args.get('context', 2)) try:
context_count = max(0, min(int(request.args.get('context', 2)), 10))
except (ValueError, TypeError):
context_count = 2
# 获取所有切片 # 获取所有切片
all_chunks = kb_manager.get_document_chunks(collection, filename) all_chunks = kb_manager.get_document_chunks(collection, filename)
@@ -1082,5 +1140,5 @@ def delete_chunks_by_source() -> Tuple[Any, int]:
"message": f"已删除 {deleted_count} 个切片" "message": f"已删除 {deleted_count} 个切片"
}) })
except Exception as e: except Exception as e:
logger.error(f"批量删除切片失败: {e}") logger.error(f"操作异常: {e}")
return jsonify({"error": str(e)}), 500 return jsonify({"error": "操作失败,请稍后重试"}), 500

View File

@@ -33,6 +33,9 @@
from flask import Blueprint, request, jsonify from flask import Blueprint, request, jsonify
from auth.gateway import require_gateway_auth, require_role from auth.gateway import require_gateway_auth, require_role
import logging
logger = logging.getLogger(__name__)
feedback_bp = Blueprint('feedback', __name__) feedback_bp = Blueprint('feedback', __name__)
@@ -95,7 +98,8 @@ def submit_feedback():
"suggestion_id": result.get('suggestion_id') "suggestion_id": result.get('suggestion_id')
}) })
except Exception as e: except Exception as e:
return jsonify({"error": str(e)}), 500 logger.error(f"反馈操作异常: {e}")
return jsonify({"error": "操作失败,请稍后重试"}), 500
@feedback_bp.route('/feedback/stats', methods=['GET']) @feedback_bp.route('/feedback/stats', methods=['GET'])
@@ -113,7 +117,8 @@ def get_feedback_stats():
"stats": stats "stats": stats
}) })
except Exception as e: except Exception as e:
return jsonify({"error": str(e)}), 500 logger.error(f"反馈操作异常: {e}")
return jsonify({"error": "操作失败,请稍后重试"}), 500
@feedback_bp.route('/feedback/list', methods=['GET']) @feedback_bp.route('/feedback/list', methods=['GET'])
@@ -141,7 +146,8 @@ def get_feedback_list():
"total": len(feedbacks) "total": len(feedbacks)
}) })
except Exception as e: except Exception as e:
return jsonify({"error": str(e)}), 500 logger.error(f"反馈操作异常: {e}")
return jsonify({"error": "操作失败,请稍后重试"}), 500
@feedback_bp.route('/reports/weekly', methods=['GET']) @feedback_bp.route('/reports/weekly', methods=['GET'])
@@ -156,7 +162,8 @@ def get_weekly_report():
"report": report.to_dict() "report": report.to_dict()
}) })
except Exception as e: except Exception as e:
return jsonify({"error": str(e)}), 500 logger.error(f"反馈操作异常: {e}")
return jsonify({"error": "操作失败,请稍后重试"}), 500
@feedback_bp.route('/reports/monthly', methods=['GET']) @feedback_bp.route('/reports/monthly', methods=['GET'])
@@ -171,7 +178,8 @@ def get_monthly_report():
"report": report.to_dict() "report": report.to_dict()
}) })
except Exception as e: except Exception as e:
return jsonify({"error": str(e)}), 500 logger.error(f"反馈操作异常: {e}")
return jsonify({"error": "操作失败,请稍后重试"}), 500
@feedback_bp.route('/faq', methods=['GET']) @feedback_bp.route('/faq', methods=['GET'])
@@ -190,7 +198,8 @@ def get_faq_list():
"total": len(faqs) "total": len(faqs)
}) })
except Exception as e: except Exception as e:
return jsonify({"error": str(e)}), 500 logger.error(f"反馈操作异常: {e}")
return jsonify({"error": "操作失败,请稍后重试"}), 500
@feedback_bp.route('/faq', methods=['POST']) @feedback_bp.route('/faq', methods=['POST'])
@@ -233,7 +242,8 @@ def create_faq():
"message": "FAQ已创建请通过 /faq/<id>/approve 接口确认后生效" "message": "FAQ已创建请通过 /faq/<id>/approve 接口确认后生效"
}) })
except Exception as e: except Exception as e:
return jsonify({"error": str(e)}), 500 logger.error(f"反馈操作异常: {e}")
return jsonify({"error": "操作失败,请稍后重试"}), 500
@feedback_bp.route('/faq/<int:faq_id>/approve', methods=['POST']) @feedback_bp.route('/faq/<int:faq_id>/approve', methods=['POST'])
@@ -276,7 +286,8 @@ def approve_faq(faq_id):
"message": "FAQ已批准并同步到知识库" "message": "FAQ已批准并同步到知识库"
}) })
except Exception as e: except Exception as e:
return jsonify({"error": str(e)}), 500 logger.error(f"反馈操作异常: {e}")
return jsonify({"error": "操作失败,请稍后重试"}), 500
@feedback_bp.route('/faq/<int:faq_id>', methods=['PUT']) @feedback_bp.route('/faq/<int:faq_id>', methods=['PUT'])
@@ -326,7 +337,8 @@ def update_faq(faq_id):
"sync_status": sync_status "sync_status": sync_status
}) })
except Exception as e: except Exception as e:
return jsonify({"error": str(e)}), 500 logger.error(f"反馈操作异常: {e}")
return jsonify({"error": "操作失败,请稍后重试"}), 500
@feedback_bp.route('/faq/<int:faq_id>', methods=['DELETE']) @feedback_bp.route('/faq/<int:faq_id>', methods=['DELETE'])
@@ -358,7 +370,8 @@ def delete_faq(faq_id):
"message": "FAQ删除成功" if deleted else "FAQ不存在" "message": "FAQ删除成功" if deleted else "FAQ不存在"
}) })
except Exception as e: except Exception as e:
return jsonify({"error": str(e)}), 500 logger.error(f"反馈操作异常: {e}")
return jsonify({"error": "操作失败,请稍后重试"}), 500
@feedback_bp.route('/faq/suggestions', methods=['GET']) @feedback_bp.route('/faq/suggestions', methods=['GET'])
@@ -378,7 +391,8 @@ def get_faq_suggestions():
"total": len(suggestions) "total": len(suggestions)
}) })
except Exception as e: except Exception as e:
return jsonify({"error": str(e)}), 500 logger.error(f"反馈操作异常: {e}")
return jsonify({"error": "操作失败,请稍后重试"}), 500
@feedback_bp.route('/faq/suggestions/<int:suggestion_id>/approve', methods=['POST']) @feedback_bp.route('/faq/suggestions/<int:suggestion_id>/approve', methods=['POST'])
@@ -413,7 +427,8 @@ def approve_faq_suggestion(suggestion_id):
else: else:
return jsonify({"error": result.get('error', '批准失败')}), 400 return jsonify({"error": result.get('error', '批准失败')}), 400
except Exception as e: except Exception as e:
return jsonify({"error": str(e)}), 500 logger.error(f"反馈操作异常: {e}")
return jsonify({"error": "操作失败,请稍后重试"}), 500
@feedback_bp.route('/faq/suggestions/<int:suggestion_id>/reject', methods=['POST']) @feedback_bp.route('/faq/suggestions/<int:suggestion_id>/reject', methods=['POST'])
@@ -429,7 +444,8 @@ def reject_faq_suggestion(suggestion_id):
"message": "FAQ建议已拒绝" if rejected else "建议不存在" "message": "FAQ建议已拒绝" if rejected else "建议不存在"
}) })
except Exception as e: except Exception as e:
return jsonify({"error": str(e)}), 500 logger.error(f"反馈操作异常: {e}")
return jsonify({"error": "操作失败,请稍后重试"}), 500
# ==================== Bad Case 分析接口 ==================== # ==================== Bad Case 分析接口 ====================
@@ -472,7 +488,8 @@ def get_bad_cases():
] ]
}) })
except Exception as e: except Exception as e:
return jsonify({"error": str(e)}), 500 logger.error(f"反馈操作异常: {e}")
return jsonify({"error": "操作失败,请稍后重试"}), 500
@feedback_bp.route('/feedback/blacklist', methods=['GET']) @feedback_bp.route('/feedback/blacklist', methods=['GET'])
@@ -497,4 +514,5 @@ def get_chunk_blacklist():
"usage": "在检索时过滤这些来源以提升回答质量" "usage": "在检索时过滤这些来源以提升回答质量"
}) })
except Exception as e: except Exception as e:
return jsonify({"error": str(e)}), 500 logger.error(f"反馈操作异常: {e}")
return jsonify({"error": "操作失败,请稍后重试"}), 500

View File

@@ -9,8 +9,11 @@
""" """
import os import os
import logging
from flask import Blueprint, send_file, jsonify, current_app from flask import Blueprint, send_file, jsonify, current_app
logger = logging.getLogger(__name__)
image_bp = Blueprint('images', __name__) image_bp = Blueprint('images', __name__)
@@ -70,7 +73,8 @@ def get_image(image_id: str):
return send_file(image_path, mimetype=mimetype) return send_file(image_path, mimetype=mimetype)
except Exception as e: except Exception as e:
return jsonify({"error": f"读取图片失败: {str(e)}"}), 500 logger.error(f"读取图片异常: {e}")
return jsonify({"error": "读取图片失败"}), 500
return jsonify({"error": "图片不存在", "image_id": image_id}), 404 return jsonify({"error": "图片不存在", "image_id": image_id}), 404
@@ -122,7 +126,8 @@ def get_image_info(image_id: str):
"url": f"/images/{image_id}" "url": f"/images/{image_id}"
}) })
except Exception as e: except Exception as e:
return jsonify({"error": f"读取图片信息失败: {str(e)}"}), 500 logger.error(f"读取图片信息异常: {e}")
return jsonify({"error": "读取图片信息失败"}), 500
return jsonify({"error": "图片不存在", "image_id": image_id}), 404 return jsonify({"error": "图片不存在", "image_id": image_id}), 404
@@ -179,7 +184,8 @@ def list_images():
}) })
except Exception as e: except Exception as e:
return jsonify({"error": f"列出图片失败: {str(e)}"}), 500 logger.error(f"列出图片异常: {e}")
return jsonify({"error": "列出图片失败"}), 500
@image_bp.route('/images/stats', methods=['GET']) @image_bp.route('/images/stats', methods=['GET'])
@@ -222,4 +228,5 @@ def image_stats():
}) })
except Exception as e: except Exception as e:
return jsonify({"error": f"获取统计信息失败: {str(e)}"}), 500 logger.error(f"获取统计信息异常: {e}")
return jsonify({"error": "获取统计信息失败"}), 500

View File

@@ -391,10 +391,11 @@ def sync_documents() -> Tuple[Any, int]:
} }
}) })
except Exception as e: except Exception as e:
logger.error(f"知识库操作异常: {e}")
results.append({ results.append({
"collection": "all", "collection": "all",
"status": "error", "status": "error",
"message": str(e) "message": "操作失败"
}) })
else: else:
# 没有 sync_service返回提示 # 没有 sync_service返回提示
@@ -462,7 +463,8 @@ def debug_scan() -> Tuple[Any, int]:
result["scanned_count"] = len(scanned) result["scanned_count"] = len(scanned)
result["scanned_ids"] = list(scanned.keys())[:10] result["scanned_ids"] = list(scanned.keys())[:10]
except Exception as e: except Exception as e:
result["scan_error"] = str(e) logger.error(f"扫描异常: {e}")
result["scan_error"] = "扫描失败"
return jsonify(result) return jsonify(result)
@@ -499,9 +501,11 @@ def reindex_collection(kb_name: str) -> Tuple[Any, int]:
from data.db import get_connection from data.db import get_connection
with get_connection("knowledge") as conn: with get_connection("knowledge") as conn:
cursor = conn.cursor() cursor = conn.cursor()
# 转义 LIKE 通配符,防止 kb_name 中的 % 或 _ 导致非预期匹配
escaped_kb = kb_name.replace('%', '\\%').replace('_', '\\_')
# 删除以 "{kb_name}/" 或 "{kb_name}\" 开头的文档哈希(兼容 Windows 和 Linux # 删除以 "{kb_name}/" 或 "{kb_name}\" 开头的文档哈希(兼容 Windows 和 Linux
cursor.execute("DELETE FROM document_hashes WHERE document_id LIKE ? OR document_id LIKE ?", cursor.execute("DELETE FROM document_hashes WHERE document_id LIKE ? ESCAPE '\\' OR document_id LIKE ? ESCAPE '\\'",
(f"{kb_name}/%", f"{kb_name}\\%")) (f"{escaped_kb}/%", f"{escaped_kb}\\%"))
deleted = cursor.rowcount deleted = cursor.rowcount
logger.info(f"已清除 {deleted} 条哈希记录: {kb_name}") logger.info(f"已清除 {deleted} 条哈希记录: {kb_name}")
except Exception as e: except Exception as e:
@@ -520,7 +524,8 @@ def reindex_collection(kb_name: str) -> Tuple[Any, int]:
"errors": result.errors "errors": result.errors
}) })
except Exception as e: except Exception as e:
return jsonify({"error": str(e)}), 500 logger.error(f"reindex 异常: {e}")
return jsonify({"error": "操作失败,请稍后重试"}), 500
else: else:
return jsonify({"error": "同步服务不可用"}), 503 return jsonify({"error": "同步服务不可用"}), 503
@@ -633,7 +638,8 @@ def deprecate_document(kb_name: str, filename: str) -> Tuple[Any, int]:
) )
return jsonify(result) return jsonify(result)
except Exception as e: except Exception as e:
return jsonify({"success": False, "error": str(e)}), 500 logger.error(f"操作异常: {e}")
return jsonify({"success": False, "error": "操作失败,请稍后重试"}), 500
@kb_bp.route('/collections/<kb_name>/documents/<path:filename>/restore', methods=['POST']) @kb_bp.route('/collections/<kb_name>/documents/<path:filename>/restore', methods=['POST'])
@@ -664,7 +670,8 @@ def restore_document(kb_name: str, filename: str) -> Tuple[Any, int]:
result = kb_manager.restore_document(kb_name, filename) result = kb_manager.restore_document(kb_name, filename)
return jsonify(result) return jsonify(result)
except Exception as e: except Exception as e:
return jsonify({"success": False, "error": str(e)}), 500 logger.error(f"操作异常: {e}")
return jsonify({"success": False, "error": "操作失败,请稍后重试"}), 500
@kb_bp.route('/collections/<kb_name>/documents/<path:filename>/versions', methods=['GET']) @kb_bp.route('/collections/<kb_name>/documents/<path:filename>/versions', methods=['GET'])
@@ -720,7 +727,8 @@ def get_document_versions(kb_name: str, filename: str) -> Tuple[Any, int]:
"total": len(versions_data) "total": len(versions_data)
}) })
except Exception as e: except Exception as e:
return jsonify({"success": False, "error": str(e)}), 500 logger.error(f"操作异常: {e}")
return jsonify({"success": False, "error": "操作失败,请稍后重试"}), 500
@kb_bp.route('/collections/<kb_name>/update-image-descriptions', methods=['POST']) @kb_bp.route('/collections/<kb_name>/update-image-descriptions', methods=['POST'])
@@ -754,7 +762,8 @@ def update_image_descriptions(kb_name: str) -> Tuple[Any, int]:
result = kb_manager.update_image_descriptions(kb_name) result = kb_manager.update_image_descriptions(kb_name)
return jsonify(result) return jsonify(result)
except Exception as e: except Exception as e:
return jsonify({"success": False, "error": str(e)}), 500 logger.error(f"操作异常: {e}")
return jsonify({"success": False, "error": "操作失败,请稍后重试"}), 500
@kb_bp.route('/collections/sync-vlm-cache', methods=['POST']) @kb_bp.route('/collections/sync-vlm-cache', methods=['POST'])

View File

@@ -33,6 +33,8 @@ def get_sessions():
} }
""" """
session_manager = current_app.config['SESSION_MANAGER'] session_manager = current_app.config['SESSION_MANAGER']
if session_manager is None:
return jsonify({"error": "会话服务不可用"}), 503
user_id = request.current_user["user_id"] user_id = request.current_user["user_id"]
sessions = session_manager.get_user_sessions(user_id, limit=20) sessions = session_manager.get_user_sessions(user_id, limit=20)
@@ -62,6 +64,8 @@ def get_history(session_id):
} }
""" """
session_manager = current_app.config['SESSION_MANAGER'] session_manager = current_app.config['SESSION_MANAGER']
if session_manager is None:
return jsonify({"error": "会话服务不可用"}), 503
user_id = request.current_user["user_id"] user_id = request.current_user["user_id"]
# 验证会话归属 # 验证会话归属
@@ -81,6 +85,8 @@ def get_history(session_id):
def delete_session(session_id): def delete_session(session_id):
"""删除会话""" """删除会话"""
session_manager = current_app.config['SESSION_MANAGER'] session_manager = current_app.config['SESSION_MANAGER']
if session_manager is None:
return jsonify({"error": "会话服务不可用"}), 503
user_id = request.current_user["user_id"] user_id = request.current_user["user_id"]
# 验证会话归属 # 验证会话归属
@@ -100,6 +106,8 @@ def delete_session(session_id):
def clear_history(session_id): def clear_history(session_id):
"""清空会话历史(保留会话)""" """清空会话历史(保留会话)"""
session_manager = current_app.config['SESSION_MANAGER'] session_manager = current_app.config['SESSION_MANAGER']
if session_manager is None:
return jsonify({"error": "会话服务不可用"}), 503
user_id = request.current_user["user_id"] user_id = request.current_user["user_id"]
# 验证会话归属 # 验证会话归属

View File

@@ -112,10 +112,11 @@ def trigger_sync() -> Tuple[Any, int]:
message="同步完成" message="同步完成"
) )
except Exception as e: except Exception as e:
logger.error(f"同步操作异常: {e}")
return error_response( return error_response(
error="SYNC_ERROR", error="SYNC_ERROR",
error_code=SYNC_ERROR, error_code=SYNC_ERROR,
message=str(e), message="操作失败",
http_status=500 http_status=500
) )
@@ -164,11 +165,12 @@ def get_sync_status() -> Tuple[Any, int]:
return jsonify(status) return jsonify(status)
except Exception as e: except Exception as e:
logger.error(f"状态查询异常: {e}")
return jsonify({ return jsonify({
"status": "failed", "status": "failed",
"status_code": INTERNAL_ERROR, "status_code": INTERNAL_ERROR,
"enabled": True, "enabled": True,
"error": str(e) "error": "操作失败"
}) })
@@ -196,10 +198,11 @@ def get_sync_history() -> Tuple[Any, int]:
history = service.get_sync_history(limit=limit) if hasattr(service, 'get_sync_history') else [] history = service.get_sync_history(limit=limit) if hasattr(service, 'get_sync_history') else []
return jsonify({"history": history}) return jsonify({"history": history})
except Exception as e: except Exception as e:
logger.error(f"同步操作异常: {e}")
return error_response( return error_response(
error="SYNC_ERROR", error="SYNC_ERROR",
error_code=SYNC_ERROR, error_code=SYNC_ERROR,
message=str(e), message="操作失败",
http_status=500 http_status=500
) )
@@ -230,10 +233,11 @@ def get_change_logs() -> Tuple[Any, int]:
changes = service.get_change_logs(limit=limit, collection=collection) if hasattr(service, 'get_change_logs') else [] changes = service.get_change_logs(limit=limit, collection=collection) if hasattr(service, 'get_change_logs') else []
return jsonify({"changes": changes}) return jsonify({"changes": changes})
except Exception as e: except Exception as e:
logger.error(f"同步操作异常: {e}")
return error_response( return error_response(
error="SYNC_ERROR", error="SYNC_ERROR",
error_code=SYNC_ERROR, error_code=SYNC_ERROR,
message=str(e), message="操作失败",
http_status=500 http_status=500
) )
@@ -275,10 +279,11 @@ def start_sync_monitor() -> Tuple[Any, int]:
else: else:
return jsonify({"status": "success", "message": "文件监控功能不可用"}) return jsonify({"status": "success", "message": "文件监控功能不可用"})
except Exception as e: except Exception as e:
logger.error(f"同步操作异常: {e}")
return error_response( return error_response(
error="SYNC_ERROR", error="SYNC_ERROR",
error_code=SYNC_ERROR, error_code=SYNC_ERROR,
message=str(e), message="操作失败",
http_status=500 http_status=500
) )
@@ -303,9 +308,10 @@ def stop_sync_monitor() -> Tuple[Any, int]:
service.stop() service.stop()
return jsonify({"status": "success", "status_code": SYNC_SUCCESS, "message": "文件监控已停止"}) return jsonify({"status": "success", "status_code": SYNC_SUCCESS, "message": "文件监控已停止"})
except Exception as e: except Exception as e:
logger.error(f"同步操作异常: {e}")
return error_response( return error_response(
error="SYNC_ERROR", error="SYNC_ERROR",
error_code=SYNC_ERROR, error_code=SYNC_ERROR,
message=str(e), message="操作失败",
http_status=500 http_status=500
) )

View File

@@ -14,6 +14,7 @@ BM25 关键词检索索引
import os import os
import pickle import pickle
import threading
import numpy as np import numpy as np
from rank_bm25 import BM25Okapi from rank_bm25 import BM25Okapi
import jieba import jieba
@@ -105,16 +106,19 @@ class BM25Index:
# ==================== 全局 BM25 索引管理器 ==================== # ==================== 全局 BM25 索引管理器 ====================
_bm25_indexer: BM25Index = None _bm25_indexer: BM25Index = None
_bm25_lock = threading.Lock() # 单例创建锁,防止多线程竞态条件
def get_bm25_indexer() -> BM25Index: def get_bm25_indexer() -> BM25Index:
""" """
获取全局 BM25 索引器实例 获取全局 BM25 索引器实例(双重检查锁定,线程安全)
Returns: Returns:
BM25Index 实例 BM25Index 实例
""" """
global _bm25_indexer global _bm25_indexer
if _bm25_indexer is None:
with _bm25_lock:
if _bm25_indexer is None: if _bm25_indexer is None:
_bm25_indexer = BM25Index() _bm25_indexer = BM25Index()
return _bm25_indexer return _bm25_indexer
@@ -122,7 +126,7 @@ def get_bm25_indexer() -> BM25Index:
def init_bm25_indexer(ids=None, documents=None, metadatas=None) -> BM25Index: def init_bm25_indexer(ids=None, documents=None, metadatas=None) -> BM25Index:
""" """
初始化 BM25 索引器并添加文档 初始化 BM25 索引器并添加文档(加锁,线程安全)
Args: Args:
ids: 文档 ID 列表 ids: 文档 ID 列表
@@ -133,6 +137,7 @@ def init_bm25_indexer(ids=None, documents=None, metadatas=None) -> BM25Index:
初始化后的 BM25Index 实例 初始化后的 BM25Index 实例
""" """
global _bm25_indexer global _bm25_indexer
with _bm25_lock:
_bm25_indexer = BM25Index() _bm25_indexer = BM25Index()
if ids and documents: if ids and documents:
_bm25_indexer.add_documents(ids, documents, metadatas or []) _bm25_indexer.add_documents(ids, documents, metadatas or [])

View File

@@ -31,6 +31,7 @@ import os
import gc import gc
import time import time
import logging import logging
import threading
import numpy as np import numpy as np
from typing import List, Tuple, Optional, Dict, Any from typing import List, Tuple, Optional, Dict, Any
@@ -50,6 +51,7 @@ except ImportError:
# 延迟导入,防止循环依赖 # 延迟导入,防止循环依赖
_engine_instance = None _engine_instance = None
_engine_lock = threading.Lock() # 单例创建锁,防止多线程竞态条件
try: try:
from config import ( from config import (
@@ -327,8 +329,10 @@ class RAGEngine:
@classmethod @classmethod
def get_instance(cls) -> 'RAGEngine': def get_instance(cls) -> 'RAGEngine':
"""获取引擎单例实例""" """获取引擎单例实例(双重检查锁定,线程安全)"""
global _engine_instance global _engine_instance
if _engine_instance is None:
with _engine_lock:
if _engine_instance is None: if _engine_instance is None:
_engine_instance = cls() _engine_instance = cls()
return _engine_instance return _engine_instance

View File

@@ -18,6 +18,7 @@
import json import json
import logging import logging
import hashlib import hashlib
import threading
from dataclasses import dataclass, asdict from dataclasses import dataclass, asdict
from typing import List, Dict, Optional, Tuple from typing import List, Dict, Optional, Tuple
@@ -503,10 +504,13 @@ class IntentAnalyzer:
# ==================== 便捷函数 ==================== # ==================== 便捷函数 ====================
_analyzer = None _analyzer = None
_analyzer_lock = threading.Lock() # 单例创建锁,防止多线程竞态条件
def get_intent_analyzer() -> IntentAnalyzer: def get_intent_analyzer() -> IntentAnalyzer:
"""获取意图分析器单例""" """获取意图分析器单例(双重检查锁定,线程安全)"""
global _analyzer global _analyzer
if _analyzer is None:
with _analyzer_lock:
if _analyzer is None: if _analyzer is None:
_analyzer = IntentAnalyzer() _analyzer = IntentAnalyzer()
return _analyzer return _analyzer

View File

@@ -237,6 +237,76 @@ def parse_json_list_from_response(content: str) -> Optional[List[dict]]:
return None return None
def extract_json_object(content: str) -> Optional[dict]:
"""
多策略从 LLM 响应中提取 JSON 对象(增强版)
在 parse_json_from_response 基础上增加 fallback 策略:
1. 先调用 parse_json_from_responsemarkdown 代码块 → 直接解析)
2. 失败后 fallback 到正则匹配最外层 {...} 块
Args:
content: LLM 返回的原始内容
Returns:
解析后的字典,全部策略失败返回 None
"""
if not content:
return None
# 策略1+2markdown 代码块提取 + 直接 json.loads
result = parse_json_from_response(content)
if result is not None and isinstance(result, dict):
return result
# 策略3fallback正则匹配最外层 JSON 对象 {...}
brace_match = re.search(r'\{[\s\S]*\}', content)
if brace_match:
try:
parsed = json.loads(brace_match.group(0))
if isinstance(parsed, dict):
return parsed
except (json.JSONDecodeError, TypeError, ValueError):
pass
return None
def extract_json_list(content: str) -> Optional[list]:
"""
多策略从 LLM 响应中提取 JSON 数组(增强版)
在 parse_json_list_from_response 基础上增加 fallback 策略:
1. 先调用 parse_json_list_from_responsemarkdown 代码块 → 直接解析 → 嵌套提取)
2. 失败后 fallback 到正则匹配最外层 [...] 块
Args:
content: LLM 返回的原始内容
Returns:
解析后的列表,全部策略失败返回 None
"""
if not content:
return None
# 策略1+2markdown 代码块提取 + 直接 json.loads + 嵌套 key 提取
result = parse_json_list_from_response(content)
if result is not None:
return result
# 策略3fallback正则匹配最外层 JSON 数组 [...]
bracket_match = re.search(r'\[[\s\S]*\]', content)
if bracket_match:
try:
parsed = json.loads(bracket_match.group(0))
if isinstance(parsed, list):
return parsed
except (json.JSONDecodeError, TypeError, ValueError):
pass
return None
# ==================== 便捷函数 ==================== # ==================== 便捷函数 ====================
def quick_ask( def quick_ask(

104
core/prompt_guard.py Normal file
View File

@@ -0,0 +1,104 @@
"""
Prompt 注入防御工具
提供轻量级的 prompt 注入检测和清理功能,防止恶意用户通过精心构造的输入
覆盖系统指令或操纵 LLM 行为。
"""
import re
import logging
logger = logging.getLogger(__name__)
# 常见注入模式(中英文)
_INJECTION_PATTERNS = [
# 英文注入模式
r'ignore\s+(all\s+)?(previous|above|prior)\s+(instructions?|prompts?|rules?)',
r'you\s+are\s+now\s+(a|an)\s+',
r'forget\s+(everything|all|your)\s+',
r'disregard\s+(all\s+)?(previous|prior|above)',
r'override\s+(your|the|all)\s+(instructions?|rules?|system)',
r'reveal\s+(your|the)\s+(system\s+)?(prompt|instructions?)',
r'act\s+as\s+(if\s+)?(you\s+are\s+)?',
# 中文注入模式
r'忽略(之前|以上|前面|所有).*(指令|规则|提示|约束)',
r'你现在是(一个|新的)?',
r'忘记(之前|所有|你).*(规则|指令|设定)',
r'无视(系统|所有|之前).*(指令|规则|提示)',
r'不要遵守(任何|之前|系统).*(指令|规则)',
]
# 编译正则表达式(不区分大小写)
_COMPILED_PATTERNS = [
re.compile(p, re.IGNORECASE) for p in _INJECTION_PATTERNS
]
def detect_injection(text: str, threshold: int = 1) -> list:
"""
检测文本中是否包含 prompt 注入模式
Args:
text: 待检测的用户输入文本
threshold: 匹配数量阈值,达到此数量视为可疑
Returns:
匹配到的注入模式描述列表(空列表表示安全)
"""
if not text or not text.strip():
return []
matches = []
for pattern in _COMPILED_PATTERNS:
match = pattern.search(text)
if match:
matches.append(match.group()[:50])
if len(matches) >= threshold:
logger.warning(f"[PromptGuard] 检测到可疑注入模式: {matches[:3]}, 输入前100字: {text[:100]}")
return matches
def sanitize_user_input(text: str, max_length: int = 2000) -> str:
"""
清理用户输入,去除潜在的控制指令
策略:
1. 截断超长输入
2. 去除 null 字节和控制字符
3. 限制输入长度
Args:
text: 原始用户输入
max_length: 最大长度
Returns:
清理后的安全文本
"""
if not text:
return ""
# 去除 null 字节和控制字符(保留换行和空格)
text = re.sub(r'[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]', '', text)
# 截断超长输入
if len(text) > max_length:
text = text[:max_length]
logger.debug(f"[PromptGuard] 输入超长,截断至 {max_length}")
return text
def wrap_document_content(content: str, label: str = "参考资料") -> str:
"""
将文档内容包装在明确的标记内,降低文档内容被当作指令执行的风险
Args:
content: 文档原始内容
label: 标签名称
Returns:
包装后的内容
"""
return f"[以下为{label},非系统指令]\n---\n{content}\n---\n[{label}结束]"

View File

@@ -16,7 +16,6 @@
import json import json
import logging import logging
import re
import threading import threading
import time import time
from concurrent.futures import ThreadPoolExecutor, as_completed from concurrent.futures import ThreadPoolExecutor, as_completed
@@ -24,7 +23,7 @@ from functools import wraps
from typing import List, Dict, Any, Optional from typing import List, Dict, Any, Optional
# 导入 LLM 工具函数 # 导入 LLM 工具函数
from core.llm_utils import call_llm from core.llm_utils import call_llm, extract_json_object
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -378,47 +377,12 @@ class AnswerGrader:
return result return result
def _extract_json(self, response: str) -> dict: def _extract_json(self, response: str) -> dict:
""" """多策略从 LLM 响应中提取 JSON 对象(使用共享工具)"""
多策略从 LLM 响应中提取 JSON 对象 result = extract_json_object(response)
if result is not None:
策略优先级: return result
1. markdown 代码块提取 ```json ... ```
2. 直接 json.loads
3. 正则匹配第一个 {...} 块
"""
if not response: if not response:
raise ValueError("LLM 返回为空") raise ValueError("LLM 返回为空")
# 策略1提取 markdown 代码块
json_match = re.search(r'```(?:json)?\s*([\s\S]*?)\s*```', response)
if json_match:
json_str = json_match.group(1).strip()
try:
result = json.loads(json_str)
if isinstance(result, dict):
return result
except json.JSONDecodeError:
pass # 继续下一策略
# 策略2直接解析整个响应
try:
result = json.loads(response.strip())
if isinstance(result, dict):
return result
except json.JSONDecodeError:
pass # 继续下一策略
# 策略3正则匹配最外层 JSON 对象
brace_match = re.search(r'\{[\s\S]*\}', response)
if brace_match:
try:
result = json.loads(brace_match.group(0))
if isinstance(result, dict):
return result
except json.JSONDecodeError:
pass
# 全部策略失败
raise json.JSONDecodeError( raise json.JSONDecodeError(
f"无法从 LLM 响应中提取有效 JSON响应前300字: {response[:300]}", f"无法从 LLM 响应中提取有效 JSON响应前300字: {response[:300]}",
response, 0 response, 0

View File

@@ -28,7 +28,7 @@ def main():
parser = argparse.ArgumentParser(description='RAG API 服务') parser = argparse.ArgumentParser(description='RAG API 服务')
parser.add_argument('--host', default='0.0.0.0', help='监听地址(默认 0.0.0.0') parser.add_argument('--host', default='0.0.0.0', help='监听地址(默认 0.0.0.0')
parser.add_argument('--port', type=int, default=5001, help='监听端口(默认 5001') parser.add_argument('--port', type=int, default=5001, help='监听端口(默认 5001')
parser.add_argument('--debug', action='store_true', default=True, help='调试模式') parser.add_argument('--debug', action='store_true', default=False, help='调试模式')
parser.add_argument('--no-debug', action='store_true', help='关闭调试模式') parser.add_argument('--no-debug', action='store_true', help='关闭调试模式')
args = parser.parse_args() args = parser.parse_args()

View File

@@ -24,6 +24,9 @@ logger = logging.getLogger(__name__)
# 大表切片阈值 # 大表切片阈值
MAX_ROWS_PER_CHUNK = 200 MAX_ROWS_PER_CHUNK = 200
# 文件大小限制
MAX_FILE_SIZE = 50 * 1024 * 1024 # 50MB
@dataclass @dataclass
class UnifiedChunk: class UnifiedChunk:
@@ -67,6 +70,13 @@ def parse_excel(
} }
""" """
filepath = Path(filepath) filepath = Path(filepath)
# 检查文件大小
if filepath.exists():
file_size = filepath.stat().st_size
if file_size > MAX_FILE_SIZE:
raise ValueError(f"文件过大: {file_size / 1024 / 1024:.1f}MB最大允许 {MAX_FILE_SIZE / 1024 / 1024:.0f}MB")
if not filepath.exists(): if not filepath.exists():
raise FileNotFoundError(f"文件不存在: {filepath}") raise FileNotFoundError(f"文件不存在: {filepath}")

View File

@@ -46,6 +46,9 @@ import logging
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
# 文件大小限制
MAX_PDF_SIZE = 100 * 1024 * 1024 # 100MB
# 支持的文件格式 # 支持的文件格式
SUPPORTED_FORMATS = { SUPPORTED_FORMATS = {
'.pdf': 'PDF 文档', '.pdf': 'PDF 文档',
@@ -176,6 +179,11 @@ def parse_with_mineru_online(
if not file_path.exists(): if not file_path.exists():
raise FileNotFoundError(f"文件不存在: {file_path}") raise FileNotFoundError(f"文件不存在: {file_path}")
# 检查文件大小
file_size = file_path.stat().st_size
if file_size > MAX_PDF_SIZE:
raise ValueError(f"文件过大: {file_size / 1024 / 1024:.1f}MB最大允许 {MAX_PDF_SIZE / 1024 / 1024:.0f}MB")
logger.info(f"使用 MinerU 在线 API 解析: {file_path.name}") logger.info(f"使用 MinerU 在线 API 解析: {file_path.name}")
# 读取文件内容 # 读取文件内容
@@ -600,6 +608,11 @@ def parse_with_mineru(
if not file_path.exists(): if not file_path.exists():
raise FileNotFoundError(f"文件不存在: {file_path}") raise FileNotFoundError(f"文件不存在: {file_path}")
# 检查文件大小
file_size = file_path.stat().st_size
if file_size > MAX_PDF_SIZE:
raise ValueError(f"文件过大: {file_size / 1024 / 1024:.1f}MB最大允许 {MAX_PDF_SIZE / 1024 / 1024:.0f}MB")
# 检查文件格式 # 检查文件格式
suffix = file_path.suffix.lower() suffix = file_path.suffix.lower()
if suffix not in SUPPORTED_FORMATS: if suffix not in SUPPORTED_FORMATS:
@@ -629,8 +642,17 @@ def parse_with_mineru(
if not mineru_exe.exists(): if not mineru_exe.exists():
mineru_exe = "mineru" # 回退到系统 PATH mineru_exe = "mineru" # 回退到系统 PATH
# 参数白名单校验,防止注入非法参数
ALLOWED_BACKENDS = {'auto', 'pipeline', 'vlm', 'vlm-sglang', 'vlm-auto-engine', 'hybrid-auto-engine', 'ocr'}
ALLOWED_LANGS = {'ch', 'en', 'ch_lite', 'en_lite', 'formula', 'table'}
if backend not in ALLOWED_BACKENDS:
backend = 'pipeline'
if lang not in ALLOWED_LANGS:
lang = 'ch'
cmd = [ cmd = [
str(mineru_exe), str(mineru_exe),
"--",
"-p", str(file_path), "-p", str(file_path),
"-o", str(output_dir), "-o", str(output_dir),
"-m", "auto", "-m", "auto",
@@ -1458,6 +1480,11 @@ def parse_with_mineru_persistent(
if not file_path.exists(): if not file_path.exists():
raise FileNotFoundError(f"文件不存在: {file_path}") raise FileNotFoundError(f"文件不存在: {file_path}")
# 检查文件大小
file_size = file_path.stat().st_size
if file_size > MAX_PDF_SIZE:
raise ValueError(f"文件过大: {file_size / 1024 / 1024:.1f}MB最大允许 {MAX_PDF_SIZE / 1024 / 1024:.0f}MB")
# 计算文件 hash用于隔离输出目录 # 计算文件 hash用于隔离输出目录
file_hash = compute_file_hash(str(file_path)) file_hash = compute_file_hash(str(file_path))
output_dir = Path(output_base) / file_hash output_dir = Path(output_base) / file_hash

View File

@@ -5,12 +5,25 @@ TXT 文本解析器
""" """
import logging import logging
import os
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
# 文件大小限制
MAX_TXT_SIZE = 20 * 1024 * 1024 # 20MB
def extract_text_from_txt(filepath): def extract_text_from_txt(filepath):
"""从TXT提取文本""" """从TXT提取文本"""
# 检查文件大小
try:
file_size = os.path.getsize(filepath)
if file_size > MAX_TXT_SIZE:
logger.error(f"TXT文件过大: {file_size / 1024 / 1024:.1f}MB")
return ""
except OSError:
return ""
try: try:
with open(filepath, 'r', encoding='utf-8') as f: with open(filepath, 'r', encoding='utf-8') as f:
return f.read() return f.read()