diff --git a/apollo/conf/constant.py b/apollo/conf/constant.py index 26343d4f4c8ff70cc143f066573fac8dca6b5551..ce4baf024e5e7a66fcf952a857b1ce2be519ff04 100644 --- a/apollo/conf/constant.py +++ b/apollo/conf/constant.py @@ -141,6 +141,7 @@ VUL_TASK_CVE_FIX_RPM_INFO_GET = "/vulnerability/task/cve-fix/rpm/get" VUL_TASK_HOTPATCH_REMOVE_INFO_GET = "/vulnerability/task/hotpatch-remove/info/get" VUL_TASK_HOTPATCH_REMOVE_RESULT_GET = "/vulnerability/task/hotpatch-remove/result/get" VUL_TASK_HOTPATCH_REMOVE_PROGRESS_GET = "/vulnerability/task/hotpatch-remove/progress/get" +VUL_EXIST_RUNNING_TASK_HOST = "/vulnerability/task/running-task/host" # route of callback VUL_TASK_CVE_FIX_CALLBACK = "/vulnerability/task/callback/cve/fix" diff --git a/apollo/database/proxy/task/base.py b/apollo/database/proxy/task/base.py index a5ddede64f4eece411af474c04d1245efa3f530f..345ab413d8dcb3dd51794b8759e8d1d3ad981d93 100644 --- a/apollo/database/proxy/task/base.py +++ b/apollo/database/proxy/task/base.py @@ -896,9 +896,7 @@ class TaskProxy(TaskMysqlProxy, TaskEsProxy): Returns: list: task id list """ - cve_fix_query = ( - self.session.query(CveFixTask.task_id).filter(CveFixTask.status == TaskStatus.RUNNING).all() - ) + cve_fix_query = self.session.query(CveFixTask.task_id).filter(CveFixTask.status == TaskStatus.RUNNING).all() task_id_list = [task.task_id for task in cve_fix_query] return task_id_list @@ -1001,9 +999,7 @@ class TaskProxy(TaskMysqlProxy, TaskEsProxy): Returns: sqlalchemy.orm.Query """ - task_query = self.session.query( - Task - ).filter(Task.task_id == task_id) + task_query = self.session.query(Task).filter(Task.task_id == task_id) return task_query def get_task_hosts(self, task_id) -> tuple: @@ -1055,3 +1051,79 @@ class TaskProxy(TaskMysqlProxy, TaskEsProxy): .all() ) return [host.host_id for host in hosts] + + def exists_running_task_hosts(self, host_ids: list) -> bool: + """ + Detects a running host + + Args: + host_ids: host id list + """ + try: + exists_running_host = self._exists_running_task_hosts(host_ids) + return SUCCEED, exists_running_host + except SQLAlchemyError as error: + LOGGER.error(error) + return DATABASE_QUERY_ERROR, False + + def _exists_running_task_hosts(self, host_ids: list) -> bool: + exists_running_host_query = self.session.query(CveFixTask).filter(CveFixTask.status == TaskStatus.RUNNING) + running_fix_hosts = ( + exists_running_host_query.filter(CveFixTask.host_id.in_(host_ids)).count() + if host_ids + else exists_running_host_query.count() + ) + # cve fix + if running_fix_hosts: + LOGGER.info("At least one host is performing fix tasks.") + return True + exists_running_host_query = self.session.query(HotpatchRemoveTask).filter( + HotpatchRemoveTask.status == TaskStatus.RUNNING + ) + running_hotpatch_remove_hosts = ( + exists_running_host_query.filter(HotpatchRemoveTask.host_id.in_(host_ids)).count() + if host_ids + else exists_running_host_query.count() + ) + # hotpatch remove + if running_hotpatch_remove_hosts: + LOGGER.info("At least one host is performing hotpatch remove.") + return True + + exists_running_host_query = self.session.query(CveRollbackTask).filter( + CveRollbackTask.status == TaskStatus.RUNNING + ) + running_rollback_hosts = ( + exists_running_host_query.filter(CveRollbackTask.host_id.in_(host_ids)).count() + if host_ids + else exists_running_host_query.count() + ) + # cve rollback + if running_rollback_hosts: + LOGGER.info("At least one host is performing a repair rollback.") + return True + exists_running_host_query = self.session.query(TaskHostRepoAssociation).filter( + TaskHostRepoAssociation.status == TaskStatus.RUNNING + ) + running_repo_set_hosts = ( + exists_running_host_query.filter(TaskHostRepoAssociation.host_id.in_(host_ids)).count() + if host_ids + else exists_running_host_query.count() + ) + # repo set + if running_repo_set_hosts: + LOGGER.info("At least one host is performing repo set.") + return True + + # cve scan + exists_running_host_query = self.session.query(Host).filter(Host.status == HostStatus.SCANNING) + running_scan_hosts = ( + exists_running_host_query.filter(Host.host_id.in_(host_ids)).count() + if host_ids + else exists_running_host_query.count() + ) + if running_scan_hosts: + LOGGER.info("At least one host is performing a vulnerability scan.") + return True + + return False diff --git a/apollo/function/schema/task.py b/apollo/function/schema/task.py index e136182a4698183c22d63255aa1331f3fc4df031..6a492be3a535ae42903be7b40970917a2cd895a7 100644 --- a/apollo/function/schema/task.py +++ b/apollo/function/schema/task.py @@ -324,6 +324,7 @@ class CveRollbackTaskInfoFilterSchema(Schema): """ filter schema of cve rollback task info getting interface """ + # search_key could be host name or host ip search_key = fields.String(required=False, validate=lambda s: 0 < len(s) <= 50) status = fields.List(fields.String(validate=validate.OneOf(TaskStatus.attribute())), required=False) @@ -333,6 +334,7 @@ class GetCveRollbackTaskInfoSchema(PaginationSchema): """ validators for parameter of /vulnerability/task/cve-rollback/info/get """ + task_id = fields.String(required=True, validate=lambda s: 0 < len(s) <= 32) filter = fields.Nested(CveRollbackTaskInfoFilterSchema, required=False) @@ -351,6 +353,7 @@ class GetCveRollbackTaskResultSchema(Schema): """ validators for parameter of /vulnerability/task/cve-rollback/result/get """ + task_id = fields.String(required=True, validate=lambda s: 0 < len(s) <= 32) @@ -392,6 +395,12 @@ class TaskCveRpmHostSchema(Schema): installed_rpm = fields.String(required=True, validate=lambda s: 0 < len(s) <= 100) +class ExistRunningTaskHostSchema(Schema): + host_ids = fields.List( + fields.Integer(required=True, validate=lambda s: s > 0), required=True, validate=lambda s: len(s) > 0 + ) + + __all__ = [ "GetTaskListSchema", "GetTaskProgressSchema", diff --git a/apollo/handler/task_handler/view.py b/apollo/handler/task_handler/view.py index 17164fb71746c228222e11f0898c9f160d3de9b8..18b4daf517788445beea4da798f1337ca513c8e1 100644 --- a/apollo/handler/task_handler/view.py +++ b/apollo/handler/task_handler/view.py @@ -41,6 +41,7 @@ from apollo.database.proxy.task.scan import ScanProxy from apollo.function.schema.host import ScanHostSchema from apollo.function.schema.task import * from apollo.function.schema.task import GetHotpatchRemoveTaskCveInfoSchema +from apollo.function.schema.task import ExistRunningTaskHostSchema from apollo.handler.task_handler.callback.cve_fix import CveFixCallback from apollo.handler.task_handler.callback.cve_rollback import CveRollbackCallback from apollo.handler.task_handler.callback.hotpatch_remove import HotpatchRemoveCallback @@ -422,10 +423,7 @@ class VulGenerateRepoTask(BaseResponse): """ task_id = str(uuid.uuid1()).replace("-", "") task_info = dict( - task_id=task_id, - task_type=TaskType.REPO_SET, - create_time=int(time.time()), - username=params["username"] + task_id=task_id, task_type=TaskType.REPO_SET, create_time=int(time.time()), username=params["username"] ) task_info.update(params) data = dict(task_id=task_id) @@ -487,7 +485,7 @@ class VulExecuteTask(BaseResponse): TaskType.CVE_FIX: "_handle_cve_fix", TaskType.REPO_SET: "_handle_repo_set", TaskType.HOTPATCH_REMOVE: "_handle_hotpatch_remove", - TaskType.CVE_ROLLBACK: "_handle_cve_rollback" + TaskType.CVE_ROLLBACK: "_handle_cve_rollback", } @staticmethod @@ -1234,3 +1232,27 @@ class VulGetTaskHost(BaseResponse): """ status_code, data = callback.get_task_hosts(params["task_id"]) return self.response(code=status_code, data=data) + + +class VulExistRunningTaskHost(BaseResponse): + """ + Restful interface for getting host running the task. + """ + + @BaseResponse.handle(schema=ExistRunningTaskHostSchema, proxy=TaskProxy) + def post(self, callback: TaskProxy, **params): + """ + Args: + callback(TaskProxy): task proxy + params(dict): e.g. {"host_ids": [1,2]} + + Returns: + dict: response body, e.g. + { + "code": 200, + "message": "", + "data": True/False + } + """ + status_code, data = callback.exists_running_task_hosts(params["host_ids"]) + return self.response(code=status_code, data=data) diff --git a/apollo/url.py b/apollo/url.py index 46bf580d47e64de5ba327c392de70e6f98336148..64f40cfdb1c01b524cdd05e77554f20e508b3cb8 100644 --- a/apollo/url.py +++ b/apollo/url.py @@ -73,6 +73,7 @@ SPECIFIC_URLS = { (cve_task_view.VulGetHotpatchRemoveTaskResult, VUL_TASK_HOTPATCH_REMOVE_RESULT_GET), (cve_task_view.VulGetHotpatchRemoveTaskProgress, VUL_TASK_HOTPATCH_REMOVE_PROGRESS_GET), (cve_task_view.VulGetTaskHost, VUL_GET_TASK_HOST), + (cve_task_view.VulExistRunningTaskHost, VUL_EXIST_RUNNING_TASK_HOST), ], "CVE_TASK_CALLBACK_URLS": [ (cve_task_view.VulCveFixTaskCallback, VUL_TASK_CVE_FIX_CALLBACK),