From 89f303b051b1bed2119bed92a3d5744bd7e059ac Mon Sep 17 00:00:00 2001
From: Nassim Tabchiche <n.tabchiche2@gmail.com>
Date: Thu, 5 Dec 2024 18:48:23 +0100
Subject: [PATCH] Add likelihood endpoint for attack paths

---
 backend/ebios_rm/models.py |  8 ++++++++
 backend/ebios_rm/views.py  | 19 +++++++++++++++++++
 2 files changed, 27 insertions(+)

diff --git a/backend/ebios_rm/models.py b/backend/ebios_rm/models.py
index e0f09a192d..2db08a5819 100644
--- a/backend/ebios_rm/models.py
+++ b/backend/ebios_rm/models.py
@@ -388,6 +388,14 @@ def save(self, *args, **kwargs):
         self.folder = self.ebios_rm_study.folder
         super().save(*args, **kwargs)
 
+    @property
+    def risk_matrix(self):
+        return self.ebios_rm_study.risk_matrix
+
+    @property
+    def parsed_matrix(self):
+        return self.risk_matrix.parse_json_translated()
+
 
 class OperationalScenario(AbstractBaseModel, FolderMixin):
     ebios_rm_study = models.ForeignKey(
diff --git a/backend/ebios_rm/views.py b/backend/ebios_rm/views.py
index 8fcaf6ce9a..0416e3adc8 100644
--- a/backend/ebios_rm/views.py
+++ b/backend/ebios_rm/views.py
@@ -109,6 +109,25 @@ def category(self, request):
 class AttackPathViewSet(BaseModelViewSet):
     model = AttackPath
 
+    @action(detail=True, name="Get risk matrix", url_path="risk-matrix")
+    def risk_matrix(self, request, pk=None):
+        attack_path = self.get_object()
+        return Response(RiskMatrixReadSerializer(attack_path.risk_matrix).data)
+
+    @method_decorator(cache_page(60 * LONG_CACHE_TTL))
+    @action(detail=True, name="Get likelihood choices")
+    def likelihood(self, request, pk):
+        attack_path: AttackPath = self.get_object()
+        undefined = dict([(-1, "--")])
+        _choices = dict(
+            zip(
+                list(range(0, 64)),
+                [x["name"] for x in attack_path.parsed_matrix["probability"]],
+            )
+        )
+        choices = undefined | _choices
+        return Response(choices)
+
 
 class OperationalScenarioViewSet(BaseModelViewSet):
     model = OperationalScenario