125 lines
4.5 KiB
Diff
125 lines
4.5 KiB
Diff
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
|
|
|