From 53943f1244dba0179b63e0e5cccd07500ca5342b Mon Sep 17 00:00:00 2001
From: "hu.sima" <simahu@yahoo.com>
Date: Mon, 25 Mar 2024 11:16:04 +0800
Subject: [PATCH] fix(api): check file ext with magic

---
 cmdb-api/Pipfile                              |  1 +
 cmdb-api/api/lib/common_setting/const.py      | 29 +++++++++
 .../api/views/common_setting/file_manage.py   | 61 ++++++++++---------
 cmdb-api/requirements.txt                     |  3 +-
 4 files changed, 64 insertions(+), 30 deletions(-)

diff --git a/cmdb-api/Pipfile b/cmdb-api/Pipfile
index 4082181..2e024a1 100644
--- a/cmdb-api/Pipfile
+++ b/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
diff --git a/cmdb-api/api/lib/common_setting/const.py b/cmdb-api/api/lib/common_setting/const.py
index f2f2f34..a68dcbf 100644
--- a/cmdb-api/api/lib/common_setting/const.py
+++ b/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',
+}
diff --git a/cmdb-api/api/views/common_setting/file_manage.py b/cmdb-api/api/views/common_setting/file_manage.py
index d66d409..fb304df 100644
--- a/cmdb-api/api/views/common_setting/file_manage.py
+++ b/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, send_from_directory
+from flask import request, abort, current_app
 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)):
-            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,
-                )
+        m_type = magic.from_buffer(file.read(2048), mime=True)
+        file.seek(0)
 
-                return self.jsonify(file_name=new_filename)
-            except Exception as e:
-                current_app.logger.error(e)
-                abort(400, ErrFormat.upload_failed.format(e))
+        if m_type == 'application/octet-stream':
+            m_type = file.mimetype
+        elif m_type == 'text/plain':
+            # 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))
diff --git a/cmdb-api/requirements.txt b/cmdb-api/requirements.txt
index f147be8..c624b59 100644
--- a/cmdb-api/requirements.txt
+++ b/cmdb-api/requirements.txt
@@ -52,4 +52,5 @@ WTForms==3.0.0
 shamir~=17.12.0
 pycryptodomex>=3.19.0
 colorama>=0.4.6
-lz4>=4.3.2
\ No newline at end of file
+lz4>=4.3.2
+python-magic==0.4.27
\ No newline at end of file