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
             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/db_migrations/0230_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', '0229_alter_finding_unique_id_from_tool'),
+    ]
+
+    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/0231_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', '0230_add_engagement_risk_acceptance'),
+    ]
+
+    operations = [
+        migrations.RunPython(set_engagement_based_on_findings, migrations.RunPython.noop),
+    ]

dojo/db_migrations/0232_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', '0231_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.all().select_related("owner").annotate(accepted_findings_count=Count("accepted_findings__id"))  # TODO: check
 
     def get_filtered_tests(
         self,

dojo/models.py

@@ -1505,10 +1505,10 @@ 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)
+    # risk_acceptance = models.ManyToManyField("Risk_Acceptance",  # TODO: remove this
+    #                                          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 +1550,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.all())  # TODO check this
         old_tests = list(Test.objects.filter(engagement=self))
         # Save the object before setting any ManyToMany relationships
         copy.save()
@@ -1564,7 +1564,7 @@ def copy(self):
         for test in old_tests:
             test.copy(engagement=copy)
         # Copy the risk_acceptances
-        for risk_acceptance in old_risk_acceptances:
+        for risk_acceptance in old_risk_acceptances:  # TODO check this
             copy.risk_acceptance.add(risk_acceptance.copy(engagement=copy))
         # Assign any tags
         copy.tags.set(old_tags)
@@ -1596,7 +1596,7 @@ def unaccepted_open_findings(self):
         return findings
 
     def accept_risks(self, accepted_risks):
-        self.risk_acceptance.add(*accepted_risks)
+        self.risk_acceptance.add(*accepted_risks)  # TODO check this
 
     @property
     def has_jira_issue(self):
@@ -2164,7 +2164,7 @@ def unaccepted_open_findings(self):
         return findings
 
     def accept_risks(self, accepted_risks):
-        self.engagement.risk_acceptance.add(*accepted_risks)
+        self.engagement.risk_acceptance.add(*accepted_risks)  # TODO check this
 
     @property
     def deduplication_algorithm(self):
@@ -3663,6 +3663,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,7 +3706,7 @@ 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_set.first().get_breadcrumbs()  # TODO check this
         bc += [{"title": str(self),
                 "url": reverse("view_risk_acceptance", args=(
                     self.engagement_set.first().product.id, self.id))}]
@@ -3714,16 +3716,16 @@ def get_breadcrumbs(self):
     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]
+    # # 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
+    #     return None
 
-    def copy(self, engagement=None):
+    def copy(self, engagement=None):  # TODO check this
         copy = _copy_model_util(self)
         # Save the necessary ManyToMany relationships
         old_notes = list(self.notes.all())
@@ -3734,7 +3736,7 @@ def copy(self, engagement=None):
         for notes in old_notes:
             copy.notes.add(notes.copy())
         # Assign any accepted findings
-        if engagement:
+        if engagement:  # TODO check this
             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)
         return copy