Guess file extension if there isnt any

Author: samppanja

main.py

@@ -4,7 +4,7 @@
 import os
 from util.misc import *
 from util.filenameGenerator import filenameGenerator
-
+import util.extensionSolver
 
 app = FastAPI()
 app.mount(configuration.IMAGE_DIR, StaticFiles(directory=configuration.UPLOAD_PATH), name=configuration.IMAGE_DIR)
@@ -16,11 +16,19 @@ async def index():
 
 @app.post("/upload")
 async def upload(file: UploadFile = File(...), authorized = Depends(checkCredentials)):
+    if len(file.filename.split('.')) == 1: # Guess the extension if there is none
+        _file = await file.read()
+        guessExt = util.extensionSolver.guessFileExtension(_file)
+        guessMime = util.extensionSolver.guessMime(_file)
+        file.content_type = guessMime
+        file.filename = file.filename + guessExt
+        await file.seek(0) # Seek back, otherwise upload will fail
     if file.content_type not in configuration.ALLOWED_CONTENT:
-        return {"message": configuration.ERROR_UNALLOWED_CONTENT}
+        return {"error": configuration.ERROR_UNALLOWED_CONTENT}
     if authorized:
         if configuration.RANDOMIZED_FILENAMES:
             _extension = file.filename.split('.')[1].lower() # Save the extension of the file for upload
+            print(_extension)
             file.filename = filenameGenerator.generateName(5) +f".{_extension}"
         await uploadFile(file)
         return {"url": configuration.BASE_URL+configuration.IMAGE_DIR+file.filename}

requirements.txt

@@ -1,3 +1,4 @@
 fastapi
 python-multipart
-uvicorn
+uvicorn
+python-magic

util/extensionSolver.py

@@ -0,0 +1,10 @@
+import magic
+import mimetypes
+
+def guessMime(file) -> str:
+    return magic.from_buffer(file, mime=True)
+
+def guessFileExtension(file) -> str:
+    mime = guessMime(file)
+    return mimetypes.guess_extension(mime)
+