Validate PE files by MZ magic bytes instead of extension

Author: Dev9269

api/routes.py

@@ -51,9 +51,7 @@ async def spa_middleware(request, call_next):
         return await call_next(request)
 
 
-def _is_pe_filename(name: str) -> bool:
-    ext = Path(name).suffix.lower()
-    return ext in (".exe", ".dll", ".sys", ".ocx", ".scr", ".cpl")
+PE_MAGIC = b"MZ"
 
 
 def _format_size(b: int) -> str:
@@ -71,15 +69,17 @@ async def health():
 
 @app.post("/api/scan")
 async def scan_file(file: UploadFile = File(...)):
-    if not file.filename or not _is_pe_filename(file.filename):
-        raise HTTPException(status_code=400, detail="File must be a .exe or .dll")
+    content = await file.read()
+    if len(content) < 2 or content[:2] != PE_MAGIC:
+        raise HTTPException(
+            status_code=400, detail="Not a valid PE file (missing MZ header)"
+        )
 
-    suffix = Path(file.filename).suffix.lower()
+    if len(content) > 1024 * 1024 * 1024:
+        raise HTTPException(status_code=400, detail="File exceeds 1GB limit")
+
+    suffix = Path(file.filename or "file.bin").suffix.lower() or ".bin"
     with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
-        content = await file.read()
-        if len(content) > 1024 * 1024 * 1024:
-            os.unlink(tmp.name)
-            raise HTTPException(status_code=400, detail="File exceeds 1GB limit")
         tmp.write(content)
         tmp_path = tmp.name