Source code for swh.osv.luigi.commit2vuln

# 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 building the commit-to-vulnerability mapping
============================================================
"""

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

import luigi

from swh.graph.luigi.compressed_graph import LocalGraph
from swh.osv.luigi import DownloadVulnerabilities
from swh.osv.luigi.reachability import ComputeConnectedComponents
from swh.osv.luigi.utils import _BaseUploadToS3
from swh.osv.luigi.vuln2commit import Vuln2CommitParquet




class Commit2VulnBv(luigi.Task):
    """Creates a BVGraph that maps each revision to the vulnerabilities that
    affect it."""

    local_graph_path = luigi.PathParameter()
    graph_name = luigi.StrParameter(default="graph")
    vulnerabilities_dir = luigi.PathParameter()
    include_vcs_matches = luigi.BoolParameter(default=True)
    use_cherrypicks = luigi.BoolParameter(default=False)

    @property
    def resources(self):
        """Declares the task uses every CPU available"""
        import socket

        hostname = socket.getfqdn()
        return {
            f"{hostname}_max_cpu": 1,
        }



    def requires(self) -> Dict[str, luigi.Task]:
        """Returns instances of :class:`swh.graph.luigi.compressed_graph.LocalGraph`,
        :class:`swh.osv.luigi.DownloadVulnerabilities`,
        :class:`swh.osv.luigi.reachability.ComputeConnectedComponents` and
        :class:`swh.osv.luigi.vuln2commit.Vuln2CommitParquet`."""
        return {
            "graph": LocalGraph(local_graph_path=self.local_graph_path),
            "sqlite": DownloadVulnerabilities(
                local_graph_path=self.local_graph_path,
                graph_name=self.graph_name,
                vulnerabilities_dir=self.vulnerabilities_dir,
            ),
            "subgraphwccs": ComputeConnectedComponents(
                local_graph_path=self.local_graph_path,
                graph_name=self.graph_name,
                vulnerabilities_dir=self.vulnerabilities_dir,
                include_vcs_matches=self.include_vcs_matches,
            ),
            "vuln2commit": Vuln2CommitParquet(
                local_graph_path=self.local_graph_path,
                graph_name=self.graph_name,
                vulnerabilities_dir=self.vulnerabilities_dir,
                include_vcs_matches=self.include_vcs_matches,
                use_cherrypicks=self.use_cherrypicks,
            ),
        }




    def output(self) -> luigi.LocalTarget:
        """BVGraph file mapping revision ids to vulnerability ids."""
        dir_name = (
            "commit2vuln_using_cherrypicks"
            if self.use_cherrypicks
            else "commit2vuln_without_cherrypicks"
        )
        return luigi.LocalTarget(self.vulnerabilities_dir / f"{dir_name}.graph")




    def run(self) -> None:
        """Runs ``build-commit2vuln-bv`` and ``build-ef``."""
        from .shell import Rust

        multimap_base = Path(self.output().path).with_suffix("")
        # fmt: off
        Rust(
            "build-commit2vuln-bv",
            "--graph", self.local_graph_path / self.graph_name,
            "--db", self.input()["sqlite"],
            "--subgraphwccs", self.input()["subgraphwccs"],
            "--vuln2commit-dir", self.input()["vuln2commit"]["dir"],
            "--multimap-out", multimap_base,
        ).run()
        Rust(
            "build-ef",
            "--graph", multimap_base,
        ).run()


        # fmt: on




class UploadCommit2VulnBvToS3(_BaseUploadToS3):
    """Uploads the commit2vuln BVGraph to S3."""

    include_vcs_matches = luigi.BoolParameter(default=True)
    use_cherrypicks = luigi.BoolParameter(default=False)



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




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

        dir_name = (
            "commit2vuln_using_cherrypicks"
            if self.use_cherrypicks
            else "commit2vuln_without_cherrypicks"
        )
        return luigi.contrib.s3.S3Target(f"{self._s3_prefix()}/{dir_name}.graph")


    def _files_to_upload(self) -> Iterator[tuple[PathLike, str]]:
        base = Path(self.input()["commit2vuln"].path).with_suffix("")
        for path in sorted(base.parent.glob(f"{base.name}.*")):
            if path.is_dir():
                continue
            yield (path, str(path.relative_to(self.vulnerabilities_dir)))