Migrate vulnerability sync to the pipeline

Signed-off-by: Keshav Priyadarshi <git@keshav.space>

Author: keshav-space

fedcode/importer.py fedcode/pipelines/sync_vulnerablecode.pyfedcode/importer.py renamed to fedcode/pipelines/sync_vulnerablecode.py

@@ -9,83 +9,104 @@
 
 import logging
 import os.path
-from dataclasses import dataclass
 from itertools import zip_longest
 
 import saneyaml
+from aboutcode.pipeline import LoopProgress
 
 from fedcode.activitypub import Activity
 from fedcode.activitypub import UpdateActivity
 from fedcode.models import Note
 from fedcode.models import Package
 from fedcode.models import Repository
-from fedcode.models import Service
 from fedcode.models import Vulnerability
+from fedcode.pipelines import FederatedCodePipeline
 from fedcode.pipes import utils
 
-logger = logging.getLogger(__name__)
 
+class SyncVulnerableCode(FederatedCodePipeline):
+    """Sync VulnerableCode data from FederatedCode git repositories."""
 
-@dataclass
-class Importer:
-    repo_obj: Repository
-    default_service: Service
+    pipeline_id = "sync_vulnerablecode"
 
-    def run(self):
-        repo = self.repo_obj.git_repo_obj
-        latest_commit_hash = repo.head.commit.hexsha
-        latest_commit = repo.commit(latest_commit_hash)
-        if self.repo_obj.last_imported_commit:
-            last_imported_commit = repo.commit(self.repo_obj.last_imported_commit)
-            diffs = last_imported_commit.diff(latest_commit)
-        else:
-            last_imported_commit = None
-            # Diff between empty trees and last_imported_commit
-            diffs = latest_commit.diff("4b825dc642cb6eb9a060e54bf8d69288fbee4904", R=True)
-
-        if repo.head.commit.hexsha == self.repo_obj.last_imported_commit:
-            logger.error("Nothing to import!")
-            return
+    @classmethod
+    def steps(cls):
+        return (
+            cls.get_git_repos,
+            cls.sync_vulnerablecode_repositories,
+        )
 
-        for diff in diffs:
-            if not diff.a_path.endswith(".yaml"):
-                continue
+    def get_git_repos(self):
+        self.git_repos = Repository.objects.all()
 
-            if diff.a_path.startswith("."):
-                continue
+    def sync_vulnerablecode_repositories(self):
+        repositories_count = self.git_repos.count()
+        self.log(f"Syncing vulnerability from {repositories_count:,d} repositories")
 
-            yaml_data_a_blob = (
-                saneyaml.load(diff.a_blob.data_stream.read()) if diff.a_blob else None
-            )
-            yaml_data_b_blob = (
-                saneyaml.load(diff.b_blob.data_stream.read()) if diff.b_blob else None
+        progress = LoopProgress(total_iterations=repositories_count, logger=self.log)
+        for repository in progress.iter(self.git_repos.iterator(chunk_size=2000)):
+            repository.git_repo_obj.remotes.origin.pull()
+            sync_vulnerabilities(
+                repository=repository,
+                logger=self.log,
             )
 
-            if os.path.split(diff.a_path)[1].startswith("VCID") or os.path.split(diff.b_path)[
-                1
-            ].startswith("VCID"):
-                vul_handler(
-                    diff.change_type,
-                    self.repo_obj,
-                    yaml_data_a_blob,
-                    yaml_data_b_blob,
-                    diff.a_path,
-                    diff.b_path,
-                )
-                continue
 
-            pkg_handler(
+def sync_vulnerabilities(repository, logger):
+    repo = repository.git_repo_obj
+    latest_commit_hash = repo.head.commit.hexsha
+    latest_commit = repo.commit(latest_commit_hash)
+    if repository.last_imported_commit:
+        last_imported_commit = repo.commit(repository.last_imported_commit)
+        diffs = last_imported_commit.diff(latest_commit)
+    else:
+        last_imported_commit = None
+        # Diff between empty trees and last_imported_commit
+        diffs = latest_commit.diff("4b825dc642cb6eb9a060e54bf8d69288fbee4904", R=True)
+
+    if repo.head.commit.hexsha == repository.last_imported_commit:
+        logger("Nothing to import!", level=logging.ERROR)
+        return
+
+    diff_count = len(diffs)
+
+    logger(f"Syncing {diff_count:,d} vulnerability scan from {repository.url}")
+    progress = LoopProgress(total_iterations=diff_count, logger=logger)
+    for diff in progress.iter(diffs):
+        if not diff.a_path.endswith(".yaml"):
+            continue
+
+        if diff.a_path.startswith("."):
+            continue
+
+        yaml_data_a_blob = saneyaml.load(diff.a_blob.data_stream.read()) if diff.a_blob else None
+        yaml_data_b_blob = saneyaml.load(diff.b_blob.data_stream.read()) if diff.b_blob else None
+
+        if os.path.split(diff.a_path)[1].startswith("VCID") or os.path.split(diff.b_path)[
+            1
+        ].startswith("VCID"):
+            vul_handler(
                 diff.change_type,
-                self.default_service,
+                repository,
                 yaml_data_a_blob,
                 yaml_data_b_blob,
+                logger,
             )
-        self.repo_obj.last_imported_commit = latest_commit_hash
-        self.repo_obj.save()
-        logger.info("The Importer run successfully")
+            continue
+
+        pkg_handler(
+            diff.change_type,
+            repository.admin,
+            yaml_data_a_blob,
+            yaml_data_b_blob,
+        )
+        break
+    repository.last_imported_commit = latest_commit_hash
+    repository.save()
+    logger("The Importer run successfully")
 
 
-def vul_handler(change_type, repo_obj, yaml_data_a_blob, yaml_data_b_blob, a_path, b_path):
+def vul_handler(change_type, repo_obj, yaml_data_a_blob, yaml_data_b_blob, logger):
     if change_type == "A":  # A for added paths
         Vulnerability.objects.get_or_create(
             id=yaml_data_b_blob.get("vulnerability_id"),
@@ -108,7 +129,7 @@ def vul_handler(change_type, repo_obj, yaml_data_a_blob, yaml_data_b_blob, a_pat
         )
         vul.delete()
     else:
-        logger.error(f"Invalid Vulnerability File")
+        logger(f"Invalid Vulnerability File", level=logging.ERROR)
 
 
 def pkg_handler(change_type, default_service, yaml_data_a_blob, yaml_data_b_blob):