Merge pull request #52 from mantiumai/zim/32

scan status polling

chirps/scan/migrations/0003_scan_progress.py

@@ -0,0 +1,17 @@
+# Generated by Django 4.2.2 on 2023-07-07 19:56
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ('scan', '0002_initial'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='scan',
+            name='progress',
+            field=models.IntegerField(default=0),
+        ),
+    ]

chirps/scan/models.py

@@ -16,6 +16,7 @@ class Scan(models.Model):
     plan = models.ForeignKey('plan.Plan', on_delete=models.CASCADE)
     target = models.ForeignKey('target.BaseTarget', on_delete=models.CASCADE)
     celery_task_id = models.CharField(max_length=256, null=True)
+    progress = models.IntegerField(default=0)
     user = models.ForeignKey(User, on_delete=models.CASCADE, null=True)
 
     def __str__(self) -> str:
@@ -55,6 +56,7 @@ class Result(models.Model):
     def __str__(self):
         return f'{self.rule.name} - {self.scan.id}'
 
+
 class Finding(models.Model):
     """Model to identify the location of a finding within a result."""
 
@@ -67,22 +69,22 @@ def __str__(self):
 
     def text(self):
         """Return the text of the finding."""
-        return self.result.text[self.offset:self.offset + self.length]
+        return self.result.text[self.offset : self.offset + self.length]
 
     def surrounding_text(self):
         """return the text of the finding, with some surrounding context."""
-        buffer = self.result.text[self.offset - 20: self.offset - 1]
+        buffer = self.result.text[self.offset - 20 : self.offset - 1]
         buffer += "<span class='text-danger'>"
         buffer += self.result.text[self.offset : self.offset + self.length]
-        buffer += "</span>"
-        buffer += self.result.text[self.offset + self.length + 1: self.offset + self.length + 19]
+        buffer += '</span>'
+        buffer += self.result.text[self.offset + self.length + 1 : self.offset + self.length + 19]
         return mark_safe(buffer)
 
     def with_highlight(self):
         """return the entire text searched by the finding's rule - highlight the finding."""
         buffer = self.result.text[0 : self.offset - 1]
         buffer += "<span class='bg-danger text-white'>"
         buffer += self.result.text[self.offset : self.offset + self.length]
-        buffer += "</span>"
-        buffer += self.result.text[self.offset + self.length + 1 : ]
+        buffer += '</span>'
+        buffer += self.result.text[self.offset + self.length + 1 :]
         return mark_safe(buffer)

chirps/scan/tasks.py

@@ -30,6 +30,8 @@ def scan_task(scan_id):
     target = BaseTarget.objects.get(id=scan.target.id)
 
     # Now that we have the derrived class, call its implementation of search()
+    total_rules = scan.plan.rules.all().count()
+    rules_run = 0
     for rule in scan.plan.rules.all():
         logger.info('Starting rule evaluation', extra={'id': rule.id})
         results = target.search(query=rule.query_string, max_results=100)
@@ -47,6 +49,11 @@ def scan_task(scan_id):
                 finding = Finding(result=result, offset=match.start(), length=match.end() - match.start())
                 finding.save()
 
+        # Update the progress counter based on the number of rules that have been evaluated
+        rules_run += 1
+        scan.progress = int(rules_run / total_rules * 100)
+        scan.save()
+
     # Persist the completion time of the scan
     scan.finished_at = timezone.now()
     scan.save()

chirps/scan/templates/scan/dashboard.html

@@ -3,7 +3,7 @@
 
 {% block content %}
 <div class="d-flex">
-  <h1 class="my-auto ml-0 mr-auto text-success">Scans</h1> 
+  <h1 class="my-auto ml-0 mr-auto text-success">Scans</h1>
   <a class="btn btn-primary my-auto ml-auto mr-0" href="{% url 'scan_create' %}">New Scan</a>
 </div>
 
@@ -27,9 +27,7 @@ <h1 class="my-auto ml-0 mr-auto text-success">Scans</h1>
       <td>{{ scan.finished_at }}</td>
       <td>{{ scan.description }}</td>
       <td>{{ scan.target }}</td>
-      <td>{{ scan.celery_task_status }}
-        {% if scan.celery_task_status == 'FAILURE' %}<br />{{scan.celery_task_output}}{% endif %}
-      </td>
+      <td><div hx-get="{% url 'scan_status' scan.id %}" hx-trigger="every 1s"></div></td>
       <td>
         <button class="btn btn-link" type="button" data-bs-toggle="collapse" data-bs-target="#collapse{{scan.id}}"
           aria-expanded="false" aria-controls="collapse{{scan.id}}"> View Details </button>

chirps/scan/urls.py

@@ -7,5 +7,6 @@
     path('', views.dashboard, name='scan_dashboard'),
     path('create/', views.create, name='scan_create'),
     path('result/<int:result_id>/', views.result_detail, name='result_detail'),
-    path('finding/<int:finding_id>/', views.finding_detail, name='finding_detail')
+    path('finding/<int:finding_id>/', views.finding_detail, name='finding_detail'),
+    path('status/<int:scan_id>/', views.status, name='scan_status'),
 ]

chirps/scan/views.py

@@ -1,6 +1,7 @@
 """Views for the scan application."""
 
 from django.contrib.auth.decorators import login_required
+from django.http import HttpResponse
 from django.shortcuts import get_object_or_404, redirect, render
 
 from .forms import ScanForm
@@ -14,12 +15,14 @@ def finding_detail(request, finding_id):
     finding = get_object_or_404(Finding, pk=finding_id, result__scan__user=request.user)
     return render(request, 'scan/finding_detail.html', {'finding': finding})
 
+
 @login_required
 def result_detail(request, result_id):
     """Render the scan result detail page."""
     result = get_object_or_404(Result, pk=result_id, scan__user=request.user)
     return render(request, 'scan/result_detail.html', {'result': result})
 
+
 @login_required
 def create(request):
     """Render the scan creation page and handle POST requests."""
@@ -68,10 +71,25 @@ def dashboard(request):
                 scan.rules[result.rule.name] = {
                     'id': result.id,
                     'rule': result.rule,
-                    'findings': Finding.objects.filter(result=result).count()
+                    'findings': Finding.objects.filter(result=result).count(),
                 }
 
         # Convert the dictionary into a list that the template can iterate on
         scan.rules = scan.rules.values()
 
     return render(request, 'scan/dashboard.html', {'scans': user_scans})
+
+
+@login_required
+def status(request, scan_id):
+    """Fetch the status of a scan job."""
+    scan = get_object_or_404(Scan, pk=scan_id, user=request.user)
+
+    # Respond with the status of the celery task and the progress percentage of the scan
+    response = f'{scan.celery_task_status()} : {scan.progress} %'
+
+    if scan.finished_at is not None:
+        # HTMX will stop polling if we return a 286
+        return HttpResponse(content=response, status=286)
+
+    return HttpResponse(content=response, status=200)