feat(knowledge): 知识源维护与 TXT/Markdown/PDF 导入
支持文本、Markdown 与非扫描 PDF 上传解析,展示提取文本与切块结果。
This commit is contained in:
1 parent
4ff934ad98
commit
7937de5d7d
23 files changed
+7545
No files matched your search
@@ -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
@@ -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
|
||||
@@ -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)}")
|
||||
|
||||
@@ -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))
|
||||
Reference in new issue
Block a user