From bed7b008c28b1f658ccc7261077958e7dd644c3e Mon Sep 17 00:00:00 2001
From: ziadhany <ziadhany2016@gmail.com>
Date: Tue, 21 May 2024 18:43:20 +0300
Subject: [PATCH] Add kev to api Access the vulnerability directly from Alias

Signed-off-by: ziadhany <ziadhany2016@gmail.com>
---
 vulnerabilities/api.py                         | 13 +++++++++++++
 vulnerabilities/improvers/vulnerability_kev.py |  3 +--
 2 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/vulnerabilities/api.py b/vulnerabilities/api.py
index f51998b87..173ec3e5e 100644
--- a/vulnerabilities/api.py
+++ b/vulnerabilities/api.py
@@ -25,6 +25,7 @@
 from rest_framework.throttling import UserRateThrottle
 
 from vulnerabilities.models import Alias
+from vulnerabilities.models import Kev
 from vulnerabilities.models import Package
 from vulnerabilities.models import Vulnerability
 from vulnerabilities.models import VulnerabilityReference
@@ -167,6 +168,12 @@ def to_representation(self, instance):
         return representation
 
 
+class KEVSerializer(serializers.ModelSerializer):
+    class Meta:
+        model = Kev
+        fields = ["date_added", "description", "required_action", "due_date", "resources_and_notes"]
+
+
 class VulnerabilitySerializer(BaseResourceSerializer):
     fixed_packages = MinimalPackageSerializer(
         many=True, source="filtered_fixed_packages", read_only=True
@@ -175,6 +182,7 @@ class VulnerabilitySerializer(BaseResourceSerializer):
 
     references = VulnerabilityReferenceSerializer(many=True, source="vulnerabilityreference_set")
     aliases = AliasSerializer(many=True, source="alias")
+    kev = KEVSerializer(read_only=True)
     weaknesses = WeaknessSerializer(many=True)
 
     def to_representation(self, instance):
@@ -183,6 +191,10 @@ def to_representation(self, instance):
         weaknesses = data.get("weaknesses", [])
         data["weaknesses"] = [weakness for weakness in weaknesses if weakness is not None]
 
+        kev = data.get("kev", None)
+        if not kev:
+            data.pop("kev")
+
         return data
 
     class Meta:
@@ -196,6 +208,7 @@ class Meta:
             "affected_packages",
             "references",
             "weaknesses",
+            "kev",
         ]
 
 
diff --git a/vulnerabilities/improvers/vulnerability_kev.py b/vulnerabilities/improvers/vulnerability_kev.py
index 8cd02d4ad..fff043e86 100644
--- a/vulnerabilities/improvers/vulnerability_kev.py
+++ b/vulnerabilities/improvers/vulnerability_kev.py
@@ -11,7 +11,6 @@
 from vulnerabilities.models import Alias
 from vulnerabilities.models import Kev
 from vulnerabilities.models import KnownRansomwareCampaignUseType
-from vulnerabilities.models import Vulnerability
 
 logger = logging.getLogger(__name__)
 
@@ -43,7 +42,7 @@ def get_inferences(self, advisory_data) -> Iterable[Inference]:
             if not alias:
                 continue
 
-            vul = Vulnerability.objects.get_or_none(aliases__alias=alias)
+            vul = alias.vulnerability
 
             if not vul:
                 continue