populate_vulnerability_summary_pipeline.py

#
# 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/aboutcode-org/vulnerablecode for support or download.
# See https://aboutcode.org for more information about nexB OSS projects.
#

import logging

from aboutcode.pipeline import LoopProgress
from django.db.models import Q

from vulnerabilities.models import Advisory
from vulnerabilities.models import Vulnerability
from vulnerabilities.pipelines import VulnerableCodePipeline


class PopulateVulnerabilitySummariesPipeline(VulnerableCodePipeline):
    """Pipeline to populate missing vulnerability summaries from advisories."""

    pipeline_id = "populate_vulnerability_summaries"

    @classmethod
    def steps(cls):
        return (cls.populate_missing_summaries,)

    def populate_missing_summaries(self):
        """Find vulnerabilities with missing summaries and populate them using advisories with the same aliases."""
        vulnerabilities_qs = Vulnerability.objects.filter(summary="")
        self.log(
            f"Processing {vulnerabilities_qs.count()} vulnerabilities without summaries",
            level=logging.INFO,
        )

        progress = LoopProgress(total_iterations=vulnerabilities_qs.count(), logger=self.log)

        vulnerabilities_to_be_updated = []

        for vulnerability in progress.iter(vulnerabilities_qs.iterator()):
            cve_alias = vulnerability.aliases.filter(alias__startswith="CVE-").first()

            if not cve_alias:
                self.log(
                    f"Vulnerability {vulnerability.vulnerability_id} has no CVE alias",
                    level=logging.DEBUG,
                )
                continue

            matching_advisories = Advisory.objects.filter(
                aliases=cve_alias, created_by="nvd_importer"
            ).exclude(summary="")

            if matching_advisories.exists():
                best_advisory = matching_advisories.order_by("-date_collected").first()
                # Note: we filtered above to only get non-empty summaries
                vulnerability.summary = best_advisory.summary
                vulnerabilities_to_be_updated.append(vulnerability)
                self.log(
                    f"Updated summary for vulnerability {vulnerability.vulnerability_id}",
                    level=logging.INFO,
                )
            else:
                self.log(f"No advisory found for alias {cve_alias}", level=logging.DEBUG)
        Vulnerability.objects.bulk_update(vulnerabilities_to_be_updated, ["summary"])
        self.log(
            f"Successfully populated {len(vulnerabilities_to_be_updated)} vulnerabilities with summary",
            level=logging.INFO,
        )
        self.log("Pipeline completed", level=logging.INFO)