container_analyzer.py

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

"""Container Analyzer agent function for NAT."""

import logging
from datetime import UTC
from datetime import datetime
from pathlib import Path

from langchain_core.messages import HumanMessage, SystemMessage
from nat.builder.builder import Builder
from nat.builder.framework_enum import LLMFrameworkEnum
from nat.builder.function_info import FunctionInfo
from nat.cli.register_workflow import register_function
from nat.data_models.component_ref import FunctionRef
from nat.data_models.function import FunctionBaseConfig
from pydantic import Field

from vuln_analysis.data_models.container_analyzer import (
    ContainerAnalyzerInput,
    ContainerAnalyzerResult,
)
from vuln_analysis.prompts.container_analyzer import PHASE2_SYSTEM_PROMPT, build_user_message
from vuln_analysis.semantic_failure import serialize_result_for_cli
from vuln_analysis.tools.filesystem_tools import configure_allowed_paths
from vuln_analysis.utils.agentic_loop import run_agentic_loop
from vuln_analysis.utils.concurrency import make_lazy_rate_limiter
from vuln_analysis.utils.container_report_metadata import ContainerReportMetadataError
from vuln_analysis.utils.container_report_metadata import build_container_report_metadata
from vuln_analysis.utils.container_report_metadata import normalize_container_report_metadata
from vuln_analysis.utils.model_errors import is_retryable_model_error
from vuln_analysis.utils.web_researcher import run_web_research

logger = logging.getLogger(__name__)

_DEFAULT_PUBLIC_PREFIXES = [
    "nvcr.io/", "docker.io/", "gcr.io/", "ghcr.io/", "registry.hub.docker.com/",
]


class ContainerAnalyzerConfig(FunctionBaseConfig, name="container_analyzer"):
    phase1_llm_name: str = Field(description="LLM for Phase 1 web research")
    phase2_llm_name: str = Field(description="LLM for Phase 2 filesystem analysis")
    web_tool_names: list[FunctionRef] = Field(description="Phase 1 tools: web_search, web_fetch")
    fs_tool_names: list[FunctionRef] = Field(description="Phase 2 tools: filesystem tools")
    phase1_max_iterations: int = Field(default=8)
    phase2_max_iterations: int = Field(default=40)
    context_window_tokens: int | None = Field(
        default=None,
        description=(
            "Optional deployment-specific context cap. When omitted, derive the limit from "
            "the configured model's metadata and hosted-endpoint profile."
        ),
    )
    wrap_up_fraction: float = Field(default=0.80)
    loop_detection_threshold: int = Field(default=2)
    llm_max_rate: int | None = Field(
        default=None,
        description="Per-function LLM rate limit (requests/second). Overrides the workflow setting.",
    )
    public_registry_prefixes: list[str] = Field(default=_DEFAULT_PUBLIC_PREFIXES)
    output_subdir: str = Field(default="container_reports")
    description: str = Field(default="Analyzes a container image and produces a security report")


@register_function(config_type=ContainerAnalyzerConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])
async def container_analyzer(config: ContainerAnalyzerConfig, builder: Builder):
    phase1_llm = await builder.get_llm(llm_name=config.phase1_llm_name, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
    phase2_llm = await builder.get_llm(llm_name=config.phase2_llm_name, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
    web_tools = await builder.get_tools(tool_names=config.web_tool_names, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
    fs_tools = await builder.get_tools(tool_names=config.fs_tool_names, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
    get_rate_limiter = make_lazy_rate_limiter(config.llm_max_rate, builder)

    async def _arun(inp: ContainerAnalyzerInput) -> ContainerAnalyzerResult:
        rate_limiter = get_rate_limiter()
        pull_ref = inp.pull_ref
        logger.info("Starting container analysis for %s", pull_ref)

        if not Path(inp.filesystem_path).is_dir():
            return ContainerAnalyzerResult(
                image_name=inp.image_name, image_tag=inp.image_tag, success=False,
                error=f"Filesystem directory not found: {inp.filesystem_path}",
            )

        image_config_content: str | None = None
        image_config_loaded = False
        if inp.image_config_path:
            if not Path(inp.image_config_path).is_file():
                return ContainerAnalyzerResult(
                    image_name=inp.image_name, image_tag=inp.image_tag, success=False,
                    error=f"Image config file not found: {inp.image_config_path}",
                )
            try:
                image_config_content = Path(inp.image_config_path).read_text(encoding="utf-8")
                image_config_loaded = True
            except Exception as exc:
                logger.warning("Failed to pre-read image config: %s", exc)

        phase1_skipped = True
        web_research_context: str | None = None

        normalized_pull_ref = pull_ref.lower()
        public_registry = any(
            normalized_pull_ref.startswith(prefix.lower())
            for prefix in config.public_registry_prefixes
        )
        should_web_research = not inp.skip_web_research and public_registry

        if should_web_research:
            phase1_skipped = False
            logger.info("Phase 1: running web research for public image")
            web_research_context = await run_web_research(
                phase1_llm, web_tools, pull_ref,
                max_iterations=config.phase1_max_iterations,
                rate_limiter=rate_limiter,
            )
        else:
            logger.info("Phase 1: skipped")

        report_metadata = build_container_report_metadata(
            analysis_date=datetime.now(UTC).date(),
            image_config_path=inp.image_config_path,
            image_config_loaded=image_config_loaded,
            public_registry=public_registry,
            skip_web_research=inp.skip_web_research,
        )

        logger.info("Phase 2: starting filesystem analysis")
        configure_allowed_paths(
            roots=[inp.filesystem_path],
            filesystem_root=inp.filesystem_path,
        )

        messages = [
            SystemMessage(content=PHASE2_SYSTEM_PROMPT),
            HumanMessage(
                content=build_user_message(
                    inp,
                    web_research_context,
                    image_config_content,
                    report_metadata=report_metadata,
                )
            ),
        ]

        try:
            content, tool_calls, llm_calls = await run_agentic_loop(
                llm=phase2_llm,
                messages=messages,
                tools=fs_tools,
                max_iterations=config.phase2_max_iterations,
                context_window_tokens=config.context_window_tokens,
                wrap_up_fraction=config.wrap_up_fraction,
                loop_threshold=config.loop_detection_threshold,
                phase_label="Phase 2",
                rate_limiter=rate_limiter,
            )
        except Exception as exc:
            logger.error("Phase 2 failed: %s", exc, exc_info=True)
            return ContainerAnalyzerResult(
                image_name=inp.image_name, image_tag=inp.image_tag, success=False, error=str(exc),
                phase1_skipped=phase1_skipped,
                retryable=is_retryable_model_error(exc),
            )

        if not content:
            return ContainerAnalyzerResult(
                image_name=inp.image_name, image_tag=inp.image_tag, success=False,
                error="No report content produced by the LLM.",
                phase1_skipped=phase1_skipped,
                tool_call_count=tool_calls, llm_call_count=llm_calls,
            )

        try:
            content = normalize_container_report_metadata(content, report_metadata)
        except ContainerReportMetadataError as exc:
            logger.error("Container report metadata validation failed: %s", exc)
            return ContainerAnalyzerResult(
                image_name=inp.image_name,
                image_tag=inp.image_tag,
                success=False,
                error=f"Container report metadata validation failed: {exc}",
                retryable=False,
                phase1_skipped=phase1_skipped,
                tool_call_count=tool_calls,
                llm_call_count=llm_calls,
            )

        logger.info(
            "Analysis complete: %s (%d chars, %d tool calls, %d LLM calls)",
            pull_ref,
            len(content),
            tool_calls,
            llm_calls,
        )
        return ContainerAnalyzerResult(
            image_name=inp.image_name, image_tag=inp.image_tag, report_content=content, success=True,
            phase1_skipped=phase1_skipped,
            tool_call_count=tool_calls, llm_call_count=llm_calls,
        )

    def _str_to_input(s: str) -> ContainerAnalyzerInput:
        return ContainerAnalyzerInput.model_validate_json(s)

    def _result_to_str(r: ContainerAnalyzerResult) -> str:
        return serialize_result_for_cli(r)

    yield FunctionInfo.from_fn(
        _arun, input_schema=ContainerAnalyzerInput, description=config.description,
        converters=[_str_to_input, _result_to_str],
    )