Refine the code and add unit tests

Signed-off-by: tdruez <[email protected]>

Author: tdruez

scanpipe/management/commands/check-compliance.py

@@ -82,19 +82,38 @@ def check_compliance(self, fail_level):
                     self.stderr.write(f" > {severity.upper()}: {len(entries)}")
                     if self.verbosity > 1:
                         self.stderr.write("   " + "\n   ".join(entries))
+
         return count > 0
 
     def check_vulnerabilities(self):
-        packages = self.project.discoveredpackages.vulnerable()
-        dependencies = self.project.discovereddependencies.vulnerable()
-        vulnerable_records = list(packages) + list(dependencies)
+        # TODO: Remove duplication with scanpipe.pipes.output.add_vulnerabilities_sheet
+        vulnerable_packages_queryset = (
+            self.project.discoveredpackages.vulnerable()
+            .only_package_url_fields(extra=["affected_by_vulnerabilities"])
+            .order_by_package_url()
+        )
+        vulnerable_dependencies_queryset = (
+            self.project.discovereddependencies.vulnerable()
+            .only_package_url_fields(extra=["affected_by_vulnerabilities"])
+            .order_by_package_url()
+        )
+
+        vulnerable_records = list(vulnerable_packages_queryset) + list(
+            vulnerable_dependencies_queryset
+        )
         count = len(vulnerable_records)
 
         if self.verbosity > 0:
             if count:
                 self.stderr.write(f"{count} vulnerable records found:")
                 for entry in vulnerable_records:
                     self.stderr.write(str(entry))
+                    vulnerability_ids = [
+                        vulnerability.get("vulnerability_id")
+                        for vulnerability in entry.affected_by_vulnerabilities
+                    ]
+                    self.stderr.write(" > " + ", ".join(vulnerability_ids))
             else:
                 self.stdout.write("No vulnerabilities found")
+
         return count > 0

scanpipe/tests/test_commands.py

@@ -1212,6 +1212,32 @@ def test_scanpipe_management_command_check_compliance(self):
         )
         self.assertEqual(expected, out_value)
 
+    def test_scanpipe_management_command_check_compliance_vulnerabilities(self):
+        project = make_project(name="my_project")
+        package1 = make_package(project, package_url="pkg:generic/[email protected]")
+
+        out = StringIO()
+        options = ["--project", project.name, "--fail-on-vulnerabilities"]
+        with self.assertRaises(SystemExit) as cm:
+            call_command("check-compliance", *options, stdout=out)
+        self.assertEqual(cm.exception.code, 0)
+        out_value = out.getvalue().strip()
+        self.assertEqual("No vulnerabilities found", out_value)
+
+        package1.update(
+            affected_by_vulnerabilities=[{"vulnerability_id": "VCID-cah8-awtr-aaad"}]
+        )
+        out = StringIO()
+        options = ["--project", project.name, "--fail-on-vulnerabilities"]
+        with self.assertRaises(SystemExit) as cm:
+            call_command("check-compliance", *options, stderr=out)
+        self.assertEqual(cm.exception.code, 1)
+        out_value = out.getvalue().strip()
+        expected = (
+            "1 vulnerable records found:\npkg:generic/[email protected]\n > VCID-cah8-awtr-aaad"
+        )
+        self.assertEqual(expected, out_value)
+
     def test_scanpipe_management_command_report(self):
         label1 = "label1"
         project1 = make_project("project1", labels=[label1])