Clean imported data after import process

vulnerabilities/importer.py

@@ -24,6 +24,7 @@
 
 import pytz
 from dateutil import parser as dateparser
+from fetchcode.vcs import VCSResponse
 from fetchcode.vcs import fetch_via_vcs
 from license_expression import Licensing
 from packageurl import PackageURL
@@ -288,6 +289,10 @@ class InvalidSPDXLicense(Exception):
     pass
 
 
+class ForkError(Exception):
+    pass
+
+
 class Importer:
     """
     An Importer collects data from various upstreams and returns corresponding AdvisoryData objects
@@ -297,7 +302,7 @@ class Importer:
     spdx_license_expression = ""
     license_url = ""
     notice = ""
-    vcs_response = None
+    vcs_response: VCSResponse = None
 
     def __init__(self):
         if not self.spdx_license_expression:
@@ -324,47 +329,18 @@ def advisory_data(self) -> Iterable[AdvisoryData]:
         raise NotImplementedError
 
     def clone(self, repo_url):
+        """
+        Clone the repo at repo_url and return the VCSResponse object
+        """
         try:
             self.vcs_response = fetch_via_vcs(repo_url)
+            return self.vcs_response
         except Exception as e:
             msg = f"Failed to fetch {repo_url} via vcs: {e}"
             logger.error(msg)
             raise ForkError(msg) from e
 
 
-class ForkError(Exception):
-    pass
-
-
-class GitImporter(Importer):
-    def __init__(self, repo_url):
-        super().__init__()
-        self.repo_url = repo_url
-        self.vcs_response = None
-
-    def __enter__(self):
-        super().__enter__()
-        self.clone()
-        return self
-
-    def __exit__(self):
-        self.vcs_response.delete()
-
-    def clone(self):
-        try:
-            self.vcs_response = fetch_via_vcs(self.repo_url)
-        except Exception as e:
-            msg = f"Failed to fetch {self.repo_url} via vcs: {e}"
-            logger.error(msg)
-            raise ForkError(msg) from e
-
-    def advisory_data(self) -> Iterable[AdvisoryData]:
-        """
-        Return AdvisoryData objects corresponding to the data being imported
-        """
-        raise NotImplementedError
-
-
 # TODO: Needs rewrite
 class OvalImporter(Importer):
     """

vulnerabilities/importers/elixir_security.py

@@ -29,7 +29,7 @@ class ElixirSecurityImporter(Importer):
 
     def advisory_data(self) -> Set[AdvisoryData]:
         try:
-            self.clone(self.repo_url)
+            self.clone(repo_url=self.repo_url)
             path = Path(self.vcs_response.dest_dir)
             vuln = path / "packages"
             for file in vuln.glob("**/*.yml"):

vulnerabilities/importers/fireeye.py

@@ -13,15 +13,15 @@
 from typing import List
 
 from vulnerabilities.importer import AdvisoryData
-from vulnerabilities.importer import GitImporter
+from vulnerabilities.importer import Importer
 from vulnerabilities.importer import Reference
 from vulnerabilities.utils import build_description
 from vulnerabilities.utils import dedupe
 
 logger = logging.getLogger(__name__)
 
 
-class FireyeImporter(GitImporter):
+class FireyeImporter(Importer):
     spdx_license_expression = "CC-BY-SA-4.0 AND MIT"
     license_url = "https://github.com/mandiant/Vulnerability-Disclosures/blob/master/README.md"
     notice = """
@@ -30,23 +30,25 @@ class FireyeImporter(GitImporter):
     1. CC BY-SA 4.0 - For CVE related information not including source code (such as PoCs)
     2. MIT - For source code contained within provided CVE information
     """
-
-    def __init__(self):
-        super().__init__(repo_url="git+https://github.com/mandiant/Vulnerability-Disclosures")
+    repo_url = "git+https://github.com/mandiant/Vulnerability-Disclosures"
 
     def advisory_data(self) -> Iterable[AdvisoryData]:
-        self.clone()
-        files = filter(
-            lambda p: p.suffix in [".md", ".MD"], Path(self.vcs_response.dest_dir).glob("**/*")
-        )
-        for file in files:
-            if Path(file).stem == "README":
-                continue
-            try:
-                with open(file) as f:
-                    yield parse_advisory_data(f.read())
-            except UnicodeError:
-                logger.error(f"Invalid file {file}")
+        try:
+            self.clone(repo_url=self.repo_url)
+            files = filter(
+                lambda p: p.suffix in [".md", ".MD"], Path(self.vcs_response.dest_dir).glob("**/*")
+            )
+            for file in files:
+                if Path(file).stem == "README":
+                    continue
+                try:
+                    with open(file) as f:
+                        yield parse_advisory_data(f.read())
+                except UnicodeError:
+                    logger.error(f"Invalid file {file}")
+        finally:
+            if self.vcs_response:
+                self.vcs_response.delete()
 
 
 def parse_advisory_data(raw_data) -> AdvisoryData:

vulnerabilities/importers/gitlab.py

@@ -25,7 +25,7 @@
 
 from vulnerabilities.importer import AdvisoryData
 from vulnerabilities.importer import AffectedPackage
-from vulnerabilities.importer import GitImporter
+from vulnerabilities.importer import Importer
 from vulnerabilities.importer import Reference
 from vulnerabilities.utils import build_description
 
@@ -48,16 +48,14 @@
 GITLAB_SCHEME_BY_PURL_TYPE = {v: k for k, v in PURL_TYPE_BY_GITLAB_SCHEME.items()}
 
 
-class GitLabAPIImporter(GitImporter):
+class GitLabAPIImporter(Importer):
     spdx_license_expression = "MIT"
     license_url = "https://gitlab.com/gitlab-org/advisories-community/-/blob/main/LICENSE"
+    repo_url = "git+https://gitlab.com/gitlab-org/advisories-community/"
 
-    def __init__(self):
-        super().__init__(repo_url="git+https://gitlab.com/gitlab-org/advisories-community/")
-
-    def advisory_data(self, _keep_clone=True) -> Iterable[AdvisoryData]:
+    def advisory_data(self, _keep_clone=False) -> Iterable[AdvisoryData]:
         try:
-            self.clone()
+            self.clone(repo_url=self.repo_url)
             base_path = Path(self.vcs_response.dest_dir)
 
             for file_path in base_path.glob("**/*.yml"):

vulnerabilities/importers/istio.py

@@ -53,17 +53,21 @@ class IstioImporter(Importer):
     repo_url = "git+https://github.com/istio/istio.io/"
 
     def advisory_data(self) -> Set[AdvisoryData]:
-        self.clone(self.repo_url)
-        path = Path(self.vcs_response.dest_dir)
-        vuln = path / "content/en/news/security/"
-        for file in vuln.glob("**/*.md"):
-            # Istio website has files with name starting with underscore, these contain metadata
-            # required for rendering the website. We're not interested in these.
-            # See also https://github.com/nexB/vulnerablecode/issues/563
-            file = str(file)
-            if file.endswith("_index.md"):
-                continue
-            yield from self.process_file(file)
+        try:
+            self.clone(repo_url=self.repo_url)
+            path = Path(self.vcs_response.dest_dir)
+            vuln = path / "content/en/news/security/"
+            for file in vuln.glob("**/*.md"):
+                # Istio website has files with name starting with underscore, these contain metadata
+                # required for rendering the website. We're not interested in these.
+                # See also https://github.com/nexB/vulnerablecode/issues/563
+                file = str(file)
+                if file.endswith("_index.md"):
+                    continue
+                yield from self.process_file(file)
+        finally:
+            if self.vcs_response:
+                self.vcs_response.delete()
 
     def process_file(self, path):
 

vulnerabilities/importers/kaybee.py

@@ -10,13 +10,13 @@
 from packageurl import PackageURL
 
 from vulnerabilities.importer import AdvisoryData
-from vulnerabilities.importer import GitImporter
+from vulnerabilities.importer import Importer
 from vulnerabilities.importer import Reference
 from vulnerabilities.utils import load_yaml
 from vulnerabilities.utils import nearest_patched_package
 
 
-class KaybeeImporter(GitImporter):
+class KaybeeImporter(Importer):
     def __enter__(self):
         super(KaybeeImporter, self).__enter__()
         self._added_files, self._updated_files = self.file_changes(

vulnerabilities/importers/mozilla.py

@@ -39,7 +39,7 @@ class MozillaImporter(Importer):
 
     def advisory_data(self) -> Iterable[AdvisoryData]:
         try:
-            self.clone(self.repo_url)
+            self.clone(repo_url=self.repo_url)
             path = Path(self.vcs_response.dest_dir)
 
             vuln = path / "announce"

vulnerabilities/importers/npm.py

@@ -36,7 +36,7 @@ class NpmImporter(Importer):
 
     def advisory_data(self) -> Iterable[AdvisoryData]:
         try:
-            self.clone(self.repo_url)
+            self.clone(repo_url=self.repo_url)
             path = Path(self.vcs_response.dest_dir)
 
             vuln = path / "vuln"

vulnerabilities/importers/retiredotnet.py

@@ -30,7 +30,7 @@ class RetireDotnetImporter(Importer):
 
     def advisory_data(self) -> Iterable[AdvisoryData]:
         try:
-            self.clone(self.repo_url)
+            self.clone(repo_url=self.repo_url)
             path = Path(self.vcs_response.dest_dir)
 
             vuln = path / "Content"

vulnerabilities/importers/ruby.py

@@ -18,14 +18,14 @@
 from univers.versions import SemverVersion
 
 from vulnerabilities.importer import AdvisoryData
-from vulnerabilities.importer import GitImporter
+from vulnerabilities.importer import Importer
 from vulnerabilities.importer import Reference
 from vulnerabilities.package_managers import RubyVersionAPI
 from vulnerabilities.utils import load_yaml
 from vulnerabilities.utils import nearest_patched_package
 
 
-class RubyImporter(GitImporter):
+class RubyImporter(Importer):
     def __enter__(self):
         super(RubyImporter, self).__enter__()
 

vulnerabilities/importers/rust.py

@@ -22,13 +22,13 @@
 from univers.versions import SemverVersion
 
 from vulnerabilities.importer import AdvisoryData
-from vulnerabilities.importer import GitImporter
+from vulnerabilities.importer import Importer
 from vulnerabilities.importer import Reference
 from vulnerabilities.package_managers import CratesVersionAPI
 from vulnerabilities.utils import nearest_patched_package
 
 
-class RustImporter(GitImporter):
+class RustImporter(Importer):
     def __enter__(self):
         super(RustImporter, self).__enter__()
 

vulnerabilities/tests/test_data_source.py

@@ -9,12 +9,22 @@
 
 import os
 import xml.etree.ElementTree as ET
+from unittest.mock import MagicMock
 from unittest.mock import patch
 
+from fetchcode.vcs import VCSResponse
 from packageurl import PackageURL
 
-from vulnerabilities.importer import GitImporter
+from vulnerabilities.importer import Importer
 from vulnerabilities.importer import OvalImporter
+from vulnerabilities.importers.elixir_security import ElixirSecurityImporter
+from vulnerabilities.importers.fireeye import FireyeImporter
+from vulnerabilities.importers.gentoo import GentooImporter
+from vulnerabilities.importers.gitlab import GitLabAPIImporter
+from vulnerabilities.importers.istio import IstioImporter
+from vulnerabilities.importers.mozilla import MozillaImporter
+from vulnerabilities.importers.npm import NpmImporter
+from vulnerabilities.importers.retiredotnet import RetireDotnetImporter
 from vulnerabilities.oval_parser import OvalParser
 
 BASE_DIR = os.path.dirname(os.path.abspath(__file__))
@@ -35,7 +45,7 @@ class MockOvalImporter(OvalImporter):
     spdx_license_expression = "FOO-BAR"
 
 
-class MockGitImporter(GitImporter):
+class MockGitImporter(Importer):
     spdx_license_expression = "FOO-BAR"
 
 
@@ -83,12 +93,42 @@ def test__collect_pkgs():
     assert found_ubuntu_pkgs == expected_ubuntu_pkgs
 
 
-def clone(self):
-    pass
+@patch("vulnerabilities.importer.fetch_via_vcs")
+def test_git_importer(mock_clone):
+    mock_clone.return_value = VCSResponse(
+        dest_dir="test",
+        vcs_type="git",
+        domain="test",
+    )
+    git_importer = MockGitImporter()
+    git_importer.clone("test-url") == VCSResponse(
+        dest_dir="test",
+        vcs_type="git",
+        domain="test",
+    )
 
 
-@patch("vulnerabilities.importer.GitImporter.clone")
-def test_git_importer(mock_clone):
-    mock_clone.return_value = clone
-    imp = MockGitImporter("test-url")
-    assert imp.repo_url == "test-url"
+def test_git_importer_clone():
+    git_importers = [
+        ElixirSecurityImporter,
+        FireyeImporter,
+        GentooImporter,
+        GitLabAPIImporter,
+        IstioImporter,
+        MozillaImporter,
+        NpmImporter,
+        RetireDotnetImporter,
+    ]
+    for git_importer in git_importers:
+        mock_function = MagicMock(
+            return_value=VCSResponse(
+                dest_dir="test",
+                vcs_type="git",
+                domain="test",
+            )
+        )
+        with patch("vulnerabilities.importer.fetch_via_vcs", mock_function) as mock_fetch:
+            with patch.object(VCSResponse, "delete") as mock_delete:
+                list(git_importer().advisory_data())
+                mock_fetch.assert_called_once()
+                mock_delete.assert_called_once()