72 lines
2.4 KiB
Python
72 lines
2.4 KiB
Python
# utils/file_handler.py
|
|
import aiofiles
|
|
import hashlib
|
|
import re
|
|
import uuid
|
|
from pathlib import Path
|
|
from fastapi import UploadFile
|
|
from typing import Tuple, Dict, Any
|
|
|
|
|
|
class FileHandler:
|
|
def __init__(self, upload_dir: str = "uploads", max_file_size: int = 50 * 1024 * 1024):
|
|
self.upload_dir = Path(upload_dir)
|
|
self.upload_dir.mkdir(exist_ok=True)
|
|
self.max_file_size = max_file_size
|
|
|
|
def _sanitize_filename(self, filename: str) -> str:
|
|
original = Path(filename or "upload.step").name
|
|
suffix = Path(original).suffix.lower()
|
|
stem = Path(original).stem or "upload"
|
|
safe_stem = re.sub(r"[^A-Za-z0-9._-]+", "_", stem).strip("._-") or "upload"
|
|
if suffix not in {".stp", ".step"}:
|
|
suffix = ".step"
|
|
return f"{safe_stem}{suffix}"
|
|
|
|
@staticmethod
|
|
def _looks_like_step(content: bytes) -> bool:
|
|
if not content:
|
|
return False
|
|
head = content[:4096].decode("utf-8", errors="ignore").upper()
|
|
return (
|
|
"ISO-10303-21" in head
|
|
or "HEADER;" in head
|
|
or "FILE_SCHEMA" in head
|
|
or "DATA;" in head
|
|
)
|
|
|
|
async def save_uploaded_file(self, file: UploadFile) -> Tuple[Path, int, Dict[str, Any]]:
|
|
"""保存上传的文件并返回安全元数据"""
|
|
content = await file.read()
|
|
|
|
if not content:
|
|
raise ValueError("上传文件为空")
|
|
if len(content) > self.max_file_size:
|
|
raise ValueError(f"上传文件过大，限制 {self.max_file_size // (1024 * 1024)}MB")
|
|
if not self._looks_like_step(content):
|
|
raise ValueError("文件内容不是有效的 STP/STEP 数据")
|
|
|
|
safe_name = self._sanitize_filename(file.filename)
|
|
unique_name = f"{uuid.uuid4().hex}_{safe_name}"
|
|
file_path = self.upload_dir / unique_name
|
|
|
|
async with aiofiles.open(file_path, "wb") as f:
|
|
await f.write(content)
|
|
|
|
metadata = {
|
|
"original_filename": file.filename,
|
|
"safe_original_name": safe_name,
|
|
"stored_filename": unique_name,
|
|
"sha256": hashlib.sha256(content).hexdigest(),
|
|
}
|
|
|
|
return file_path, len(content), metadata
|
|
|
|
def cleanup_file(self, file_path: Path):
|
|
"""清理文件"""
|
|
try:
|
|
if file_path.exists():
|
|
file_path.unlink()
|
|
except Exception as e:
|
|
print(f"文件清理失败: {e}")
|