WIP

Signed-off-by: Haiko Schol <[email protected]>

Author: haikoschol

vulnerabilities/admin.py

@@ -25,11 +25,11 @@
 
 from vulnerabilities.models import (
     ImpactedPackage,
+    Importer,
     Package,
-    PackageReference,
     ResolvedPackage,
     Vulnerability,
-    VulnerabilityReference
+    VulnerabilityReference,
 )
 
 
@@ -58,6 +58,6 @@ class ResolvedPackageAdmin(admin.ModelAdmin):
     pass
 
 
-@admin.register(PackageReference)
-class PackageReferenceAdmin(admin.ModelAdmin):
+@admin.register(Importer)
+class ImporterAdmin(admin.ModelAdmin):
     pass

vulnerabilities/api.py

@@ -22,29 +22,15 @@
 #  Visit https://github.com/nexB/vulnerablecode/ for support and download.
 
 from rest_framework import serializers
-from rest_framework import status
 from rest_framework import viewsets
-from rest_framework.response import Response
 
 from packageurl import PackageURL
 
 from vulnerabilities.models import Package
-from vulnerabilities.models import PackageReference
 from vulnerabilities.models import Vulnerability
 from vulnerabilities.models import VulnerabilityReference
 
 
-class PackageReferenceSerializer(serializers.ModelSerializer):
-    class Meta:
-        model = PackageReference
-        fields = [
-            'repository',
-            'platform',
-            'name',
-            'version',
-        ]
-
-
 class VulnerabilityReferenceSerializer(serializers.ModelSerializer):
     class Meta:
         model = VulnerabilityReference
@@ -69,7 +55,6 @@ class Meta:
 
 class PackageSerializer(serializers.ModelSerializer):
     vulnerabilities = VulnerabilitySerializer(many=True)
-    references = PackageReferenceSerializer(source='packagereference_set', many=True)
 
     class Meta:
         model = Package
@@ -78,7 +63,6 @@ class Meta:
             'version',
             'package_url',
             'vulnerabilities',
-            'references',
         ]
 
 

vulnerabilities/data_dump.py

@@ -23,7 +23,6 @@
 
 from vulnerabilities.models import ImpactedPackage
 from vulnerabilities.models import Package
-from vulnerabilities.models import PackageReference
 from vulnerabilities.models import ResolvedPackage
 from vulnerabilities.models import Vulnerability
 from vulnerabilities.models import VulnerabilityReference

vulnerabilities/data_source.py

@@ -20,61 +20,101 @@
 #  VulnerableCode is a free software code scanning tool from nexB Inc. and others.
 #  Visit https://github.com/nexB/vulnerablecode/ for support and download.
 
-from attr import dataclass
 from datetime import datetime
 from typing import Any
 from typing import Mapping
+from typing import Optional
 from typing import Sequence
+from typing import Set
+import dataclasses
 
+from packageurl import PackageURL
 
-@dataclass
+
+@dataclasses.dataclass
 class DataSource:
     """
     TODO
     """
     batch_size: int
-    cutoff_date: datetime
-    config: Mapping[str, Any]
+    cutoff_date: Optional[datetime] = None
+    config: Optional[Mapping[str, Any]] = dataclasses.field(default_factory=dict)
 
     def __enter__(self):
         """
         Subclasses acquire per-run resources, such as network connections, file downloads, etc. here.
         """
-        pass
+        return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         """
         Subclasses release per-run resources acquired in __enter__() here.
         """
         pass
 
-    def __next__(self):
+    def new_records(self):
         """
-        Subclasses return batch_size sized batches of VulnerabilityInfo objects
+        Subclasses return batch_size sized batches of VulnerabilityInfo objects that have been added to the data source
+        since self.cutoff_date.
         """
-        pass
+        raise StopIteration
+
+    def updated_records(self):
+        """
+        Subclasses return batch_size sized batches of VulnerabilityInfo objects that have been modified since
+        self.cutoff_date.
+
+        NOTE: Data sources that do now enable detection of changes to existing records vs added records must only
+              implement this method, not new_records(). The ImportRunner relies on this contract to decide between
+              insert and update operations.
+        """
+        raise StopIteration
 
 
 # The following data classes express the contract between data sources and the import runner.
-# Data sources are expected to be usable as context managers and generators, yielding
-# batches of VulnerabilityInfo sequences.
+# Data sources are expected to be usable as context managers and generators, yielding batches of VulnerabilityInfo
+# sequences.
 
-@dataclass
+@dataclasses.dataclass(frozen=True)
 class Package:
     name: str
-    namespace: str
     type: str
     version: str
-    qualifiers: str
-    subpath: str
-    references: Sequence[str]
+    namespace: str = ''
+    qualifiers: str = ''
+    subpath: str = ''
+
+    @property
+    def package_url(self) -> str:
+        """
+        Return a compact Package URL "purl" string.
+        """
+        purl = PackageURL(
+            self.type, self.namespace, self.name,
+            self.version, self.qualifiers, self.subpath
+        )
+        return str(purl)
 
 
-@dataclass
+@dataclasses.dataclass
 class VulnerabilityInfo:
-    cve_id: str
     summary: str
-    affected_packages: Sequence[Package]
-    unaffected_packages: Sequence[Package]
-    fixed_packages: Sequence[Package]
-    references: Sequence[str]
+    impacted_packages: Sequence[Package]
+    resolved_packages: Sequence[Package] = dataclasses.field(default_factory=list)
+    references: Sequence[str] = dataclasses.field(default_factory=list)
+    cve_id: Optional[str] = None
+
+    _impacted_purls: Set[str] = dataclasses.field(init=False, compare=False, default=None)
+    _resolved_purls: Set[str] = dataclasses.field(init=False, compare=False, default=None)
+
+    @property
+    def impacted_package_urls(self):
+        if self._impacted_purls is None:
+            self._impacted_purls = {p.package_url for p in self.impacted_packages}
+        return self._impacted_purls
+
+    @property
+    def resolved_package_urls(self):
+        if self._resolved_purls is None:
+            self._resolved_purls = {p.package_url for p in self.resolved_packages}
+        return self._resolved_purls

vulnerabilities/import_runner.py

@@ -23,29 +23,189 @@
 
 import datetime
 import logging
+from typing import List
+from typing import Mapping
+from typing import Sequence
+from typing import Set
+from typing import Tuple
 
+from django.db import IntegrityError
+from django.db import transaction
+
+from vulnerabilities import models
+from vulnerabilities.data_source import Package
+from vulnerabilities.data_source import VulnerabilityInfo
 
 logger = logging.getLogger(__name__)
 
 
-# TODO This really should use asyncio for network and database, but sadly the Django ORM won't allow it.
 class ImportRunner:
+    """
+    The ImportRunner is responsible for inserting and updating data about vulnerabilities and
+    affected/unaffected/fixed packages in the database. The two main goals for the implementation are correctness and
+    efficiency.
+
+    Correctness:
+        - There must be no duplicates in the database (should be enforced by the schema).
+        - No valid data from the data source must be skipped or truncated.
 
+    Efficiency:
+        - Bulk inserts should be used whenever possible.
+        - Checking whether a record already exists should be kept to a minimum
+          (the data source should know this instead).
+        - All update and select operations must use indexed columns.
+    """
     def __init__(self, importer, batch_size=None):
         self.importer = importer
         self.batch_size = batch_size
 
-    def run(self, cutoff_date=None):
+    def run(self, cutoff_date=None) -> None:
+        """
+        Create a data source for the given importer and store the data retrieved in the database.
+
+        NB: Data sources provide two kinds of records; vulnerabilities and packages. Vulnerabilities are potentially
+        shared across many packages, from the same data source and from different data sources. For example, a
+        vulnerability in the Linux kernel is mentioned by advisories from all Linux distributions that package this
+        kernel version.
+        """
         logger.debug(f'Starting import for {self.importer.name}.')
         data_source = self.importer.make_data_source(cutoff_date=cutoff_date, batch_size=self.batch_size)
+
         with data_source as ds:
-            for batch in ds:
-                # TODO
-                # Check if any Vulnerability or Package from this batch already exists in the DB
-                # If not: Bulk insert everything
-                # If yes: Update existing ones and bulk insert the rest
-                pass
+            for batch in ds.new_records():
+                impacted, resolved = _collect_packages_from_batch(batch)
+                impacted, resolved = _bulk_insert_packages(impacted, resolved)
+
+                vulnerabilities = _insert_vulnerabilities_and_references(batch)
+
+                _bulk_insert_impacted_and_resolved_packages(batch, vulnerabilities, impacted, resolved)
 
             self.importer.last_run = datetime.datetime.utcnow()
             self.importer.save()
-            logger.debug(f'Successfully finished import for {self.importer.name}.')
+
+        logger.debug(f'Successfully finished import for {self.importer.name}.')
+
+
+def _bulk_insert_impacted_and_resolved_packages(
+    batch: Sequence[VulnerabilityInfo],
+    vulnerability_models: Set[models.Vulnerability],
+    impacted_package_models: Mapping[str, models.Package],
+    resolved_package_models: Mapping[str, models.Package],
+) -> None:
+
+    impacted_refs: List[models.ImpactedPackage] = []
+    resolved_refs: List[models.ResolvedPackage] = []
+
+    for vuln_info in batch:
+        vuln_model = _vuln_info_to_vuln_model(vuln_info, vulnerability_models)
+        vulnerability_models.remove(vuln_model)  # minor optimization
+
+        for impacted_package in vuln_info.impacted_packages:
+            ip = models.ImpactedPackage(
+                vulnerability=vuln_model,
+                package=impacted_package_models[impacted_package.package_url]
+            )
+            impacted_refs.append(ip)
+
+        for resolved_package in vuln_info.resolved_packages:
+            ip = models.ResolvedPackage(
+                vulnerability=vuln_model,
+                package=resolved_package_models[resolved_package.package_url]
+            )
+            resolved_refs.append(ip)
+
+    models.ImpactedPackage.objects.bulk_create(impacted_refs)
+    models.ResolvedPackage.objects.bulk_create(resolved_refs)
+
+
+@transaction.atomic
+def _insert_vulnerabilities_and_references(batch: Sequence[VulnerabilityInfo]) -> Set[models.Vulnerability]:
+    """
+    TODO Consider refactoring to use bulk_create() and avoid get_or_create() when possible.
+    """
+    vulnerabilities = set()
+
+    for vuln_info in batch:
+        vuln: models.Vulnerability
+
+        if vuln_info.cve_id:
+            vuln, created = models.Vulnerability.objects.get_or_create(cve_id=vuln_info.cve_id)
+            if created:
+                vuln.summary = vuln_info.summary
+                vuln.save()
+        else:
+            # FIXME Currently there is no way to check whether a vulnerability without a CVE ID already exists in the
+            # FIXME database.
+            vuln = models.Vulnerability.objects.create(summary=vuln_info.summary)
+
+        vulnerabilities.add(vuln)
+
+        for url in vuln_info.references:
+            try:
+                models.VulnerabilityReference.objects.create(vulnerability=vuln, url=url)
+            except IntegrityError:
+                # This vulnerability reference already exists, nothing to do.
+                # TODO Find a more efficient way to do this rather than trying and ignoring any errors.
+                pass
+
+    return vulnerabilities
+
+
+def _vuln_info_to_vuln_model(
+        vuln_info: VulnerabilityInfo,
+        vulnerability_models: Set[models.Vulnerability]
+) -> models.Vulnerability:
+
+    for v in vulnerability_models:
+        if vuln_info.cve_id and vuln_info.cve_id == v.cve_id:
+            return v
+
+        if vuln_info.summary == v.summary:
+            return v
+
+    raise RuntimeError(f'No Vulnerability model object found for this VulnerabilityInfo: {vuln_info.summary}')
+
+
+def _collect_packages_from_batch(batch: Sequence[VulnerabilityInfo]) -> Tuple[Set[Package], Set[Package]]:
+    impacted, resolved = set(), set()
+
+    for vuln_info in batch:
+        impacted.update(vuln_info.impacted_packages)
+        resolved.update(vuln_info.resolved_packages)
+
+    return impacted, resolved
+
+
+def _bulk_insert_packages(
+        impacted: Set[Package],
+        resolved: Set[Package]
+) -> Tuple[Mapping[str, Package], Mapping[str, Package]]:
+
+    pkg_models = [_to_package_model(p) for p in impacted.union(resolved)]
+    pkg_models = models.Package.objects.bulk_create(pkg_models)
+
+    impacted_purls = {p.package_url for p in impacted}
+    resolved_purls = {p.package_url for p in resolved}
+
+    impacted_models, resolved_models = {}, {}
+
+    for pkg_model in pkg_models:
+        purl = pkg_model.package_url
+
+        if purl in impacted_purls:
+            impacted_models[purl] = pkg_model
+        elif purl in resolved_purls:
+            resolved_models[purl] = pkg_model
+
+    return impacted_models, resolved_models
+
+
+def _to_package_model(pkg: Package) -> models.Package:
+    return models.Package(
+        name=pkg.name,
+        type=pkg.type,
+        version=pkg.version,
+        namespace=pkg.namespace,
+        qualifiers=pkg.qualifiers,
+        subpath=pkg.subpath,
+    )