Skip to content
This repository has been archived by the owner on Apr 20, 2022. It is now read-only.

when return_queue is 0, do not return queue data. return_queue defaul… #119

Merged
merged 1 commit into from
Nov 15, 2021
Merged

when return_queue is 0, do not return queue data. return_queue defaul… #119

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
61 changes: 32 additions & 29 deletions apiserver/report/handler/heartbeat_handler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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')
Expand All @@ -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()
Expand Down Expand Up @@ -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):
Expand Down