From 2dbc352d9870049fa0f9226e015e5909007355fe Mon Sep 17 00:00:00 2001
From: rearcher <123781007@qq.com>
Date: Mon, 2 Dec 2024 19:48:20 +0800
Subject: [PATCH] fix upload filed error, optimize the generated rollback task
 information

---
 apollo/database/proxy/task/cve_rollback.py |  4 +-
 apollo/handler/cve_handler/view.py         | 45 +++++++++++++++++-----
 2 files changed, 38 insertions(+), 11 deletions(-)

diff --git a/apollo/database/proxy/task/cve_rollback.py b/apollo/database/proxy/task/cve_rollback.py
index 0525602..2e761c9 100644
--- a/apollo/database/proxy/task/cve_rollback.py
+++ b/apollo/database/proxy/task/cve_rollback.py
@@ -117,8 +117,8 @@ class CveRollbackTaskProxy(TaskProxy):
         host_num = cve_fix_task_info.host_num
 
         if lang.startswith("en"):
-            task_name = "ROLLBACK_TASK: %s" % fix_task_name
-            description = "ORIGIN_TASK_DESCRIPTION: %s" % fix_task_description
+            task_name = "Rollback task: %s" % fix_task_name
+            description = "Origin task description: %s" % fix_task_description
         else:
             task_name = "回滚: %s" % fix_task_name
             description = "原CVE修复任务描述: %s" % fix_task_description
diff --git a/apollo/handler/cve_handler/view.py b/apollo/handler/cve_handler/view.py
index 200cc0d..25c5d7c 100644
--- a/apollo/handler/cve_handler/view.py
+++ b/apollo/handler/cve_handler/view.py
@@ -20,10 +20,12 @@ import glob
 import os
 import shutil
 import time
+import uuid
 from collections import defaultdict
 from typing import List, Optional
 
-from flask import g
+from flask import g, request
+from werkzeug.utils import secure_filename
 from vulcanus.database.helper import judge_return_code
 from vulcanus.log.log import LOGGER
 from vulcanus.restful.resp.state import (
@@ -430,7 +432,36 @@ class VulGetCveTaskHost(BaseResponse):
         return self.response(code=status_code, data=result)
 
 
-class VulUploadAdvisory(BaseResponse):
+class FileUpload:
+    @classmethod
+    def _upload_file(cls, save_path, file_key="file"):
+        """
+        upload file to save_path
+        Args:
+            save_path (str): path the file to be saved
+            file_key (str): body key for the file
+
+        Returns:
+            int: verify status code
+            str: file_path
+            str: file_name
+        """
+
+        file_name = ""
+        file = request.files.get(file_key)
+        if file is None or not file.filename:
+            return PARAM_ERROR, "", file_name
+        username = g.username
+        filename = secure_filename(file.filename)
+        file_name = str(uuid.uuid4()) + "." + filename.rsplit('.', 1)[-1]
+        if not os.path.exists(os.path.join(save_path, username)):
+            os.makedirs(os.path.join(save_path, username))
+        file_path = os.path.join(save_path, username, file_name)
+        file.save(file_path)
+        return SUCCEED, file_path, file_name
+
+
+class VulUploadAdvisory(BaseResponse, FileUpload):
     """
     Restful interface for importing security advisory xml (compressed files or single file)
     """
@@ -442,13 +473,11 @@ class VulUploadAdvisory(BaseResponse):
             int: status code
         """
         save_path = FILE_UPLOAD_PATH
-        status, username, file_name = self.verify_upload_request(save_path)
+        status, file_path, file_name = self._upload_file(save_path)
 
         if status != SUCCEED:
             return status
 
-        file_path = os.path.join(save_path, username, file_name)
-
         suffix = file_name.split('.')[-1]
         if suffix == "xml":
             status_code = self._save_single_advisory(proxy, file_path)
@@ -548,7 +577,7 @@ class VulUploadAdvisory(BaseResponse):
         return self.response(code=self._handle(callback))
 
 
-class VulUploadUnaffected(BaseResponse):
+class VulUploadUnaffected(BaseResponse, FileUpload):
     """
     Restful interface for importing unaffected cve xml (compressed files or single file)
     """
@@ -560,13 +589,11 @@ class VulUploadUnaffected(BaseResponse):
             int: status code
         """
         save_path = FILE_UPLOAD_PATH
-        status, username, file_name = self.verify_upload_request(save_path)
+        status, file_path, file_name = self._upload_file(save_path)
 
         if status != SUCCEED:
             return status
 
-        file_path = os.path.join(save_path, username, file_name)
-
         suffix = file_name.split('.')[-1]
         if suffix == "xml":
             status_code = self._save_unaffected_cve(proxy, file_path)
-- 
Gitee