129 lines
3.8 KiB
Python

import os
import shutil
import uuid
from pathlib import Path
from typing import Optional
from fastapi import UploadFile, HTTPException, status
from app.core.config import settings
def validate_file_size(file: UploadFile, max_size: int = settings.MAX_UPLOAD_SIZE) -> None:
"""Validate the size of an uploaded file."""
# Move cursor to the end to get size, then back to the beginning
file.file.seek(0, os.SEEK_END)
file_size = file.file.tell()
file.file.seek(0)
if file_size > max_size:
raise HTTPException(
status_code=status.HTTP_413_REQUEST_ENTITY_TOO_LARGE,
detail=f"File size too large. Maximum size is {max_size / (1024 * 1024):.1f} MB.",
)
def save_upload_file(
file: UploadFile,
destination_dir: Path,
allowed_extensions: Optional[list[str]] = None,
filename: Optional[str] = None,
) -> str:
"""Save an uploaded file to a destination directory."""
# Create directory if it doesn't exist
destination_dir.mkdir(parents=True, exist_ok=True)
# Validate file size
validate_file_size(file)
# Get file extension and validate if needed
if file.filename:
ext = file.filename.split(".")[-1].lower()
if allowed_extensions and ext not in allowed_extensions:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=f"File type not allowed. Allowed types: {', '.join(allowed_extensions)}",
)
else:
ext = "bin" # Default extension if no filename
# Generate a unique filename if not provided
if not filename:
filename = f"{uuid.uuid4()}.{ext}"
elif not filename.endswith(f".{ext}"):
filename = f"{filename}.{ext}"
# Full path to save the file
file_path = destination_dir / filename
# Save the file
with open(file_path, "wb") as buffer:
shutil.copyfileobj(file.file, buffer)
return str(file_path)
def save_deposit_proof(user_id: int, file: UploadFile) -> str:
"""Save a deposit proof image."""
# Define allowed extensions for deposit proofs
allowed_extensions = ["jpg", "jpeg", "png", "pdf"]
# Define the destination directory
destination_dir = settings.DEPOSIT_PROOFS_DIR
# Generate a filename with user_id for organization
filename = f"deposit_proof_{user_id}_{uuid.uuid4()}"
# Save the file
file_path = save_upload_file(
file=file,
destination_dir=destination_dir,
allowed_extensions=allowed_extensions,
filename=filename,
)
return file_path
def save_kyc_document(user_id: int, document_type: str, file: UploadFile) -> str:
"""Save a KYC document."""
# Define allowed extensions for KYC documents
allowed_extensions = ["jpg", "jpeg", "png", "pdf"]
# Define the destination directory
destination_dir = settings.KYC_UPLOAD_DIR
# Generate a filename with user_id and document_type for organization
filename = f"kyc_{user_id}_{document_type}_{uuid.uuid4()}"
# Save the file
file_path = save_upload_file(
file=file,
destination_dir=destination_dir,
allowed_extensions=allowed_extensions,
filename=filename,
)
return file_path
def save_bot_image(bot_id: int, file: UploadFile) -> str:
"""Save a bot image."""
# Define allowed extensions for bot images
allowed_extensions = ["jpg", "jpeg", "png"]
# Define the destination directory
destination_dir = Path(settings.UPLOAD_DIR) / "bot_images"
# Generate a filename with bot_id for organization
filename = f"bot_{bot_id}_{uuid.uuid4()}"
# Save the file
file_path = save_upload_file(
file=file,
destination_dir=destination_dir,
allowed_extensions=allowed_extensions,
filename=filename,
)
return file_path