From 5d699207efd9164bef54b10e31c2082980e37f38 Mon Sep 17 00:00:00 2001
From: rabbitali <shusheng.wen@outlook.com>
Date: Thu, 1 Jun 2023 15:53:34 +0800
Subject: [PATCH] update hot patch status related operation support
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 apollo/database/proxy/cve.py   | 20 ++++++++++++++++----
 apollo/database/proxy/host.py  |  8 +++++---
 apollo/database/proxy/task.py  |  7 +++++--
 apollo/database/table.py       |  3 ++-
 apollo/function/schema/cve.py  |  2 ++
 apollo/function/schema/host.py |  1 +
 apollo/function/schema/task.py |  4 +++-
 7 files changed, 34 insertions(+), 11 deletions(-)

diff --git a/apollo/database/proxy/cve.py b/apollo/database/proxy/cve.py
index 9dc96ae..13a1ae6 100644
--- a/apollo/database/proxy/cve.py
+++ b/apollo/database/proxy/cve.py
@@ -224,7 +224,10 @@ class CveMysqlProxy(MysqlProxy):
         Returns:
             set
         """
-        filters = {CveHostAssociation.fixed == filter_dict.get("fixed", False)}
+        # when fixed does not have a value, the query data is not meaningful
+        # the default query is unfixed CVE information 
+        fixed = filter_dict.get("fixed", False)
+        filters = {CveHostAssociation.fixed == fixed}
         if not filter_dict:
             return filters
 
@@ -235,6 +238,12 @@ class CveMysqlProxy(MysqlProxy):
             filters.add(Host.host_group_name.in_(filter_dict["host_group"]))
         if filter_dict.get("repo"):
             filters.add(Host.repo_name.in_(filter_dict["repo"]))
+        if filter_dict.get("hp_status"):
+            filters.add(CveHostAssociation.hp_status.in_(filter_dict["hp_status"]))
+        if filter_dict.get("hotpatch") and fixed is True:
+            filters.add(CveHostAssociation.fixed_by_hp.in_(filter_dict["hotpatch"]))
+        elif filter_dict.get("hotpatch") and fixed is False:
+            filters.add(CveHostAssociation.support_hp.in_(filter_dict["hotpatch"]))
         return filters
 
     def _query_cve_hosts(self, username, cve_id, filters):
@@ -250,7 +259,8 @@ class CveMysqlProxy(MysqlProxy):
         """
         cve_query = self.session.query(Host.host_id, Host.host_name, Host.host_ip, Host.host_group_name,
                                        Host.repo_name, Host.last_scan, CveHostAssociation.support_hp,
-                                       CveHostAssociation.fixed, CveHostAssociation.fixed_by_hp) \
+                                       CveHostAssociation.fixed, CveHostAssociation.fixed_by_hp,
+                                       CveHostAssociation.hp_status ) \
             .join(CveHostAssociation, Host.host_id == CveHostAssociation.host_id) \
             .filter(Host.user == username, CveHostAssociation.cve_id == cve_id) \
             .filter(*filters)
@@ -268,7 +278,8 @@ class CveMysqlProxy(MysqlProxy):
                 "host_group": row.host_group_name,
                 "repo": row.repo_name,
                 "last_scan": row.last_scan,
-                "hotpatch": row.fixed_by_hp if row.fixed is True else row.support_hp
+                "hotpatch": row.fixed_by_hp if row.fixed is True else row.support_hp,
+                "hp_status": row.hp_status
             }
             result.append(host_info)
         return result
@@ -382,7 +393,8 @@ class CveMysqlProxy(MysqlProxy):
             sqlalchemy.orm.query.Query
         """
         cve_query = self.session.query(CveHostAssociation.cve_id, Host.host_id, Host.host_name, Host.host_ip, 
-        CveHostAssociation.support_hp, CveHostAssociation.fixed_by_hp, CveHostAssociation.fixed) \
+                                       CveHostAssociation.support_hp, CveHostAssociation.fixed_by_hp, 
+                                       CveHostAssociation.fixed) \
             .join(CveHostAssociation, Host.host_id == CveHostAssociation.host_id) \
             .filter(CveHostAssociation.cve_id.in_(cve_list)) \
             .filter(*filters)
diff --git a/apollo/database/proxy/host.py b/apollo/database/proxy/host.py
index a9431a9..3fdf97b 100644
--- a/apollo/database/proxy/host.py
+++ b/apollo/database/proxy/host.py
@@ -525,7 +525,8 @@ class HostProxy(HostMysqlProxy, CveEsProxy):
                 "%" + filter_dict["cve_id"] + "%"))
         if filter_dict.get("severity"):
             filters.add(Cve.severity.in_(filter_dict["severity"]))
-
+        if filter_dict.get("hp_status"):
+            filters.add(CveHostAssociation.hp_status.in_(filter_dict["hp_status"]))
         if filter_dict.get("hotpatch") and fixed is True:
             filters.add(CveHostAssociation.fixed_by_hp.in_(filter_dict["hotpatch"]))
         elif filter_dict.get("hotpatch") and fixed is False:
@@ -548,7 +549,7 @@ class HostProxy(HostMysqlProxy, CveEsProxy):
         """
         host_cve_query = self.session.query(CveHostAssociation.cve_id, Cve.publish_time, Cve.severity, Cve.cvss_score,
                                             CveHostAssociation.fixed, CveHostAssociation.support_hp,
-                                            CveHostAssociation.fixed_by_hp) \
+                                            CveHostAssociation.fixed_by_hp, CveHostAssociation.hp_status) \
             .select_from(CveHostAssociation) \
             .outerjoin(Cve, CveHostAssociation.cve_id == Cve.cve_id) \
             .outerjoin(Host, Host.host_id == CveHostAssociation.host_id) \
@@ -577,7 +578,8 @@ class HostProxy(HostMysqlProxy, CveEsProxy):
                 "severity": row.severity,
                 "description": description_dict[cve_id] if description_dict.get(cve_id) else "",
                 "cvss_score": row.cvss_score,
-                "hotpatch": row.fixed_by_hp if row.fixed is True else row.support_hp
+                "hotpatch": row.fixed_by_hp if row.fixed is True else row.support_hp,
+                "hp_status": row.hp_status
             }
             result.append(cve_info)
         return result
diff --git a/apollo/database/proxy/task.py b/apollo/database/proxy/task.py
index e660f02..ac15485 100644
--- a/apollo/database/proxy/task.py
+++ b/apollo/database/proxy/task.py
@@ -265,7 +265,8 @@ class TaskMysqlProxy(MysqlProxy):
                 "affected": True,
                 "fixed": True,
                 "fixed_by_hp": fix_cve.get("fixed_by_hp"),
-                "support_hp": None
+                "support_hp": None,
+                "hp_status": fix_cve.get("hp_status")
             }
 
         self.session.query(CveHostAssociation) \
@@ -1397,6 +1398,7 @@ class TaskMysqlProxy(MysqlProxy):
             "task_name": basic_task.task_name,
             "task_type": basic_task.task_type,
             "check_items": basic_task.check_items.split(',') if basic_task.check_items else [],
+            "accepted": basic_task.accepted,
             "total_hosts": [],
             "tasks": []
         }
@@ -1423,7 +1425,7 @@ class TaskMysqlProxy(MysqlProxy):
         Returns:
             sqlalchemy.orm.Query
         """
-        task_query = self.session.query(Task.task_id, Task.task_name, Task.task_type, Task.check_items) \
+        task_query = self.session.query(Task.task_id, Task.task_name, Task.task_type, Task.check_items, Task.accepted) \
             .filter(Task.task_id == task_id)
         return task_query
 
@@ -2606,6 +2608,7 @@ class TaskProxy(TaskMysqlProxy, TaskEsProxy):
                     "auto_reboot": True,
                     "create_time": 1,
                     "check_items": "",
+                    "accepted": True
                     "info": [
                         {
                             "cve_id": "cve1",
diff --git a/apollo/database/table.py b/apollo/database/table.py
index 251aaeb..33f4380 100644
--- a/apollo/database/table.py
+++ b/apollo/database/table.py
@@ -37,7 +37,7 @@ class CveHostAssociation(Base, MyBase):
     fixed = Column(Boolean)
     support_hp = Column(Boolean, default=None)
     fixed_by_hp = Column(Boolean, default=None)
-
+    hp_status = Column(String(20))
 
 class CveAffectedPkgs(Base, MyBase):
     """
@@ -144,6 +144,7 @@ class Task(Base, MyBase):
     create_time = Column(Integer)
     host_num = Column(Integer)
     check_items = Column(String(32))
+    accepted = Column(Boolean, default=False)
 
     username = Column(String(40), ForeignKey('user.username'))
 
diff --git a/apollo/function/schema/cve.py b/apollo/function/schema/cve.py
index 635e5eb..6584941 100644
--- a/apollo/function/schema/cve.py
+++ b/apollo/function/schema/cve.py
@@ -61,6 +61,8 @@ class CveHostFilterSchema(Schema):
     repo = fields.List(fields.String(
         validate=lambda s: len(s) != 0), required=False)
     fixed = fields.Boolean(required=True, validate=validate.OneOf([True, False]))
+    hotpatch = fields.List(fields.Boolean(validate=validate.OneOf([True, False])), required=False)
+    hp_status = fields.List(fields.String(validate=validate.OneOf(["ACCEPTED", "ACTIVED"])), required=False)
 
 
 class GetCveHostsSchema(Schema):
diff --git a/apollo/function/schema/host.py b/apollo/function/schema/host.py
index 84dcfbe..a0cc4b5 100644
--- a/apollo/function/schema/host.py
+++ b/apollo/function/schema/host.py
@@ -93,6 +93,7 @@ class HostCvesFilterSchema(Schema):
     hotpatch = fields.List(fields.Boolean(
         validate=validate.OneOf([True, False])), required=False)
     fixed = fields.Boolean(validate=validate.OneOf([True, False]))
+    hp_status = fields.List(fields.String(validate=validate.OneOf(["ACCEPTED", "ACTIVED"])), required=False)
 
 
 class GetHostCvesSchema(Schema):
diff --git a/apollo/function/schema/task.py b/apollo/function/schema/task.py
index 472fd53..415c2ca 100644
--- a/apollo/function/schema/task.py
+++ b/apollo/function/schema/task.py
@@ -89,7 +89,8 @@ class GenerateCveTaskSchema(Schema):
     task_name = fields.String(required=True, validate=lambda s: len(s) != 0)
     description = fields.String(
         required=True, validate=lambda s: 0 < len(s) <= 50)
-    auto_reboot = fields.Boolean(required=False, default=True)
+    auto_reboot = fields.Boolean(required=False, default=False)
+    accepted = fields.Boolean(required=True, validate=validate.OneOf([True, False]))
     check_items = fields.String(required=False, validate=lambda s: 0 < len(s) <= 32)
     info = fields.List(fields.Nested(CveInfoDictSchema), required=True, validate=lambda s: len(s) > 0)
 
@@ -226,6 +227,7 @@ class InstallPcakageInfoSchema(Schema):
 class FixedCveInfoSchema(Schema):
     cve_id = fields.String(required=True, validate=lambda s: len(s) != 0)
     fixed_by_hp = fields.Boolean(required=True, validate=validate.OneOf([True, False]))
+    hp_status = fields.String(validate=validate.OneOf(["ACCEPTED", "ACTIVED"]), required=False)
 
 
 class CveScanCallbackSchema(Schema):
-- 
