Refactor export command

- Improve memroy usage of main querysets
- Do not leak internal ids in serialized data
- Work towards reusing serializers

Signed-off-by: Philippe Ombredanne <[email protected]>

Author: pombredanne

vulnerabilities/management/commands/export.py

@@ -7,127 +7,175 @@
 # See https://aboutcode.org for more information about nexB OSS projects.
 #
 import logging
-import os
-from hashlib import sha512
+from itertools import groupby
 from pathlib import Path
 
 import saneyaml
 from django.core.management.base import BaseCommand
 from django.core.management.base import CommandError
 from packageurl import PackageURL
 
+from aboutcode import hashid
 from vulnerabilities.models import Package
 
 logger = logging.getLogger(__name__)
 
 
+def serialize_severity(sev):
+    # inlines refs
+    ref = sev.reference
+    sevref = {
+        "url": ref.url,
+        "reference_type": ref.reference_type,
+        "reference_id": ref.reference_id,
+    }
+
+    return {
+        "score": sev.value,
+        "scoring_system": sev.scoring_system,
+        "scoring_elements": sev.scoring_elements,
+        "published_at": sev.published_at,
+        "reference": sevref,
+    }
+
+
+def serialize_vulnerability(vuln):
+    """
+    Return a plain data mapping seralized from ``vuln`` Vulnerability instance.
+    """
+    aliases = list(vuln.aliases.values_list("alias", flat=True))
+    severities = [serialize_severity(sev) for sev in vuln.severities]
+    weaknesses = [wkns.cwe for wkns in vuln.weaknesses.all()]
+
+    references = list(
+        vuln.references.values(
+            "url",
+            "reference_type",
+            "reference_id",
+        )
+    )
+
+    return {
+        "vulnerability_id": vuln.vcid,
+        "aliases": aliases,
+        "summary": vuln.summary,
+        "severities": severities,
+        "weaknesses": weaknesses,
+        "references": references,
+    }
+
+
 class Command(BaseCommand):
-    help = "export vulnerablecode data"
+    help = """Export vulnerability and package data as YAML for use in FederatedCode
+
+    This command exports the data in a tree of directories and YAML files designed such that
+    it is possible to access directly a vulnerability data file by only knowing its VCID, and that
+    it is possible to access directly the package data files by only knowing its PURL.
+    """
 
     def add_arguments(self, parser):
-        parser.add_argument("path")
+        parser.add_argument(
+            "path",
+            help="Path to a directory where to export data.",
+        )
 
     def handle(self, *args, **options):
-        if options["path"]:
-            git_path = Path(options["path"])
-            if not git_path.is_dir():
-                raise CommandError("Please enter a valid path")
+        if path := options["path"]:
+            base_path = Path(path)
 
-            self.export_data(git_path)
+        if not path or not base_path.is_dir():
+            raise CommandError("Enter a valid directory path")
 
-        self.stdout.write(self.style.SUCCESS("Successfully exported vulnerablecode data"))
+        self.stdout.write("Exporting vulnerablecode Package and Vulnerability data.")
+        self.export_data(base_path)
+        self.stdout.write(self.style.SUCCESS(f"Successfully exported data to {base_path}."))
 
-    def export_data(self, git_path):
+    def export_data(self, base_path: Path):
         """
-        export vulnerablecode data
-        by running `python manage.py export /path/vulnerablecode-data`
+        Export vulnerablecode data to ``base_path``.`
         """
-        self.stdout.write("Exporting vulnerablecode data")
-
-        ecosystems = [pkg.type for pkg in Package.objects.distinct("type")]
-
-        for ecosystem in ecosystems:
-            package_files = {}  # {"package path": "data" }
-            vul_files = {}  # {"vulnerability path": "data" }
-
-            for purl in (
-                Package.objects.filter(type=ecosystem)
-                .prefetch_related("vulnerabilities")
-                .paginated()
-            ):
-                purl_without_version = PackageURL(
-                    type=purl.type,
-                    namespace=purl.namespace,
-                    name=purl.name,
-                )
-
-                # ./aboutcode-packages-ed5/maven/org.apache.log4j/log4j-core/versions/vulnerabilities.yml
-                pkg_filepath = (
-                    f"./aboutcode-packages-{get_purl_hash(purl_without_version)}/{purl.type}/{purl.namespace}/{purl.name}"
-                    f"/versions/vulnerabilities.yml"
-                )
-
-                package_data = {
-                    "purl": str(purl),
-                    "affected_by_vulnerabilities": [
-                        vuln.vulnerability_id for vuln in purl.affected_by
-                    ],
-                    "fixing_vulnerabilities": [vuln.vulnerability_id for vuln in purl.fixing],
-                }
-
-                if pkg_filepath in package_files:
-                    package_files[pkg_filepath]["versions"].append(package_data)
-                else:
-                    package_files[pkg_filepath] = {
-                        "package": str(purl_without_version),
-                        "versions": [package_data],
+        i = 0
+        seen_vcid = set()
+
+        for i, (purl_without_version, package_versions) in enumerate(packages_by_type_ns_name(), 1):
+            pkg_version = None
+            try:
+                package_urls = []
+                package_vulnerabilities = []
+                for pkg_version in package_versions:
+                    purl = pkg_version.package_url
+                    package_urls.append(purl)
+                    package_data = {
+                        "purl": purl,
+                        "affected_by_vulnerabilities": list(
+                            pkg_version.affected_by.values_list("vulnerability_id", flat=True)
+                        ),
+                        "fixing_vulnerabilities": list(
+                            pkg_version.fixing.values_list("vulnerability_id", flat=True)
+                        ),
                     }
+                    package_vulnerabilities.append(package_data)
 
-                for vul in purl.vulnerabilities.all():
-                    vulnerability_id = vul.vulnerability_id
-                    # ./aboutcode-vulnerabilities-12/34/VCID-1223-3434-34343/VCID-1223-3434-34343.yml
-                    vul_filepath = (
-                        f"./aboutcode-vulnerabilities-{vulnerability_id[5:7]}/{vulnerability_id[10:12]}"
-                        f"/{vulnerability_id}/{vulnerability_id}.yml"
-                    )
-                    vul_files[vul_filepath] = {
-                        "vulnerability_id": vul.vulnerability_id,
-                        "aliases": [alias.alias for alias in vul.get_aliases],
-                        "summary": vul.summary,
-                        "severities": [severity for severity in vul.severities.values()],
-                        "references": [ref for ref in vul.references.values()],
-                        "weaknesses": [
-                            "CWE-" + str(weakness["cwe_id"]) for weakness in vul.weaknesses.values()
-                        ],
-                    }
+                    for vuln in pkg_version.vulnerabilities.all():
+                        vcid = vuln.vulnerability_id
+                        # do not write twice the same file
+                        if vcid in seen_vcid:
+                            continue
+
+                        seen_vcid.add(vcid)
+                        vulnerability = serialize_vulnerability(vuln)
+                        vpath = hashid.get_vcid_yml_file_path(vcid)
+                        write_file(base_path=base_path, file_path=vpath, data=vulnerability)
+                        if (lv := len(seen_vcid)) % 100 == 0:
+                            self.stdout.write(f"Processed {lv} vulnerabilities. Last VCID: {vcid}")
+
+                ppath = hashid.get_package_purls_yml_file_path(purl)
+                write_file(base_path=base_path, file_path=ppath, data=package_urls)
 
-            for items in [package_files, vul_files]:
-                for filepath, data in items.items():
-                    create_file(filepath, git_path, data)
+                pvpath = hashid.get_package_vulnerabilities_yml_file_path(purl)
+                write_file(base_path=base_path, file_path=pvpath, data=package_vulnerabilities)
 
-            self.stdout.write(f"Successfully exported {ecosystem} data")
+                if i % 100 == 0:
+                    self.stdout.write(f"Processed {i} package. Last PURL: {purl_without_version}")
 
+            except Exception as e:
+                raise Exception(f"Failed to process Package: {pkg_version}") from e
 
-def create_file(filepath, git_path, data):
+        self.stdout.write(f"Exported data for: {i} package and {len(seen_vcid)} vulnerabilities.")
+
+
+def by_purl_type_ns_name(package):
     """
-    Check if the directories exist if it doesn't exist create a new one then Create the file
-    ./aboutcode-vulnerabilities-12/34/VCID-1223-3434-34343/VCID-1223-3434-34343.yml
-    ./aboutcode-packages-ed5/maven/org.apache.log4j/log4j-core/versions/vulnerabilities.yml
-    ./aboutcode-packages-ed5/maven/org.apache.log4j/log4j-core/versions/1.2.3/vulnerabilities.yml
+    Key function to sort packages by type, namespace and name
     """
-    filepath = git_path.joinpath(filepath)
-    dirname = os.path.dirname(filepath)
-    os.makedirs(dirname, exist_ok=True)
-    data = saneyaml.dump(data)
-    with open(filepath, encoding="utf-8", mode="w") as f:
-        f.write(data)
+    return package.type, package.namespace, package.name
 
 
-def get_purl_hash(purl: PackageURL, length: int = 3) -> str:
+def packages_by_type_ns_name():
+    """
+    Return a two-level iterator over all Packages grouped-by package, ignoring version.
+    """
+    qs = (
+        Package.objects.order_by("type", "namespace", "name", "version")
+        .prefetch_related(
+            "vulnerabilities",
+            "vulnerabilities__references",
+            "vulnerabilities__weaknesses",
+            "vulnerabilities__references__vulnerabilityseverity_set",
+        )
+        .paginated()
+    )
+
+    for tp_ns_name, packages in groupby(qs, key=by_purl_type_ns_name):
+        yield PackageURL(*tp_ns_name), packages
+
+
+def write_file(base_path: Path, file_path: Path, data: dict):
     """
-    Return a short lower cased hash of a purl.
-    https://github.com/nexB/purldb/pull/235/files#diff-a1fd023bd42d73f56019d540f38be711255403547add15108540d70f9948dd40R154
+    Write the ``data`` as YAML to the ``file_path`` in the ``base_path`` root directory.
+    Create directories in the path as needed.
     """
-    purl_bytes = str(purl).encode("utf-8")
-    short_hash = sha512(purl_bytes).hexdigest()[:length]
-    return short_hash.lower()
+    write_to = base_path / file_path
+    write_to.parent.mkdir(parents=True, exist_ok=True)
+    with open(write_to, encoding="utf-8", mode="w") as f:
+        f.write(saneyaml.dump(data))

vulnerabilities/tests/test_data/export_command/aboutcode-packages-1ccd/generic/nginx/test/purls.yml

@@ -0,0 +1 @@
+- pkg:generic/nginx/test@2

vulnerabilities/tests/test_data/export_command/aboutcode-packages-1ccd/generic/nginx/test/vulnerabilities.yml

@@ -0,0 +1,4 @@
+- purl: pkg:generic/nginx/test@2
+  affected_by_vulnerabilities:
+    - VCID-pst6-b358-aaap
+  fixing_vulnerabilities: []

vulnerabilities/tests/test_data/export_command/aboutcode-vulnerabilities/ps/VCID-pst6-b358-aaap.yml

@@ -0,0 +1,19 @@
+vulnerability_id: VCID-pst6-b358-aaap
+aliases:
+  - CVE-xxx-xxx-xx
+summary: test-vuln
+severities:
+  - score: '7.0'
+    scoring_system: cvssv3_vector
+    scoring_elements: CVSS:3.0/AV:N/AC:L/PR:N/UI:N/S:U/C:N/I:N/A:H
+    published_at:
+    reference:
+      url: https://..
+      reference_type:
+      reference_id: fake
+weaknesses:
+  - CWE-15
+references:
+  - url: https://..
+    reference_type:
+    reference_id: fake