Source code for swh.vulns.naive_server

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

import concurrent.futures
import functools
import json
from pathlib import Path
import sys
from typing import Callable, Dict, Iterator, List, TypeVar

import grpc

from swh.model.exceptions import ValidationError
from swh.model.swhids import CoreSWHID

from .grpc.swhvulns_pb2 import AffectingVulnerabilitiesRequest, DetectedVulnerability
from .grpc.swhvulns_pb2_grpc import (
    VulnerabilityServiceServicer,
    add_VulnerabilityServiceServicer_to_server,
)

F = TypeVar("F", bound=Callable)




def log_errors(f: F) -> F:
    @functools.wraps(f)
    def newf(*args, **kwargs):
        try:
            return f(*args, **kwargs)
        except KeyboardInterrupt:
            raise
        except Exception:
            import traceback

            traceback.print_exc()
            raise

    return newf  # type: ignore[return-value]





class NaiveVulnerabilityServicer(VulnerabilityServiceServicer):
    """The simplest implementation of :class:`VulnerabilityServiceServicer`, which just returns
    :class:`DetectedVulnerability` from an in-memory list."""

    vulnerabilities: Dict[CoreSWHID, List[DetectedVulnerability]] = {}

    def __init__(self, vulnerabilities: List[DetectedVulnerability]):
        self.vulnerabilities = {}
        for vuln in vulnerabilities:
            self.vulnerabilities.setdefault(
                CoreSWHID.from_string(vuln.swhid), []
            ).append(vuln)



    @log_errors
    def AffectingVulnerabilities(
        self, request: AffectingVulnerabilitiesRequest, context
    ) -> Iterator[DetectedVulnerability]:
        swhids = []
        for string_swhid in request.swhid:
            try:
                swhids.append(CoreSWHID.from_string(string_swhid))
            except ValidationError:
                context.set_code(grpc.StatusCode.INVALID_ARGUMENT)
                context.set_details(f"Invalid SWHID: {string_swhid}")
                return
        for swhid in swhids:
            yield from self.vulnerabilities.get(swhid) or []






def run(address: str, database: List[DetectedVulnerability]) -> None:
    """
    Starts :class:`NaiveVulnerabilityServicer` and never returns.

    ``address`` should be ``<hostname>:<port>``, eg. ``[::]:12345``
    """
    servicer = NaiveVulnerabilityServicer(database)

    server = grpc.server(concurrent.futures.ThreadPoolExecutor(max_workers=10))
    add_VulnerabilityServiceServicer_to_server(servicer, server)
    server.add_insecure_port(address)
    server.start()
    print(f"Server started, listening on {address}")
    server.wait_for_termination()





def main():
    try:
        _, address, database_path = sys.argv
    except ValueError:
        print(f"Syntax: {sys.argv[0]} [::]:<port> database.json", file=sys.stderr)
        print()
        print(
            "Serves a read-only list of DetectedVulnerability (see :file:`proto/swhvulns.proto`) from a JSON file"
        )
        exit(1)

    try:
        database = json.loads(Path(database_path).read_text())
    except FileNotFoundError:
        print(f"{database_path} does not exist", file=sys.stderr)
        exit(2)

    try:
        run(address, database)
    except KeyboardInterrupt:
        pass



if __name__ == "__main__":
    main()