Try to improve performance by adding pagination

Fix filename for export files
Add multiple parameterizes for create_sub_path test .
Add new format for exporting vulnerablecode-data
Add a test
Fix export test with yaml format
Change the export format from json to yaml
Add test for export command
Add test for write_vuln_data function
Edit export.py , Fix missing attribute in vuln_data
Export vulnerablecode-data
Add new format for exporting vulnerablecode-data
Add a test
Fix export test with yaml format
Change the export format from json to yaml
Add test for export command
Add test for write_vuln_data function
Edit export.py , Fix missing attribute in vuln_data
Export vulnerablecode-data
Add new format for exporting vulnerablecode-data
Add a test
Fix export test with yaml format
Change the export format from json to yaml
Add test for export command
Add test for write_vuln_data function
Edit export.py , Fix missing attribute in vuln_data
Export vulnerablecode-data

Signed-off-by: ziadhany <ziadhany2016@gmail.com>

Author: ziadhany

vulnerabilities/management/commands/export.py

@@ -0,0 +1,138 @@
+#
+# Copyright (c) nexB Inc. and others. All rights reserved.
+# VulnerableCode is a trademark of nexB Inc.
+# SPDX-License-Identifier: Apache-2.0
+# See http://www.apache.org/licenses/LICENSE-2.0 for the license text.
+# See https://github.com/nexB/vulnerablecode for support or download.
+# See https://aboutcode.org for more information about nexB OSS projects.
+#
+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 vulnerabilities.models import Package
+
+
+class Command(BaseCommand):
+    help = "export vulnerablecode data"
+
+    def add_arguments(self, parser):
+        parser.add_argument("path")
+
+    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")
+
+            self.export_data(git_path)
+
+        self.stdout.write(self.style.SUCCESS("Successfully exported vulnerablecode data"))
+
+    def export_data(self, git_path):
+        """
+        export vulnerablecode data
+        by run `python manage.py export /path/vulnerablecode-data`
+        """
+        self.stdout.write(f"Exporting vulnerablecode data")
+
+        for purl in Package.objects.all().paginated():
+            package_dir = create_sub_paths(git_path, purl.type, purl.namespace, purl.name)
+            pck_filepath = package_dir.joinpath(f"{purl.type}-{purl.name}.yaml")
+            write_package_data(pck_filepath, purl)
+
+            for vul in purl.vulnerabilities.all():
+                vul_filepath = package_dir.joinpath(f"{vul.vulnerability_id}.yaml")
+                write_vul_data(vul_filepath, vul)
+
+
+def write_vul_data(filepath, vul):
+    """
+     write the vulnerability data in a file with a path like this
+    `path/ecosystem/package/VCID-XXXX-XXXX-XXXX` and create the directories if it doesn't exist
+    """
+
+    vul_data = saneyaml.dump(
+        {
+            "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()
+            ],
+        }
+    )
+
+    with open(filepath, encoding="utf-8", mode="w") as f:
+        f.write(vul_data)
+
+
+def write_package_data(filepath, purl):
+    """
+     write the pacakge data in a file with a path like this
+    `path/ecosystem/package/package_type-package_name` and create the directories if it doesn't exist
+    and if the pacakge file doesn't exist we create the new header
+    """
+    purl_object = PackageURL.from_string(purl.package_url)
+    purl_without_version = PackageURL(
+        type=purl_object.type,
+        namespace=purl_object.namespace,
+        name=purl_object.name,
+    )
+
+    package_data_with_header = saneyaml.dump(
+        {
+            "pacakge": str(purl_without_version),
+            "versions": [
+                {
+                    "purl": str(purl_object),
+                    "affected_by_vulnerabilities": [
+                        vuln.vulnerability_id for vuln in purl.affected_by
+                    ],
+                    "fixing_vulnerabilities": [vuln.vulnerability_id for vuln in purl.fixing],
+                }
+            ],
+        }
+    )
+
+    package_data = {
+        "purl": str(purl_object),
+        "affected_by_vulnerabilities": [vuln.vulnerability_id for vuln in purl.affected_by],
+        "fixing_vulnerabilities": [vuln.vulnerability_id for vuln in purl.fixing],
+    }
+
+    if Path(filepath).is_file():
+        with open(filepath, "r") as f:
+            old_yaml = saneyaml.load(f)
+            old_yaml["versions"].append(package_data)
+
+        if old_yaml:
+            with open(filepath, "w") as f:
+                f.write(saneyaml.dump(old_yaml))
+    else:
+        with open(filepath, encoding="utf-8", mode="w") as f:
+            f.write(package_data_with_header)
+
+
+def create_sub_paths(git_path, purl_type, purl_namespace, purl_name):
+    """
+    create the directories if it doesn't exist : `path/purl_type/purl_namespace/purl_name`
+    """
+    ecosystem_dir = git_path.joinpath(purl_type)
+    if not ecosystem_dir.is_dir():
+        ecosystem_dir.mkdir()
+
+    namespace_dir = ecosystem_dir.joinpath(purl_namespace)
+    if not namespace_dir.is_dir():
+        namespace_dir.mkdir()
+
+    package_dir = namespace_dir.joinpath(purl_name)
+    if not package_dir.is_dir():
+        package_dir.mkdir()
+
+    return package_dir

vulnerabilities/tests/test_export.py

@@ -0,0 +1,113 @@
+import os
+from io import StringIO
+from pathlib import Path
+from unittest import TestCase
+
+import pytest
+import saneyaml
+from django.core.management import call_command
+from django.core.management.base import CommandError
+
+from vulnerabilities.management.commands.export import create_sub_paths
+from vulnerabilities.management.commands.export import write_package_data
+from vulnerabilities.management.commands.export import write_vul_data
+from vulnerabilities.models import Package
+from vulnerabilities.models import PackageRelatedVulnerability
+from vulnerabilities.models import Vulnerability
+
+
+@pytest.mark.parametrize(
+    "purl_type,purl_namespace,purl_name",
+    [
+        ("generic", "", "nginx"),
+        ("github", "package-url", "purl-spec"),
+        ("pypi", "", "django"),
+    ],
+)
+def test_create_sub_paths(tmp_path, purl_type, purl_namespace, purl_name):
+    create_sub_paths(tmp_path, purl_type, purl_namespace, purl_name)
+
+    ecosystem_dir = os.path.join(tmp_path, purl_type)
+    assert os.path.isdir(ecosystem_dir)
+
+    namespace_dir = os.path.join(ecosystem_dir, purl_namespace)
+    assert os.path.isdir(namespace_dir)
+
+    name_dir = os.path.join(namespace_dir, purl_name)
+    assert os.path.isdir(name_dir)
+
+
+@pytest.fixture
+def package(db):
+    return Package.objects.create(
+        type="generic", namespace="nginx", name="test", version="2", qualifiers={}, subpath=""
+    )
+
+
+@pytest.fixture
+def vulnerability(db):
+    return Vulnerability.objects.create(
+        vulnerability_id="VCID-pst6-b358-aaap",
+        summary="test-vuln",
+    )
+
+
+@pytest.fixture
+def package_related_vulnerability(db, package, vulnerability):
+    PackageRelatedVulnerability.objects.create(
+        package=package,
+        vulnerability=vulnerability,
+        fix=False,
+    )
+    return package
+
+
+@pytest.mark.django_db
+def test_write_vul_data(tmp_path, vulnerability):
+    expected_data = {
+        "vulnerability_id": "VCID-pst6-b358-aaap",
+        "aliases": [],
+        "summary": "test-vuln",
+        "severities": [],
+        "references": [],
+        "weaknesses": [],
+    }
+    pck_filepath = os.path.join(tmp_path, "filename")
+    write_vul_data(pck_filepath, vulnerability)
+    assert os.path.isfile(pck_filepath)
+    assert Path(pck_filepath).read_text() == saneyaml.dump(expected_data)
+
+
+@pytest.mark.django_db
+def test_write_package_data(tmp_path, package_related_vulnerability):
+    expected_data = {
+        "pacakge": "pkg:generic/nginx/test",
+        "versions": [
+            {
+                "purl": "pkg:generic/nginx/test@2",
+                "affected_by_vulnerabilities": ["VCID-pst6-b358-aaap"],
+                "fixing_vulnerabilities": [],
+            }
+        ],
+    }
+    pck_filepath = os.path.join(tmp_path, "filename")
+    write_package_data(pck_filepath, package_related_vulnerability)
+    assert os.path.isfile(pck_filepath)
+    assert Path(pck_filepath).read_text() == saneyaml.dump(expected_data)
+
+
+class TestExportCommand(TestCase):
+    def test_missing_path(self):
+        with pytest.raises(CommandError) as cm:
+            call_command("export", stdout=StringIO())
+
+        err = str(cm)
+        assert "Error: the following arguments are required: path" in err
+
+    def test_bad_path_fail_error(self):
+        buf = StringIO()
+        with pytest.raises(CommandError) as cm:
+            call_command("export", "/bad path", stdout=buf)
+
+        err = str(cm)
+        assert "Please enter a valid path" in err