From e7e9871111a67d1aee5b7a7d68029b13894f8fae Mon Sep 17 00:00:00 2001 From: rabbitali 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