Fix ehancement pipelines

Signed-off-by: Tushar Goel <tushar.goel.dav@gmail.com>

Author: TG1999

vulnerabilities/pipelines/v2_improvers/enhance_with_exploitdb.py

@@ -22,6 +22,7 @@
 from vulnerabilities.models import AdvisoryReference
 from vulnerabilities.models import AdvisoryV2
 from vulnerabilities.pipelines import VulnerableCodePipeline
+from vulnerabilities.utils import relate_aliases_with_advisories
 
 
 class ExploitDBImproverPipeline(VulnerableCodePipeline):
@@ -83,17 +84,7 @@ def add_vulnerability_exploit(row, logger):
     if not aliases:
         return 0
 
-    for raw_alias in aliases:
-        try:
-            if alias := AdvisoryAlias.objects.get(alias=raw_alias):
-                for adv in alias.advisories.all():
-                    advisories.add(adv)
-            else:
-                advs = AdvisoryV2.objects.filter(advisory_id=raw_alias).latest_per_avid()
-                for adv in advs:
-                    advisories.add(adv)
-        except AdvisoryAlias.DoesNotExist:
-            continue
+    advisories = relate_aliases_with_advisories(aliases)
 
     if not advisories:
         logger(f"No advisory found for aliases {aliases}")

vulnerabilities/pipelines/v2_improvers/enhance_with_github_poc.py

@@ -17,6 +17,7 @@
 from vulnerabilities.models import AdvisoryPOC
 from vulnerabilities.models import AdvisoryV2
 from vulnerabilities.pipelines import VulnerableCodePipeline
+from vulnerabilities.utils import relate_aliases_with_advisories
 
 
 class GithubPocsImproverPipeline(VulnerableCodePipeline):
@@ -61,18 +62,7 @@ def collect_and_store_exploits(self):
 
             filename = file_path.stem.strip()
 
-            advisories = set()
-            try:
-                if alias := AdvisoryAlias.objects.get(alias=filename):
-                    for adv in alias.advisories.all():
-                        advisories.add(adv)
-                else:
-                    advs = AdvisoryV2.objects.filter(advisory_id=filename).latest_per_avid()
-                    for adv in advs:
-                        advisories.add(adv)
-            except AdvisoryAlias.DoesNotExist:
-                self.log(f"Advisory {filename} not found.")
-                continue
+            advisories = relate_aliases_with_advisories([filename])
 
             for advisory in advisories:
                 for exploit_data in exploits_data:

vulnerabilities/pipelines/v2_improvers/enhance_with_kev.py

@@ -17,6 +17,7 @@
 from vulnerabilities.models import AdvisoryExploit
 from vulnerabilities.models import AdvisoryV2
 from vulnerabilities.pipelines import VulnerableCodePipeline
+from vulnerabilities.utils import relate_aliases_with_advisories
 
 
 class VulnerabilityKevPipeline(VulnerableCodePipeline):
@@ -72,18 +73,7 @@ def add_vulnerability_exploit(kev_vul, logger):
     if not cve_id:
         return 0
 
-    advisories = set()
-    try:
-        if alias := AdvisoryAlias.objects.get(alias=cve_id):
-            for adv in alias.advisories.all():
-                advisories.add(adv)
-        else:
-            advs = AdvisoryV2.objects.filter(advisory_id=cve_id).latest_per_avid()
-            for adv in advs:
-                advisories.add(adv)
-    except AdvisoryAlias.DoesNotExist:
-        logger(f"No advisory found for aliases {cve_id}")
-        return 0
+    advisories = relate_aliases_with_advisories([cve_id])
 
     for advisory in advisories:
         AdvisoryExploit.objects.update_or_create(

vulnerabilities/pipelines/v2_improvers/enhance_with_metasploit.py

@@ -19,6 +19,7 @@
 from vulnerabilities.models import AdvisoryExploit
 from vulnerabilities.models import AdvisoryV2
 from vulnerabilities.pipelines import VulnerableCodePipeline
+from vulnerabilities.utils import relate_aliases_with_advisories
 
 
 class MetasploitImproverPipeline(VulnerableCodePipeline):
@@ -77,17 +78,9 @@ def add_advisory_exploit(record, logger):
     if not interesting_references:
         return 0
 
-    for ref in interesting_references:
-        try:
-            if alias := AdvisoryAlias.objects.get(alias=ref):
-                for adv in alias.advisories.all():
-                    advisories.add(adv)
-            else:
-                advs = AdvisoryV2.objects.filter(advisory_id=ref).latest_per_avid()
-                for adv in advs:
-                    advisories.add(adv)
-        except AdvisoryAlias.DoesNotExist:
-            continue
+    advisories = []
+
+    advisories = relate_aliases_with_advisories(interesting_references)
 
     if not advisories:
         logger(f"No advisories found for aliases {interesting_references}")

vulnerabilities/tests/test_utils.py

@@ -13,6 +13,7 @@
 from django.test import TestCase
 from fetchcode.package_versions import PackageVersion
 from packageurl import PackageURL
+import pytest
 from univers.version_constraint import VersionConstraint
 from univers.version_range import GemVersionRange
 from univers.version_range import VersionRange
@@ -24,12 +25,12 @@
 from vulnerabilities.importer import PackageCommitPatchData
 from vulnerabilities.importer import PatchData
 from vulnerabilities.importer import VulnerabilitySeverity
-from vulnerabilities.models import AdvisoryV2
+from vulnerabilities.models import AdvisoryAlias, AdvisoryV2
 from vulnerabilities.pipelines import insert_advisory_v2
 from vulnerabilities.references import XsaReferenceV2
 from vulnerabilities.references import ZbxReferenceV2
 from vulnerabilities.tests.pipelines import TestLogger
-from vulnerabilities.utils import AffectedPackage
+from vulnerabilities.utils import AffectedPackage, relate_aliases_with_advisories
 from vulnerabilities.utils import get_item
 from vulnerabilities.utils import get_severity_range
 from vulnerabilities.utils import nearest_patched_package
@@ -262,3 +263,156 @@ def test_content_id_from_adv_data_roundtrip_are_same(self):
         id_from_roundtrip = utils.compute_content_id_v2(adv_roundtrip)
 
         self.assertEqual(id_from_data, id_from_roundtrip)
+
+
+@pytest.mark.django_db
+def test_returns_advisories_from_alias():
+    advisory = AdvisoryV2.objects.create(
+        advisory_id="GHSA-123",
+        datasource_id="ds",
+        pipeline_id="ds_importer_v2",
+        unique_content_id="i3giu",
+        url="https://test.com",
+        avid="GHSA-123",
+        is_latest=True,
+        _all_impacts_unfurled_at="2025-01-01",
+    )
+
+    alias = AdvisoryAlias.objects.create(alias="CVE-2025-123")
+    alias.advisories.add(advisory)
+
+    result = relate_aliases_with_advisories(["CVE-2025-123"])
+
+    assert result == {advisory}
+
+
+@pytest.mark.django_db
+def test_returns_latest_advisory_from_advisory_id():
+    advisory = AdvisoryV2.objects.create(
+        advisory_id="GHSA-456",
+        avid="GHSA-456",
+        is_latest=True,
+        _all_impacts_unfurled_at="2025-01-01",
+        datasource_id="ds",
+        pipeline_id="ds_importer_v2",
+        unique_content_id="i3giu2w",
+        url="https://test.com",
+    )
+
+    result = relate_aliases_with_advisories(["GHSA-456"])
+
+    assert advisory in result
+    assert len(result) == 1
+
+
+@pytest.mark.django_db
+def test_handles_mixed_aliases_and_advisory_ids():
+    alias_adv = AdvisoryV2.objects.create(
+        advisory_id="GHSA-1",
+        avid="GHSA-1",
+        is_latest=True,
+        _all_impacts_unfurled_at="2025-01-01",
+        datasource_id="ds",
+        pipeline_id="ds_importer_v2",
+        unique_content_id="i3gswiu",
+        url="https://test.com",
+    )
+
+    advisory_id_adv = AdvisoryV2.objects.create(
+        advisory_id="GHSA-2",
+        avid="GHSA-2",
+        is_latest=True,
+        _all_impacts_unfurled_at="2025-01-01",
+        datasource_id="ds",
+        pipeline_id="ds_importer_v2",
+        unique_content_id="i3giu2ww",
+        url="https://test.com",
+    )
+
+    alias = AdvisoryAlias.objects.create(alias="CVE-2025-1")
+    alias.advisories.add(alias_adv)
+
+    result = relate_aliases_with_advisories(
+        ["CVE-2025-1", "GHSA-2"]
+    )
+
+    assert result == {alias_adv, advisory_id_adv}
+
+
+@pytest.mark.django_db
+def test_excludes_non_latest_alias_advisories():
+    latest = AdvisoryV2.objects.create(
+        advisory_id="GHSA-1",
+        avid="GHSA-1",
+        is_latest=True,
+        _all_impacts_unfurled_at="2025-01-01",
+        datasource_id="ds",
+        pipeline_id="ds_importer_v2",
+        unique_content_id="i3giu",
+        url="https://test.com",
+    )
+
+    old = AdvisoryV2.objects.create(
+        advisory_id="GHSA-1-old",
+        avid="GHSA-1",
+        is_latest=False,
+        _all_impacts_unfurled_at="2025-01-01",
+        datasource_id="ds",
+        pipeline_id="ds_importer_v2",
+        unique_content_id="i3giu2w",
+        url="https://test.com",
+    )
+
+    alias = AdvisoryAlias.objects.create(alias="CVE-2025-1")
+    alias.advisories.add(latest, old)
+
+    result = relate_aliases_with_advisories(["CVE-2025-1"])
+
+    assert result == {latest}
+
+
+@pytest.mark.django_db
+def test_excludes_unfurled_advisories():
+    advisory = AdvisoryV2.objects.create(
+        advisory_id="GHSA-1",
+        avid="GHSA-1",
+        is_latest=True,
+        _all_impacts_unfurled_at=None,
+        datasource_id="ds",
+        pipeline_id="ds_importer_v2",
+        unique_content_id="i3giu",
+        url="https://test.com",
+    )
+
+    alias = AdvisoryAlias.objects.create(alias="CVE-2025-1")
+    alias.advisories.add(advisory)
+
+    result = relate_aliases_with_advisories(["CVE-2025-1"])
+
+    assert result == set()
+
+
+@pytest.mark.django_db
+def test_deduplicates_results():
+    advisory = AdvisoryV2.objects.create(
+        advisory_id="GHSA-1",
+        avid="GHSA-1",
+        is_latest=True,
+        _all_impacts_unfurled_at="2025-01-01",
+        datasource_id="ds",
+        pipeline_id="ds_importer_v2",
+        unique_content_id="i3giu",
+        url="https://test.com",
+    )
+
+    alias1 = AdvisoryAlias.objects.create(alias="CVE-1")
+    alias2 = AdvisoryAlias.objects.create(alias="CVE-2")
+
+    alias1.advisories.add(advisory)
+    alias2.advisories.add(advisory)
+
+    result = relate_aliases_with_advisories(
+        ["CVE-1", "CVE-2"]
+    )
+
+    assert result == {advisory}

vulnerabilities/utils.py

@@ -1055,3 +1055,31 @@ def get_advisories_from_groups(groups, include_ssvc_trees=False):
     "gem",
     "conan",
 ]
+
+
+def relate_aliases_with_advisories(aliases):
+    from vulnerabilities.models import AdvisoryAlias
+    from vulnerabilities.models import AdvisoryV2
+
+    aliases = set(aliases)
+
+    found_aliases = AdvisoryAlias.objects.filter(alias__in=aliases).prefetch_related("advisories")
+
+    found_alias_values = set(found_aliases.values_list("alias", flat=True))
+
+    alias_advisories = AdvisoryV2.objects.filter(
+        aliases__alias__in=found_alias_values,
+        is_latest=True,
+        _all_impacts_unfurled_at__isnull=False,
+    ).distinct()
+
+    missing_aliases = aliases - found_alias_values
+
+    advisory_id_advisories = AdvisoryV2.objects.filter(
+        advisory_id__in=missing_aliases,
+        _all_impacts_unfurled_at__isnull=False,
+    ).latest_per_avid()
+
+    advisories = set(alias_advisories)
+    advisories.update(advisory_id_advisories)
+    return advisories