check_run_script.py

import logging

from typing import Any

import litellm

from litellm import Choices
from litellm import completion
from litellm import get_supported_openai_params
from litellm import ModelResponse
from pydantic import BaseModel

from linux_mcp_server.config import CONFIG
from linux_mcp_server.config import ReasoningEffort
from linux_mcp_server.utils import StrEnum


# LiteLLM has bugs where the custom_llm_provider is not properly
# passed back into query functions for things like "model supports reasoning".
# Without this, each such failure spits out a debug message to the terminal
# https://github.com/BerriAI/litellm/issues/23879
litellm.suppress_debug_info = True


logger = logging.getLogger("linux-mcp-server")


def get_model() -> str:
    if CONFIG.gatekeeper.model:
        return CONFIG.gatekeeper.model
    else:
        raise ValueError("To use run_script tools, you must set LINUX_MCP_GATEKEEPER__MODEL")


READONLY_INSTRUCTION = """
 * The script does not modify the system in any way, and only does read-only operations.
"""

READONLY_RESULT = """
 * `MODIFIES_SYSTEM`: The script modifies the system in some way.
"""

READONLY_USE_CASES = "get the status of a system, or diagnose a problem"
READWRITE_USE_CASES = "modify system settings, or fix a problem"

PROMPT = """\
The user is using a language model connected to the linux-mcp-server
Model Context Protocol server, and are using it to do one or more of the
following: {use_cases}.

You are a gatekeeper that checks if commands and scripts generated by the
language model are correct and correspond to the user's intent.

# General policies

 * Software can only be installed from repositories that are already
   configured on the system, such as apt or DNF repositories.
   No new repositories may be added.
 * Except for installing software from pre-configured repositories,
   nothing may be downloaded from the internet.

# Your Task

In response to the user's prompt, the language model has generated the
following script of type {script_type}.

<START_OF_SCRIPT>
{script}
<END_OF_SCRIPT>

and has described what the script does with the following text:

<START_OF_DESCRIPTION>
{description}
<END_OF_DESCRIPTION>

Your job is to check that:

{readonly_instruction}* The script does what is described, and does not do anything unrelated.
* It conforms to the general policies outlined above.
* The script is clear and simple and written in an obvious, expected
  manner.
* The script does not contain malicious code or introduce security vulnerabilities
  on the system.

It is *OK* if the script does not consider every possible failure mode, as long
as it fails cleanly. The model may have previously investigated the system
and found out things like the operating system version and configuration
file paths.

The description of the script may reference querying or modifying a specific host;
this is *OK* - the MCP server will take care of executing the script on the
appropriate host. The script does not need to, for example, use SSH to connect to
the target host.

Return a JSON object with the following keys, `status` (required) and `detail` (optional).

status should be one of the following:

* `OK`: The script does what is described and safe to run.
* `BAD_DESCRIPTION`: The script does not do what is described.
* `POLICY`: The script violates one or more of the policies listed under "General Policies" above.
{readonly_result}* `UNCLEAR`: The script is unclear or difficult to understand.
* `DANGEROUS`: The script is likely to cause damage to the system
* `MALICIOUS`: The script contains malicious code or would introduce security vulnerabilities.

If the status is not OK, the detail should be short 1-3 sentence description of
what is wrong with the script. Be specific to allow the language model to correct
the problem.

If the script seems buggy but does not fall into any of the categories above, return
a status of `OK`.

Example response (respond with only the JSON object - do not wrap it in a code block).
{{ "status": "DANGEROUS", "detail": "The critical example file `/etc/example.conf` will be deleted if an error occurs in the first call to sed." }}
"""


class GatekeeperStatus(StrEnum):
    OK = "OK"
    BAD_DESCRIPTION = "BAD_DESCRIPTION"
    POLICY = "POLICY"
    MODIFIES_SYSTEM = "MODIFIES_SYSTEM"
    UNCLEAR = "UNCLEAR"
    DANGEROUS = "DANGEROUS"
    MALICIOUS = "MALICIOUS"


class GatekeeperResult(BaseModel):
    status: GatekeeperStatus
    detail: str = ""

    @property
    def description(self) -> str:
        match self.status:
            case GatekeeperStatus.OK:
                return "OK"
            case GatekeeperStatus.BAD_DESCRIPTION:
                return f"Bad description: {self.detail}"
            case GatekeeperStatus.POLICY:
                return f"Policy violation: {self.detail}"
            case GatekeeperStatus.MODIFIES_SYSTEM:
                return f"Script modifies the system and readonly is true: {self.detail}"
            case GatekeeperStatus.UNCLEAR:
                return f"Unclear script: {self.detail}"
            case GatekeeperStatus.DANGEROUS:
                return f"Dangerous script: {self.detail}"
            case GatekeeperStatus.MALICIOUS:
                # We don't provide detail here to make it harder for a malicious model
                # to figure out workarounds
                return "Possibly malicious script: not allowed"
            case _:  # pragma: no cover
                raise ValueError(f"Unknown status: {self.status}")

    @classmethod
    def parse_from_description(cls, description: str) -> "GatekeeperResult":
        """Parse a GatekeeperResult from a description string.

        Args:
            description: A description string like "OK" or "Policy violation: details..."

        Returns:
            A GatekeeperResult with the parsed status and detail

        Raises:
            ValueError: If the description prefix doesn't match any known status
        """
        _prefix_to_status = {
            "OK": GatekeeperStatus.OK,
            "Bad description": GatekeeperStatus.BAD_DESCRIPTION,
            "Policy violation": GatekeeperStatus.POLICY,
            "Script modifies the system and readonly is true": GatekeeperStatus.MODIFIES_SYSTEM,
            "Unclear script": GatekeeperStatus.UNCLEAR,
            "Dangerous script": GatekeeperStatus.DANGEROUS,
            "Possibly malicious script": GatekeeperStatus.MALICIOUS,
        }

        prefix = description.split(":")[0]
        status = _prefix_to_status.get(prefix, None)
        if status is None:
            raise ValueError(f"Unknown description prefix: {description}")

        if status == GatekeeperStatus.OK:
            return cls(status=status, detail="")
        else:
            detail = description[len(prefix) :].lstrip(": ").strip()
            return cls(status=status, detail=detail)


def _build_completion_kwargs():
    extra_kwargs: dict[str, Any] = {}
    model = get_model()

    structured_output = CONFIG.gatekeeper.structured_output
    if structured_output is None:
        params = get_supported_openai_params(model=model)
        structured_output = params is not None and "response_format" in params

    if structured_output:
        extra_kwargs["response_format"] = GatekeeperResult

    reasoning_effort = CONFIG.gatekeeper.reasoning_effort
    if reasoning_effort is not None:
        if model.startswith("openrouter/"):
            if reasoning_effort == ReasoningEffort.NONE:
                extra_kwargs["reasoning"] = {"enabled": False}
            else:
                extra_kwargs["reasoning"] = {"enabled": True, "effort": reasoning_effort.value}
        else:
            extra_kwargs["reasoning_effort"] = reasoning_effort.value

    if model.startswith("openrouter/"):
        provider: dict[str, Any] = {
            "require_parameters": True,
        }
        extra_kwargs["provider"] = provider
        if CONFIG.gatekeeper.quantization:
            provider["quantizations"] = [CONFIG.gatekeeper.quantization]

    if CONFIG.gatekeeper.template_kwargs:
        extra_kwargs["chat_template_kwargs"] = CONFIG.gatekeeper.template_kwargs

    return extra_kwargs


def check_run_script(description: str, script_type: str, script: str, *, readonly: bool) -> GatekeeperResult:
    # Check that the script does what is described
    if "start_of_script" in script.lower() or "end_of_script" in script.lower():
        return GatekeeperResult(
            status=GatekeeperStatus.MALICIOUS, detail="Script contains 'start_of_script' or 'end_of_script'"
        )

    if "start_of_description" in script.lower() or "end_of_description" in script.lower():
        return GatekeeperResult(
            status=GatekeeperStatus.MALICIOUS, detail="Script contains 'start_of_description' or 'end_of_description'"
        )

    prompt = PROMPT.format(
        script_type=script_type,
        script=script,
        description=description,
        use_cases=READONLY_USE_CASES if readonly else READWRITE_USE_CASES,
        readonly_instruction=READONLY_INSTRUCTION if readonly else "",
        readonly_result=READONLY_RESULT if readonly else "",
    )

    messages = [{"role": "user", "content": prompt}]

    extra_kwargs = _build_completion_kwargs()

    response = completion(
        model=get_model(),
        messages=messages,
        temperature=CONFIG.gatekeeper.temperature,
        **extra_kwargs,
    )
    assert isinstance(response, ModelResponse)
    assert isinstance(response.choices[0], Choices)
    response_text = (response.choices[0].message.content or "").strip()

    logger.info(f"Gatekeeper response: {response_text}")

    return GatekeeperResult.model_validate_json(response_text)