Source code for swh.osv.luigi.vuln_reports

# Copyright (C) 2026  The Software Heritage developers
# See the AUTHORS file at the top-level directory of this distribution
# License: GNU General Public License version 3, or any later version
# See top-level LICENSE file for more information

"""
Luigi tasks for downloading and processing vulnerability reports
================================================================
"""

from os import PathLike
from pathlib import Path
from typing import Dict, Iterator, Tuple

import luigi

from swh.graph.luigi.compressed_graph import LocalGraph

from .utils import _BaseUploadToS3

ALL_ZIP_URL = "https://osv-vulnerabilities.storage.googleapis.com/all.zip"




class DownloadNvdTarballs(luigi.Task):
    """Downloads NVD tarballs from Google Cloud Storage."""

    vulnerabilities_dir = luigi.PathParameter()



    def output(self) -> luigi.LocalTarget:
        """``nvd_cve/.stamp`` file indicating NVD tarballs have been downloaded."""
        return luigi.LocalTarget(self.vulnerabilities_dir / "nvd_cve" / ".stamp")




    def run(self) -> None:
        """Downloads NVD tarballs from GCS and creates a stamp file."""
        from ..download_nvd_osv import download_to_tarballs

        self.vulnerabilities_dir.mkdir(exist_ok=True, parents=True)
        nvd_dir = self.vulnerabilities_dir / "nvd_cve"
        download_to_tarballs(nvd_dir)
        Path(self.output().path).touch()






class DownloadOsvZip(luigi.Task):
    """Downloads the all.zip file from the OSV vulnerability feed."""

    vulnerabilities_dir = luigi.PathParameter()



    def output(self) -> luigi.LocalTarget:
        """``all.zip`` file containing all OSV vulnerabilities."""
        return luigi.LocalTarget(self.vulnerabilities_dir / "all.zip")




    def run(self) -> None:
        """Downloads all.zip from the OSV vulnerability feed."""
        import urllib.request

        self.vulnerabilities_dir.mkdir(exist_ok=True, parents=True)
        tmp_path = self.vulnerabilities_dir / "all.zip.tmp"
        urllib.request.urlretrieve(ALL_ZIP_URL, tmp_path)
        tmp_path.rename(self.output().path)






class DownloadVulnerabilities(luigi.Task):
    """Downloads vulnerability reports and imports them into a SQLite database."""

    local_graph_path = luigi.PathParameter()
    graph_name = luigi.StrParameter(default="graph")
    vulnerabilities_dir = luigi.PathParameter()



    def requires(self) -> Dict[str, luigi.Task]:
        """Returns instances of :class:`swh.graph.luigi.compressed_graph.LocalGraph`,
        :class:`DownloadNvdTarballs` and :class:`DownloadOsvZip`."""
        return {
            "graph": LocalGraph(local_graph_path=self.local_graph_path),
            "nvd": DownloadNvdTarballs(
                vulnerabilities_dir=self.vulnerabilities_dir,
            ),
            "all_zip": DownloadOsvZip(
                vulnerabilities_dir=self.vulnerabilities_dir,
            ),
        }




    def output(self) -> luigi.LocalTarget:
        """``all.sqlite`` file containing all vulnerability data."""
        return luigi.LocalTarget(self.vulnerabilities_dir / "all.sqlite")




    def run(self) -> None:
        """Downloads vulnerability reports and imports them into a SQLite database."""
        import sqlite3

        from ..to_sqlite import import_osv_tarballs, import_osv_zip, init_db

        nvd_dir = Path(self.input()["nvd"].path).parent
        all_zip_path = Path(self.input()["all_zip"].path)

        tmp_path = self.vulnerabilities_dir / "all.sqlite.tmp"
        with sqlite3.connect(tmp_path) as conn:
            init_db(conn)
            import_osv_tarballs(conn, nvd_dir)
            import_osv_zip(conn, all_zip_path)

        tmp_path.rename(self.output().path)






class UploadVulnerabilitiesSqliteToS3(_BaseUploadToS3):
    """Uploads the vulnerabilities SQLite database to S3."""



    def requires(self) -> Dict[str, luigi.Task]:
        """Returns an instance of :class:`DownloadVulnerabilities`."""
        return {
            "sqlite": DownloadVulnerabilities(
                local_graph_path=self.local_graph_path,
                graph_name=self.graph_name,
                vulnerabilities_dir=self.vulnerabilities_dir,
            ),
        }




    def output(self) -> luigi.Target:
        """S3 target for the uploaded all.sqlite file."""
        import luigi.contrib.s3

        return luigi.contrib.s3.S3Target(f"{self._s3_prefix()}/all.sqlite")


    def _files_to_upload(self) -> Iterator[Tuple[PathLike, str]]:
        """Yields (local_path, s3_key) pairs for the sqlite database."""
        yield (self.input()["sqlite"].path, "all.sqlite")