update model + migration

Author: kiblik

dojo/api_v2/serializers.py

@@ -1578,10 +1578,10 @@ def get_engagement(self, obj):
         )
 
     def validate(self, data):
-        def validate_findings_have_same_engagement(finding_objects: list[Finding]):
+        def validate_findings_have_same_engagement(finding_objects: list[Finding]):  # TODO: check
             engagements = finding_objects.values_list("test__engagement__id", flat=True).distinct().count()
             if engagements > 1:
-                msg = "You are not permitted to add findings from multiple engagements"
+                msg = "You are not permitted to add findings from multiple engagements"  # TODO: same is missing for UI
                 raise PermissionDenied(msg)
 
         findings = data.get("accepted_findings", [])

dojo/api_v2/views.py

@@ -704,7 +704,7 @@ def get_queryset(self):
         return (
             get_authorized_risk_acceptances(Permissions.Risk_Acceptance)
             .prefetch_related(
-                "notes", "engagement_set", "owner", "accepted_findings",
+                "notes", "engagement_set", "owner", "accepted_findings",  # TODO: Check this
             )
             .distinct()
         )

dojo/db_migrations/0231_add_engagement_risk_acceptance.py

@@ -0,0 +1,19 @@
+# Generated by Django 5.1.8 on 2025-05-01 12:54
+
+import django.db.models.deletion
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('dojo', '0230_alter_jira_instance_accepted_mapping_resolution_and_more'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='risk_acceptance',
+            name='engagement',
+            field=models.ForeignKey(blank=True, editable=False, null=True, on_delete=django.db.models.deletion.CASCADE, to='dojo.engagement'),
+        ),
+    ]

dojo/db_migrations/0232_set_risk_acceptance_engagement.py

@@ -0,0 +1,28 @@
+# Generated by Django 5.1.8 on 2025-05-01 12:59
+
+import django.db.models.deletion
+from django.db import migrations, models
+import logging
+
+logger = logging.getLogger(__name__)
+
+def set_engagement_based_on_findings(apps, schema_editor):
+    Engagement = apps.get_model('dojo', 'Engagement')
+    RiskAcceptance = apps.get_model('dojo', 'Risk_Acceptance')
+    through_model = Engagement.risk_acceptance.through
+
+    for rel in through_model.objects.all():
+        ra = RiskAcceptance.objects.get(pk=rel.risk_acceptance_id)
+        ra.engagement_id = rel.engagement_id
+        ra.save()
+            
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('dojo', '0231_add_engagement_risk_acceptance'),
+    ]
+
+    operations = [
+        migrations.RunPython(set_engagement_based_on_findings, migrations.RunPython.noop),
+    ]

dojo/db_migrations/0233_alter_risk_acceptance_engagement.py

@@ -0,0 +1,23 @@
+# Generated by Django 5.1.8 on 2025-05-01 12:59
+
+import django.db.models.deletion
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('dojo', '0232_set_risk_acceptance_engagement'),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name='risk_acceptance',
+            name='engagement',
+            field=models.ForeignKey(editable=False, on_delete=django.db.models.deletion.CASCADE, to='dojo.engagement'),
+        ),
+        migrations.RemoveField(
+            model_name='engagement',
+            name='risk_acceptance',
+        ),
+    ]

dojo/engagement/views.py

@@ -425,7 +425,7 @@ def get_template(self):
         return "dojo/view_eng.html"
 
     def get_risks_accepted(self, eng):
-        return eng.risk_acceptance.all().select_related("owner").annotate(accepted_findings_count=Count("accepted_findings__id"))
+        return eng.risk_acceptance_set.all().select_related("owner").annotate(accepted_findings_count=Count("accepted_findings__id"))
 
     def get_filtered_tests(
         self,

dojo/models.py

@@ -1505,10 +1505,6 @@ class Engagement(models.Model):
                                 default="threat_model", editable=False)
     tmodel_path = models.CharField(max_length=1000, default="none",
                                    editable=False, blank=True, null=True)
-    risk_acceptance = models.ManyToManyField("Risk_Acceptance",
-                                             default=None,
-                                             editable=False,
-                                             blank=True)
     done_testing = models.BooleanField(default=False, editable=False)
     engagement_type = models.CharField(editable=True, max_length=30, default="Interactive",
                                        null=True,
@@ -1550,7 +1546,7 @@ def copy(self):
         old_notes = list(self.notes.all())
         old_files = list(self.files.all())
         old_tags = list(self.tags.all())
-        old_risk_acceptances = list(self.risk_acceptance.all())
+        old_risk_acceptances = list(self.risk_acceptance_set.all())
         old_tests = list(Test.objects.filter(engagement=self))
         # Save the object before setting any ManyToMany relationships
         copy.save()
@@ -1565,7 +1561,7 @@ def copy(self):
             test.copy(engagement=copy)
         # Copy the risk_acceptances
         for risk_acceptance in old_risk_acceptances:
-            copy.risk_acceptance.add(risk_acceptance.copy(engagement=copy))
+            risk_acceptance.copy(engagement=copy)
         # Assign any tags
         copy.tags.set(old_tags)
 
@@ -1595,9 +1591,6 @@ def unaccepted_open_findings(self):
 
         return findings
 
-    def accept_risks(self, accepted_risks):
-        self.risk_acceptance.add(*accepted_risks)
-
     @property
     def has_jira_issue(self):
         import dojo.jira_link.helper as jira_helper
@@ -2163,9 +2156,6 @@ def unaccepted_open_findings(self):
 
         return findings
 
-    def accept_risks(self, accepted_risks):
-        self.engagement.risk_acceptance.add(*accepted_risks)
-
     @property
     def deduplication_algorithm(self):
         deduplicationAlgorithm = settings.DEDUPE_ALGO_LEGACY
@@ -3663,6 +3653,8 @@ class Risk_Acceptance(models.Model):
 
     name = models.CharField(max_length=300, null=False, blank=False, help_text=_("Descriptive name which in the future may also be used to group risk acceptances together across engagements and products"))
 
+    engagement = models.ForeignKey(Engagement, editable=False, blank=False, null=False, on_delete=models.CASCADE)
+
     accepted_findings = models.ManyToManyField(Finding)
 
     recommendation = models.CharField(choices=TREATMENT_CHOICES, max_length=2, null=False, default=TREATMENT_FIX, help_text=_("Recommendation from the security team."), verbose_name=_("Security Recommendation"))
@@ -3704,26 +3696,17 @@ def name_and_expiration_info(self):
         return str(self.name) + (" (expired " if self.is_expired else " (expires ") + (timezone.localtime(self.expiration_date).strftime("%b %d, %Y") if self.expiration_date else "Never") + ")"
 
     def get_breadcrumbs(self):
-        bc = self.engagement_set.first().get_breadcrumbs()
+        bc = self.engagement.get_breadcrumbs()
         bc += [{"title": str(self),
                 "url": reverse("view_risk_acceptance", args=(
-                    self.engagement_set.first().product.id, self.id))}]
+                    self.engagement.product.id, self.id))}]
         return bc
 
     @property
     def is_expired(self):
         return self.expiration_date_handled is not None
 
-    # relationship is many to many, but we use it as one-to-many
-    @property
-    def engagement(self):
-        engs = self.engagement_set.all()
-        if engs:
-            return engs[0]
-
-        return None
-
-    def copy(self, engagement=None):
+    def copy(self, engagement):
         copy = _copy_model_util(self)
         # Save the necessary ManyToMany relationships
         old_notes = list(self.notes.all())
@@ -3734,9 +3717,10 @@ def copy(self, engagement=None):
         for notes in old_notes:
             copy.notes.add(notes.copy())
         # Assign any accepted findings
-        if engagement:
-            new_accepted_findings = Finding.objects.filter(test__engagement=engagement, hash_code__in=old_accepted_findings_hash_codes, risk_accepted=True).distinct()
-            copy.accepted_findings.set(new_accepted_findings)
+        new_accepted_findings = Finding.objects.filter(test__engagement=engagement, hash_code__in=old_accepted_findings_hash_codes, risk_accepted=True).distinct()
+        copy.accepted_findings.set(new_accepted_findings)
+        copy.engagement = engagement
+        copy.save()
         return copy
 
 

dojo/risk_acceptance/api.py

@@ -52,7 +52,6 @@ def accept_risks(self, request, pk=None):
         base_findings = model.unaccepted_open_findings
         owner = request.user
         accepted = _accept_risks(accepted_risks, base_findings, owner)
-        model.accept_risks(accepted)
         result = RiskAcceptanceSerializer(instance=accepted, many=True)
         return Response(status=status.HTTP_201_CREATED, data=result.data)
 
@@ -75,7 +74,6 @@ def accept_risks(self, request):
         for engagement in get_authorized_engagements(Permissions.Engagement_View):
             base_findings = engagement.unaccepted_open_findings
             accepted = _accept_risks(accepted_risks, base_findings, owner)
-            engagement.accept_risks(accepted)
             accepted_result.extend(accepted)
         result = RiskAcceptanceSerializer(instance=accepted_result, many=True)
         return Response(status=201, data=result.data)
@@ -95,7 +93,8 @@ def _accept_risks(accepted_risks: list[AcceptedRisk], base_findings: QuerySet, o
             acceptance = Risk_Acceptance.objects.create(owner=owner, name=name[:100],
                                                         decision=Risk_Acceptance.TREATMENT_ACCEPT,
                                                         decision_details=risk.justification,
-                                                        accepted_by=risk.accepted_by[:200])
+                                                        accepted_by=risk.accepted_by[:200],
+                                                        engagement=findings[0].test.engagement)  # TODO: Add validator that all findings are from this Eng to Model
             acceptance.accepted_findings.set(findings)
             findings.update(risk_accepted=True, active=False)
             acceptance.save()

dojo/risk_acceptance/helper.py

@@ -294,7 +294,7 @@ def get_almost_expired_risk_acceptances_to_handle(heads_up_days):
 
 def prefetch_for_expiration(risk_acceptances):
     return risk_acceptances.prefetch_related("accepted_findings", "accepted_findings__jira_issue",
-                                                "engagement_set",
+                                                "engagement_set",  # TODO: check this
                                                 "engagement__jira_project",
                                                 "engagement__jira_project__jira_instance",
                                              )