From af818b3d88cfb1b4a2e4907811af672149e57240 Mon Sep 17 00:00:00 2001 From: luzhongyang Date: Mon, 15 Nov 2021 12:44:09 +0800 Subject: [PATCH 1/2] when return_queue is 0, do not return queue data. return_queue default is 1 --- apiserver/report/handler/heartbeat_handler.py | 61 ++++++++++--------- 1 file changed, 32 insertions(+), 29 deletions(-) diff --git a/apiserver/report/handler/heartbeat_handler.py b/apiserver/report/handler/heartbeat_handler.py index 04e1145..da4216a 100644 --- a/apiserver/report/handler/heartbeat_handler.py +++ b/apiserver/report/handler/heartbeat_handler.py @@ -31,6 +31,7 @@ def __init__(self): self.report_queue = None self.method_queue = None self.replay_queue = None + self.return_queue = None def parse(self): self.cpu = self.detail.get('cpu') @@ -39,6 +40,7 @@ def parse(self): self.report_queue = self.detail.get('reportQueue', 0) self.method_queue = self.detail.get('methodQueue', 0) self.replay_queue = self.detail.get('replayQueue', 0) + self.return_queue = self.detail.get('replayQueue', 1) def has_permission(self): self.agent = IastAgent.objects.filter(id=self.agent_id, user=self.user_id).first() @@ -75,40 +77,41 @@ def save_heartbeat(self): agent=self.agent) def get_result(self, msg=None): - try: - project_agents = IastAgent.objects.values('id').filter(bind_project_id=self.agent.bind_project_id) - if project_agents is None: - logger.info(_('There is no probe under the project')) + if self.return_queue == 1: + try: + project_agents = IastAgent.objects.values('id').filter(bind_project_id=self.agent.bind_project_id) + if project_agents is None: + logger.info(_('There is no probe under the project')) - replay_queryset = IastReplayQueue.objects.values( - 'id', 'relation_id', 'uri', 'method', 'scheme', 'header', 'params', 'body', 'replay_type' - ).filter(agent_id__in=project_agents, state=const.WAITING)[:10] - if len(replay_queryset) == 0: - logger.info(_('Replay request does not exist')) + replay_queryset = IastReplayQueue.objects.values( + 'id', 'relation_id', 'uri', 'method', 'scheme', 'header', 'params', 'body', 'replay_type' + ).filter(agent_id__in=project_agents, state=const.WAITING)[:10] + if len(replay_queryset) == 0: + logger.info(_('Replay request does not exist')) - success_ids, success_vul_ids, failure_ids, failure_vul_ids, replay_requests = [], [], [], [], [] - for replay_request in replay_queryset: - if replay_request['uri']: - replay_requests.append(replay_request) - success_ids.append(replay_request['id']) - if replay_request['replay_type'] == const.VUL_REPLAY: - success_vul_ids.append(replay_request['relation_id']) - else: - failure_ids.append(replay_request['id']) - if replay_request['replay_type'] == const.VUL_REPLAY: - failure_vul_ids.append(replay_request['relation_id']) + success_ids, success_vul_ids, failure_ids, failure_vul_ids, replay_requests = [], [], [], [], [] + for replay_request in replay_queryset: + if replay_request['uri']: + replay_requests.append(replay_request) + success_ids.append(replay_request['id']) + if replay_request['replay_type'] == const.VUL_REPLAY: + success_vul_ids.append(replay_request['relation_id']) + else: + failure_ids.append(replay_request['id']) + if replay_request['replay_type'] == const.VUL_REPLAY: + failure_vul_ids.append(replay_request['relation_id']) - timestamp = int(time.time()) - IastReplayQueue.objects.filter(id__in=success_ids).update(update_time=timestamp, state=const.SOLVING) - IastReplayQueue.objects.filter(id__in=failure_ids).update(update_time=timestamp, state=const.SOLVED) + timestamp = int(time.time()) + IastReplayQueue.objects.filter(id__in=success_ids).update(update_time=timestamp, state=const.SOLVING) + IastReplayQueue.objects.filter(id__in=failure_ids).update(update_time=timestamp, state=const.SOLVED) - IastVulnerabilityModel.objects.filter(id__in=success_vul_ids).update(latest_time=timestamp, status_id=2) - IastVulnerabilityModel.objects.filter(id__in=failure_vul_ids).update(latest_time=timestamp, status_id=1) - logger.info(_('Reproduction request issued successfully')) + IastVulnerabilityModel.objects.filter(id__in=success_vul_ids).update(latest_time=timestamp, status_id=2) + IastVulnerabilityModel.objects.filter(id__in=failure_vul_ids).update(latest_time=timestamp, status_id=1) + logger.info(_('Reproduction request issued successfully')) - return replay_requests - except Exception as e: - logger.info(_('Replay request query failed, reason: {}').format(e)) + return replay_requests + except Exception as e: + logger.info(_('Replay request query failed, reason: {}').format(e)) return list() def save(self): From fb32144aab29f6008f7757941ab7192c4e0643d1 Mon Sep 17 00:00:00 2001 From: luzhongyang Date: Mon, 15 Nov 2021 14:12:48 +0800 Subject: [PATCH 2/2] fix bug --- apiserver/report/handler/heartbeat_handler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/apiserver/report/handler/heartbeat_handler.py b/apiserver/report/handler/heartbeat_handler.py index da4216a..587aa37 100644 --- a/apiserver/report/handler/heartbeat_handler.py +++ b/apiserver/report/handler/heartbeat_handler.py @@ -40,7 +40,7 @@ def parse(self): self.report_queue = self.detail.get('reportQueue', 0) self.method_queue = self.detail.get('methodQueue', 0) self.replay_queue = self.detail.get('replayQueue', 0) - self.return_queue = self.detail.get('replayQueue', 1) + self.return_queue = self.detail.get('returnQueue', 1) def has_permission(self): self.agent = IastAgent.objects.filter(id=self.agent_id, user=self.user_id).first()