20260131

back/main.py

@@ -7,6 +7,7 @@ from sqlalchemy.orm import foreign
 import os
 import uuid
 import json
+import filetype
 from datetime import datetime
 
 app = Flask(__name__)
@@ -24,7 +25,7 @@ db = SQLAlchemy(app)
 # 全局配置变量
 NEED_AUDIT = True
 FILE_SIZE_LIMIT_MB = 10.0
-FILE_FORMATS = ["png", "jpg", "jpeg", "gif", "webp"]
+FILE_FORMATS = ["image/png", "image/jpeg", "image/gif", "image/webp"]
 
 def now_time():
     return datetime.now()
@@ -165,7 +166,7 @@ def load_config():
                         else:
                             parsed = raw
                         if isinstance(parsed, list):
-                            FILE_FORMATS = [str(x).strip().lstrip('.').lower() for x in parsed if str(x).strip()]
+                            FILE_FORMATS = [str(x).strip().lower() for x in parsed if str(x).strip()]
                     except Exception:
                         pass
         except Exception as e:
@@ -211,17 +212,6 @@ def save_hashtags(tag_type, target_id, hashtopic):
         )
         db.session.add(new_tag)
 
-def detect_image_type(header: bytes):
-    if header.startswith(b"\xFF\xD8\xFF"):
-        return "jpeg"
-    if header.startswith(b"\x89PNG\r\n\x1a\n"):
-        return "png"
-    if header.startswith(b"GIF87a") or header.startswith(b"GIF89a"):
-        return "gif"
-    if header.startswith(b"RIFF") and len(header) >= 12 and header[8:12] == b"WEBP":
-        return "webp"
-    return None
-    
 # --- 用户普通api端点 ---
 @app.route('/api/settings', methods=['GET'])
 def get_settings():
@@ -512,24 +502,15 @@ def upload_pic():
                 return jsonify({"code": 2006, "data": "上传的图片超出限制大小"})
 
         ext = os.path.splitext(file.filename)[1].lstrip('.').lower()
-        if not ext or (FILE_FORMATS and ext not in FILE_FORMATS):
-            return jsonify({"code": 2007, "data": "上传的文件类型不支持"})
-
-        header = file.read(512)
+        kind = filetype.guess(file.read(5120))
         file.seek(0)
-        detected = detect_image_type(header)
-        ext_map = {
-            'jpg': 'jpeg',
-            'jpeg': 'jpeg',
-            'png': 'png',
-            'gif': 'gif',
-            'webp': 'webp',
-        }
-        expected = ext_map.get(ext)
-        if not expected or detected != expected:
+        detected_mime = kind.mime if kind else None
+        if not detected_mime or (FILE_FORMATS and detected_mime.lower() not in FILE_FORMATS):
             return jsonify({"code": 2007, "data": "上传的文件类型不支持"})
 
-        filename = f"{uuid.uuid4().hex}.{ext}"
+        if not ext and kind:
+            ext = kind.extension
+        filename = f"{uuid.uuid4().hex}.{ext}" if ext else uuid.uuid4().hex
         filepath = os.path.join(IMG_DIR, filename)
         file.save(filepath)