Patched code sample

import os
import requests
from pydantic import BaseModel, Field

# This Pydantic model defines the schema for data the LLM can provide.
# Notice that connection parameters like 'es_url' or 'api_key' are NOT
# included. The LLM can only control operational parameters like the query.
class FixedElasticsearchMemoryInput(BaseModel):
    query: str = Field(description="The query to search for in the memory.")
    content_to_store: str | None = Field(
        default=None, description="Content to store in the memory."
    )

class FixedElasticsearchMemoryTool:
    """
    A representation of the fixed elasticsearch_memory tool.

    The fix prevents the Server-Side Request Forgery (SSRF) vulnerability
    by removing connection parameters (es_url, cloud_id, api_key) from the
    LLM-controllable input schema. These parameters are now configured
    securely during the tool's initialization using environment variables,
    ensuring the LLM cannot redirect requests to an arbitrary server.
    """
    name = "elasticsearch_memory"
    description = "Interacts with Elasticsearch memory to store or retrieve information."
    model = FixedElasticsearchMemoryInput

    def __init__(self):
        # FIX: Connection parameters are loaded securely from the operator's
        # environment during initialization. They are no longer part of the
        # runtime input from the LLM.
        self.es_url = os.environ.get("ELASTICSEARCH_URL")
        api_key = os.environ.get("ELASTICSEARCH_API_KEY")

        if not self.es_url or not api_key:
            raise ValueError(
                "Operator must set ELASTICSEARCH_URL and ELASTICSEARCH_API_KEY "
                "environment variables."
            )

        # The connection client is configured once with the secure parameters.
        self.session = requests.Session()
        self.session.headers.update({"Authorization": f"ApiKey {api_key}"})

    def run(self, **kwargs):
        # The 'kwargs' from the LLM are validated against the restricted
        # FixedElasticsearchMemoryInput model. Any attempt to pass 'es_url'
        # would be ignored.
        try:
            inputs = self.model(**kwargs)
        except Exception as e:
            return f"Invalid input: {e}"

        # FIX: The request is always sent to the pre-configured 'self.es_url',
        # preventing the LLM from specifying a malicious destination.
        if inputs.query:
            print(f"Executing search on pre-configured host: {self.es_url}")
            # In a real implementation, a request would be made, e.g.:
            # response = self.session.post(f"{self.es_url}/_search", json=...)
            return f"Successfully queried memory for '{inputs.query}'."

        return "No operation was performed."