refactor signal

Signed-off-by: tdruez <tdruez@aboutcode.org>

Author: tdruez

vulnerabilities/triage/admin.py

@@ -13,6 +13,7 @@
 from dje.admin import DataspacedAdmin
 from dje.admin import dejacode_site
 from dje.list_display import AsLink
+from vulnerabilities.triage.engine import evaluate_ruleset
 from vulnerabilities.triage.forms import AnalysisPresetForm
 from vulnerabilities.triage.forms import TriageRulesetForm
 from vulnerabilities.triage.models import AnalysisPreset
@@ -130,6 +131,21 @@ def get_changes_details(self, form):
         ]
         return super().get_changes_details(form)
 
+    def save_model(self, request, obj, form, change):
+        super().save_model(request, obj, form, change)
+
+        if not obj.enabled:
+            # Records with an open Request are kept so re-enabling the ruleset reconnects to
+            # it instead of opening a duplicate Request.
+            obj.triage_records.filter(request__isnull=True).delete()
+            return
+
+        if not change:
+            return
+
+        for assignment in obj.product_triage_rulesets.select_related("product"):
+            evaluate_ruleset(ruleset=obj, product=assignment.product)
+
     def get_form(self, request, obj=None, change=False, **kwargs):
         kwargs["fields"] = [
             "name",

vulnerabilities/triage/engine.py

@@ -9,6 +9,7 @@
 from django.apps import apps
 from django.utils import timezone
 
+from vulnerabilities.triage.models import ProductTriageRuleset
 from vulnerabilities.triage.models import TriageRecord
 from vulnerabilities.triage.rules import RULE_REGISTRY
 from vulnerabilities.triage.rules import rule_parameters_from_config
@@ -225,3 +226,13 @@ def evaluate_ruleset(ruleset, product, apply_preset=True):
         matched_rules_per_vulnerability_id=matched_rules_per_vulnerability_id,
         apply_preset=apply_preset,
     )
+
+
+def reevaluate_product_rulesets(product, apply_preset=True):
+    """Re-evaluate all enabled triage rulesets currently assigned to the product."""
+    assignments = ProductTriageRuleset.objects.filter(
+        product=product, ruleset__enabled=True
+    ).select_related("ruleset", "ruleset__analysis_preset")
+
+    for assignment in assignments:
+        evaluate_ruleset(ruleset=assignment.ruleset, product=product, apply_preset=apply_preset)

vulnerabilities/triage/management/commands/create_triage_rulesets.py

@@ -139,9 +139,18 @@ def add_arguments(self, parser):
             action="store_true",
             help=(
                 "Delete all existing triage rulesets in the dataspace before recreating them."
-                " This also removes all associated product assignments and triage records."
+                " This also removes all associated product assignments and triage records,"
+                " and every product using these rulesets will need to be reassigned manually"
+                " afterward."
             ),
         )
+        parser.add_argument(
+            "--noinput",
+            "--no-input",
+            action="store_false",
+            dest="interactive",
+            help="Tells Django to NOT prompt the user for input of any kind.",
+        )
 
     def handle(self, *args, **options):
         dataspace_name = options["dataspace"]
@@ -152,62 +161,68 @@ def handle(self, *args, **options):
             raise CommandError(f'Dataspace "{dataspace_name}" does not exist.')
 
         if options["reset"]:
+            if options["interactive"]:
+                confirm = input(
+                    "You have requested a reset of the triage rulesets in dataspace"
+                    f' "{dataspace_name}".\n'
+                    "This will delete all triage rulesets and analysis presets in this"
+                    " dataspace, along with every associated product assignment and triage"
+                    " record. Every product using these rulesets will need to be reassigned"
+                    " manually afterward.\n"
+                    "Are you sure you want to do this?\n\n"
+                    "    Type 'yes' to continue, or 'no' to cancel: "
+                )
+            else:
+                confirm = "yes"
+
+            if confirm != "yes":
+                self.stdout.write("Reset cancelled.")
+                return
+
             deleted_rulesets, _ = TriageRuleset.objects.filter(dataspace=dataspace).delete()
             deleted_presets, _ = AnalysisPreset.objects.filter(dataspace=dataspace).delete()
             self.stdout.write(
                 f"  Deleted {deleted_rulesets} existing ruleset(s) and {deleted_presets} preset(s)."
             )
+        elif TriageRuleset.objects.filter(dataspace=dataspace).exists():
+            raise CommandError(
+                f'Dataspace "{dataspace_name}" already has triage rulesets.'
+                " Use --reset to delete and recreate them."
+            )
 
-        ruleset_created_count = 0
         for ruleset_data in REFERENCE_RULESETS:
-            _, created = TriageRuleset.objects.get_or_create(
+            TriageRuleset.objects.create(
                 dataspace=dataspace,
                 name=ruleset_data["name"],
-                defaults={
-                    "description": ruleset_data["description"],
-                    "recommended_action": ruleset_data["recommended_action"],
-                    "precedence": ruleset_data["precedence"],
-                    "rules_config": ruleset_data["rules_config"],
-                    "enabled": True,
-                },
+                description=ruleset_data["description"],
+                recommended_action=ruleset_data["recommended_action"],
+                precedence=ruleset_data["precedence"],
+                rules_config=ruleset_data["rules_config"],
+                enabled=True,
             )
-            if created:
-                ruleset_created_count += 1
-                self.stdout.write(f"  Created: {ruleset_data['name']}")
-            else:
-                self.stdout.write(f"  Already exists: {ruleset_data['name']}")
+            self.stdout.write(f"  Created: {ruleset_data['name']}")
 
-        preset_created_count = 0
         for preset_data in REFERENCE_PRESETS:
-            ruleset_name = preset_data["ruleset_name"]
-            preset, preset_created = AnalysisPreset.objects.get_or_create(
+            preset = AnalysisPreset.objects.create(
                 dataspace=dataspace,
                 name=preset_data["name"],
-                defaults={
-                    "description": preset_data.get("description", ""),
-                    "state": preset_data.get("state", ""),
-                    "justification": preset_data.get("justification", ""),
-                    "responses": preset_data.get("responses"),
-                    "detail": preset_data.get("detail", ""),
-                },
+                description=preset_data.get("description", ""),
+                state=preset_data.get("state", ""),
+                justification=preset_data.get("justification", ""),
+                responses=preset_data.get("responses"),
+                detail=preset_data.get("detail", ""),
             )
-            if preset_created:
-                preset_created_count += 1
-                self.stdout.write(f"  Created preset: {preset_data['name']}")
-            else:
-                self.stdout.write(f"  Already exists: {preset_data['name']}")
-            try:
-                ruleset = TriageRuleset.objects.get(dataspace=dataspace, name=ruleset_name)
-                if ruleset.analysis_preset_id != preset.pk:
-                    ruleset.analysis_preset = preset
-                    ruleset.save(update_fields=["analysis_preset"])
-                    self.stdout.write(f"    Linked preset to ruleset: {ruleset_name}")
-            except TriageRuleset.DoesNotExist:
-                self.stdout.write(f"    Ruleset not found: {ruleset_name}")
+            self.stdout.write(f"  Created preset: {preset_data['name']}")
+
+            ruleset_name = preset_data["ruleset_name"]
+            ruleset = TriageRuleset.objects.get(dataspace=dataspace, name=ruleset_name)
+            ruleset.analysis_preset = preset
+            ruleset.save(update_fields=["analysis_preset"])
+            self.stdout.write(f"    Linked preset to ruleset: {ruleset_name}")
 
         self.stdout.write(
             self.style.SUCCESS(
-                f"{ruleset_created_count} ruleset(s) and {preset_created_count} preset(s)"
+                f"{len(REFERENCE_RULESETS)} ruleset(s) and {len(REFERENCE_PRESETS)} preset(s)"
                 f" created in dataspace '{dataspace_name}'."
             )
         )

vulnerabilities/triage/signals.py

@@ -12,36 +12,10 @@
 
 from vulnerabilities.triage.engine import delete_preset_analyses_for_product
 from vulnerabilities.triage.engine import evaluate_ruleset
-from vulnerabilities.triage.models import ProductTriageRuleset
+from vulnerabilities.triage.engine import reevaluate_product_rulesets
 from vulnerabilities.triage.models import TriageRecord
 
 
-def reevaluate_product_rulesets(product, apply_preset=True):
-    """Re-evaluate all enabled triage rulesets currently assigned to the product."""
-    assignments = ProductTriageRuleset.objects.filter(
-        product=product, ruleset__enabled=True
-    ).select_related("ruleset", "ruleset__analysis_preset")
-
-    for assignment in assignments:
-        evaluate_ruleset(ruleset=assignment.ruleset, product=product, apply_preset=apply_preset)
-
-
-@receiver(post_save, sender="vulnerabilities_triage.TriageRuleset")
-def reevaluate_or_delete_on_ruleset_save(sender, instance, created, **kwargs):
-    """Re-evaluate assigned products on config change; delete records when disabled."""
-    if not instance.enabled:
-        # Records with an open Request are kept so re-enabling the ruleset reconnects to it
-        # instead of opening a duplicate Request.
-        instance.triage_records.filter(request__isnull=True).delete()
-        return
-
-    if created:
-        return
-
-    for assignment in instance.product_triage_rulesets.select_related("product"):
-        evaluate_ruleset(ruleset=instance, product=assignment.product)
-
-
 @receiver(post_save, sender="vulnerabilities_triage.ProductTriageRuleset")
 def evaluate_on_assign(sender, instance, created, **kwargs):
     """Evaluate the ruleset against the product as soon as it is assigned."""