respect system settings product grading enabled

Author: valentijnscholten

dojo/importers/default_importer.py

@@ -1,6 +1,6 @@
 import logging
 
-from celery import chord
+from celery import chord, group
 from django.core.files.uploadedfile import TemporaryUploadedFile
 from django.core.serializers import serialize
 from django.db.models.query_utils import Q
@@ -15,6 +15,7 @@
 from dojo.models import (
     Engagement,
     Finding,
+    System_Settings,
     Test,
     Test_Import,
 )
@@ -274,8 +275,13 @@ def process_findings(
                 if batch_full or is_final:
                     # Launch chord with current batch of signatures
                     product = self.test.engagement.product
-                    calculate_grade_signature = utils.calculate_grade_signature(product)
-                    chord(post_processing_task_signatures)(calculate_grade_signature)
+                    system_settings = System_Settings.objects.get()
+                    if system_settings.enable_product_grade:
+                        calculate_grade_signature = utils.calculate_grade_signature(product)
+                        chord(post_processing_task_signatures)(calculate_grade_signature)
+                    # If product grading is disabled, just run the post-processing tasks without the grade calculation callback
+                    elif post_processing_task_signatures:
+                        group(post_processing_task_signatures).apply_async()
 
                     logger.debug(f"Launched chord with {len(post_processing_task_signatures)} tasks (batch #{current_batch_number}, size: {len(post_processing_task_signatures)})")
 
@@ -305,7 +311,9 @@ def process_findings(
 
         # Always perform an initial grading, even though it might get overwritten later.
         product = self.test.engagement.product
-        calculate_grade(product)
+        system_settings = System_Settings.objects.get()
+        if system_settings.enable_product_grade:
+            calculate_grade(product)
 
         sync = kwargs.get("sync", True)
         if not sync:
@@ -394,7 +402,9 @@ def close_old_findings(
         # Calculate grade once after all findings have been closed
         if old_findings:
             product = self.test.engagement.product
-            calculate_grade(product)
+            system_settings = System_Settings.objects.get()
+            if system_settings.enable_product_grade:
+                calculate_grade(product)
 
         return old_findings
 

dojo/importers/default_reimporter.py

@@ -1,6 +1,6 @@
 import logging
 
-from celery import chord
+from celery import chord, group
 from django.core.files.uploadedfile import TemporaryUploadedFile
 from django.core.serializers import serialize
 from django.db.models.query_utils import Q
@@ -15,6 +15,7 @@
     Development_Environment,
     Finding,
     Notes,
+    System_Settings,
     Test,
     Test_Import,
 )
@@ -279,8 +280,13 @@ def process_findings(
                 if batch_full or is_final:
                     # Launch chord with current batch of signatures
                     product = self.test.engagement.product
-                    calculate_grade_signature = utils.calculate_grade_signature(product)
-                    chord(post_processing_task_signatures)(calculate_grade_signature)
+                    system_settings = System_Settings.objects.get()
+                    if system_settings.enable_product_grade:
+                        calculate_grade_signature = utils.calculate_grade_signature(product)
+                        chord(post_processing_task_signatures)(calculate_grade_signature)
+                    # If product grading is disabled, just run the post-processing tasks without the grade calculation callback
+                    elif post_processing_task_signatures:
+                        group(post_processing_task_signatures).apply_async()
 
                     logger.debug(f"Launched chord with {len(post_processing_task_signatures)} tasks (batch #{current_batch_number}, size: {len(post_processing_task_signatures)})")
 
@@ -306,7 +312,9 @@ def process_findings(
 
         # Synchronous tasks were already executed during processing, just calculate grade
         product = self.test.engagement.product
-        calculate_grade(product)
+        system_settings = System_Settings.objects.get()
+        if system_settings.enable_product_grade:
+            calculate_grade(product)
 
         # Process the results and return them back
         return self.process_results(**kwargs)
@@ -353,7 +361,9 @@ def close_old_findings(
         # Calculate grade once after all findings have been closed
         if mitigated_findings:
             product = self.test.engagement.product
-            calculate_grade(product)
+            system_settings = System_Settings.objects.get()
+            if system_settings.enable_product_grade:
+                calculate_grade(product)
 
         return mitigated_findings
 

unittests/test_importers_performance.py

@@ -176,12 +176,12 @@ def import_reimport_performance(self, expected_num_queries1, expected_num_async_
     # def test_import_reimport_reimport_performance_async(self, mock):
     def test_import_reimport_reimport_performance_async(self):
         self.import_reimport_performance(
-            expected_num_queries1=715,
-            expected_num_async_tasks1=13,
-            expected_num_queries2=612,
-            expected_num_async_tasks2=25,
-            expected_num_queries3=294,
-            expected_num_async_tasks3=23,
+            expected_num_queries1=713,
+            expected_num_async_tasks1=10,
+            expected_num_queries2=610,
+            expected_num_async_tasks2=22,
+            expected_num_queries3=292,
+            expected_num_async_tasks3=20,
         )
 
     # @patch("dojo.decorators.we_want_async", return_value=False)
@@ -199,11 +199,11 @@ def test_import_reimport_reimport_performance_no_async(self):
         testuser.usercontactinfo.save()
         self.import_reimport_performance(
             expected_num_queries1=713,
-            expected_num_async_tasks1=11,
+            expected_num_async_tasks1=10,
             expected_num_queries2=615,
-            expected_num_async_tasks2=23,
+            expected_num_async_tasks2=22,
             expected_num_queries3=297,
-            expected_num_async_tasks3=21,
+            expected_num_async_tasks3=20,
         )
 
     # @patch("dojo.decorators.we_want_async", return_value=False)