Avoid duplication of package type inference in github importer

Signed-off-by: Shivam Sandbhor <shivam.sandbhor@gmail.com>

Author: sbs2001

vulnerabilities/importers/github.py

@@ -97,14 +97,6 @@ def __init__(self, *args, **kwargs):
         except KeyError:
             raise GitHubTokenError("Environment variable GH_TOKEN is missing")
 
-        self.ecosytem_type = {
-            "RUBYGEMS": "gem",
-            "NUGET": "nuget",
-            "PIP": "pypi",
-            "MAVEN": "maven",
-            "COMPOSER": "composer",
-        }
-
     def __enter__(self):
         self.advisories = self.fetch()
 
@@ -177,7 +169,7 @@ def process_response(self) -> List[Advisory]:
         adv_list = []
         for ecosystem in self.advisories:
             self.set_version_api(ecosystem)
-            pkg_type = self.ecosytem_type[ecosystem]
+            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"]

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'):