106 lines
3.9 KiB
Diff
106 lines
3.9 KiB
Diff
From a032e1e0b11365a0dc5d725fd234771cd53c0858 Mon Sep 17 00:00:00 2001
|
|
From: gongzt <gong_zhengtang@163.com>
|
|
Date: Fri, 2 Jun 2023 14:29:57 +0800
|
|
Subject: [PATCH] Repair Host cve verification is not performed in a generation task
|
|
MIME-Version: 1.0
|
|
Content-Type: text/plain; charset=UTF-8
|
|
Content-Transfer-Encoding: 8bit
|
|
|
|
---
|
|
apollo/database/proxy/task.py | 37 +++++++++++++++++++++++++++++
|
|
apollo/handler/task_handler/view.py | 20 ++++++++++++++++
|
|
2 files changed, 57 insertions(+)
|
|
|
|
diff --git a/apollo/database/proxy/task.py b/apollo/database/proxy/task.py
|
|
index e660f02..edba161 100644
|
|
--- a/apollo/database/proxy/task.py
|
|
+++ b/apollo/database/proxy/task.py
|
|
@@ -3208,3 +3208,40 @@ class TaskProxy(TaskMysqlProxy, TaskEsProxy):
|
|
|
|
# insert task id and username into es
|
|
self._init_task_in_es(task_id, data["username"])
|
|
+
|
|
+ def validate_cves(self, cve_id: list) -> bool:
|
|
+ """
|
|
+ Verifying cve validity
|
|
+
|
|
+ Args:
|
|
+ cve_id: id of the cve to be validate
|
|
+
|
|
+ Returns:
|
|
+ bool: A return of true indicates that the validation passed
|
|
+ """
|
|
+
|
|
+ try:
|
|
+ exists_cve_count = self.session.query(CveHostAssociation).filter(
|
|
+ CveHostAssociation.cve_id.in_(cve_id)).count()
|
|
+
|
|
+ return True if exists_cve_count == len(cve_id) else False
|
|
+ except SQLAlchemyError as error:
|
|
+ LOGGER.error(error)
|
|
+ return False
|
|
+
|
|
+ def validate_hosts(self, host_id: list) -> bool:
|
|
+ """
|
|
+ Verifying host validity
|
|
+
|
|
+ Args:
|
|
+ host_id: id of the host to be validate
|
|
+
|
|
+ Returns:
|
|
+ bool: A return of true indicates that the validation passed
|
|
+ """
|
|
+ try:
|
|
+ exists_host_count = self.session.query(Host).filter(Host.host_id.in_(host_id)).count()
|
|
+ return True if exists_host_count == len(host_id) else False
|
|
+ except SQLAlchemyError as error:
|
|
+ LOGGER.error(error)
|
|
+ return False
|
|
diff --git a/apollo/handler/task_handler/view.py b/apollo/handler/task_handler/view.py
|
|
index 214053c..314f7bb 100644
|
|
--- a/apollo/handler/task_handler/view.py
|
|
+++ b/apollo/handler/task_handler/view.py
|
|
@@ -287,6 +287,14 @@ class VulGenerateCveTask(BaseResponse):
|
|
"task_id": "id1"
|
|
}
|
|
"""
|
|
+ host_ids = [host["host_id"] for hosts in params["info"] for host in hosts["host_info"]]
|
|
+ if not callback.validate_hosts(host_id=list(set(host_ids))):
|
|
+ return self.response(code=PARAM_ERROR)
|
|
+
|
|
+ cve_ids = [cve["cve_id"] for cve in params["info"]]
|
|
+ if not callback.validate_cves(cve_id=list(set(cve_ids))):
|
|
+ return self.response(code=PARAM_ERROR)
|
|
+
|
|
status_code, data = self._handle(callback, params)
|
|
return self.response(code=status_code, data=data)
|
|
|
|
@@ -488,6 +496,10 @@ class VulGenerateRepoTask(BaseResponse):
|
|
"task_id": "1"
|
|
}
|
|
"""
|
|
+ host_ids = [host["host_id"] for host in params["info"]]
|
|
+ if not callback.validate_hosts(host_id=list(set(host_ids))):
|
|
+ return self.response(code=PARAM_ERROR)
|
|
+
|
|
status_code, data = self._handle(callback, params)
|
|
return self.response(code=status_code, data=data)
|
|
|
|
@@ -836,6 +848,14 @@ class VulGenerateCveRollback(BaseResponse):
|
|
"task_id": "1"
|
|
}
|
|
"""
|
|
+ host_ids = [host["host_id"] for host in params["info"]]
|
|
+ if not callback.validate_hosts(host_id=list(set(host_ids))):
|
|
+ return self.response(code=PARAM_ERROR)
|
|
+
|
|
+ cve_ids = [cve["cve_id"] for host in params["info"] for cve in host["cves"]]
|
|
+ if not callback.validate_cves(cve_id=list(set(cve_ids))):
|
|
+ return self.response(code=PARAM_ERROR)
|
|
+
|
|
status_code, data = self._handle(callback, params)
|
|
return self.response(code=status_code, data=data)
|
|
|
|
--
|
|
Gitee
|