DjangoAdmin_FastAPI_EleVue/modules/upload/service.py

import asyncio
import io
import logging
import mimetypes
import os
import random
import re
import time

from config import FASTAPI_TEMP_PATH, FASTAPI_IMAGE_URL, FASTAPI_UPLOAD_DIR
from config.app import UPLOAD_ALLOWED_EXTS, UPLOAD_MAX_SIZE_MB
from core import R
from utils.file import mkdir, verify_file_magic

logger = logging.getLogger(__name__)


def _parse_allowed_extensions(exts):
    """解析允许的扩展名参数，统一转为带点的列表格式

    支持列表或字符串输入，字符串支持逗号/竖线分隔
    """
    if not exts:
        return []

    if isinstance(exts, list):
        return [ext if ext.startswith('.') else f'.{ext}' for ext in exts]

    if isinstance(exts, str):
        if '|' in exts:
            parts = exts.split('|')
        elif ',' in exts:
            parts = exts.split(',')
        else:
            parts = [exts]

        result = []
        for part in parts:
            part = part.strip().lower()
            if part:
                if not part.startswith('.'):
                    part = f'.{part}'
                result.append(part)
        return result

    return []


def _is_safe_filename(filename: str) -> bool:
    """校验文件名安全性，防止路径遍历攻击

    拒绝包含路径分隔符或 .. 的文件名，确保不会跳出目标目录。
    """
    if not filename:
        return False
    # 规范化后检测路径遍历
    normalized = os.path.normpath(filename)
    if '..' in normalized.split(os.sep):
        return False
    # 文件名不应包含路径分隔符
    basename = os.path.basename(normalized)
    return basename == normalized or basename == filename


def _guess_mime_type(ext_with_dot):
    """根据扩展名推断MIME类型"""
    mime_map = {
        '.jpg': 'image/jpeg', '.jpeg': 'image/jpeg', '.png': 'image/png',
        '.gif': 'image/gif', '.bmp': 'image/bmp', '.webp': 'image/webp',
        '.svg': 'image/svg+xml', '.ico': 'image/x-icon',
        '.txt': 'text/plain', '.pdf': 'application/pdf',
        '.doc': 'application/msword',
        '.docx': 'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
        '.xls': 'application/vnd.ms-excel',
        '.xlsx': 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
        '.ppt': 'application/vnd.ms-powerpoint',
        '.pptx': 'application/vnd.openxmlformats-officedocument.presentationml.presentation',
        '.mp3': 'audio/mpeg', '.mp4': 'video/mp4',
        '.zip': 'application/zip', '.rar': 'application/x-rar-compressed',
        '.7z': 'application/x-7z-compressed',
    }
    return mime_map.get(ext_with_dot, 'application/octet-stream')


def _get_mime_type(filename, ext_with_dot):
    """获取文件的MIME类型"""
    mime_type, _ = mimetypes.guess_type(filename)
    if not mime_type:
        mime_type = _guess_mime_type(ext_with_dot)
    return mime_type


def _generate_filename(ext_with_dot):
    """生成唯一文件名：年月日时分秒 + 5位随机数 + 扩展名"""
    timestamp = time.strftime('%Y%m%d%H%M%S')
    random_num = '%05d' % random.randint(0, 99999)
    return f"{timestamp}{random_num}{ext_with_dot}"


# SVG中需要移除的危险标签和属性
_SVG_DANGEROUS_TAGS = re.compile(
    r'<\s*/?(script|iframe|object|embed|applet|form|input|button|textarea|select|link|meta|base)'
    r'\b[^>]*>',
    re.IGNORECASE | re.DOTALL,
)

_SVG_DANGEROUS_ATTRS = re.compile(
    r'\s+'
    r'(on\w+'          # on* 事件属性：onclick, onload, onerror 等
    r'|href\s*=\s*["\']?\s*javascript'  # javascript: 链接
    r'|xlink:href\s*=\s*["\']?\s*javascript)'
    r'[^>]*>',
    re.IGNORECASE,
)


def _sanitize_svg(content: bytes) -> bytes:
    """清洗SVG内容，移除可能导致XSS攻击的危险标签和属性"""
    try:
        text = content.decode('utf-8')
    except UnicodeDecodeError:
        return content

    # 移除危险标签
    text = _SVG_DANGEROUS_TAGS.sub('', text)
    # 移除危险属性
    text = _SVG_DANGEROUS_ATTRS.sub('>', text)
    return text.encode('utf-8')


def _save_upload(content, original_name, file_ext_with_dot):
    """同步校验文件内容并落盘，返回 R 结果（在线程池中执行）"""
    file_size = len(content)

    # 校验文件大小
    max_size_bytes = UPLOAD_MAX_SIZE_MB * 1024 * 1024
    if file_size > max_size_bytes:
        return R.failed(msg=f"文件大小超过限制，最大允许 {UPLOAD_MAX_SIZE_MB}MB")

    # 校验魔术字节
    is_valid, magic_msg = verify_file_magic(io.BytesIO(content), file_ext_with_dot)
    if not is_valid:
        return R.failed(msg=magic_msg)

    # 生成文件名和路径
    generated_name = _generate_filename(file_ext_with_dot)
    date_dir = time.strftime('%Y%m%d')
    save_dir = os.path.join(FASTAPI_TEMP_PATH, date_dir)
    mkdir(save_dir)
    file_path = os.path.join(save_dir, generated_name)

    # 保存文件
    if file_ext_with_dot == '.svg':
        content = _sanitize_svg(content)
    with open(file_path, 'wb') as f:
        f.write(content)

    # 生成访问信息
    relative_path = file_path.replace(FASTAPI_UPLOAD_DIR, "").replace("\\", "/")
    file_url = FASTAPI_IMAGE_URL + relative_path
    file_ext = file_ext_with_dot[1:] if file_ext_with_dot.startswith('.') else file_ext_with_dot
    mime_type = _get_mime_type(original_name, file_ext_with_dot)

    result = {
        'originalName': original_name,
        'fileExtension': file_ext,
        'fileType': mime_type,
        'fileSize': file_size,
        'fileName': generated_name,
        'filePath': relative_path,
        'fileUrl': file_url,
    }

    return R.ok(msg="上传成功", data=result)


async def upload_image(file):
    """上传图片

    Args:
        file: FastAPI UploadFile 对象

    Returns:
        R: 上传结果，包含文件信息
    """
    try:
        if not file or not file.filename:
            return R.failed(msg="请选择要上传的文件")

        # 获取文件基本信息
        original_name = file.filename

        # 校验文件名安全性，防止路径遍历
        if not _is_safe_filename(original_name):
            return R.failed(msg="文件名不合法")

        file_ext_with_dot = os.path.splitext(original_name)[1].lower()

        # 验证图片扩展名
        allowed_image_exts = ['.jpg', '.jpeg', '.png', '.gif', '.bmp', '.webp', '.svg', '.ico']
        if file_ext_with_dot not in allowed_image_exts:
            return R.failed(msg="文件格式不支持，仅支持图片格式")

        # 读取文件内容
        content = await file.read()

        # 同步校验与落盘放入线程池，避免阻塞事件循环
        return await asyncio.to_thread(_save_upload, content, original_name, file_ext_with_dot)

    except Exception as e:
        logger.error(f"图片上传异常: {str(e)}")
        return R.failed(msg=f"上传失败: {str(e)}")


async def upload_file(file, file_exts=None):
    """上传文件（兼容图片和各类文件）

    Args:
        file: FastAPI UploadFile 对象
        file_exts: 允许的文件扩展名，默认使用 UPLOAD_ALLOWED_EXTS 全局配置

    Returns:
        R: 上传结果，包含文件信息
    """
    try:
        if not file or not file.filename:
            return R.failed(msg="请选择要上传的文件")

        # 获取文件基本信息
        original_name = file.filename

        # 校验文件名安全性，防止路径遍历
        if not _is_safe_filename(original_name):
            return R.failed(msg="文件名不合法")

        file_ext_with_dot = os.path.splitext(original_name)[1].lower()
        file_ext = file_ext_with_dot[1:] if file_ext_with_dot.startswith('.') else ''

        # 验证文件扩展名
        allowed_exts = _parse_allowed_extensions(file_exts or UPLOAD_ALLOWED_EXTS)
        if allowed_exts and file_ext_with_dot not in allowed_exts:
            return R.failed(msg=f"文件格式不支持，仅支持 {', '.join(allowed_exts)} 格式")

        # 读取文件内容
        content = await file.read()

        # 同步校验与落盘放入线程池，避免阻塞事件循环
        return await asyncio.to_thread(_save_upload, content, original_name, file_ext_with_dot)

    except Exception as e:
        logger.error(f"文件上传异常: {str(e)}")
        return R.failed(msg=f"上传失败: {str(e)}")