new extensions checker magically

Author: JexPY

api/app/services/helpers/detectFileExtension.py

@@ -0,0 +1,12 @@
+import magic
+import mimetypes
+from pathlib import Path
+from fastapi import HTTPException,status
+
+def magic_extensions(file_path: Path):
+    mime = magic.Magic(mime=True)
+    detectedExtension = mimetypes.guess_all_extensions(mime.from_buffer(open(file_path, "rb").read(2048)))
+    if len(detectedExtension) > 0:
+        return detectedExtension[0]
+    else:
+        raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail='File extension not detected')

api/app/services/serveUploadedFiles.py

@@ -1,6 +1,5 @@
 import os
 import shutil
-import imghdr
 import _thread
 import copy
 from typing import List
@@ -12,27 +11,40 @@
 from fastapi import UploadFile,HTTPException,status
 from .storage.googleCloud import upload_image_file_to_google_storage,upload_video_file_to_google_storage
 from .storage.s3 import upload_image_file_to_s3_storage,upload_video_file_to_s3_storage
-from ffmpeg import probe
+from .helpers.detectFileExtension import magic_extensions
+import sys
 
-def save_upload_file_tmp(upload_file: UploadFile) -> Path:
+def save_upload_file_tmp(upload_file: None, raw_data_file = None) -> Path:
     try:
-        suffix = Path(upload_file.filename).suffix
-        with NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
-            shutil.copyfileobj(upload_file.file, tmp)
-            tmp_path = Path(tmp.name)
+        if raw_data_file:
+            with NamedTemporaryFile(delete=False, suffix=None) as tmp:
+                    shutil.copyfileobj(raw_data_file.raw, tmp)
+
+        else:
+            with NamedTemporaryFile(delete=False, suffix=None) as tmp:
+                shutil.copyfileobj(upload_file.file, tmp)
+
+        extension = magic_extensions(Path(tmp.name))
+        final_temp_file = tmp.name + extension
+        os.rename(Path(tmp.name), final_temp_file)
+
     except:
         raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail='Impossible to manipulate the file')
     finally:
-        upload_file.file.close()
-    return tmp_path
+        if upload_file:
+            upload_file.file.close()
+        else:
+            raw_data_file.close()
+    return Path(final_temp_file), extension
 
 
-def handle_upload_image_file(thumbnail, upload_file: UploadFile):
+def handle_upload_image_file(thumbnail, upload_file: None, raw_data_file = None):
     try:
-        tmp_path = save_upload_file_tmp(upload_file)
-        if imghdr.what(tmp_path):
+        tmp_path, file_extension = save_upload_file_tmp(upload_file, raw_data_file)
+
+        if file_extension[1:] in os.environ.get('IMAGE_AllOWED_FILE_FORMAT').split(','):
 
-            imagePaths = resize_image(tmp_path, imghdr.what(tmp_path), thumbnail, os.environ.get('IMAGE_CONVERTING_PREFERED_FORMAT'))
+            imagePaths = resize_image(tmp_path, file_extension[1:], thumbnail, os.environ.get('IMAGE_CONVERTING_PREFERED_FORMAT'))
 
             if os.environ.get('PREFERED_STORAGE') == 'google':
                 _thread.start_new_thread(upload_image_file_to_google_storage, (copy.deepcopy(imagePaths),))
@@ -56,11 +68,10 @@ def handle_upload_image_file(thumbnail, upload_file: UploadFile):
         tmp_path.unlink()  # Delete the temp file
 
 
-def handle_multiple_image_file_uploads(FILES: List[UploadFile], workers: int):
+def handle_multiple_image_file_uploads(FILES: List[UploadFile], workers: int, thumbnail: bool):
     # We can use a with statement to ensure threads are cleaned up promptly
     with concurrent.futures.ThreadPoolExecutor(max_workers = workers) as executor:
-        # Start the load operations and mark each future with its FILES thumbnail is default for this function
-        future_to_url = {executor.submit(handle_upload_image_file, True, eachFile): eachFile for eachFile in FILES}
+        future_to_url = {executor.submit(handle_upload_image_file, eachFile, thumbnail): eachFile for eachFile in FILES}
         result = []
         for future in concurrent.futures.as_completed(future_to_url):
             try:
@@ -70,15 +81,12 @@ def handle_multiple_image_file_uploads(FILES: List[UploadFile], workers: int):
         return result
 
 
-def handle_upload_video_file(optimize, upload_file: UploadFile):
+def handle_upload_video_file(optimize, upload_file: None, raw_data_file = None):
+
     try:
-        tmp_path = save_upload_file_tmp(upload_file)
-        videoFileCheck = probe(tmp_path).get('format')
-        if videoFileCheck.get('format_name'):
-            # Checks for video file type also is possible to restict for only mp4 format by checking major_brand
-            # videoFileCheck.get('tags').get('major_brand') == 'mp42'
-            if os.environ.get('VIDEO_AllOWED_FILE_FORMAT') in videoFileCheck.get('format_name').split(','):
+        tmp_path, file_extension = save_upload_file_tmp(upload_file, raw_data_file)
 
+        if file_extension[1:] in os.environ.get('VIDEO_AllOWED_FILE_FORMAT').split(','):
                 videoPaths = video_file_FFMPEG(tmp_path, optimize)
 
                 if os.environ.get('PREFERED_STORAGE') == 'google':
@@ -100,8 +108,6 @@ def handle_upload_video_file(optimize, upload_file: UploadFile):
                 videoPaths['storage'] = os.environ.get('PREFERED_STORAGE')
 
                 return videoPaths
-            else:
-                raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail='Format not granted')
         else:
             raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail='Not valid format')
     except: