Merge pull request #430 from simontigers/common_file_ext_check

fix(api): check file ext with magic

cmdb-api/Pipfile

@@ -66,6 +66,7 @@ hvac = "==2.0.0"
 colorama = ">=0.4.6"
 pycryptodomex = ">=3.19.0"
 lz4 = ">=4.3.2"
+python-magic = "==0.4.27"
 
 [dev-packages]
 # Testing

cmdb-api/api/lib/common_setting/const.py

@@ -35,3 +35,32 @@ AuthCommonConfigAutoRedirect = 'auto_redirect'
 class TestType(BaseEnum):
     Connect = 'connect'
     Login = 'login'
+
+
+MIMEExtMap = {
+    'application/vnd.openxmlformats-officedocument.wordprocessingml.document': '.docx',
+    'application/msword': '.doc',
+    'application/vnd.ms-word.document.macroEnabled.12': '.docm',
+    'application/vnd.ms-excel': '.xls',
+    'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet': '.xlsx',
+    'application/vnd.ms-excel.sheet.macroEnabled.12': '.xlsm',
+    'application/vnd.ms-powerpoint': '.ppt',
+    'application/vnd.openxmlformats-officedocument.presentationml.presentation': '.pptx',
+    'application/vnd.ms-powerpoint.presentation.macroEnabled.12': '.pptm',
+    'application/zip': '.zip',
+    'application/x-7z-compressed': '.7z',
+    'application/json': '.json',
+    'application/pdf': '.pdf',
+    'image/png': '.png',
+    'image/bmp': '.bmp',
+    'image/prs.btif': '.btif',
+    'image/gif': '.gif',
+    'image/jpeg': '.jpg',
+    'image/tiff': '.tif',
+    'image/vnd.microsoft.icon': '.ico',
+    'image/webp': '.webp',
+    'image/svg+xml': '.svg',
+    'image/vnd.adobe.photoshop': '.psd',
+    'text/plain': '.txt',
+    'text/csv': '.csv',
+}

cmdb-api/api/views/common_setting/file_manage.py

@@ -1,10 +1,10 @@
 # -*- coding:utf-8 -*-
-import os
+from flask import request, abort, current_app
-
-from flask import request, abort, current_app, send_from_directory
 from werkzeug.utils import secure_filename
 import lz4.frame
+import magic
 
+from api.lib.common_setting.const import MIMEExtMap
 from api.lib.common_setting.resp_format import ErrFormat
 from api.lib.common_setting.upload_file import allowed_file, generate_new_file_name, CommonFileCRUD
 from api.resource import APIView
@@ -45,32 +45,35 @@ class PostFileView(APIView):
 
         if not file:
             abort(400, ErrFormat.file_is_required)
-        extension = file.mimetype.split('/')[-1]
-        if '+' in extension:
-            extension = file.filename.split('.')[-1]
-        if file.filename == '':
-            filename = f'.{extension}'
-        else:
-            if extension not in file.filename:
-                filename = file.filename + f".{extension}"
-            else:
-                filename = file.filename
 
-        if allowed_file(filename, current_app.config.get('ALLOWED_EXTENSIONS', ALLOWED_EXTENSIONS)):
+        m_type = magic.from_buffer(file.read(2048), mime=True)
-            new_filename = generate_new_file_name(filename)
+        file.seek(0)
-            new_filename = secure_filename(new_filename)
-            file_content = file.read()
-            compressed_data = lz4.frame.compress(file_content)
-            try:
-                CommonFileCRUD.add_file(
-                    origin_name=filename,
-                    file_name=new_filename,
-                    binary=compressed_data,
-                )
 
-                return self.jsonify(file_name=new_filename)
+        if m_type == 'application/octet-stream':
-            except Exception as e:
+            m_type = file.mimetype
-                current_app.logger.error(e)
+        elif m_type == 'text/plain':
-                abort(400, ErrFormat.upload_failed.format(e))
+            # https://github.com/ahupp/python-magic/issues/193
+            m_type = m_type if file.mimetype == m_type else file.mimetype
 
-        abort(400, ErrFormat.file_type_not_allowed.format(filename))
+        extension = MIMEExtMap.get(m_type, None)
+
+        if extension is None:
+            abort(400, f"不支持的文件类型: {m_type}")
+
+        filename = file.filename if file.filename and file.filename.endswith(extension) else file.filename + extension
+
+        new_filename = generate_new_file_name(filename)
+        new_filename = secure_filename(new_filename)
+        file_content = file.read()
+        compressed_data = lz4.frame.compress(file_content)
+        try:
+            CommonFileCRUD.add_file(
+                origin_name=filename,
+                file_name=new_filename,
+                binary=compressed_data,
+            )
+
+            return self.jsonify(file_name=new_filename)
+        except Exception as e:
+            current_app.logger.error(e)
+            abort(400, ErrFormat.upload_failed.format(e))

cmdb-api/requirements.txt

@@ -53,3 +53,4 @@ shamir~=17.12.0
 pycryptodomex>=3.19.0
 colorama>=0.4.6
 lz4>=4.3.2
+python-magic==0.4.27