From e7e9871111a67d1aee5b7a7d68029b13894f8fae Mon Sep 17 00:00:00 2001
From: rabbitali <wenxin32@foxmail.com>
Date: Wed, 13 Dec 2023 10:11:22 +0800
Subject: [PATCH] add rollback task execution method and fix cve scan
 callback error

---
 zeus/conf/constant.py                 |   1 +
 zeus/function/verify/vulnerability.py |  18 ++--
 zeus/url.py                           |   1 +
 zeus/vulnerability_manage/view.py     | 135 +++++++++++++++++++++++++-
 4 files changed, 143 insertions(+), 12 deletions(-)

diff --git a/zeus/conf/constant.py b/zeus/conf/constant.py
index 994dd90..1370d6e 100644
--- a/zeus/conf/constant.py
+++ b/zeus/conf/constant.py
@@ -32,6 +32,7 @@ CERES_HOST_INFO = "aops-ceres collect --host '%s'"
 CERES_CVE_REPO_SET = "aops-ceres apollo --set-repo '%s'"
 CERES_CVE_SCAN = "aops-ceres apollo --scan '%s'"
 CERES_CVE_FIX = "aops-ceres apollo --fix '%s'"
+CERES_CVE_ROLLBACK = "aops-ceres apollo --rollback '%s'"
 CERES_HOTPATCH_REMOVE = "aops-ceres apollo --remove-hotpatch '%s'"
 CERES_SYNC_CONF = "aops-ceres sync --conf '%s'"
 CERES_OBJECT_FILE_CONF = "aops-ceres ragdoll --list '%s'"
diff --git a/zeus/function/verify/vulnerability.py b/zeus/function/verify/vulnerability.py
index 07875e0..ff25c8d 100644
--- a/zeus/function/verify/vulnerability.py
+++ b/zeus/function/verify/vulnerability.py
@@ -86,14 +86,12 @@ class CveFixSchema(TaskGeneralSchema):
     fix_type = fields.String(validate=validate.OneOf(["hotpatch", "coldpatch"]), required=True)
 
 
-class CveRollbackSingleInfoSchema(Schema):
-    cve_id = fields.String(validate=lambda s: len(s) > 0)
-    hotpatch = fields.Boolean(validate=validate.OneOf([True, False]))
-
-
-class CveRollbackTask(Schema):
+class CveRollbackTaskSchema(Schema):
     host_id = fields.Integer(required=True, validate=lambda s: s > 0)
-    cves = fields.List(fields.Nested(CveRollbackSingleInfoSchema()), required=True)
+    installed_rpm = fields.String(required=True, validate=lambda s: 100 >= len(s) > 0)
+    target_rpm = fields.String(required=True, validate=lambda s: 100 >= len(s) > 0)
+    dnf_event_start = fields.Integer(allow_none=True, required=True, validate=lambda s: s > 0)
+    dnf_event_end = fields.Integer(allow_none=True, required=True, validate=lambda s: s > 0)
 
 
 class CveRollbackSchema(TaskGeneralSchema):
@@ -101,10 +99,12 @@ class CveRollbackSchema(TaskGeneralSchema):
     validator for cve rollback
     """
 
-    tasks = fields.List(fields.Nested(CveRollbackTask()), required=True, validate=lambda s: len(s) > 0)
+    tasks = fields.List(fields.Nested(CveRollbackTaskSchema()), required=True, validate=lambda s: len(s) > 0)
+    fix_task_id = fields.String(required=True, validate=lambda s: len(s) > 0)
+    rollback_type = fields.String(validate=validate.OneOf(["hotpatch", "coldpatch"]), required=True)
 
     class Meta:
-        fields = ("tasks", "task_id", "task_name", "total_hosts", "task_type", "callback")
+        exclude = ("total_hosts",)
 
 
 class HotpatchRemoveTask(Schema):
diff --git a/zeus/url.py b/zeus/url.py
index ad8cec9..5f00ef9 100644
--- a/zeus/url.py
+++ b/zeus/url.py
@@ -101,6 +101,7 @@ SPECIFIC_URLS = {
     ],
     'CVE_URLS': [
         (vulnerability_view.ExecuteRepoSetTask, EXECUTE_REPO_SET),
+        (vulnerability_view.ExecuteCveRollbackTask, EXECUTE_CVE_ROLLBACK),
         (vulnerability_view.ExecuteCveScanTask, EXECUTE_CVE_SCAN),
         (vulnerability_view.ExecuteCveFixTask, EXECUTE_CVE_FIX),
         (vulnerability_view.ExecuteHotpatchRemoveTask, EXECUTE_HOTPATCH_REMOVE),
diff --git a/zeus/vulnerability_manage/view.py b/zeus/vulnerability_manage/view.py
index be52e23..37ab633 100644
--- a/zeus/vulnerability_manage/view.py
+++ b/zeus/vulnerability_manage/view.py
@@ -26,6 +26,7 @@ from zeus.conf import configuration
 from zeus.conf.constant import (
     CERES_CVE_FIX,
     CERES_CVE_REPO_SET,
+    CERES_CVE_ROLLBACK,
     CERES_HOTPATCH_REMOVE,
     CERES_CVE_SCAN,
     CveTaskStatus,
@@ -34,7 +35,13 @@ from zeus.conf.constant import (
 from zeus.database.proxy.host import HostProxy
 from zeus.database.table import Host
 from zeus.function.model import ClientConnectArgs
-from zeus.function.verify.vulnerability import CveFixSchema, CveScanSchema, HotpatchRemoveSchema, RepoSetSchema
+from zeus.function.verify.vulnerability import (
+    CveFixSchema,
+    CveRollbackSchema,
+    CveScanSchema,
+    HotpatchRemoveSchema,
+    RepoSetSchema,
+)
 from zeus.host_manager.ssh import execute_command_and_parse_its_result
 
 
@@ -283,7 +290,7 @@ class ExecuteCveScanTask(BaseResponse, BaseExcuteTask):
             CERES_CVE_SCAN % json.dumps({"check_items": self._check_items}),
         )
         if status != state.SUCCEED:
-            request_body["status"] = CveTaskStatus.FAIL
+            request_body.update({"status":CveTaskStatus.FAIL, "reboot":False})
         else:
             request_body.update(json.loads(cve_scan_result))
 
@@ -500,7 +507,7 @@ class ExecuteHotpatchRemoveTask(BaseResponse, BaseExcuteTask):
             ),
             command,
         )
-        
+
         if status == state.SUCCEED:
             request_body.update(json.loads(hotpatch_remove_result))
         else:
@@ -552,3 +559,125 @@ class ExecuteHotpatchRemoveTask(BaseResponse, BaseExcuteTask):
         ]
         threading.Thread(target=lambda: gevent.joinall(wait_execute_tasks)).start()
         return self.response(code=state.SUCCEED)
+
+
+class ExecuteCveRollbackTask(BaseResponse, BaseExcuteTask):
+    """
+    Interface for cve rollback.
+    Restful API: POST
+    """
+
+    def _execute_task(self, host_info: dict, task_info: dict) -> None:
+        """
+        Execute cve rollback task
+
+        Args:
+            host_info(dict): e.g
+                {
+                    "host_id": 1,
+                    "host_ip": "127.0.0.1",
+                    "host_name": "test_host",
+                    "ssh_port": 22,
+                    "ssh_user": "root",
+                    "pkey": "RSA-KEY-string",
+                }
+            task_info (dict): e.g
+                {
+                    "host_id": "id1",
+                    "check_items":[],
+                    "rollback_type": "hotpatch",
+                    "installed_kernel": "kernel-5.1.10",
+                    "target_kernel": "kernel-5.1.9",
+                    "dnf_event_start": 1,
+                    "dnf_event_end": 2,
+                }
+        Returns:
+            None
+        """
+        request_body = {
+            "execution_time": int(time.time()),
+            "task_id": self._task_id,
+            "host_id": host_info.get("host_id"),
+            "host_ip": host_info.get("host_ip"),
+            "host_name": host_info.get("host_name"),
+        }
+
+        task_info.pop("host_id")
+        command = CERES_CVE_ROLLBACK % json.dumps(task_info)
+        status, cve_rollback_result = execute_command_and_parse_its_result(
+            ClientConnectArgs(
+                host_info.get("host_ip"),
+                host_info.get("ssh_port"),
+                host_info.get("ssh_user"),
+                host_info.get("pkey"),
+                60 * 10,
+            ),
+            command,
+        )
+        if status != state.SUCCEED:
+            request_body.update(
+                {
+                    "status": CveTaskStatus.FAIL,
+                    "log": cve_rollback_result,
+                    "check_items": [
+                        {"item": item, "result": CveTaskStatus.FAIL} for item in task_info.get("check_items")
+                    ],
+                }
+            )
+        else:
+            request_body.update(json.loads(cve_rollback_result))
+
+        url = f'http://{configuration.apollo.get("IP")}:{ configuration.apollo.get("PORT")}{self._callback_url}'
+        self.get_response("post", url, request_body, self._header, timeout=10)
+
+    @BaseResponse.handle(schema=CveRollbackSchema)
+    def post(self, **params) -> Response:
+        """
+        execute cve rollback task
+
+        Args:
+            params (dict): e.g
+                {
+                    "task_id": "c6714973c9b342a380fd01fdf7f90ef5",
+                    "task_name": "cve rollback task",
+                    "fix_task_id": "string",
+                    "task_type": "cve rollback",
+                    "rollback_type": "coldpatch",
+                    "check_items": ["network"],
+                    "tasks": [
+                        {
+                            "host_id": 74,
+                            "installed_rpm": "kernel-5.1.10",
+                            "target_rpm": "kernel-5.1.9",
+                            "dnf_event_start": 1,
+                            "dnf_event_end": 2
+                        }
+                    ],
+                    "callback": "/vulnerability/task/callback/cve/rollback"
+                }
+        Returns:
+            response body
+        """
+        total_host = [task_info["host_id"] for task_info in params.get("tasks")]
+        status_code, host_infos = query_host_basic_info(total_host, params.get('username'))
+        if status_code != state.SUCCEED:
+            return self.response(code=status_code)
+        # parse args
+        self._task_id = params.get("task_id")
+        self._task_name = params.get("task_name")
+        self._task_type = params.get("task_type")
+        self._header["local_account"] = params.get("username")
+        self._callback_url = params.get('callback')
+        # Execute task
+        tasks = generate_tasks(
+            params.get('tasks'),
+            host_infos,
+            **{
+                "check_items": params.get('check_items'),
+                "rollback_type": params.get('rollback_type'),
+            },
+        )
+        threading.Thread(
+            target=lambda: gevent.joinall([gevent.spawn(self._execute_task, *task) for task in tasks])
+        ).start()
+        return self.response(code=state.SUCCEED)
-- 
2.33.0

