cve_justification_node.py

# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import logging
from textwrap import dedent

from morpheus_llm.llm import LLMLambdaNode
from morpheus_llm.llm import LLMNode
from morpheus_llm.llm.nodes.llm_generate_node import LLMGenerateNode
from morpheus_llm.llm.nodes.prompt_template_node import PromptTemplateNode
from morpheus_llm.llm.services.llm_service import LLMClient

logger = logging.getLogger(__name__)


class CVEJustifyNode(LLMNode):
    """
    A node to classify the results of the summary node.
    """

    JUSTIFICATION_LABEL_COL_NAME = "justification_label"
    JUSTIFICATION_REASON_COL_NAME = "justification"
    AFFECTED_STATUS_COL_NAME = "affected_status"

    JUSTIFICATION_PROMPT = dedent("""
        The summary provided below (delimited with XML tags), generated by the software agent named Sherlock, evaluates
        a specific CVE (Common Vulnerabilities and Exposures) against the backdrop of a software package or environment
        information. This information may include a Software Bill of Materials (SBOM), source code, and dependency
        documentation among others related to environments such as Docker containers. Sherlock's role is to analyze this
        data to ascertain if the CVE impacts the software environment, determining the necessity for a patch.

        Your task is to review Sherlock's analysis and classify the situation into one of the following categories based
        on the described scenario:
        - false_positive: The association between the software package and the CVE is incorrect, either due to an
        erroneous package identification or a mismatched CVE.
        - code_not_present: The software product is unaffected as it does not contain the code or library that harbors
        the vulnerability.
        - code_not_reachable: During runtime, the vulnerable code is not executed.
        - requires_configuration: The exploitability of this issue depends on whether a specific configuration option is
        enabled or disabled. In this case, the configuration is set in a manner that prevents exploitability.
        - requires_dependency: Exploitability is contingent upon a missing dependency.
        - requires_environment: A specific environment, absent in this case, is required for exploitability.
        - compiler_protected: Exploitability hinges on the setting or unsetting of a compiler flag. In this case, the
        flag is set in a manner that prevents exploitability.
        - runtime_protected: Mechanisms are in place that prevent exploits during runtime.
        - perimeter_protected: Protective measures block attacks at the physical, logical, or network perimeters.
        - mitigating_control_protected: Implemented controls mitigate the likelihood or impact of the vulnerability.
        - uncertain: Not enough information to determine the package's exploitability status.
        - vulnerable: the package is actually vulnerable to the CVE and needs to be patched.

        Response only with the category name on the first line, and reasoning on the second line.
        <summary>{{summary}}</summary>
        """).strip("\n")

    # Map raw justification labels optimized for LLM acccuracy to final labels expected by downstream systems
    RAW_TO_FINAL_JUSTIFICATION_MAP = {
        "false_positive": "false_positive",
        "code_not_present": "code_not_present",
        "code_not_reachable": "code_not_reachable",
        "requires_configuration": "requires_configuration",
        "requires_dependency": "requires_dependency",
        "requires_environment": "requires_environment",
        "compiler_protected": "protected_by_compiler",
        "runtime_protected": "protected_at_runtime",
        "perimeter_protected": "protected_at_perimeter",
        "mitigating_control_protected": "protected_by_mitigating_control",
        "uncertain": "uncertain",
        "vulnerable": "vulnerable"
    }

    JUSTIFICATION_TO_AFFECTED_STATUS_MAP = {
        "false_positive": "FALSE",
        "code_not_present": "FALSE",
        "code_not_reachable": "FALSE",
        "requires_configuration": "FALSE",
        "requires_dependency": "FALSE",
        "requires_environment": "FALSE",
        "protected_by_compiler": "FALSE",
        "protected_at_runtime": "FALSE",
        "protected_at_perimeter": "FALSE",
        "protected_by_mitigating_control": "FALSE",
        "uncertain": "UNKNOWN",
        "vulnerable": "TRUE"
    }

    def __init__(self, *, llm_client: LLMClient):
        """
        Initialize the CVEJustificationNode with a selected model.

        Parameters
        ----------
        llm_client : LLMClient
            The LLM client to use for generating the justification.
        """
        super().__init__()

        async def _strip_summaries(summaries: list[str]) -> list[str]:
            return [summary.strip() for summary in summaries]

        self.add_node('stripped_summaries', inputs=['summaries'], node=LLMLambdaNode(_strip_summaries))

        self.add_node("justification_prompt",
                      inputs=[("/stripped_summaries", "summary")],
                      node=PromptTemplateNode(template=self.JUSTIFICATION_PROMPT, template_format='jinja'))

        self.add_node("justify", inputs=["/justification_prompt"], node=LLMGenerateNode(llm_client=llm_client))

        async def _parse_justification(justifications: list[str]) -> dict[str, list[str]]:
            labels: list[str] = []
            reasons: list[str] = []
            affected_status: list[str] = []

            split_justifications = [j.split('\n') for j in justifications]

            for j in split_justifications:

                if len(j) < 2:
                    raise ValueError(
                        f"Invalid justification format: {j}. Must be the label and the reason separated by a newline")

                justification_label_raw = j[0].strip()
                justification_label = self.RAW_TO_FINAL_JUSTIFICATION_MAP.get(justification_label_raw,
                                                                              justification_label_raw)
                labels.append(justification_label)
                reasons.append('\n'.join(j[1:]).strip())
                try:
                    affected_status.append(self.JUSTIFICATION_TO_AFFECTED_STATUS_MAP[justification_label])
                except KeyError:
                    logger.error("Invalid justification label: '%s', setting affected_status='UNKNOWN'",
                                 justification_label)
                    affected_status.append("UNKNOWN")

            return {
                self.JUSTIFICATION_LABEL_COL_NAME: labels,
                self.JUSTIFICATION_REASON_COL_NAME: reasons,
                self.AFFECTED_STATUS_COL_NAME: affected_status
            }

        self.add_node("parse_justification",
                      inputs=['/justify'],
                      node=LLMLambdaNode(_parse_justification),
                      is_output=True)