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 Vulnerability
        Analysis for Container Security, 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. Vulnerability Analysis for Container Security'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 Vulnerability Analysis for Container Security'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)