input.py

# SPDX-FileCopyrightText: Copyright (c) 2024-2026, 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.

from __future__ import annotations

import os
import warnings

from pydantic import BaseModel, ConfigDict, Field, model_validator

from vuln_analysis.data_models.vulnerability_id import VulnerabilityId


def _cache_default_from_env() -> bool:
    """Match standalone pipeline: CACHE_ENABLED=false disables on-disk report caching."""
    raw = os.environ.get("CACHE_ENABLED")
    if raw is None:
        return True
    return str(raw).strip().lower() not in ("0", "false", "no", "off")


class CVEEntry(BaseModel):
    """A single vulnerability identifier with optional hints for the pipeline."""
    model_config = ConfigDict(extra="forbid")

    vuln_id: VulnerabilityId
    package_hint: str | None = None
    version_hint: str | None = None

    @property
    def is_ghsa(self) -> bool:
        return self.vuln_id.startswith("GHSA-")


class PipelineInput(BaseModel):
    """Everything needed to run the full 4-agent vulnerability analysis pipeline."""
    model_config = ConfigDict(extra="forbid")

    vulns: list[CVEEntry] = Field(min_length=1)
    image_name: str
    image_tag: str
    image_digest: str | None = None
    filesystem_path: str = ""
    image_config_path: str | None = None
    skip_web_research: bool = False
    output_dir: str | None = None
    cache: bool = Field(default_factory=_cache_default_from_env)

    cache_cve_researcher: bool | None = None
    cache_container_analyzer: bool | None = None
    cache_exploitability_analyzer: bool | None = None
    cache_vex_categorizer: bool | None = None

    # Opaque, client-supplied correlation data echoed verbatim into the terminal webhook's
    # ``context`` (see server/webhooks/). Kept a raw ``dict`` (not a typed submodel) so values
    # round-trip type-preserving -- integers stay JSON numbers, etc. The pipeline attaches no
    # meaning to its keys.
    callback_context: dict | None = None

    @model_validator(mode="before")
    @classmethod
    def _discard_legacy_model(cls, value):
        if not isinstance(value, dict) or "model" not in value:
            return value

        warnings.warn(
            "PipelineInput.model is deprecated and ignored; configure agent models through workflow LLM settings.",
            DeprecationWarning,
            stacklevel=2,
        )
        value = value.copy()
        value.pop("model")
        return value

    @model_validator(mode="after")
    def _check_input(self):
        seen: set[str] = set()
        for entry in self.vulns:
            if entry.vuln_id in seen:
                raise ValueError(f"Duplicate vuln_id: {entry.vuln_id}")
            seen.add(entry.vuln_id)

        if not self.image_name:
            raise ValueError("image_name must be provided.")
        if not self.image_tag:
            raise ValueError("image_tag must be provided.")
        if _final_path_component_has_tag(self.image_name):
            raise ValueError("image_name must not include a tag when image_tag is provided.")
        return self

    _AGENT_CACHE_FIELDS = {
        "cve_researcher": "cache_cve_researcher",
        "container_analyzer": "cache_container_analyzer",
        "exploitability_analyzer": "cache_exploitability_analyzer",
        "vex_categorizer": "cache_vex_categorizer",
    }

    def use_cache_for(self, agent: str) -> bool:
        """Resolve the effective cache setting for *agent*.

        Per-agent overrides take precedence when set.  If the global
        ``cache`` flag is ``False`` it always wins (disables everything).
        """
        if not self.cache:
            return False
        field_name = self._AGENT_CACHE_FIELDS.get(agent)
        if field_name is None:
            raise ValueError(f"Unknown agent: {agent!r}")
        override = getattr(self, field_name)
        return override if override is not None else self.cache

    @property
    def identifier(self) -> str:
        return self.vulns[0].vuln_id


def _final_path_component_has_tag(image_name: str) -> bool:
    """Return True when the repository name already contains a tag suffix."""
    return ":" in image_name.rsplit("/", 1)[-1]