feat(knowledge): 知识源维护与 TXT/Markdown/PDF 导入

支持文本、Markdown 与非扫描 PDF 上传解析,展示提取文本与切块结果。
This commit is contained in:
WangShen007 committed 2026-09-29 20:35:00 +08:00
1 parent 4ff934ad98
commit 7937de5d7d
23 files changed
+7545

No files matched your search

+281
View File
@@ -0,0 +1,281 @@
"""
Copyright 2024-2026 ReplyPier
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.
"""
from fastapi import APIRouter, UploadFile, File, HTTPException, Depends, status, Header, Request, Query
from fastapi.responses import FileResponse, JSONResponse
from typing import Optional, Tuple
import uuid
import os
from datetime import datetime
from urllib.parse import quote
from sqlalchemy.orm import Session
from app.core.config import settings
from app.core.s3 import upload_file_to_s3, get_s3_signed_url
from app.core.local_file_urls import verify_local_signature
from app.core.logger import get_logger
from app.core.security import verify_conversation_token
from app.models import User
from app.core.auth import get_current_user
from app.database import get_db
logger = get_logger(__name__)
router = APIRouter()
# Allowed file types and size limits
ALLOWED_IMAGE_TYPES = {'image/jpeg', 'image/jpg', 'image/png', 'image/gif', 'image/webp', 'image/svg+xml'}
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
MAX_FILE_SIZE = 10 * 1024 * 1024 # 10MB
MAX_IMAGE_SIZE = 5 * 1024 * 1024 # 5MB
async def get_current_user_or_widget(
request: Request,
db: Session = Depends(get_db),
authorization: Optional[str] = Header(None),
x_conversation_token: Optional[str] = Header(None)
) -> dict:
"""
Get current user from JWT token, cookies, or widget conversation token
Returns dict with user info or widget info
"""
# Try to get user from cookies or Authorization header (admin/agent users)
try:
# Call get_current_user with the request - it handles both cookies and headers
user = await get_current_user(request, None, db)
if user:
return {
"type": "user",
"user_id": str(user.id),
"org_id": str(user.organization_id)
}
except HTTPException:
# User auth failed, continue to try other methods
pass
except Exception as e:
logger.error(f"Error authenticating user: {str(e)}")
# Try to get widget info from conversation token
if x_conversation_token:
try:
payload = verify_conversation_token(x_conversation_token)
if payload:
return {
"type": "widget",
"widget_id": payload.get("widget_id"),
"org_id": payload.get("org_id"),
"customer_id": payload.get("customer_id")
}
except Exception as e:
logger.error(f"Error verifying conversation token: {str(e)}")
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="Authentication required"
)
def _not_found() -> HTTPException:
return HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="File not found")
def attachment_key(file_path: str) -> Tuple[str, str]:
"""Parse a download path into (org_id, storage key), or 404.
Uploads only ever land at `chat_attachments/<org_id>/<uuid>.<ext>` (see
FileUploadService.upload_file), so anything else is not a servable
attachment. Matching that shape exactly rules out traversal: no `..`
segment can survive a three-part match on a UUID directory.
"""
# Accept the stored form ("/uploads/chat_attachments/...") and the bare key.
clean_path = file_path.lstrip('/')
if clean_path.startswith('uploads/'):
clean_path = clean_path[len('uploads/'):]
parts = clean_path.split('/')
if len(parts) != 3 or parts[0] != 'chat_attachments' or not parts[2] or parts[2] in ('.', '..'):
raise _not_found()
try:
path_org_id = str(uuid.UUID(parts[1]))
except ValueError:
raise _not_found()
return path_org_id, '/'.join(('chat_attachments', path_org_id, parts[2]))
def authorized_attachment_key(file_path: str, org_id: Optional[str]) -> str:
"""Resolve a download path to a storage key the caller's org may read."""
if not org_id:
raise _not_found()
path_org_id, key = attachment_key(file_path)
# 404 rather than 403: whether another org's attachment exists is itself
# something the caller shouldn't learn.
if path_org_id != str(org_id):
raise _not_found()
return key
# Visitor-uploaded files (SVG included) are served from the API origin, so one
# opened directly must not run script there. Only scripts are blocked: a
# sandbox CSP would also stop the browser's PDF viewer.
ATTACHMENT_HEADERS = {
"Content-Security-Policy": "script-src 'none'",
"X-Content-Type-Options": "nosniff",
"Cache-Control": "private, max-age=300",
}
def _inline_attachment(path: str, filename: str) -> FileResponse:
return FileResponse(path, content_disposition_type="inline", filename=filename,
headers=ATTACHMENT_HEADERS)
@router.get("/local/{file_path:path}")
async def serve_signed_local_file(
file_path: str,
exp: Optional[int] = Query(None),
sig: Optional[str] = Query(None),
):
"""Serve a local-storage attachment through a link from sign_local_key.
The signature is the authorization, as with an S3 presigned URL, because
the widget loads attachments with plain <img>/<a> tags.
"""
_, key = attachment_key(file_path)
if not verify_local_signature(key, exp, sig):
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="Invalid or expired link")
path = os.path.join("uploads", key)
if not os.path.isfile(path):
raise _not_found()
return _inline_attachment(path, os.path.basename(key))
@router.get("/download/{file_path:path}")
async def download_file(
file_path: str,
request: Request,
db: Session = Depends(get_db),
authorization: Optional[str] = Header(None),
x_conversation_token: Optional[str] = Header(None)
):
"""
Download/serve a file from AWS S3 or local storage
Usage: GET /api/v1/files/download/chat_attachments/org-id/filename.jpg
"""
try:
from fastapi.responses import StreamingResponse
from app.core.s3 import get_s3_client
# Authentication is required, and the result is what authorizes the
# read below. This used to swallow the failure and serve the file
# anyway, making every org's attachments world-readable by path.
auth_info = await get_current_user_or_widget(
request, db, authorization, x_conversation_token)
storage_key = authorized_attachment_key(file_path, auth_info.get("org_id"))
if not settings.S3_FILE_STORAGE:
local_file_path = os.path.join("uploads", storage_key)
if not os.path.isfile(local_file_path):
raise _not_found()
return _inline_attachment(local_file_path, os.path.basename(storage_key))
# For AWS S3 - stream file from S3
try:
s3_client = get_s3_client()
except Exception as client_err:
logger.error(f"Failed to create S3 client: {str(client_err)}", exc_info=True)
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail="Failed to create S3 client"
)
# S3 keys don't include the /uploads/ prefix
s3_key = storage_key
bucket = settings.S3_BUCKET
# Check if the file exists
try:
s3_client.head_object(Bucket=bucket, Key=s3_key)
except s3_client.exceptions.NoSuchKey:
logger.warning(f"File not found in S3: {s3_key}")
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail="File not found"
)
except Exception as head_err:
logger.error(f"Error checking file in S3: {str(head_err)}", exc_info=True)
# Get the object from S3
try:
response = s3_client.get_object(Bucket=bucket, Key=s3_key)
except s3_client.exceptions.NoSuchKey:
logger.warning(f"File not found in S3: {s3_key}")
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail="File not found"
)
except Exception as s3_err:
logger.error(f"S3 error: {str(s3_err)}", exc_info=True)
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail="Failed to retrieve file from S3"
)
# Stream the file back to client
def generate():
while True:
chunk = response['Body'].read(8192)
if not chunk:
break
yield chunk
# Determine content type
content_type = response.get('ContentType', 'application/octet-stream')
# Extract filename from the validated key
filename = storage_key.split('/')[-1]
logger.info(f"Streaming file from S3: {filename}, content-type: {content_type}")
return StreamingResponse(
generate(),
media_type=content_type,
headers={
**ATTACHMENT_HEADERS,
"Content-Disposition": f"inline; filename*=utf-8''{quote(filename)}",
}
)
except HTTPException:
raise
except Exception as e:
logger.error(f"Error downloading file {file_path}: {str(e)}", exc_info=True)
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail="Failed to download file"
)
File diff suppressed because it is too large. Load diff
+117
View File
@@ -0,0 +1,117 @@
"""
Copyright 2024-2026 ReplyPier
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.
Storage backend shim shared by upload endpoints: S3 when configured, else the
local uploads/ directory. Only PUBLIC_UPLOAD_FOLDERS are static-mounted under
/api/v1/uploads; private files are reached through signed links.
"""
import os
from typing import Optional
import aiofiles
from app.core.config import settings
from app.core.local_file_urls import sign_local_key
from app.core.s3 import (
delete_file_from_s3,
download_file_from_s3,
get_s3_signed_url,
upload_file_to_s3,
)
async def store_upload(content: bytes, folder: str, file_name: str, content_type: str) -> str:
"""Persist bytes and return the URL to store. S3 mode returns the raw S3
URL, local mode the {API_V1_STR}/uploads/... path; either way, pass it
through resolve_public_url before handing it to a client."""
if settings.S3_FILE_STORAGE:
return await upload_file_to_s3(content, folder, file_name, content_type=content_type)
upload_dir = os.path.join("uploads", folder)
os.makedirs(upload_dir, exist_ok=True)
async with aiofiles.open(os.path.join(upload_dir, file_name), "wb") as f:
await f.write(content)
return f"{settings.API_V1_STR}/uploads/{folder}/{file_name}"
# Local-storage folders anyone may read, static-mounted at {API_V1_STR}/uploads/<folder>:
# agent avatars and agents' profile pictures shown to anonymous widget visitors,
# and help-center logos and article images. Everything else under uploads/ (chat
# attachments, PDF imports awaiting a worker) is private.
PUBLIC_UPLOAD_FOLDERS = ("agents", "user", "help_center")
# Private folders handed to clients as short-lived signed links, the local
# counterpart of an S3 presigned URL.
SIGNED_UPLOAD_FOLDERS = ("chat_attachments",)
def upload_key(stored: str) -> str:
"""Key relative to `uploads/` for a stored local upload URL.
Tolerates every shape the app has ever written: the current
`{API_V1_STR}/uploads/...`, the legacy bare `/uploads/...` that
agent photos, profile pictures and chat attachments carried, and a plain
relative key with no prefix at all.
"""
relative = stored.lstrip("/")
api_prefix = f"{settings.API_V1_STR.strip('/')}/"
if relative.startswith(api_prefix):
relative = relative[len(api_prefix):]
if relative.startswith("uploads/"):
relative = relative[len("uploads/"):]
return relative
def local_upload_path(stored: str) -> str:
"""Filesystem path under `uploads/` for a stored local upload URL."""
return os.path.join("uploads", upload_key(stored))
def signed_upload_key(stored: Optional[str]) -> Optional[str]:
"""The key of a stored local upload that is only served through a signed
link, or None for anything else (S3 and other absolute URLs, public files)."""
if not stored or "://" in stored or stored.startswith(("data:", "blob:")):
return None
key = upload_key(stored)
return key if key.split("/", 1)[0] in SIGNED_UPLOAD_FOLDERS else None
async def load_upload(stored: str) -> bytes:
"""Read back a previously stored upload by the URL/path store_upload
returned. Used by workers running in a different container than the API
(shared local uploads mount or S3)."""
if stored.startswith("http"):
return await download_file_from_s3(stored)
async with aiofiles.open(local_upload_path(stored), "rb") as f:
return await f.read()
async def delete_upload(stored: str) -> None:
"""Best-effort removal of a stored upload (temp import files)."""
try:
if stored.startswith("http"):
await delete_file_from_s3(stored)
return
os.remove(local_upload_path(stored))
except Exception:
pass
async def resolve_public_url(stored_url: str) -> str:
"""Client-facing URL for a stored upload: presigned for S3 objects, a
signed link for private local files, verbatim otherwise."""
if stored_url and settings.S3_FILE_STORAGE and stored_url.startswith("http"):
return await get_s3_signed_url(stored_url)
key = signed_upload_key(stored_url)
return sign_local_key(key) if key else stored_url
+384
View File
@@ -0,0 +1,384 @@
"""
Copyright 2024-2026 ReplyPier
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)}")
+117
View File
@@ -0,0 +1,117 @@
"""
Copyright 2024-2026 ReplyPier
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.
Hardening for user-uploaded images (help-center logo + article images). Every
uploaded raster is:
- size-capped (bytes),
- checked for a real, decodable image (not a polyglot / renamed file),
- guarded against decompression bombs (pixel-count cap read from the header
BEFORE the full image is decoded),
- re-encoded from decoded pixels — which strips EXIF/metadata, trailing bytes
and any appended payload, and normalises the stored format.
SVG is intentionally NOT accepted here: it is active markup and would be stored
XSS on the public help center. Callers that previously took SVG rasterise on the
client (the cropper) or drop it.
"""
import io
from fastapi import HTTPException
from PIL import Image
# Refuse absurd pixel counts before decoding — a few-KB file can claim to be
# tens of thousands of pixels per side and exhaust memory on decode.
_MAX_PIXELS = 40_000_000 # 40 MP
Image.MAX_IMAGE_PIXELS = _MAX_PIXELS
# PIL format -> (content type, extension) for the formats we re-encode to.
_PNG = ("image/png", ".png")
_JPEG = ("image/jpeg", ".jpg")
_FORMAT_CONTENT_TYPE = {
"PNG": "image/png",
"JPEG": "image/jpeg",
"WEBP": "image/webp",
"GIF": "image/gif",
}
def _open_validated(content: bytes, *, max_bytes: int, max_dim: int) -> Image.Image:
"""Open bytes as an image, enforcing byte size, decodability and a pixel /
dimension cap (bomb guard) — all before the pixels are fully decoded."""
if not content:
raise HTTPException(status_code=400, detail="Empty image file.")
if len(content) > max_bytes:
raise HTTPException(status_code=400, detail=f"Image must be {max_bytes // (1024 * 1024)} MB or smaller.")
try:
img = Image.open(io.BytesIO(content))
width, height = img.size # header only; no full decode yet
except HTTPException:
raise
except Exception:
raise HTTPException(status_code=400, detail="Invalid or corrupt image file.")
if width <= 0 or height <= 0:
raise HTTPException(status_code=400, detail="Invalid image dimensions.")
if width > max_dim or height > max_dim or width * height > _MAX_PIXELS:
raise HTTPException(
status_code=400,
detail=f"Image is too large — keep it within {max_dim}×{max_dim} pixels.",
)
try:
img.load() # dimensions are now known-safe, so decoding is bounded
except Exception:
raise HTTPException(status_code=400, detail="Invalid or corrupt image file.")
return img
def sanitize_image(
content: bytes,
*,
allowed_content_types: set[str],
max_bytes: int,
max_dim: int,
fit: int | None = None,
) -> tuple[bytes, str, str]:
"""Validate and re-encode an uploaded raster image.
Returns (safe_bytes, content_type, extension). Raises HTTPException(400) on
anything that is not a clean, in-policy raster image. `fit`, when given,
downscales the longest side to at most `fit` px (aspect preserved).
JPEG re-encodes to JPEG (keeps photos small); everything else re-encodes to
PNG (keeps transparency). Either path strips metadata and any non-pixel
payload because the output is built from decoded pixels only.
"""
img = _open_validated(content, max_bytes=max_bytes, max_dim=max_dim)
detected = _FORMAT_CONTENT_TYPE.get((img.format or "").upper())
if detected is None or detected not in allowed_content_types:
raise HTTPException(
status_code=400,
detail="Unsupported image type. Use PNG, JPEG or WebP.",
)
if fit and max(img.size) > fit:
img.thumbnail((fit, fit), Image.LANCZOS)
out = io.BytesIO()
if detected == "image/jpeg":
img.convert("RGB").save(out, format="JPEG", quality=85, optimize=True)
content_type, ext = _JPEG
else:
img.convert("RGBA").save(out, format="PNG", optimize=True)
content_type, ext = _PNG
return out.getvalue(), content_type, ext
@@ -0,0 +1,512 @@
"""
Copyright 2024-2026 ReplyPier
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 pytest
import base64
import os
import tempfile
from unittest.mock import Mock, patch, AsyncMock, MagicMock
from uuid import uuid4
from app.services.file_upload_service import (
FileUploadService,
ALLOWED_IMAGE_TYPES,
ALLOWED_DOCUMENT_TYPES,
ALLOWED_FILE_TYPES,
MAX_FILE_SIZE,
MAX_IMAGE_SIZE
)
class TestFileUploadServiceValidation:
"""Test file validation methods"""
def test_validate_file_type_allowed_image(self):
"""Test validation of allowed image types"""
for content_type in ALLOWED_IMAGE_TYPES:
assert FileUploadService.validate_file_type(content_type) is True
def test_validate_file_type_allowed_document(self):
"""Test validation of allowed document types"""
for content_type in ALLOWED_DOCUMENT_TYPES:
assert FileUploadService.validate_file_type(content_type) is True
def test_validate_file_type_disallowed(self):
"""Test validation of disallowed file types"""
disallowed_types = [
'video/mp4',
'audio/mp3',
'application/x-executable',
'text/html',
'application/javascript'
]
for content_type in disallowed_types:
assert FileUploadService.validate_file_type(content_type) is False
def test_validate_file_size_image_within_limit(self):
"""Test image file size validation within limits"""
# Test image under image size limit
image_size = MAX_IMAGE_SIZE - 1000 # 1KB under limit
is_valid, error = FileUploadService.validate_file_size(image_size, 'image/jpeg')
assert is_valid is True
assert error is None
def test_validate_file_size_image_exceeds_image_limit(self):
"""Test image file size validation exceeding image limit"""
# Test image over image size limit
image_size = MAX_IMAGE_SIZE + 1000 # 1KB over image limit
is_valid, error = FileUploadService.validate_file_size(image_size, 'image/jpeg')
assert is_valid is False
assert "Image file size exceeds maximum" in error
assert "5.0MB" in error
def test_validate_file_size_document_within_limit(self):
"""Test document file size validation within limits"""
# Test document under general size limit
doc_size = MAX_FILE_SIZE - 1000 # 1KB under limit
is_valid, error = FileUploadService.validate_file_size(doc_size, 'application/pdf')
assert is_valid is True
assert error is None
def test_validate_file_size_document_exceeds_limit(self):
"""Test document file size validation exceeding limit"""
# Test document over general size limit
doc_size = MAX_FILE_SIZE + 1000 # 1KB over limit
is_valid, error = FileUploadService.validate_file_size(doc_size, 'application/pdf')
assert is_valid is False
assert "File size exceeds maximum" in error
assert "10.0MB" in error
def test_validate_file_size_image_exceeds_general_limit(self):
"""Test image file size validation exceeding general limit"""
# Test image over general size limit (but under image limit is irrelevant)
image_size = MAX_FILE_SIZE + 1000 # 1KB over general limit
is_valid, error = FileUploadService.validate_file_size(image_size, 'image/png')
assert is_valid is False
# The service checks image size limit first, so we get image-specific error
assert "Image file size exceeds maximum" in error
assert "5.0MB" in error
class TestFileUploadServiceUpload:
"""Test file upload functionality"""
@pytest.fixture
def sample_file_data(self):
"""Create sample file data for testing"""
content = b"This is a test file content"
encoded_content = base64.b64encode(content).decode('utf-8')
return {
'content': encoded_content,
'filename': 'test_file.txt',
'content_type': 'text/plain',
'size': len(content)
}
@pytest.fixture
def sample_image_data(self):
"""Create sample image data for testing"""
# Simple fake image content
content = b"fake_image_data_" * 100 # Small image
encoded_content = base64.b64encode(content).decode('utf-8')
return {
'content': encoded_content,
'filename': 'test_image.jpg',
'content_type': 'image/jpeg',
'size': len(content)
}
@pytest.fixture
def large_file_data(self):
"""Create large file data that exceeds size limits"""
content = b"x" * (MAX_FILE_SIZE + 1000) # Exceeds max file size
encoded_content = base64.b64encode(content).decode('utf-8')
return {
'content': encoded_content,
'filename': 'large_file.pdf',
'content_type': 'application/pdf',
'size': len(content)
}
@pytest.mark.asyncio
async def test_upload_file_s3_success(self, sample_file_data):
"""Test successful file upload to S3"""
org_id = str(uuid4())
user_id = str(uuid4())
with patch('app.services.file_upload_service.settings') as mock_settings, \
patch('app.services.file_upload_service.upload_file_to_s3') as mock_s3_upload, \
patch('app.services.file_upload_service.resolve_public_url') as mock_signed_url, \
patch('uuid.uuid4') as mock_uuid:
# Configure mocks
mock_settings.S3_FILE_STORAGE = True
mock_uuid.return_value = Mock()
mock_uuid.return_value.__str__ = Mock(return_value='test-uuid-123')
mock_s3_upload.return_value = 'https://s3.amazonaws.com/bucket/chat_attachments/org123/test-uuid-123.txt'
mock_signed_url.return_value = 'https://s3.amazonaws.com/bucket/signed-url'
result = await FileUploadService.upload_file(
file_data=sample_file_data,
org_id=org_id,
user_id=user_id
)
# Verify S3 upload was called
mock_s3_upload.assert_called_once()
call_args = mock_s3_upload.call_args
assert call_args[1]['folder'] == f'chat_attachments/{org_id}'
assert call_args[1]['filename'] == 'test-uuid-123.txt'
assert call_args[1]['content_type'] == 'text/plain'
# Verify signed URL was generated
mock_signed_url.assert_called_once()
# Verify result
assert result['file_url'] == 'https://s3.amazonaws.com/bucket/chat_attachments/org123/test-uuid-123.txt'
assert result['signed_url'] == 'https://s3.amazonaws.com/bucket/signed-url'
assert result['filename'] == 'test_file.txt'
assert result['content_type'] == 'text/plain'
assert result['size'] == len(base64.b64decode(sample_file_data['content']))
@pytest.mark.asyncio
async def test_upload_file_local_success(self, sample_file_data):
"""Test successful file upload to local storage"""
org_id = str(uuid4())
customer_id = str(uuid4())
with patch('app.services.file_upload_service.settings') as mock_settings, \
patch('uuid.uuid4') as mock_uuid, \
patch('os.makedirs') as mock_makedirs, \
patch('builtins.open', create=True) as mock_open:
# Configure mocks
mock_settings.S3_FILE_STORAGE = False
mock_uuid.return_value = Mock()
mock_uuid.return_value.__str__ = Mock(return_value='test-uuid-456')
# Mock file operations
mock_file = MagicMock()
mock_open.return_value.__enter__.return_value = mock_file
result = await FileUploadService.upload_file(
file_data=sample_file_data,
org_id=org_id,
customer_id=customer_id
)
# Verify directory creation
mock_makedirs.assert_called_once_with(f'uploads/chat_attachments/{org_id}', exist_ok=True)
# Verify file write
mock_open.assert_called_once_with(f'uploads/chat_attachments/{org_id}/test-uuid-456.txt', 'wb')
mock_file.write.assert_called_once()
# Verify result
assert result['file_url'] == f'/uploads/chat_attachments/{org_id}/test-uuid-456.txt'
# Local attachments are never static-served: clients get a signed link.
assert result['signed_url'].startswith(
f'/api/v1/files/local/chat_attachments/{org_id}/test-uuid-456.txt?exp=')
assert result['filename'] == 'test_file.txt'
assert result['content_type'] == 'text/plain'
@pytest.mark.asyncio
async def test_upload_file_missing_content(self):
"""Test upload with missing content"""
file_data = {
'filename': 'test.txt',
'content_type': 'text/plain'
# Missing 'content'
}
with pytest.raises(ValueError) as exc_info:
await FileUploadService.upload_file(
file_data=file_data,
org_id=str(uuid4())
)
assert "Missing required file data" in str(exc_info.value)
@pytest.mark.asyncio
async def test_upload_file_missing_filename(self):
"""Test upload with missing filename"""
file_data = {
'content': base64.b64encode(b"test content").decode('utf-8'),
'content_type': 'text/plain'
# Missing 'filename'
}
with pytest.raises(ValueError) as exc_info:
await FileUploadService.upload_file(
file_data=file_data,
org_id=str(uuid4())
)
assert "Missing required file data" in str(exc_info.value)
@pytest.mark.asyncio
async def test_upload_file_invalid_type(self, sample_file_data):
"""Test upload with invalid file type"""
# Modify to invalid type
sample_file_data['content_type'] = 'video/mp4'
sample_file_data['filename'] = 'test.mp4'
with pytest.raises(ValueError) as exc_info:
await FileUploadService.upload_file(
file_data=sample_file_data,
org_id=str(uuid4())
)
# Error should indicate the file type is not allowed
assert "not allowed" in str(exc_info.value) or "File type" in str(exc_info.value)
@pytest.mark.asyncio
async def test_upload_file_invalid_base64(self):
"""Test upload with invalid base64 content"""
file_data = {
'content': 'Z', # Invalid base64 - single character that can't be decoded properly
'filename': 'test.txt',
'content_type': 'text/plain'
}
with pytest.raises(ValueError) as exc_info:
await FileUploadService.upload_file(
file_data=file_data,
org_id=str(uuid4())
)
assert "Invalid file content encoding" in str(exc_info.value)
@pytest.mark.asyncio
async def test_upload_file_size_exceeds_limit(self, large_file_data):
"""Test upload with file size exceeding limits"""
with pytest.raises(ValueError) as exc_info:
await FileUploadService.upload_file(
file_data=large_file_data,
org_id=str(uuid4())
)
assert "File size exceeds maximum" in str(exc_info.value)
@pytest.mark.asyncio
async def test_upload_file_image_size_exceeds_limit(self):
"""Test upload with image size exceeding image limits"""
# Create large image that exceeds image size limit
content = b"x" * (MAX_IMAGE_SIZE + 1000)
encoded_content = base64.b64encode(content).decode('utf-8')
file_data = {
'content': encoded_content,
'filename': 'large_image.jpg',
'content_type': 'image/jpeg'
}
with pytest.raises(ValueError) as exc_info:
await FileUploadService.upload_file(
file_data=file_data,
org_id=str(uuid4())
)
assert "Image file size exceeds maximum" in str(exc_info.value)
@pytest.mark.asyncio
async def test_upload_file_s3_upload_failure(self, sample_file_data):
"""Test handling of S3 upload failure"""
with patch('app.services.file_upload_service.settings') as mock_settings, \
patch('app.services.file_upload_service.upload_file_to_s3') as mock_s3_upload:
mock_settings.S3_FILE_STORAGE = True
mock_s3_upload.side_effect = Exception("S3 upload failed")
with pytest.raises(ValueError) as exc_info:
await FileUploadService.upload_file(
file_data=sample_file_data,
org_id=str(uuid4())
)
assert "Failed to upload file" in str(exc_info.value)
@pytest.mark.asyncio
async def test_upload_file_local_write_failure(self, sample_file_data):
"""Test handling of local file write failure"""
with patch('app.services.file_upload_service.settings') as mock_settings, \
patch('os.makedirs') as mock_makedirs, \
patch('builtins.open', create=True) as mock_open:
mock_settings.S3_FILE_STORAGE = False
mock_open.side_effect = OSError("Permission denied")
with pytest.raises(ValueError) as exc_info:
await FileUploadService.upload_file(
file_data=sample_file_data,
org_id=str(uuid4())
)
assert "Failed to upload file" in str(exc_info.value)
@pytest.mark.asyncio
async def test_upload_file_default_content_type(self):
"""Test upload with default content type when not provided"""
content = b"test content"
file_data = {
'content': base64.b64encode(content).decode('utf-8'),
'filename': 'test.txt'
# No content_type provided - will default to 'application/octet-stream' which is not allowed
}
# Should raise error because default content type is not allowed
with pytest.raises(ValueError) as exc_info:
await FileUploadService.upload_file(
file_data=file_data,
org_id=str(uuid4())
)
# Error could be extension mismatch or file type not allowed
error_str = str(exc_info.value).lower()
assert "not allowed" in error_str or "does not match" in error_str or "file type" in error_str
@pytest.mark.asyncio
async def test_upload_file_folder_structure_user(self, sample_file_data):
"""Test folder structure for user uploads"""
org_id = str(uuid4())
user_id = str(uuid4())
with patch('app.services.file_upload_service.settings') as mock_settings, \
patch('app.services.file_upload_service.upload_file_to_s3') as mock_s3_upload, \
patch('app.services.file_upload_service.resolve_public_url') as mock_signed_url:
mock_settings.S3_FILE_STORAGE = True
mock_s3_upload.return_value = 'https://s3.amazonaws.com/bucket/file.txt'
mock_signed_url.return_value = 'https://s3.amazonaws.com/bucket/signed'
await FileUploadService.upload_file(
file_data=sample_file_data,
org_id=org_id,
user_id=user_id
)
# Verify folder structure
call_args = mock_s3_upload.call_args
assert call_args[1]['folder'] == f'chat_attachments/{org_id}'
@pytest.mark.asyncio
async def test_upload_file_folder_structure_customer(self, sample_file_data):
"""Test folder structure for customer uploads"""
org_id = str(uuid4())
customer_id = str(uuid4())
with patch('app.services.file_upload_service.settings') as mock_settings, \
patch('app.services.file_upload_service.upload_file_to_s3') as mock_s3_upload, \
patch('app.services.file_upload_service.resolve_public_url') as mock_signed_url:
mock_settings.S3_FILE_STORAGE = True
mock_s3_upload.return_value = 'https://s3.amazonaws.com/bucket/file.txt'
mock_signed_url.return_value = 'https://s3.amazonaws.com/bucket/signed'
await FileUploadService.upload_file(
file_data=sample_file_data,
org_id=org_id,
customer_id=customer_id
)
# Verify folder structure (same as user for now)
call_args = mock_s3_upload.call_args
assert call_args[1]['folder'] == f'chat_attachments/{org_id}'
@pytest.mark.asyncio
async def test_upload_file_unique_filename_generation(self, sample_file_data):
"""Test that unique filenames are generated"""
org_id = str(uuid4())
with patch('app.services.file_upload_service.settings') as mock_settings, \
patch('app.services.file_upload_service.upload_file_to_s3') as mock_s3_upload, \
patch('app.services.file_upload_service.resolve_public_url') as mock_signed_url, \
patch('uuid.uuid4') as mock_uuid:
mock_settings.S3_FILE_STORAGE = True
mock_s3_upload.return_value = 'https://s3.amazonaws.com/bucket/file.txt'
mock_signed_url.return_value = 'https://s3.amazonaws.com/bucket/signed'
# Mock UUID generation
mock_uuid.return_value = Mock()
mock_uuid.return_value.__str__ = Mock(return_value='unique-id-789')
await FileUploadService.upload_file(
file_data=sample_file_data,
org_id=org_id
)
# Verify unique filename was generated
call_args = mock_s3_upload.call_args
assert call_args[1]['filename'] == 'unique-id-789.txt'
@pytest.mark.asyncio
async def test_upload_file_no_extension(self):
"""Test upload with filename that has no extension"""
content = b"test content"
file_data = {
'content': base64.b64encode(content).decode('utf-8'),
'filename': 'test_file_no_extension',
'content_type': 'text/plain'
}
# Should raise error because file has no extension
with pytest.raises(ValueError) as exc_info:
await FileUploadService.upload_file(
file_data=file_data,
org_id=str(uuid4())
)
assert "must have an extension" in str(exc_info.value) or "extension" in str(exc_info.value).lower()
class TestFileUploadServiceConstants:
"""Test service constants and configurations"""
def test_allowed_file_types_combination(self):
"""Test that ALLOWED_FILE_TYPES is correct combination"""
expected_types = ALLOWED_IMAGE_TYPES | ALLOWED_DOCUMENT_TYPES
assert ALLOWED_FILE_TYPES == expected_types
def test_file_size_constants(self):
"""Test file size constants are reasonable"""
assert MAX_FILE_SIZE == 10 * 1024 * 1024 # 10MB
assert MAX_IMAGE_SIZE == 5 * 1024 * 1024 # 5MB
assert MAX_IMAGE_SIZE < MAX_FILE_SIZE # Image limit should be smaller
def test_allowed_image_types_content(self):
"""Test allowed image types include common formats (SVG excluded for XSS safety)"""
expected_image_types = {
'image/jpeg', 'image/jpg', 'image/png',
'image/gif', 'image/webp'
}
assert ALLOWED_IMAGE_TYPES == expected_image_types
# Verify SVG is explicitly NOT included
assert 'image/svg+xml' not in ALLOWED_IMAGE_TYPES
def test_allowed_document_types_content(self):
"""Test allowed document types include common formats"""
expected_doc_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'
}
assert ALLOWED_DOCUMENT_TYPES == expected_doc_types
@@ -0,0 +1,94 @@
"""
Copyright 2024-2026 ReplyPier
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 io
import pytest
from fastapi import HTTPException
from PIL import Image
from app.services.image_security import sanitize_image
RASTER = {"image/png", "image/jpeg", "image/webp"}
def _img_bytes(fmt: str, size=(64, 48), color=(200, 30, 30)) -> bytes:
buf = io.BytesIO()
Image.new("RGB", size, color).save(buf, format=fmt)
return buf.getvalue()
def test_png_reencoded_to_png():
out, ctype, ext = sanitize_image(
_img_bytes("PNG"), allowed_content_types=RASTER, max_bytes=2_000_000, max_dim=4000
)
assert ctype == "image/png" and ext == ".png"
assert Image.open(io.BytesIO(out)).format == "PNG"
def test_jpeg_stays_jpeg():
out, ctype, ext = sanitize_image(
_img_bytes("JPEG"), allowed_content_types=RASTER, max_bytes=2_000_000, max_dim=4000
)
assert ctype == "image/jpeg" and ext == ".jpg"
assert Image.open(io.BytesIO(out)).format == "JPEG"
def test_fit_downscales_longest_side():
out, _c, _e = sanitize_image(
_img_bytes("PNG", size=(1200, 600)),
allowed_content_types=RASTER,
max_bytes=2_000_000,
max_dim=4000,
fit=512,
)
w, h = Image.open(io.BytesIO(out)).size
assert max(w, h) == 512 and (w, h) == (512, 256)
def test_rejects_non_image():
with pytest.raises(HTTPException) as e:
sanitize_image(b"not an image", allowed_content_types=RASTER, max_bytes=2_000_000, max_dim=4000)
assert e.value.status_code == 400
def test_rejects_oversized_bytes():
with pytest.raises(HTTPException) as e:
sanitize_image(_img_bytes("PNG"), allowed_content_types=RASTER, max_bytes=10, max_dim=4000)
assert e.value.status_code == 400
def test_rejects_dimensions_over_cap():
with pytest.raises(HTTPException) as e:
sanitize_image(
_img_bytes("PNG", size=(300, 300)),
allowed_content_types=RASTER,
max_bytes=2_000_000,
max_dim=100,
)
assert e.value.status_code == 400
def test_rejects_disallowed_format():
# GIF is a valid image but not in the allowlist here → rejected.
with pytest.raises(HTTPException) as e:
sanitize_image(_img_bytes("GIF"), allowed_content_types=RASTER, max_bytes=2_000_000, max_dim=4000)
assert e.value.status_code == 400
def test_empty_rejected():
with pytest.raises(HTTPException):
sanitize_image(b"", allowed_content_types=RASTER, max_bytes=2_000_000, max_dim=4000)
@@ -0,0 +1,88 @@
"""
Copyright 2024-2026 ReplyPier
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 io
import os
from uuid import uuid4
import pytest
from fastapi import UploadFile
from PIL import Image
from app.api.agent import save_file
from app.api.users import save_upload_file
from app.core.config import settings
from app.services.file_storage import local_upload_path, store_upload
MOUNT = f"{settings.API_V1_STR}/uploads"
def _png_upload(name: str = "a.png") -> UploadFile:
buf = io.BytesIO()
Image.new("RGB", (2, 2), "red").save(buf, format="PNG")
buf.seek(0)
return UploadFile(filename=name, file=buf, headers={"content-type": "image/png"})
class TestLocalUploadPath:
"""The stored-URL → filesystem-path direction has to read every shape the
app has written, because rows predating the normalisation are still bare."""
def test_current_prefixed_shape(self):
assert local_upload_path(f"{MOUNT}/agents/org/a.png") == os.path.join(
"uploads", "agents/org/a.png")
def test_legacy_bare_shape(self):
# What agent photos and profile pictures carried before the migration.
assert local_upload_path("/uploads/agents/org/a.png") == os.path.join(
"uploads", "agents/org/a.png")
def test_relative_key(self):
assert local_upload_path("agents/org/a.png") == os.path.join(
"uploads", "agents/org/a.png")
def test_never_doubles_the_uploads_segment(self):
for stored in (f"{MOUNT}/x.png", "/uploads/x.png", "uploads/x.png", "x.png"):
assert local_upload_path(stored) == os.path.join("uploads", "x.png"), stored
class TestWritersAgree:
"""Every local writer must emit the same prefix, or the frontend resolves
the path against the API origin and 404s off the static mount."""
@pytest.mark.asyncio
async def test_store_upload_prefixes(self, tmp_path, monkeypatch):
monkeypatch.chdir(tmp_path)
monkeypatch.setattr(settings, "S3_FILE_STORAGE", False)
stored = await store_upload(b"x", "help_center", "a.png", "image/png")
assert stored.startswith(f"{MOUNT}/")
@pytest.mark.asyncio
async def test_agent_photo_prefixes_and_round_trips(self, tmp_path, monkeypatch):
monkeypatch.chdir(tmp_path)
monkeypatch.setattr(settings, "S3_FILE_STORAGE", False)
stored = await save_file(_png_upload(), uuid4())
assert stored.startswith(f"{MOUNT}/")
# The delete path reads it back through local_upload_path.
assert os.path.exists(local_upload_path(stored))
@pytest.mark.asyncio
async def test_profile_pic_prefixes_and_round_trips(self, tmp_path, monkeypatch):
monkeypatch.chdir(tmp_path)
monkeypatch.setattr(settings, "S3_FILE_STORAGE", False)
stored = await save_upload_file(_png_upload(), str(uuid4()), str(uuid4()))
assert stored.startswith(f"{MOUNT}/")
assert os.path.exists(local_upload_path(stored))