Add new format for exporting vulnerablecode-data

Add a test

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

Author: ziadhany

vulnerabilities/management/commands/export.py

@@ -8,13 +8,12 @@
 #
 from pathlib import Path
 
-import yaml
+import saneyaml
 from django.core.management.base import BaseCommand
 from django.core.management.base import CommandError
 from packageurl import PackageURL
 
-from vulnerabilities.models import PackageRelatedVulnerability
-from vulnerabilities.models import Vulnerability
+from vulnerabilities.models import Package
 
 
 class Command(BaseCommand):
@@ -40,45 +39,100 @@ def export_data(self, git_path):
         """
         self.stdout.write(f"Exporting vulnerablecode data")
 
-        vulnerabilities = Vulnerability.objects.all()
-        for vuln in vulnerabilities.iterator():
-            vuln_data = {
-                "vulnerability_id": vuln.vulnerability_id,
-                "aliases": [alias.alias for alias in vuln.get_aliases],
-                "summary": vuln.summary,
-                "affected_purls": vuln.get_affected_purls(),
-                "fixed_purl": vuln.get_fixing_purls(),
-                "severities": [severity for severity in vuln.severities.values()],
-                "references": [ref for ref in vuln.references.values()],
-                "weaknesses": [
-                    "CWE-" + str(weakness["cwe_id"]) for weakness in vuln.weaknesses.values()
-                ],
-            }
-
-            relate_package = PackageRelatedVulnerability.objects.filter(
-                vulnerability_id=vuln.id
-            ).first()
-            if relate_package:
-                purl = PackageURL.from_string(relate_package.package.purl)
-                write_vuln_data(git_path, purl.type, purl.name, vuln.vulnerability_id, vuln_data)
-            else:
-                print(
-                    f"Can't create a vulnerability file without at least one purl related to it {vuln.id}"
-                )
-
-
-def write_vuln_data(git_path, ecosystem, package, vuln_id, vuln_data):
+        for purl in Package.objects.paginated():
+            package_dir = create_sub_paths(git_path, purl.type, purl.namespace, purl.name)
+            pck_filepath = package_dir.joinpath(f"{purl.type}-{purl.namespace}-{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 vulnerablecode-code data in a file with a path like this
+     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
     """
-    ecosystem_dir = git_path.joinpath(ecosystem)
+
+    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/ecosystem/package/`
+    """
+    ecosystem_dir = git_path.joinpath(purl_type)
     if not ecosystem_dir.is_dir():
         ecosystem_dir.mkdir()
 
-    package_dir = ecosystem_dir.joinpath(package)
+    package_dir = ecosystem_dir.joinpath(purl_namespace)
+    if not package_dir.is_dir():
+        package_dir.mkdir()
+
+    package_dir = ecosystem_dir.joinpath(purl_name)
     if not package_dir.is_dir():
         package_dir.mkdir()
 
-    with open(package_dir.joinpath(f"{vuln_id}.yaml"), mode="w", encoding="utf-8") as f:
-        f.write(yaml.dump(vuln_data))
+    return package_dir

vulnerabilities/tests/test_export.py

@@ -4,108 +4,94 @@
 from unittest import TestCase
 
 import pytest
-import yaml
+import saneyaml
 from django.core.management import call_command
 from django.core.management.base import CommandError
 
-from vulnerabilities.management.commands.export import write_vuln_data
+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(
-    "ecosystem,package,vuln_id,vuln_data",
+    "purl_type,purl_namespace,purl_name",
     [
-        (
-            "pypi",
-            "django",
-            "VCID-5gjv-fh72-aaah",
-            """{
-    'vulnerability_id': 'VCID-5gjv-fh72-aaah',
-    'aliases': [
-        'CVE-2023-31047',
-        'GHSA-r3xc-prgr-mg9p'
-    ],
-    'summary': 'Django bypasses validation when using one form field to upload multiple files',
-    'affected_purls': [
-        'pkg:pypi/django@3.2',
-    ],
-    'fixed_purl': [
-        'pkg:pypi/django@3.2.19',
-        'pkg:pypi/django@4.1.9',
-        'pkg:pypi/django@4.2.1'
-    ],
-    'severities': [
-        {
-            'id': 28067,
-            'reference_id': 176651,
-            'scoring_system': 'cvssv3.1_qr',
-            'value': 'CRITICAL',
-            'scoring_elements': ''
-        }
-    ],
-    'references': [
-        {
-            'id': 176644,
-            'url': 'https://docs.djangoproject.com/en/4.2/releases/security/',
-            'reference_id': ''
-        },
-
-    ],
-    'weaknesses': []
-}""",
-        ),
-        (
-            "pypi",
-            "aiohttp",
-            "VCID-e5jg-e5bm-aaaf",
-            """{
-    'vulnerability_id': 'VCID-e5jg-e5bm-aaaf',
-    'aliases': [
-        'CVE-2018-1000519',
-        'GHSA-fpwp-69xv-c67f'
-    ],
-    'summary': 'Moderate severity vulnerability that affects aiohttp-session',
-    'affected_purls': [
-        'pkg:pypi/aiohttp@0.1',
-        'pkg:pypi/aiohttp-session@2.3.0'
-    ],
-    'fixed_purl': [
-        'pkg:pypi/aiohttp@3.0.0b0',
-        'pkg:pypi/aiohttp-session@2.4.0'
-    ],
-    'severities': [
-        {
-            'id': 19999,
-            'reference_id': 143396,
-            'scoring_system': 'cvssv3.1_qr',
-            'value': 'MODERATE',
-            'scoring_elements': ''
-        }
-    ],
-    'references': [
-        {
-            'id': 143397,
-            'url': 'https://github.com/aio-libs/aiohttp-session/blob/master/aiohttp_session/redis_storage.py#L60',
-            'reference_id': ''
-        }
-    ],
-    'weaknesses': []
-}""",
-        ),
+        ("generic", "", "nginx"),
     ],
 )
-def test_write_vuln_data(tmp_path, ecosystem, package, vuln_id, vuln_data):
-    write_vuln_data(tmp_path, ecosystem, package, vuln_id, vuln_data)
+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, ecosystem)
+    ecosystem_dir = os.path.join(tmp_path, purl_type)
     assert os.path.isdir(ecosystem_dir)
 
-    package_dir = os.path.join(ecosystem_dir, package)
-    assert os.path.isdir(package_dir)
+    namespace_dir = os.path.join(tmp_path, purl_namespace)
+    assert os.path.isdir(namespace_dir)
+
+    name_dir = os.path.join(ecosystem_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)
 
-    vuln_file = os.path.join(package_dir, vuln_id + ".yaml")
-    assert os.path.isfile(vuln_file)
 
-    assert Path(vuln_file).read_text() == yaml.dump(vuln_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):