Extract API logic to utils from Gitlab datasource #1903

* Separate Gitlab API handling logic from vulntotal Gitlab datasource to utils file

Signed-off-by: Michael Ehab Mikhail <[email protected]>

Author: michaelehab

vulntotal/datasources/gitlab.py

@@ -19,6 +19,8 @@
 from fetchcode import fetch
 from packageurl import PackageURL
 
+from vulntotal.datasources.gitlab_api import fetch_gitlab_advisories_for_purl
+from vulntotal.datasources.gitlab_api import fetch_yaml
 from vulntotal.validator import DataSource
 from vulntotal.validator import VendorData
 from vulntotal.vulntotal_utils import gitlab_constraints_satisfied
@@ -40,18 +42,12 @@ def datasource_advisory(self, purl) -> Iterable[VendorData]:
         Yields:
             VendorData instance containing the advisory information for the package.
         """
-        package_slug = get_package_slug(purl)
-        directory_files = fetch_directory_contents(package_slug)
-        if not directory_files:
-            path = self.supported_ecosystem()[purl.type]
-            casesensitive_package_slug = get_casesensitive_slug(path, package_slug)
-            directory_files = fetch_directory_contents(casesensitive_package_slug)
+        advisories = fetch_gitlab_advisories_for_purl(
+            purl, self.supported_ecosystem(), get_casesensitive_slug
+        )
 
-        if directory_files:
-            yml_files = [file for file in directory_files if file["name"].endswith(".yml")]
-
-            interesting_advisories = parse_interesting_advisories(yml_files, purl)
-            return interesting_advisories
+        if advisories:
+            return parse_interesting_advisories(advisories, purl)
 
     @classmethod
     def supported_ecosystem(cls):
@@ -67,21 +63,6 @@ def supported_ecosystem(cls):
         }
 
 
-def fetch_directory_contents(package_slug):
-    url = f"https://gitlab.com/api/v4/projects/12006272/repository/tree?path={package_slug}"
-    response = requests.get(url)
-    if response.status_code == 200:
-        return response.json()
-
-
-def fetch_yaml(file_path):
-    response = requests.get(
-        f"https://gitlab.com/gitlab-org/security-products/gemnasium-db/-/raw/master/{file_path}"
-    )
-    if response.status_code == 200:
-        return response.text
-
-
 def get_package_slug(purl):
     """
     Constructs a package slug from a given purl.
@@ -163,27 +144,25 @@ def get_casesensitive_slug(path, package_slug):
         has_next = paginated_tree["pageInfo"]["hasNextPage"]
 
 
-def parse_interesting_advisories(yml_files, purl) -> Iterable[VendorData]:
+def parse_interesting_advisories(advisories, purl) -> Iterable[VendorData]:
     """
     Parses advisories from YAML files in a given location that match a given version.
 
     Parameters:
-        yml_files: An array having the paths of yml files to parse.
+        advisories: A list of advisory dictionaries fetched from the GitLab API.
         purl: PURL for the advisory.
 
     Yields:
         VendorData instance containing the advisory information for the package.
     """
     version = purl.version
 
-    for file in yml_files:
-        yml_data = fetch_yaml(file["path"])
-        gitlab_advisory = saneyaml.load(yml_data)
-        affected_range = gitlab_advisory["affected_range"]
+    for advisory in advisories:
+        affected_range = advisory.get("affected_range")
         if gitlab_constraints_satisfied(affected_range, version):
             yield VendorData(
                 purl=PackageURL(purl.type, purl.namespace, purl.name),
-                aliases=gitlab_advisory["identifiers"],
+                aliases=advisory.get("identifiers", []),
                 affected_versions=[affected_range],
-                fixed_versions=gitlab_advisory["fixed_versions"],
+                fixed_versions=advisory.get("fixed_versions", []),
             )

vulntotal/datasources/gitlab_api.py

@@ -0,0 +1,56 @@
+import requests
+import saneyaml
+
+
+def fetch_directory_contents(package_slug):
+    url = f"https://gitlab.com/api/v4/projects/12006272/repository/tree?path={package_slug}"
+    response = requests.get(url)
+    if response.status_code == 200:
+        return response.json()
+    return []
+
+
+def fetch_yaml(file_path):
+    response = requests.get(
+        f"https://gitlab.com/gitlab-org/security-products/gemnasium-db/-/raw/master/{file_path}"
+    )
+    if response.status_code == 200:
+        return response.text
+    return None
+
+
+def get_package_slug(purl, supported_ecosystem):
+    if purl.type not in supported_ecosystem:
+        return
+    ecosystem = supported_ecosystem[purl.type]
+    package_name = purl.name
+    if purl.type in ("maven", "composer", "golang"):
+        package_name = f"{purl.namespace}/{purl.name}"
+    return f"{ecosystem}/{package_name}"
+
+
+def get_directory_yml_files(purl, supported_ecosystem, get_casesensitive_slug):
+    package_slug = get_package_slug(purl, supported_ecosystem)
+    directory_files = fetch_directory_contents(package_slug)
+    if not directory_files:
+        path = supported_ecosystem[purl.type]
+        casesensitive_package_slug = get_casesensitive_slug(path, package_slug)
+        directory_files = fetch_directory_contents(casesensitive_package_slug)
+    if not directory_files:
+        return []
+    return [file for file in directory_files if file["name"].endswith(".yml")]
+
+
+def fetch_gitlab_advisories_for_purl(purl, supported_ecosystem, get_casesensitive_slug):
+    yml_files = get_directory_yml_files(purl, supported_ecosystem, get_casesensitive_slug)
+
+    advisories = []
+    for file in yml_files:
+        yml_data = fetch_yaml(file["path"])
+        if yml_data:
+            advisories.append(saneyaml.load(yml_data))
+    return advisories
+
+
+def get_estimated_advisories_count(purl, supported_ecosystem, get_casesensitive_slug):
+    return len(get_directory_yml_files(purl, supported_ecosystem, get_casesensitive_slug))

vulntotal/tests/test_data/gitlab/parsed_advisory-expected.json

@@ -1,51 +1,26 @@
 [
   {
-    "purl": "pkg:generic/namespace/test",
-    "affected_versions": [
-      "<=2.7.1"
-    ],
-    "fixed_versions": [
-      "2.7.2"
-    ],
-    "aliases": [
-      "CVE-2014-1402"
-    ]
+    "purl": "pkg:pypi/namespace/test",
+    "affected_versions": ["<=2.7.1"],
+    "fixed_versions": ["2.7.2"],
+    "aliases": ["CVE-2014-1402"]
   },
   {
-    "purl": "pkg:generic/namespace/test",
-    "affected_versions": [
-      "<2.8.1"
-    ],
-    "fixed_versions": [
-      "2.8.1"
-    ],
-    "aliases": [
-      "GHSA-hj2j-77xm-mc5v",
-      "CVE-2016-10745"
-    ]
+    "purl": "pkg:pypi/namespace/test",
+    "affected_versions": ["<2.8.1"],
+    "fixed_versions": ["2.8.1"],
+    "aliases": ["GHSA-hj2j-77xm-mc5v", "CVE-2016-10745"]
   },
   {
-    "purl": "pkg:generic/namespace/test",
-    "affected_versions": [
-      "<2.10.1"
-    ],
-    "fixed_versions": [
-      "2.10.1"
-    ],
-    "aliases": [
-      "CVE-2019-10906"
-    ]
+    "purl": "pkg:pypi/namespace/test",
+    "affected_versions": ["<2.10.1"],
+    "fixed_versions": ["2.10.1"],
+    "aliases": ["CVE-2019-10906"]
   },
   {
-    "purl": "pkg:generic/namespace/test",
-    "affected_versions": [
-      "<2.11.3"
-    ],
-    "fixed_versions": [
-      "2.11.3"
-    ],
-    "aliases": [
-      "CVE-2020-28493"
-    ]
+    "purl": "pkg:pypi/namespace/test",
+    "affected_versions": ["<2.11.3"],
+    "fixed_versions": ["2.11.3"],
+    "aliases": ["CVE-2020-28493"]
   }
-]
+]

vulntotal/tests/test_gitlab.py

@@ -14,6 +14,7 @@
 
 from vulnerabilities.tests import util_tests
 from vulntotal.datasources import gitlab
+from vulntotal.datasources import gitlab_api
 
 
 class TestGitlab(testcase.FileBasedTesting):
@@ -28,12 +29,26 @@ def test_generate_package_advisory_url(self):
             "pkg:composer/bolt/[email protected]",
             "pkg:nuget/[email protected]",
         ]
-        results = [gitlab.get_package_slug(PackageURL.from_string(purl)) for purl in purls]
+        supported_ecosystem = gitlab.GitlabDataSource.supported_ecosystem()
+        results = [
+            gitlab_api.get_package_slug(PackageURL.from_string(purl), supported_ecosystem)
+            for purl in purls
+        ]
         expected_file = self.get_test_loc("package_advisory_url-expected.json", must_exist=False)
         util_tests.check_results_against_json(results, expected_file)
 
-    @mock.patch("vulntotal.datasources.gitlab.fetch_yaml")
-    def test_parse_interesting_advisories(self, mock_fetch_yaml):
+    @mock.patch("vulntotal.datasources.gitlab_api.fetch_yaml")
+    @mock.patch("vulntotal.datasources.gitlab_api.fetch_directory_contents")
+    def test_parse_interesting_advisories(self, mock_fetch_directory_contents, mock_fetch_yaml):
+        # Mock the directory contents response
+        mock_fetch_directory_contents.return_value = [
+            {"name": "CVE-2014-1402.yml", "path": "path/to/CVE-2014-1402.yml"},
+            {"name": "CVE-2016-10745.yml", "path": "path/to/CVE-2016-10745.yml"},
+            {"name": "CVE-2019-10906.yml", "path": "path/to/CVE-2019-10906.yml"},
+            {"name": "CVE-2019-8341.yml", "path": "path/to/CVE-2019-8341.yml"},
+            {"name": "CVE-2020-28493.yml", "path": "path/to/CVE-2020-28493.yml"},
+        ]
+
         # Mock the yaml file responses
         advisory_folder = (
             Path(__file__)
@@ -51,17 +66,15 @@ def test_parse_interesting_advisories(self, mock_fetch_yaml):
 
         mock_fetch_yaml.side_effect = yaml_files
 
-        purl = PackageURL("generic", "namespace", "test", "0.1.1")
+        purl = PackageURL("pypi", "namespace", "test", "0.1.1")
 
-        yml_files = [
-            {"name": "CVE-2014-1402.yml", "path": "path/to/CVE-2014-1402.yml"},
-            {"name": "CVE-2016-10745.yml", "path": "path/to/CVE-2016-10745.yml"},
-            {"name": "CVE-2019-10906.yml", "path": "path/to/CVE-2019-10906.yml"},
-            {"name": "CVE-2019-8341.yml", "path": "path/to/CVE-2019-8341.yml"},
-            {"name": "CVE-2020-28493.yml", "path": "path/to/CVE-2020-28493.yml"},
-        ]
+        supported_ecosystem = gitlab.GitlabDataSource.supported_ecosystem()
+
+        advisories = gitlab_api.fetch_gitlab_advisories_for_purl(
+            purl, supported_ecosystem, gitlab.get_casesensitive_slug
+        )
 
-        results = [adv.to_dict() for adv in gitlab.parse_interesting_advisories(yml_files, purl)]
+        results = [adv.to_dict() for adv in gitlab.parse_interesting_advisories(advisories, purl)]
 
         expected_file = self.get_test_loc("parsed_advisory-expected.json", must_exist=False)
         util_tests.check_results_against_json(results, expected_file)