Merge pull request #291 from sbs2001/improve_github_importer

Improve GitHub importer

Author: sbs2001

vulnerabilities/importer_yielder.py

@@ -179,7 +179,7 @@
         'data_source': 'GitHubAPIDataSource',
         'data_source_cfg': {
             'endpoint': 'https://api.github.com/graphql',
-            'ecosystems': ['MAVEN', 'NUGET', 'COMPOSER']
+            'ecosystems': ['MAVEN', 'NUGET', 'COMPOSER', 'PIP', 'RUBYGEMS']
         }
     },
     {

vulnerabilities/importers/github.py

@@ -42,6 +42,8 @@
 from vulnerabilities.package_managers import MavenVersionAPI
 from vulnerabilities.package_managers import NugetVersionAPI
 from vulnerabilities.package_managers import ComposerVersionAPI
+from vulnerabilities.package_managers import PypiVersionAPI
+from vulnerabilities.package_managers import RubyVersionAPI
 
 # set of all possible values of first '%s' = {'MAVEN','COMPOSER', 'NUGET'}
 # second '%s' is interesting, it will have the value '' for the first request,
@@ -132,6 +134,8 @@ def set_version_api(self, ecosystem: str) -> None:
             "MAVEN": MavenVersionAPI,
             "NUGET": NugetVersionAPI,
             "COMPOSER": ComposerVersionAPI,
+            "PIP": PypiVersionAPI,
+            "RUBYGEMS": RubyVersionAPI,
         }
         versioner = versioners.get(ecosystem)
         if versioner:
@@ -147,13 +151,13 @@ def process_name(ecosystem: str, pkg_name: str) -> Optional[Tuple[Optional[str],
             ns, name = artifact_comps
             return ns, name
 
-        if ecosystem == "NUGET":
-            return None, pkg_name
-
         if ecosystem == "COMPOSER":
             vendor, name = pkg_name.split("/")
             return vendor, name
 
+        if ecosystem == "NUGET" or ecosystem == "PIP" or ecosystem == "RUBYGEMS":
+            return None, pkg_name
+
     def collect_packages(self, ecosystem):
         packages = set()
         for page in self.advisories[ecosystem]:
@@ -165,30 +169,29 @@ def process_response(self) -> List[Advisory]:
         adv_list = []
         for ecosystem in self.advisories:
             self.set_version_api(ecosystem)
-            pkg_type = ecosystem.lower()
+            pkg_type = self.version_api.package_type
             for resp_page in self.advisories[ecosystem]:
                 for adv in resp_page["data"]["securityVulnerabilities"]["edges"]:
                     name = adv["node"]["package"]["name"]
 
                     if self.process_name(ecosystem, name):
                         ns, pkg_name = self.process_name(ecosystem, name)
+                        aff_range = adv["node"]["vulnerableVersionRange"]
+                        aff_vers, unaff_vers = self.categorize_versions(
+                            aff_range, self.version_api.get(name)
+                        )
+                        affected_purls = {
+                            PackageURL(name=pkg_name, namespace=ns, version=version, type=pkg_type)
+                            for version in aff_vers
+                        }
+
+                        unaffected_purls = {
+                            PackageURL(name=pkg_name, namespace=ns, version=version, type=pkg_type)
+                            for version in unaff_vers
+                        }
                     else:
-                        continue
-                    aff_range = adv["node"]["vulnerableVersionRange"]
-                    aff_vers, unaff_vers = self.categorize_versions(
-                        aff_range, self.version_api.get(name)
-                    )
-                    affected_purls = {
-                        PackageURL(name=pkg_name, namespace=ns,
-                                   version=version, type=pkg_type)
-                        for version in aff_vers
-                    }
-
-                    unaffected_purls = {
-                        PackageURL(name=pkg_name, namespace=ns,
-                                   version=version, type=pkg_type)
-                        for version in unaff_vers
-                    }
+                        affected_purls = set()
+                        unaffected_purls = set()
 
                     cve_ids = set()
                     vuln_references = []
@@ -199,12 +202,13 @@ def process_response(self) -> List[Advisory]:
                             cve_ids.add(vuln["value"])
 
                         elif vuln["type"] == "GHSA":
-                            ghsa = vuln['value']
-                            vuln_references.append(Reference(
-                                reference_id=ghsa,
-                                url="https://github.com/advisories/{}".format(
-                                    ghsa)
-                            ))
+                            ghsa = vuln["value"]
+                            vuln_references.append(
+                                Reference(
+                                    reference_id=ghsa,
+                                    url="https://github.com/advisories/{}".format(ghsa),
+                                )
+                            )
 
                     for cve_id in cve_ids:
                         adv_list.append(
@@ -219,8 +223,9 @@ def process_response(self) -> List[Advisory]:
         return adv_list
 
     @staticmethod
-    def categorize_versions(version_range: str, all_versions: Set[str]) -> Tuple[Set[str], Set[str]]:  # nopep8
+    def categorize_versions(
+        version_range: str, all_versions: Set[str]
+    ) -> Tuple[Set[str], Set[str]]:  # nopep8
         version_range = RangeSpecifier(version_range)
-        affected_versions = {
-            version for version in all_versions if version in version_range}
+        affected_versions = {version for version in all_versions if version in version_range}
         return (affected_versions, all_versions - affected_versions)

vulnerabilities/package_managers.py

@@ -41,6 +41,9 @@ def get(self, package_name: str) -> Set[str]:
 
 
 class LaunchpadVersionAPI(VersionAPI):
+
+    package_type = "deb"
+
     async def load_api(self, pkg_set):
         async with ClientSession(raise_for_status=True) as session:
             await asyncio.gather(
@@ -75,6 +78,9 @@ async def set_api(self, pkg, session):
 
 
 class PypiVersionAPI(VersionAPI):
+
+    package_type = "pypi"
+
     async def load_api(self, pkg_set):
         async with ClientSession(raise_for_status=True) as session:
             await asyncio.gather(
@@ -96,6 +102,9 @@ async def fetch(self, pkg, session):
 
 
 class CratesVersionAPI(VersionAPI):
+
+    package_type = "cargo"
+
     async def load_api(self, pkg_set):
         async with ClientSession(raise_for_status=True) as session:
             await asyncio.gather(
@@ -114,6 +123,9 @@ async def fetch(self, pkg, session):
 
 
 class RubyVersionAPI(VersionAPI):
+
+    package_type = "gem"
+
     async def load_api(self, pkg_set):
         async with ClientSession(raise_for_status=True) as session:
             await asyncio.gather(
@@ -135,6 +147,9 @@ async def fetch(self, pkg, session):
 
 
 class NpmVersionAPI(VersionAPI):
+
+    package_type = "npm"
+
     async def load_api(self, pkg_set):
         async with ClientSession(raise_for_status=True) as session:
             await asyncio.gather(
@@ -156,6 +171,9 @@ async def fetch(self, pkg, session):
 
 
 class DebianVersionAPI(VersionAPI):
+
+    package_type = "deb"
+
     async def load_api(self, pkg_set):
         # Need to set the headers, because the Debian API upgrades
         # the connection to HTTP 2.0
@@ -189,6 +207,9 @@ async def set_api(self, pkg, session, retry_count=5):
 
 
 class MavenVersionAPI(VersionAPI):
+
+    package_type = "maven"
+
     async def load_api(self, pkg_set):
         async with ClientSession(raise_for_status=True) as session:
             await asyncio.gather(
@@ -242,6 +263,9 @@ def extract_versions(xml_response: ET.ElementTree) -> Set[str]:
 
 
 class NugetVersionAPI(VersionAPI):
+
+    package_type = "nuget"
+
     async def load_api(self, pkg_set):
         async with ClientSession(raise_for_status=True) as session:
             await asyncio.gather(
@@ -274,6 +298,9 @@ def extract_versions(resp: dict) -> Set[str]:
 
 
 class ComposerVersionAPI(VersionAPI):
+
+    package_type = "composer"
+
     async def load_api(self, pkg_set):
         async with ClientSession(raise_for_status=True) as session:
             await asyncio.gather(
@@ -304,6 +331,9 @@ def extract_versions(resp: dict, pkg_name: str) -> Set[str]:
 
 
 class GitHubTagsAPI(VersionAPI):
+
+    package_type = "github"
+
     async def load_api(self, repo_set):
         async with ClientSession(raise_for_status=True) as session:
             await asyncio.gather(

vulnerabilities/tests/test_github.py

@@ -303,6 +303,7 @@ def test_process_response(self):
         ]
 
         mock_version_api = MagicMock()
+        mock_version_api.package_type = "maven"
         mock_version_api.get = lambda x: {'1.2.0', '9.0.2'}
         with patch('vulnerabilities.importers.github.MavenVersionAPI', return_value=mock_version_api):  # nopep8
             with patch('vulnerabilities.importers.github.GitHubAPIDataSource.set_api'):