aops-apollo/0002-suitable-for-2003-sp3.patch
gongzt 61be517e12 优化cve关联查询性能及适配20.03-sp3版本
(cherry picked from commit a89ce8b1bc704397a0f94744e9d6acd1cdab1f90)
2023-10-19 15:56:12 +08:00

248 lines
11 KiB
Diff

From ad8982cafbe849707a2a07fa3f329f377b5aea33 Mon Sep 17 00:00:00 2001
From: gongzt <gong_zhengtang@163.com>
Date: Tue, 17 Oct 2023 09:50:16 +0800
Subject: Fixed errors in 20.03-sp3, such as task progress, cve repair task, and host cve query
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
---
apollo/database/proxy/cve.py | 2 +-
apollo/database/proxy/host.py | 136 +++++++++++++---------------------
apollo/database/proxy/task.py | 2 +-
3 files changed, 54 insertions(+), 86 deletions(-)
diff --git a/apollo/database/proxy/cve.py b/apollo/database/proxy/cve.py
index 4125894..257083a 100644
--- a/apollo/database/proxy/cve.py
+++ b/apollo/database/proxy/cve.py
@@ -485,7 +485,7 @@ class CveMysqlProxy(MysqlProxy):
CveHostAssociation.installed_rpm,
CveHostAssociation.available_rpm,
)
- .join(CveHostAssociation, Host.host_id == CveHostAssociation.host_id)
+ .join(Host, Host.host_id == CveHostAssociation.host_id)
.filter(CveHostAssociation.cve_id.in_(cve_list))
.filter(*filters)
.all()
diff --git a/apollo/database/proxy/host.py b/apollo/database/proxy/host.py
index 4b75f1c..b3cabb6 100644
--- a/apollo/database/proxy/host.py
+++ b/apollo/database/proxy/host.py
@@ -511,31 +511,27 @@ class HostProxy(HostMysqlProxy, CveEsProxy):
result = {"total_count": 0, "total_page": 1, "result": []}
host_id = data["host_id"]
- filters = self._get_host_cve_filters(data.get("filter", {}))
- host_cve_query = self._query_host_cve(data["username"], host_id, filters).all()
+ host_cve_query = self._query_host_cve(data["username"], host_id, data.get("filter", {}))
- total_count = len(host_cve_query)
+ total_count = host_cve_query.count()
if not total_count:
return SUCCEED, result
- cve_info_list, cve_packages_dict = self._preprocess_cve_list_query(host_cve_query)
sort_column = data['sort'] if "sort" in data else "cve_id"
direction, page, per_page = data.get('direction'), data.get('page'), data.get('per_page')
- processed_cve_list, total_page = self._sort_and_page_host_cve_info(
- cve_info_list, page, per_page, sort_column, direction
- )
- description_dict = self._get_cve_description(list(cve_packages_dict.keys()))
- result['result'] = self._add_additional_info_to_cve_list(
- processed_cve_list, description_dict, cve_packages_dict
- )
+ host_cve_list, total_page = sort_and_page(host_cve_query, sort_column, direction, per_page, page)
+
+ cve_id_list = [cve.cve_id for cve in host_cve_list]
+ description_dict = self._get_cve_description(cve_id_list)
+ result['result'] = self._add_additional_info_to_cve_list(host_cve_list, description_dict)
result['total_page'] = total_page
result['total_count'] = total_count
return SUCCEED, result
@staticmethod
- def _get_host_cve_filters(filter_dict):
+ def _get_host_cve_filters(filter_dict, cve_affected_pkg_subquery):
"""
Generate filters to filter host's CVEs
@@ -563,7 +559,7 @@ class HostProxy(HostMysqlProxy, CveEsProxy):
filters.add(
or_(
CveHostAssociation.cve_id.like("%" + filter_dict["search_key"] + "%"),
- CveAffectedPkgs.package.like("%" + filter_dict["search_key"] + "%"),
+ cve_affected_pkg_subquery.c.package.like("%" + filter_dict["search_key"] + "%"),
)
)
if filter_dict.get("severity"):
@@ -573,107 +569,79 @@ class HostProxy(HostMysqlProxy, CveEsProxy):
filters.add(CveHostAssociation.affected == filter_dict["affected"])
return filters
- def _query_host_cve(self, username: str, host_id: int, filters: set):
+ def _query_host_cve(self, username: str, host_id: int, filters_dict: dict):
"""
query needed host CVEs info
Args:
username (str): user name of the request
host_id (int): host id
- filters (set): filter given by user
+ filter_dict(dict): filter dict to filter host's CVEs, e.g.
+ {
+ "cve_id": "",
+ "severity": ["high", "unknown"],
+ "affected": True,
+ "package": "vim",
+ "fixed": False // The default is false if the field is null.
+ }
Returns:
sqlalchemy.orm.query.Query
"""
+ cve_affected_pkg_subquery = (
+ self.session.query(
+ CveAffectedPkgs.cve_id,
+ func.group_concat(func.distinct(CveAffectedPkgs.package), SEPARATOR=",").label("package"),
+ )
+ .group_by(CveAffectedPkgs.cve_id)
+ .distinct()
+ .subquery()
+ )
+
+ filters = self._get_host_cve_filters(filters_dict, cve_affected_pkg_subquery)
+
host_cve_query = (
self.session.query(
CveHostAssociation.cve_id,
case([(Cve.publish_time == None, "")], else_=Cve.publish_time).label("publish_time"),
case([(Cve.severity == None, "")], else_=Cve.severity).label("severity"),
case([(Cve.cvss_score == None, "")], else_=Cve.cvss_score).label("cvss_score"),
- case([(CveAffectedPkgs.package == None, "")], else_=CveAffectedPkgs.package).label("package"),
+ case(
+ [(cve_affected_pkg_subquery.c.package == None, "")], else_=cve_affected_pkg_subquery.c.package
+ ).label("package"),
)
.select_from(CveHostAssociation)
.outerjoin(Cve, CveHostAssociation.cve_id == Cve.cve_id)
- .outerjoin(CveAffectedPkgs, CveAffectedPkgs.cve_id == CveHostAssociation.cve_id)
+ .outerjoin(cve_affected_pkg_subquery, cve_affected_pkg_subquery.c.cve_id == CveHostAssociation.cve_id)
.outerjoin(Host, Host.host_id == CveHostAssociation.host_id)
.filter(CveHostAssociation.host_id == host_id, Host.user == username)
.filter(*filters)
- ).group_by(CveHostAssociation.cve_id, CveAffectedPkgs.package)
+ ).group_by(CveHostAssociation.cve_id, cve_affected_pkg_subquery.c.package)
return host_cve_query
@staticmethod
- def _preprocess_cve_list_query(cve_list_query: List[dict]) -> Tuple[List[dict], dict]:
- """
- get each cve's source package set and deduplication cve info list
-
- Args:
- cve_list_query(list): cve info list(two rows may have same cve id
- with different source package)
- Returns:
- list: list of cve info without package and description.
- dict: key is cve id, value is cve affected source package set. e.g
- { "cve-xxxx-xxxx": {"vim"}, "cve-xxxx-xxxx":{"kernel"} }
- """
- cve_packages_dict, cve_info_list = defaultdict(set), []
-
- for row in cve_list_query:
- if row.cve_id not in cve_packages_dict:
- cve_info_list.append(
- {
- "cve_id": row.cve_id,
- "publish_time": row.publish_time,
- "severity": row.severity,
- "cvss_score": row.cvss_score,
- }
- )
- cve_packages_dict[row.cve_id].add(row.package)
-
- return cve_info_list, cve_packages_dict
-
- @staticmethod
- def _sort_and_page_host_cve_info(cve_info_list, page, per_page, sort_column=None, direction="desc"):
- """
- sort and page cve info list
-
- Args:
- cve_info_list (list): cve info list. not empty.
- page(int)
- per_page(int): number of record per page
- sort_column(str): the column that sort based on
- direction(str): desc or asc
-
- Returns:
- query result, total number of pages
- """
- total_page = 1
- total_count = len(cve_info_list)
-
- if sort_column:
- cve_info_list.sort(
- key=lambda cve_info: cve_info[sort_column], reverse=True if direction == "desc" else False
- )
-
- if page and per_page:
- total_page = math.ceil(total_count / per_page)
- return cve_info_list[per_page * (page - 1) : per_page * page], total_page
-
- return cve_info_list, total_page
-
- @staticmethod
- def _add_additional_info_to_cve_list(cve_info_list: list, description_dict: dict, cve_packages_dict: list) -> list:
+ def _add_additional_info_to_cve_list(host_cve_list: list, description_dict: dict) -> list:
"""
add description and affected source packages for each cve
Args:
- cve_info_list:
+ host_cve_list:
description_dict (dict): key is cve's id, value is cve's description
- cve_pkgs_dict (dict): key is cve's id, value is cve's affected packages set
Returns:
list
"""
- for cve_info in cve_info_list:
- cve_id = cve_info["cve_id"]
- cve_info["description"] = description_dict[cve_id] if description_dict.get(cve_id) else ""
- cve_info["package"] = ",".join(cve_packages_dict[cve_id])
- return cve_info_list
+ host_cve_info_list = []
+ for host_cve in host_cve_list:
+ description = description_dict[host_cve.cve_id] if description_dict.get(host_cve.cve_id) else ""
+ host_cve_info_list.append(
+ {
+ "cve_id": host_cve.cve_id,
+ "publish_time": host_cve.publish_time,
+ "severity": host_cve.severity,
+ "cvss_score": host_cve.cvss_score,
+ "package": host_cve.package,
+ "description": description,
+ }
+ )
+
+ return host_cve_info_list
diff --git a/apollo/database/proxy/task.py b/apollo/database/proxy/task.py
index e7c3b4e..de151b2 100644
--- a/apollo/database/proxy/task.py
+++ b/apollo/database/proxy/task.py
@@ -1291,7 +1291,7 @@ class TaskMysqlProxy(MysqlProxy):
status = TaskStatus.UNKNOWN
else:
status = TaskStatus.SUCCEED
- result[cve_id] = {"progress": row.total - row.running - row.none, "status": status}
+ result[cve_id] = {"progress": int(row.total - row.running - row.none), "status": status}
succeed_list = list(result.keys())
fail_list = list(set(cve_list) - set(succeed_list))
--
Gitee