Skip to content

Commit

Permalink
fix task names and celery beat errors
Browse files Browse the repository at this point in the history
  • Loading branch information
@yogeshojha
yogeshojha committed Jul 20, 2024
1 parent 8c9187c commit 5667893
Show file tree
Hide file tree
Showing 2 changed files with 66 additions and 49 deletions.
31 changes: 30 additions & 1 deletion web/reNgine/common_func.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1028,4 +1028,33 @@ def parse_llm_vulnerability_report(report):
except Exception as e:
return data

return data
return data


def create_scan_object(host_id, engine_id, initiated_by_id=None):
'''
create task with pending status so that celery task will execute when
threads are free
Args:
host_id: int: id of Domain model
engine_id: int: id of EngineType model
initiated_by_id: int : id of User model (Optional)
'''
# get current time
current_scan_time = timezone.now()
# fetch engine and domain object
engine = EngineType.objects.get(pk=engine_id)
domain = Domain.objects.get(pk=host_id)
scan = ScanHistory()
scan.scan_status = INITIATED_TASK
scan.domain = domain
scan.scan_type = engine
scan.start_scan_date = current_scan_time
if initiated_by_id:
user = User.objects.get(pk=initiated_by_id)
scan.initiated_by = user
scan.save()
# save last scan date for domain model
domain.start_scan_date = current_scan_time
domain.save()
return scan.id
84 changes: 36 additions & 48 deletions web/startScan/views.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -271,9 +271,9 @@ def start_scan_ui(request, slug, domain_id):

# Create ScanHistory object
scan_history_id = create_scan_object(
domain_id,
engine_id,
request.user
host_id=domain_id,
engine_id=engine_id,
initiated_by_id=request.user.id
)
scan = ScanHistory.objects.get(pk=scan_history_id)

Expand All @@ -286,7 +286,8 @@ def start_scan_ui(request, slug, domain_id):
'results_dir': '/usr/src/scan_results',
'imported_subdomains': subdomains_in,
'out_of_scope_subdomains': subdomains_out,
'url_filter': filterPath
'url_filter': filterPath,
'initiated_by_id': request.user.id
}
initiate_scan.apply_async(kwargs=kwargs)
scan.save()
Expand Down Expand Up @@ -328,9 +329,9 @@ def start_multiple_scan(request, slug):
for domain_id in list_of_domains.split(","):
# Start the celery task
scan_history_id = create_scan_object(
domain_id,
engine_id,
request.user
host_id=domain_id,
engine_id=engine_id,
initiated_by_id=request.user.id
)
# domain = get_object_or_404(Domain, id=domain_id)

Expand All @@ -340,6 +341,7 @@ def start_multiple_scan(request, slug):
'engine_id': engine_id,
'scan_type': LIVE_SCAN,
'results_dir': '/usr/src/scan_results',
'initiated_by_id': request.user.id
# TODO: Add this to multiple scan view
# 'imported_subdomains': subdomains_in,
# 'out_of_scope_subdomains': subdomains_out
Expand Down Expand Up @@ -535,12 +537,15 @@ def schedule_scan(request, host_id, slug):
'scan_history_id': 1,
'scan_type': SCHEDULED_SCAN,
'imported_subdomains': subdomains_in,
'out_of_scope_subdomains': subdomains_out
'out_of_scope_subdomains': subdomains_out,
'initiated_by_id': request.user.id
}
PeriodicTask.objects.create(interval=schedule,
name=task_name,
task='reNgine.tasks.initiate_scan',
kwargs=json.dumps(kwargs))
PeriodicTask.objects.create(
interval=schedule,
name=task_name,
task='initiate_scan',
kwargs=json.dumps(kwargs)
)
elif scheduled_mode == 'clocked':
schedule_time = request.POST['scheduled_time']
clock, _ = ClockedSchedule.objects.get_or_create(
Expand All @@ -551,13 +556,16 @@ def schedule_scan(request, host_id, slug):
'engine_id': engine.id,
'scan_type': SCHEDULED_SCAN,
'imported_subdomains': subdomains_in,
'out_of_scope_subdomains': subdomains_out
'out_of_scope_subdomains': subdomains_out,
'initiated_by_id': request.user.id
}
PeriodicTask.objects.create(clocked=clock,
one_off=True,
name=task_name,
task='reNgine.tasks.initiate_scan',
kwargs=json.dumps(kwargs))
PeriodicTask.objects.create(
clocked=clock,
one_off=True,
name=task_name,
task='initiate_scan',
kwargs=json.dumps(kwargs)
)
messages.add_message(
request,
messages.INFO,
Expand Down Expand Up @@ -629,29 +637,6 @@ def change_vuln_status(request, id):
return HttpResponse('')


def create_scan_object(host_id, engine_id, initiated_by):
'''
create task with pending status so that celery task will execute when
threads are free
'''
# get current time
current_scan_time = timezone.now()
# fetch engine and domain object
engine = EngineType.objects.get(pk=engine_id)
domain = Domain.objects.get(pk=host_id)
scan = ScanHistory()
scan.scan_status = INITIATED_TASK
scan.domain = domain
scan.scan_type = engine
scan.start_scan_date = current_scan_time
scan.initiated_by = initiated_by
scan.save()
# save last scan date for domain model
domain.start_scan_date = current_scan_time
domain.save()
return scan.id


@has_permission_decorator(PERM_MODIFY_SYSTEM_CONFIGURATIONS, redirect_url=FOUR_OH_FOUR_URL)
def delete_all_scan_results(request):
if request.method == 'POST':
Expand Down Expand Up @@ -694,9 +679,9 @@ def start_organization_scan(request, id, slug):
# Start Celery task for each organization's domains
for domain in organization.get_domains():
scan_history_id = create_scan_object(
domain.id,
engine_id,
request.user
host_id=domain.id,
engine_id=engine_id,
initiated_by_id=request.user.id
)
scan = ScanHistory.objects.get(pk=scan_history_id)

Expand All @@ -706,6 +691,7 @@ def start_organization_scan(request, id, slug):
'engine_id': engine_id,
'scan_type': LIVE_SCAN,
'results_dir': '/usr/src/scan_results',
'initiated_by_id': request.user.id,
# TODO: Add this to multiple scan view
# 'imported_subdomains': subdomains_in,
# 'out_of_scope_subdomains': subdomains_out
Expand Down Expand Up @@ -773,12 +759,13 @@ def schedule_organization_scan(request, slug, id):
'engine_id': engine.id,
'scan_history_id': 0,
'scan_type': SCHEDULED_SCAN,
'imported_subdomains': None
'imported_subdomains': None,
'initiated_by_id': request.user.id
})
PeriodicTask.objects.create(
interval=schedule,
name=task_name,
task='reNgine.tasks.initiate_scan',
task='initiate_scan',
kwargs=_kwargs
)

Expand All @@ -793,12 +780,13 @@ def schedule_organization_scan(request, slug, id):
'engine_id': engine.id,
'scan_history_id': 0,
'scan_type': LIVE_SCAN,
'imported_subdomains': None
'imported_subdomains': None,
'initiated_by_id': request.user.id
})
PeriodicTask.objects.create(clocked=clock,
one_off=True,
name=task_name,
task='reNgine.tasks.initiate_scan',
task='initiate_scan',
kwargs=_kwargs
)

Expand Down

0 comments on commit 5667893

Please sign in to comment.