feat(knowledge): 知识源维护与 TXT/Markdown/PDF 导入
支持文本、Markdown 与非扫描 PDF 上传解析，展示提取文本与切块结果。
This commit is contained in:
WangShen007 committed
1 parent
4388e1d203
commit
53825b0cfb
23 files changed
+7545
No files matched your search
@@ -0,0 +1,384 @@
|
||||
"""
|
||||
Copyright 2024-2026 ChatterMate
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import os
|
||||
import uuid
|
||||
from typing import Dict, Optional, List, Set
|
||||
from app.core.config import settings
|
||||
from app.core.s3 import upload_file_to_s3
|
||||
from app.services.file_storage import resolve_public_url
|
||||
from app.core.logger import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
# Allowed file types and size limits
|
||||
ALLOWED_IMAGE_TYPES = {'image/jpeg', 'image/jpg', 'image/png', 'image/gif', 'image/webp'}
|
||||
ALLOWED_DOCUMENT_TYPES = {'application/pdf', 'application/msword',
|
||||
'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
|
||||
'text/plain', 'text/csv', 'application/vnd.ms-excel',
|
||||
'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet'}
|
||||
ALLOWED_FILE_TYPES = ALLOWED_IMAGE_TYPES | ALLOWED_DOCUMENT_TYPES
|
||||
|
||||
# Security: SVG removed from default allowed types due to XSS risks
|
||||
# Can be enabled per-agent if needed with proper sanitization
|
||||
|
||||
MAX_FILE_SIZE = 10 * 1024 * 1024 # 10MB
|
||||
MAX_IMAGE_SIZE = 5 * 1024 * 1024 # 5MB
|
||||
|
||||
# Magic byte signatures for file type validation
|
||||
# This prevents MIME type spoofing attacks
|
||||
MAGIC_BYTES = {
|
||||
# Images
|
||||
'image/jpeg': [
|
||||
(b'\xff\xd8\xff\xe0', 0), # JPEG JFIF
|
||||
(b'\xff\xd8\xff\xe1', 0), # JPEG EXIF
|
||||
(b'\xff\xd8\xff\xe2', 0), # JPEG ICC
|
||||
(b'\xff\xd8\xff\xe8', 0), # JPEG SPIFF
|
||||
(b'\xff\xd8\xff\xdb', 0), # JPEG DQT
|
||||
(b'\xff\xd8\xff\xee', 0), # JPEG Adobe
|
||||
],
|
||||
'image/jpg': [
|
||||
(b'\xff\xd8\xff\xe0', 0),
|
||||
(b'\xff\xd8\xff\xe1', 0),
|
||||
(b'\xff\xd8\xff\xe2', 0),
|
||||
(b'\xff\xd8\xff\xe8', 0),
|
||||
(b'\xff\xd8\xff\xdb', 0),
|
||||
(b'\xff\xd8\xff\xee', 0),
|
||||
],
|
||||
'image/png': [
|
||||
(b'\x89PNG\r\n\x1a\n', 0),
|
||||
],
|
||||
'image/gif': [
|
||||
(b'GIF87a', 0),
|
||||
(b'GIF89a', 0),
|
||||
],
|
||||
'image/webp': [
|
||||
(b'RIFF', 0), # WebP starts with RIFF, then has WEBP at offset 8
|
||||
],
|
||||
# Documents
|
||||
'application/pdf': [
|
||||
(b'%PDF-', 0),
|
||||
],
|
||||
'application/msword': [
|
||||
(b'\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1', 0), # OLE Compound File (DOC)
|
||||
],
|
||||
'application/vnd.openxmlformats-officedocument.wordprocessingml.document': [
|
||||
(b'PK\x03\x04', 0), # ZIP-based (DOCX)
|
||||
],
|
||||
'application/vnd.ms-excel': [
|
||||
(b'\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1', 0), # OLE Compound File (XLS)
|
||||
],
|
||||
'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet': [
|
||||
(b'PK\x03\x04', 0), # ZIP-based (XLSX)
|
||||
],
|
||||
'text/plain': [], # Text files don't have magic bytes, validated by content
|
||||
'text/csv': [], # CSV files don't have magic bytes, validated by content
|
||||
}
|
||||
|
||||
# File extension to MIME type mapping for validation
|
||||
EXTENSION_TO_MIME = {
|
||||
'.jpg': 'image/jpeg',
|
||||
'.jpeg': 'image/jpeg',
|
||||
'.png': 'image/png',
|
||||
'.gif': 'image/gif',
|
||||
'.webp': 'image/webp',
|
||||
'.pdf': 'application/pdf',
|
||||
'.doc': 'application/msword',
|
||||
'.docx': 'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
|
||||
'.txt': 'text/plain',
|
||||
'.csv': 'text/csv',
|
||||
'.xls': 'application/vnd.ms-excel',
|
||||
'.xlsx': 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
|
||||
}
|
||||
|
||||
# Category labels for display
|
||||
FILE_TYPE_CATEGORIES = {
|
||||
'images': {'image/jpeg', 'image/jpg', 'image/png', 'image/gif', 'image/webp'},
|
||||
'documents': {'application/pdf'},
|
||||
'office': {'application/msword', 'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
|
||||
'application/vnd.ms-excel', 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet'},
|
||||
'text': {'text/plain', 'text/csv'},
|
||||
}
|
||||
|
||||
|
||||
class FileUploadService:
|
||||
"""Service for handling file uploads to S3 or local storage with enhanced security"""
|
||||
|
||||
@staticmethod
|
||||
def validate_magic_bytes(file_content: bytes, claimed_content_type: str) -> tuple[bool, Optional[str]]:
|
||||
"""
|
||||
Validate file content against magic byte signatures to prevent MIME type spoofing.
|
||||
|
||||
Args:
|
||||
file_content: The raw binary content of the file
|
||||
claimed_content_type: The MIME type claimed by the client
|
||||
|
||||
Returns:
|
||||
(is_valid, error_message)
|
||||
"""
|
||||
if claimed_content_type not in MAGIC_BYTES:
|
||||
logger.warning(f"Unknown content type for magic byte validation: {claimed_content_type}")
|
||||
return False, f"File type '{claimed_content_type}' is not allowed"
|
||||
|
||||
signatures = MAGIC_BYTES[claimed_content_type]
|
||||
|
||||
# Text files don't have magic bytes - validate as text content
|
||||
if not signatures:
|
||||
if claimed_content_type in {'text/plain', 'text/csv'}:
|
||||
# Check if content is valid text (not binary)
|
||||
try:
|
||||
# Try to decode as UTF-8, allowing for some common encodings
|
||||
sample = file_content[:8192] # Check first 8KB
|
||||
sample.decode('utf-8')
|
||||
return True, None
|
||||
except UnicodeDecodeError:
|
||||
try:
|
||||
sample.decode('latin-1')
|
||||
return True, None
|
||||
except UnicodeDecodeError:
|
||||
return False, "File content does not appear to be valid text"
|
||||
return True, None
|
||||
|
||||
# Check for WebP special case (has WEBP at offset 8)
|
||||
if claimed_content_type == 'image/webp':
|
||||
if len(file_content) >= 12:
|
||||
if file_content[:4] == b'RIFF' and file_content[8:12] == b'WEBP':
|
||||
return True, None
|
||||
return False, "File content does not match WebP format signature"
|
||||
|
||||
# Check against all possible signatures for this MIME type
|
||||
for signature, offset in signatures:
|
||||
end_offset = offset + len(signature)
|
||||
if len(file_content) >= end_offset:
|
||||
if file_content[offset:end_offset] == signature:
|
||||
return True, None
|
||||
|
||||
return False, f"File content does not match the claimed type '{claimed_content_type}'. The file may be corrupted or mislabeled."
|
||||
|
||||
@staticmethod
|
||||
def validate_extension(filename: str, claimed_content_type: str) -> tuple[bool, Optional[str]]:
|
||||
"""
|
||||
Validate that file extension matches the claimed content type.
|
||||
|
||||
Args:
|
||||
filename: The name of the file
|
||||
claimed_content_type: The MIME type claimed by the client
|
||||
|
||||
Returns:
|
||||
(is_valid, error_message)
|
||||
"""
|
||||
ext = os.path.splitext(filename.lower())[1]
|
||||
|
||||
if not ext:
|
||||
return False, "File must have an extension"
|
||||
|
||||
if ext not in EXTENSION_TO_MIME:
|
||||
return False, f"File extension '{ext}' is not allowed. Allowed extensions: {', '.join(sorted(EXTENSION_TO_MIME.keys()))}"
|
||||
|
||||
expected_mime = EXTENSION_TO_MIME[ext]
|
||||
|
||||
# Handle jpg/jpeg equivalence
|
||||
if claimed_content_type in {'image/jpeg', 'image/jpg'} and expected_mime in {'image/jpeg', 'image/jpg'}:
|
||||
return True, None
|
||||
|
||||
if expected_mime != claimed_content_type:
|
||||
return False, f"File extension '{ext}' does not match content type '{claimed_content_type}'"
|
||||
|
||||
return True, None
|
||||
|
||||
@staticmethod
|
||||
def validate_file_type(content_type: str, allowed_types: Optional[Set[str]] = None) -> bool:
|
||||
"""
|
||||
Validate if the file type is allowed.
|
||||
|
||||
Args:
|
||||
content_type: The MIME type to validate
|
||||
allowed_types: Optional set of allowed MIME types. If None, uses default ALLOWED_FILE_TYPES.
|
||||
"""
|
||||
types_to_check = allowed_types if allowed_types is not None else ALLOWED_FILE_TYPES
|
||||
return content_type in types_to_check
|
||||
|
||||
@staticmethod
|
||||
def get_allowed_types_for_agent(agent_allowed_types: Optional[List[str]] = None) -> Set[str]:
|
||||
"""
|
||||
Get the set of allowed file types for an agent.
|
||||
|
||||
Args:
|
||||
agent_allowed_types: List of category names like ['images', 'documents', 'office', 'text']
|
||||
If None or empty, returns all default allowed types.
|
||||
|
||||
Returns:
|
||||
Set of MIME types that are allowed
|
||||
"""
|
||||
if not agent_allowed_types:
|
||||
return ALLOWED_FILE_TYPES
|
||||
|
||||
allowed = set()
|
||||
for category in agent_allowed_types:
|
||||
category_lower = category.lower()
|
||||
if category_lower in FILE_TYPE_CATEGORIES:
|
||||
allowed.update(FILE_TYPE_CATEGORIES[category_lower])
|
||||
|
||||
# If no valid categories found, return default
|
||||
if not allowed:
|
||||
return ALLOWED_FILE_TYPES
|
||||
|
||||
return allowed
|
||||
|
||||
@staticmethod
|
||||
def validate_file_size(file_size: int, content_type: str) -> tuple[bool, Optional[str]]:
|
||||
"""
|
||||
Validate file size based on content type
|
||||
Returns: (is_valid, error_message)
|
||||
"""
|
||||
if content_type in ALLOWED_IMAGE_TYPES and file_size > MAX_IMAGE_SIZE:
|
||||
return False, f"Image file size exceeds maximum allowed size of {MAX_IMAGE_SIZE / (1024*1024):.1f}MB"
|
||||
|
||||
if file_size > MAX_FILE_SIZE:
|
||||
return False, f"File size exceeds maximum allowed size of {MAX_FILE_SIZE / (1024*1024):.1f}MB"
|
||||
|
||||
return True, None
|
||||
|
||||
@staticmethod
|
||||
def get_friendly_allowed_types_message(allowed_types: Optional[Set[str]] = None) -> str:
|
||||
"""Generate a user-friendly message listing allowed file types."""
|
||||
types_to_use = allowed_types if allowed_types is not None else ALLOWED_FILE_TYPES
|
||||
|
||||
extensions = []
|
||||
for ext, mime in EXTENSION_TO_MIME.items():
|
||||
if mime in types_to_use:
|
||||
extensions.append(ext.upper().replace('.', ''))
|
||||
|
||||
return ', '.join(sorted(set(extensions)))
|
||||
|
||||
@staticmethod
|
||||
async def upload_file(
|
||||
file_data: dict,
|
||||
org_id: str,
|
||||
customer_id: Optional[str] = None,
|
||||
user_id: Optional[str] = None,
|
||||
allowed_types: Optional[Set[str]] = None
|
||||
) -> Dict[str, any]:
|
||||
"""
|
||||
Upload a file to S3 or local storage with enhanced security validation.
|
||||
|
||||
Args:
|
||||
file_data: Dict with 'content' (base64), 'filename', 'content_type', 'size'
|
||||
org_id: Organization ID
|
||||
customer_id: Optional customer ID (for widget users)
|
||||
user_id: Optional user ID (for authenticated users)
|
||||
allowed_types: Optional set of allowed MIME types for this upload
|
||||
|
||||
Returns:
|
||||
Dict with 'file_url', 'signed_url', 'filename', 'content_type', 'size'
|
||||
|
||||
Raises:
|
||||
ValueError: If validation fails or upload fails
|
||||
"""
|
||||
try:
|
||||
filename = file_data.get('filename', '')
|
||||
content_type = file_data.get('content_type', 'application/octet-stream')
|
||||
|
||||
# Validate file data
|
||||
if not file_data.get('content') or not filename:
|
||||
raise ValueError("Missing required file data: content and filename")
|
||||
|
||||
# Determine which types are allowed
|
||||
types_allowed = allowed_types if allowed_types is not None else ALLOWED_FILE_TYPES
|
||||
friendly_types = FileUploadService.get_friendly_allowed_types_message(types_allowed)
|
||||
|
||||
# Step 1: Validate extension matches claimed content type
|
||||
is_valid, error = FileUploadService.validate_extension(filename, content_type)
|
||||
if not is_valid:
|
||||
logger.warning(f"Extension validation failed for {filename}: {error}")
|
||||
raise ValueError(error)
|
||||
|
||||
# Step 2: Validate file type is allowed
|
||||
if not FileUploadService.validate_file_type(content_type, types_allowed):
|
||||
raise ValueError(f"File type '{content_type}' is not allowed. Allowed types: {friendly_types}")
|
||||
|
||||
# Step 3: Decode base64 content
|
||||
try:
|
||||
file_content = base64.b64decode(file_data['content'])
|
||||
except Exception as e:
|
||||
logger.error(f"Error decoding base64 content: {str(e)}")
|
||||
raise ValueError("Invalid file content encoding")
|
||||
|
||||
file_size = len(file_content)
|
||||
|
||||
# Step 4: Validate file size
|
||||
is_valid, error_message = FileUploadService.validate_file_size(file_size, content_type)
|
||||
if not is_valid:
|
||||
raise ValueError(error_message)
|
||||
|
||||
# Step 5: SECURITY - Validate magic bytes to prevent MIME spoofing
|
||||
is_valid, magic_error = FileUploadService.validate_magic_bytes(file_content, content_type)
|
||||
if not is_valid:
|
||||
logger.warning(f"Magic byte validation failed for {filename}: {magic_error}")
|
||||
raise ValueError(f"Security check failed: {magic_error}")
|
||||
|
||||
logger.info(f"File validation passed for {filename}: type={content_type}, size={file_size}")
|
||||
|
||||
# Generate unique filename with sanitized extension
|
||||
file_extension = os.path.splitext(filename)[1].lower()
|
||||
# Additional sanitization: only allow known extensions
|
||||
if file_extension not in EXTENSION_TO_MIME:
|
||||
raise ValueError(f"File extension '{file_extension}' is not allowed")
|
||||
|
||||
unique_filename = f"{uuid.uuid4()}{file_extension}"
|
||||
|
||||
# Determine folder based on auth type
|
||||
folder = f"chat_attachments/{org_id}"
|
||||
|
||||
# Upload to S3 if enabled, otherwise save locally
|
||||
if settings.S3_FILE_STORAGE:
|
||||
file_url = await upload_file_to_s3(
|
||||
file_content=file_content,
|
||||
folder=folder,
|
||||
filename=unique_filename,
|
||||
content_type=content_type
|
||||
)
|
||||
else:
|
||||
# Save locally
|
||||
upload_dir = os.path.join("uploads", folder)
|
||||
os.makedirs(upload_dir, exist_ok=True)
|
||||
|
||||
file_path = os.path.join(upload_dir, unique_filename)
|
||||
with open(file_path, "wb") as f:
|
||||
f.write(file_content)
|
||||
|
||||
file_url = f"/uploads/{folder}/{unique_filename}"
|
||||
|
||||
signed_url = await resolve_public_url(file_url)
|
||||
|
||||
logger.info(f"Successfully uploaded file: {filename} -> {file_url}")
|
||||
|
||||
return {
|
||||
'file_url': file_url,
|
||||
'signed_url': signed_url,
|
||||
'filename': filename,
|
||||
'content_type': content_type,
|
||||
'size': file_size
|
||||
}
|
||||
|
||||
except ValueError:
|
||||
# Re-raise validation errors
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Error in FileUploadService.upload_file: {str(e)}", exc_info=True)
|
||||
raise ValueError(f"Failed to upload file: {str(e)}")
|
||||
|
||||
Reference in new issue
Block a user