Add AffectedPackageV2 for advisory v2

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

Author: keshav-space

vulnerabilities/importer.py

@@ -11,11 +11,8 @@
 import datetime
 import functools
 import logging
-import os
-import shutil
 import traceback
 import xml.etree.ElementTree as ET
-from pathlib import Path
 from typing import Iterable
 from typing import List
 from typing import Mapping
@@ -339,6 +336,88 @@ def from_dict(cls, affected_pkg: dict):
         )
 
 
+@functools.total_ordering
+@dataclasses.dataclass(eq=True)
+class AffectedPackageV2:
+    """
+    Relate a Package URL with a range of affected versions and fixed versions.
+    The Package URL must *not* have a version.
+    AffectedPackage must contain either ``affected_version_range`` or ``fixed_version_range``.
+    """
+
+    package: PackageURL
+    affected_version_range: Optional[VersionRange] = None
+    fixed_version_range: Optional[VersionRange] = None
+
+    def __post_init__(self):
+        if self.package.version:
+            raise ValueError(f"Affected Package URL {self.package!r} cannot have a version.")
+
+        if not (self.affected_version_range or self.fixed_version_range):
+            raise ValueError(
+                f"Affected Package {self.package!r} should have either fixed version range or an "
+                "affected version range."
+            )
+
+    def __lt__(self, other):
+        if not isinstance(other, AffectedPackage):
+            return NotImplemented
+        return self._cmp_key() < other._cmp_key()
+
+    # TODO: Add cache
+    def _cmp_key(self):
+        return (
+            str(self.package),
+            str(self.affected_version_range or ""),
+            str(self.fixed_version_range or ""),
+        )
+
+    def to_dict(self):
+        """Return a serializable dict that can be converted back using self.from_dict"""
+
+        affected_version_range = (
+            str(self.affected_version_range) if self.affected_version_range else None
+        )
+        fixed_version_range = str(self.fixed_version_range) if self.fixed_version_range else None
+        return {
+            "package": purl_to_dict(self.package),
+            "affected_version_range": affected_version_range,
+            "fixed_version_range": fixed_version_range,
+        }
+
+    @classmethod
+    def from_dict(cls, affected_pkg: dict):
+        """Return an AffectedPackage object from dict generated by self.to_dict"""
+
+        package = PackageURL(**affected_pkg["package"])
+        affected_version_range = None
+        fixed_version_range = None
+        affected_range = affected_pkg["affected_version_range"]
+        fixed_range = affected_pkg["fixed_version_range"]
+
+        try:
+            affected_version_range = VersionRange.from_string(affected_range)
+            fixed_version_range = VersionRange.from_string(fixed_range)
+        except:
+            tb = traceback.format_exc()
+            logger.error(
+                f"Cannot create AffectedPackage with invalid or unknown range: {affected_pkg!r} with error: {tb!r}"
+            )
+            return
+
+        if not fixed_version_range and not affected_version_range:
+            logger.error(
+                f"Cannot create AffectedPackage without fixed or affected range: {affected_pkg!r}"
+            )
+            return
+
+        return cls(
+            package=package,
+            affected_version_range=affected_version_range,
+            fixed_version_range=fixed_version_range,
+        )
+
+
 @dataclasses.dataclass(order=True)
 class AdvisoryData:
     """
@@ -355,7 +434,9 @@ class AdvisoryData:
     advisory_id: str = ""
     aliases: List[str] = dataclasses.field(default_factory=list)
     summary: Optional[str] = ""
-    affected_packages: List[AffectedPackage] = dataclasses.field(default_factory=list)
+    affected_packages: Union[List[AffectedPackage], List[AffectedPackageV2]] = dataclasses.field(
+        default_factory=list
+    )
     references: List[Reference] = dataclasses.field(default_factory=list)
     references_v2: List[ReferenceV2] = dataclasses.field(default_factory=list)
     date_published: Optional[datetime.datetime] = None
@@ -392,13 +473,19 @@ def to_dict(self):
     @classmethod
     def from_dict(cls, advisory_data):
         date_published = advisory_data["date_published"]
+        affected_packages = advisory_data["affected_packages"]
+        affected_package_cls = AffectedPackage
+        if affected_packages:
+            affected_package_cls = (
+                AffectedPackageV2
+                if "fixed_version_range" in affected_packages[0]
+                else AffectedPackage
+            )
         transformed = {
             "aliases": advisory_data["aliases"],
             "summary": advisory_data["summary"],
             "affected_packages": [
-                AffectedPackage.from_dict(pkg)
-                for pkg in advisory_data["affected_packages"]
-                if pkg is not None
+                affected_package_cls.from_dict(pkg) for pkg in affected_packages if pkg is not None
             ],
             "references": [Reference.from_dict(ref) for ref in advisory_data["references"]],
             "date_published": datetime.datetime.fromisoformat(date_published)