Remove functions from API and add in model properties

Signed-off-by: Tushar Goel <[email protected]>

Author: TG1999

vulnerabilities/api.py

@@ -222,19 +222,20 @@ class PackageSerializer(BaseResourceSerializer):
     Lookup software package using Package URLs
     """
 
-    next_non_vulnerable_version = serializers.SerializerMethodField("get_next_non_vulnerable")
+    next_non_vulnerable_version = serializers.CharField(read_only=True)
+    latest_non_vulnerable_version = serializers.CharField(read_only=True)
 
-    def get_next_non_vulnerable(self, package):
-        next_non_vulnerable = package.fixed_package_details.get("next_non_vulnerable", None)
-        if next_non_vulnerable:
-            return next_non_vulnerable.version
+    # def get_next_non_vulnerable(self, package):
+    #     next_non_vulnerable = package.fixed_package_details.get("next_non_vulnerable", None)
+    #     if next_non_vulnerable:
+    #         return next_non_vulnerable.version
 
-    latest_non_vulnerable_version = serializers.SerializerMethodField("get_latest_non_vulnerable")
+    # latest_non_vulnerable_version = serializers.SerializerMethodField("get_latest_non_vulnerable")
 
-    def get_latest_non_vulnerable(self, package):
-        latest_non_vulnerable = package.fixed_package_details.get("latest_non_vulnerable", None)
-        if latest_non_vulnerable:
-            return latest_non_vulnerable.version
+    # def get_latest_non_vulnerable(self, package):
+    #     latest_non_vulnerable = package.fixed_package_details.get("latest_non_vulnerable", None)
+    #     if latest_non_vulnerable:
+    #         return latest_non_vulnerable.version
 
     purl = serializers.CharField(source="package_url")
 

vulnerabilities/models.py

@@ -710,6 +710,47 @@ def version_class(self):
     @cached_property
     def current_version(self):
         return self.version_class(self.version)
+    
+    @property
+    def next_non_vulnerable_version(self):
+        """
+        Return the version string of the next non-vulnerable package version.
+        """
+        next_non_vulnerable, _ = self.get_non_vulnerable_versions()
+        return next_non_vulnerable.version if next_non_vulnerable else None
+
+    @property
+    def latest_non_vulnerable_version(self):
+        """
+        Return the version string of the latest non-vulnerable package version.
+        """
+        _, latest_non_vulnerable = self.get_non_vulnerable_versions()
+        return latest_non_vulnerable.version if latest_non_vulnerable else None
+
+    def get_non_vulnerable_versions(self):
+        """
+        Return a tuple of the next and latest non-vulnerable versions as PackageURL objects.
+        Return a tuple of (None, None) if there is no non-vulnerable version.
+        """
+        non_vulnerable_versions = Package.objects.get_fixed_by_package_versions(
+            self, fix=False
+        ).only_non_vulnerable()
+        sorted_versions = self.sort_by_version(non_vulnerable_versions)
+
+        later_non_vulnerable_versions = [
+            non_vuln_ver
+            for non_vuln_ver in sorted_versions
+            if self.version_class(non_vuln_ver.version) > self.current_version
+        ]
+
+        if later_non_vulnerable_versions:
+            sorted_versions = self.sort_by_version(later_non_vulnerable_versions)
+            next_non_vulnerable_version = sorted_versions[0]
+            latest_non_vulnerable_version = sorted_versions[-1]
+
+            return next_non_vulnerable_version, latest_non_vulnerable_version
+
+        return None, None
 
     @property
     def fixed_package_details(self):

vulnerabilities/tests/test_api.py

@@ -454,22 +454,22 @@ def test_api_packages_all_num_queries(self):
             ]
 
     def test_api_packages_single_num_queries(self):
-        with self.assertNumQueries(10):
+        with self.assertNumQueries(8):
             self.csrf_client.get(f"/api/packages/{self.pkg_2_14_0_rc1.id}", format="json")
 
     def test_api_packages_single_with_purl_in_query_num_queries(self):
-        with self.assertNumQueries(11):
+        with self.assertNumQueries(9):
             self.csrf_client.get(f"/api/packages/?purl={self.pkg_2_14_0_rc1.purl}", format="json")
 
     def test_api_packages_single_with_purl_no_version_in_query_num_queries(self):
-        with self.assertNumQueries(88):
+        with self.assertNumQueries(68):
             self.csrf_client.get(
                 f"/api/packages/?purl=pkg:maven/com.fasterxml.jackson.core/jackson-databind",
                 format="json",
             )
 
     def test_api_packages_bulk_search(self):
-        with self.assertNumQueries(63):
+        with self.assertNumQueries(49):
             packages = [self.pkg_2_12_6, self.pkg_2_12_6_1, self.pkg_2_13_1]
             purls = [p.purl for p in packages]
 
@@ -482,7 +482,7 @@ def test_api_packages_bulk_search(self):
             ).json()
 
     def test_api_packages_with_lookup(self):
-        with self.assertNumQueries(18):
+        with self.assertNumQueries(16):
             data = {"purl": self.pkg_2_12_6.purl}
 
             resp = self.csrf_client.post(
@@ -492,7 +492,7 @@ def test_api_packages_with_lookup(self):
             ).json()
 
     def test_api_packages_bulk_lookup(self):
-        with self.assertNumQueries(63):
+        with self.assertNumQueries(49):
             packages = [self.pkg_2_12_6, self.pkg_2_12_6_1, self.pkg_2_13_1]
             purls = [p.purl for p in packages]