distinct_mcp/files.py

"""Reading and writing files inside one directory an operator nominated.

This is the dangerous server in the package, so the argument for it is set out
in full. Everything else here answers from memory; this one touches a
volunteer's disk, and a volunteer is a person who lent a machine to strangers.

**There is no default root and there never is one.** :func:`file_mcp_server`
takes a directory and refuses to build without one. It does not fall back to
the working directory, the home directory, a temporary directory or anything
else, because every one of those would turn "the operator forgot to configure
it" into "the model can read the operator's files". A refusal at install time
is a message an operator reads; a default is a capability nobody granted. The
filesystem root and the home directory itself are also refused: nominating
either is a configuration mistake rather than an intention, and the cost of
being wrong about that is the whole disk.

**Containment is decided on the resolved path, never on the joined one.** The
requested path is joined to the root and then passed through
``os.path.realpath``, and it is the *result* of that resolution that has to sit
inside the root. Doing it the other way round is the classic hole: a lexical
check on ``root/notes.txt`` passes, and then ``open`` follows the symbolic link
that ``notes.txt`` happens to be and reads ``/etc/shadow``. Resolution before
the check, not after, is the whole difference between the two.

**Symbolic links inside the root are refused rather than followed.** After the
containment check, every component below the root is examined with ``lstat``
and a link at any of them fails the call, even a link whose target is inside
the root. A link that resolves inside the root today can be repointed outside
it tomorrow, and this server would then be resolving an attacker's path on
every call. Refusing the whole shape is cheaper to reason about than deciding
which links are benign. The containment check is the load-bearing one and this
is defence in depth: on Windows a directory junction is not reported as a link
by ``S_ISLNK`` at all, and there the realpath comparison is what catches it.

**Traversal is refused by shape as well as by resolution.** A path containing
``..`` is rejected before anything touches the filesystem, along with absolute
paths, drive letters, UNC prefixes, backslashes, null bytes and the Windows
device names (``NUL`` and friends open a device rather than a file even when a
directory prefix is present). The resolution check would catch most of these on
its own. They are refused separately because the model gets a message naming
the actual problem, and because a check that never fires cannot be trusted to
be the only one.

**Both directions are capped, and the two caps are not symmetric.** A read that
runs past :data:`MAX_READ_BYTES` returns the leading bytes and says
``truncated``, because a partial answer plus an honest flag is more use than a
refusal and the model can tell the user the file was longer. A write that runs
past :data:`MAX_WRITE_BYTES` is refused outright, because truncating a write
means writing a corrupted file and calling it success. The caps are also sized
against the broker rather than chosen freely: ``JobPolicy`` defaults to 16,384
input bytes and 131,072 output bytes, a JSON string escapes a control character
to six bytes, and a schema that advertises a longer limit than the broker will
accept is a schema that lies to the model.

**Only UTF-8 text.** A file that is not valid UTF-8 is refused rather than
decoded with replacement characters. Replacement would put the mangled remains
of a binary file into a small model's context, where it is both useless and
expensive.

**What this server cannot enforce, stated plainly.**

* *It bounds where, not what.* If the operator nominates a directory holding
  their private keys, the model may read their private keys. Choosing the
  directory is the security decision and this code cannot make it.
* *There is a race it cannot close.* Between the resolution check and the
  ``open`` that follows it, anything with write access to the root can replace
  a component with a symbolic link. Closing that properly needs the whole path
  walked with ``O_NOFOLLOW`` and per-directory file descriptors, which POSIX
  supports and Windows does not, so this code does not pretend to. The window
  matters only to an attacker who can already write inside the nominated
  directory, and it is stated here rather than left for someone to discover.
* *File contents are untrusted input.* Text read from a file reaches the
  model's prompt exactly like an MCP tool description does, and a file
  containing instructions is a prompt injection. The harness should fence a
  file's contents the way it already fences tool results. Nothing in this
  module can do that for it.
* *The root is fixed when the server is built.* A deployment wanting a
  per-run directory has to build and install one server per run, which is an
  action its own operator-side code takes. Nothing a run says can move the
  root, which is the property worth keeping.
"""

from __future__ import annotations

import os
import stat
import tempfile
from collections.abc import Mapping
from pathlib import Path
from typing import Any

from distinct_tools.core import ToolInputError
from distinct_tools.mcp import InProcessMcpServer, McpError, McpServerManifest, McpToolDefinition

from .common import flag_argument, reject_unknown, structured

#: Longest read returned in one call. Sized against ``JobPolicy``'s default
#: 131,072 output bytes: the worst case for JSON escaping is six bytes per
#: source byte, so 16,384 bytes of pathological text still encodes inside the
#: cap, and ordinary text costs about one byte per byte.
MAX_READ_BYTES = 16_384

#: Longest write accepted in one call. Sized against ``JobPolicy``'s default
#: 16,384 input bytes, which the arguments must fit *after* JSON encoding, so
#: this sits at half of it. Ordinary text at this size encodes to roughly 8,300
#: bytes and fits; text dense in control characters does not, and the broker
#: refuses it as invalid input rather than writing a truncated file.
MAX_WRITE_BYTES = 8_192

#: Entries returned by one listing.
MAX_ENTRIES = 200

#: Files the root may hold before a *new* file is refused. A model in a loop
#: must not be able to fill a volunteer's disk one small file at a time.
MAX_FILES = 500

MAX_PATH_CHARACTERS = 200
MAX_PATH_COMPONENTS = 8

#: Names that open a device rather than a file on Windows, whatever directory
#: precedes them. Refused on every platform so behaviour does not depend on
#: where the worker happens to be running.
_DEVICE_NAMES = frozenset(
    {"con", "prn", "aux", "nul"}
    | {f"com{index}" for index in range(1, 10)}
    | {f"lpt{index}" for index in range(1, 10)}
)

#: The manifest an operator would approve to install this server. No hosts: a
#: file server that could also reach the network is an exfiltration tool, and
#: these two capabilities should never arrive in one approval.
FILE_MANIFEST = McpServerManifest(slug="files", version="1", timeout_seconds=10.0)


class FileVault:
    """One operator-nominated directory, and every check that guards it.

    Separate from the handlers so the path rules can be tested directly rather
    than only through a broker, and so a caller can hold one vault while
    publishing it through more than one surface.
    """

    def __init__(
        self,
        root: object,
        *,
        max_read_bytes: int = MAX_READ_BYTES,
        max_write_bytes: int = MAX_WRITE_BYTES,
        max_files: int = MAX_FILES,
    ) -> None:
        self.root = _resolve_root(root)
        self.max_read_bytes = _bound("max_read_bytes", max_read_bytes, MAX_READ_BYTES)
        self.max_write_bytes = _bound("max_write_bytes", max_write_bytes, MAX_WRITE_BYTES)
        self.max_files = _bound("max_files", max_files, MAX_FILES)

    # -- path rules -------------------------------------------------------

    def resolve(self, raw: object, *, allow_root: bool = False) -> Path:
        """Return the real path a request may touch, or refuse and say why."""

        relative = _relative_parts(raw, allow_empty=allow_root)
        joined = self.root.joinpath(*relative)

        # Resolution first. ``realpath`` follows every symbolic link it finds,
        # including ones in the middle of the path, and tolerates a final
        # component that does not exist yet, which is what a write needs.
        try:
            resolved = Path(os.path.realpath(joined))
        except OSError as exc:  # pragma: no cover - platform dependent
            raise ToolInputError(f"path could not be resolved: {type(exc).__name__}") from exc

        # Containment second, on the resolved value. This is the check that
        # actually decides the question.
        if resolved != self.root and self.root not in resolved.parents:
            raise ToolInputError(
                "path leaves the folder this worker was given; only paths inside it can be used"
            )

        # Defence in depth third. See the module docstring for why a link is
        # refused even when it points back inside the root.
        current = self.root
        for part in relative:
            current = current / part
            try:
                info = current.lstat()
            except FileNotFoundError:
                break
            except OSError as exc:
                raise ToolInputError(f"path could not be read: {type(exc).__name__}") from exc
            if stat.S_ISLNK(info.st_mode):
                raise ToolInputError(
                    "this path goes through a shortcut (a symbolic link) and this tool does "
                    "not follow shortcuts"
                )
        return resolved

    def existing_file(self, raw: object) -> Path:
        path = self.resolve(raw)
        if not path.exists():
            raise ToolInputError(
                "there is no file at that path; call list_files to see what is there"
            )
        if not path.is_file():
            raise ToolInputError("that path is a folder, not a file")
        return path

    def existing_directory(self, raw: object) -> Path:
        path = self.resolve(raw, allow_root=True)
        if not path.is_dir():
            raise ToolInputError("there is no folder at that path")
        return path

    def relative_name(self, path: Path) -> str:
        """The path as the model should see it: relative, forward slashes."""

        return path.relative_to(self.root).as_posix() or "."

    def count_files(self) -> int:
        """Count files under the root, stopping once the cap is passed.

        ``os.walk`` does not follow directory symbolic links, so a link cannot
        inflate or deflate this count by pointing somewhere else.
        """

        total = 0
        for _, _, filenames in os.walk(self.root):
            total += len(filenames)
            if total > self.max_files:
                return total
        return total


def _bound(name: str, value: object, ceiling: int) -> int:
    if type(value) is not int or not 1 <= value <= ceiling:
        raise McpError(f"{name} must be an integer between 1 and {ceiling}")
    return value


def _resolve_root(value: object) -> Path:
    """Resolve and vet the operator's directory, or refuse to build at all."""

    if not isinstance(value, str | os.PathLike) or not str(value).strip():
        raise McpError(
            "a file server needs a directory nominated by the operator; there is no default, "
            "because a default would be a capability nobody granted"
        )
    try:
        resolved = Path(os.path.realpath(Path(value).expanduser()))
    except (OSError, RuntimeError) as exc:
        raise McpError(f"the nominated directory could not be resolved: {type(exc).__name__}") from exc
    if not resolved.is_dir():
        raise McpError(f"the nominated directory does not exist or is not a directory: {resolved}")
    if resolved == Path(resolved.anchor):
        raise McpError("the filesystem root is not a workspace; nominate a specific directory")
    try:
        home = Path(os.path.realpath(Path.home()))
    except (OSError, RuntimeError):  # pragma: no cover - no home directory
        home = None
    if home is not None and resolved == home:
        raise McpError(
            "the home directory itself is too broad; nominate a specific directory inside it"
        )
    return resolved


def _relative_parts(raw: object, *, allow_empty: bool) -> tuple[str, ...]:
    """Turn a requested path into safe components, or refuse with the reason."""

    if raw is None and allow_empty:
        return ()
    if not isinstance(raw, str):
        raise ToolInputError("path must be a string")
    text = raw.strip()
    if not text:
        if allow_empty:
            return ()
        raise ToolInputError("path must be a non-empty string")
    if len(text) > MAX_PATH_CHARACTERS:
        raise ToolInputError(f"path must be at most {MAX_PATH_CHARACTERS} characters")
    if "\x00" in text:
        raise ToolInputError("path must not contain a null byte")
    if "\\" in text:
        raise ToolInputError("path must use forward slashes")
    candidate = Path(text)
    if candidate.is_absolute() or candidate.anchor:
        raise ToolInputError("path must be relative to the folder this worker was given")
    parts = tuple(part for part in candidate.parts if part not in (".",))
    if any(part == ".." for part in parts):
        raise ToolInputError("path must not contain '..'; it cannot leave the folder")
    if len(parts) > MAX_PATH_COMPONENTS:
        raise ToolInputError(f"path may be at most {MAX_PATH_COMPONENTS} folders deep")
    for part in parts:
        if part.split(".")[0].lower() in _DEVICE_NAMES:
            raise ToolInputError(f"'{part}' is a reserved device name and cannot be used")
    if not parts and not allow_empty:
        raise ToolInputError("path must name a file")
    return parts


def _trim_partial_utf8(data: bytes) -> bytes:
    """Drop a multi-byte character the read cap cut in half.

    Without this, truncating at a byte boundary inside a character would make
    an ordinary long file look like a file that is not text at all.
    """

    index = len(data) - 1
    steps = 0
    while index >= 0 and steps < 4:
        byte = data[index]
        if byte & 0b1100_0000 == 0b1000_0000:
            index -= 1
            steps += 1
            continue
        if byte & 0b1000_0000 == 0:
            return data
        if byte & 0b1111_0000 == 0b1111_0000:
            width = 4
        elif byte & 0b1110_0000 == 0b1110_0000:
            width = 3
        else:
            width = 2
        return data if index + width <= len(data) else data[:index]
    return data


def _read_handler(vault: FileVault):
    def read_file(arguments: Mapping[str, Any]) -> dict[str, Any]:
        reject_unknown(arguments, frozenset({"path"}))
        path = vault.existing_file(arguments.get("path"))
        try:
            size = path.stat().st_size
            with open(path, "rb") as handle:
                data = handle.read(vault.max_read_bytes)
        except OSError as exc:
            raise ToolInputError(f"the file could not be read: {type(exc).__name__}") from exc
        truncated = size > len(data)
        if truncated:
            data = _trim_partial_utf8(data)
        try:
            text = data.decode("utf-8")
        except UnicodeDecodeError as exc:
            raise ToolInputError(
                "that file is not UTF-8 text, and this tool reads text files only"
            ) from exc
        return structured(
            {
                "path": vault.relative_name(path),
                "text": text,
                "bytes_returned": len(data),
                "file_bytes": size,
                "truncated": truncated,
            }
        )

    return read_file


def _write_handler(vault: FileVault):
    def write_file(arguments: Mapping[str, Any]) -> dict[str, Any]:
        reject_unknown(arguments, frozenset({"path", "text", "overwrite"}))
        # Checked here rather than through the shared helper because an empty
        # string is a legitimate thing to write and a missing argument is not,
        # and the helper cannot tell those two apart.
        text = arguments.get("text")
        if not isinstance(text, str):
            raise ToolInputError("text must be a string; pass an empty string for an empty file")
        if len(text) > MAX_WRITE_BYTES:
            raise ToolInputError(f"text must be at most {MAX_WRITE_BYTES:,} characters")
        # Measured in bytes as well, because the cap exists to bound what lands
        # on a volunteer's disk and one character is up to four bytes of it.
        payload = text.encode("utf-8")
        if len(payload) > vault.max_write_bytes:
            raise ToolInputError(
                f"text must be at most {vault.max_write_bytes:,} bytes when encoded; "
                "write a shorter file rather than a truncated one"
            )
        overwrite = flag_argument(arguments, "overwrite")
        path = vault.resolve(arguments.get("path"))
        existed = path.exists()
        if existed:
            if not path.is_file():
                raise ToolInputError("that path is a folder, not a file")
            if not overwrite:
                raise ToolInputError(
                    "a file already exists at that path. Call again with overwrite set to true "
                    "if you really mean to replace it, or choose another name"
                )
        elif not path.parent.is_dir():
            # No directory creation on purpose: the operator nominated one
            # folder, and a model building a tree inside it is more surface
            # for no benefit anybody asked for.
            raise ToolInputError(
                "the folder for that path does not exist, and this tool does not create folders"
            )
        elif vault.count_files() >= vault.max_files:
            raise ToolInputError(
                f"the folder already holds {vault.max_files:,} files, which is the limit"
            )

        # Written to a temporary file in the same folder and moved into place,
        # so a failure halfway through leaves the previous contents intact
        # rather than an empty file where the user's work used to be.
        handle = None
        try:
            handle = tempfile.NamedTemporaryFile(
                dir=path.parent, prefix=".write-", suffix=".tmp", delete=False
            )
            with handle:
                handle.write(payload)
                handle.flush()
                os.fsync(handle.fileno())
            os.replace(handle.name, path)
        except OSError as exc:
            if handle is not None:
                try:
                    os.unlink(handle.name)
                except OSError:
                    pass
            raise ToolInputError(f"the file could not be written: {type(exc).__name__}") from exc
        return structured(
            {
                "path": vault.relative_name(path),
                "bytes_written": len(payload),
                "replaced": existed,
            }
        )

    return write_file


def _list_handler(vault: FileVault):
    def list_files(arguments: Mapping[str, Any]) -> dict[str, Any]:
        reject_unknown(arguments, frozenset({"path"}))
        directory = vault.existing_directory(arguments.get("path"))
        entries: list[dict[str, Any]] = []
        try:
            with os.scandir(directory) as scan:
                for item in scan:
                    if item.is_symlink():
                        kind, size = "link", None
                    elif item.is_dir(follow_symlinks=False):
                        kind, size = "folder", None
                    elif item.is_file(follow_symlinks=False):
                        kind = "file"
                        try:
                            size = item.stat(follow_symlinks=False).st_size
                        except OSError:
                            size = None
                    else:
                        kind, size = "other", None
                    entries.append({"name": item.name, "kind": kind, "bytes": size})
        except OSError as exc:
            raise ToolInputError(f"the folder could not be listed: {type(exc).__name__}") from exc
        entries.sort(key=lambda entry: str(entry["name"]))
        total = len(entries)
        return structured(
            {
                "path": vault.relative_name(directory),
                "entries": entries[:MAX_ENTRIES],
                "count": min(total, MAX_ENTRIES),
                "total": total,
                "truncated": total > MAX_ENTRIES,
                "empty": total == 0,
            }
        )

    return list_files


def _definitions(vault: FileVault) -> tuple[tuple[McpToolDefinition, Any], ...]:
    read_limit = f"{vault.max_read_bytes:,}"
    write_limit = f"{vault.max_write_bytes:,}"
    return (
        (
            McpToolDefinition(
                name="read_file",
                description=(
                    "Read a text file from the one folder this worker has been given and return "
                    "what it says. Give the path relative to that folder, such as notes.txt or "
                    "reports/march.txt. Use it when you have been asked about a file by name, or "
                    f"after list_files has told you what is there. It returns at most {read_limit} "
                    "bytes and tells you when a file was longer, and it reads plain text only. It "
                    "cannot reach anything outside that folder and does not follow shortcuts out "
                    "of it. Do not call it twice for the same path: the answer will be the same."
                ),
                input_schema={
                    "type": "object",
                    "additionalProperties": False,
                    "required": ["path"],
                    "properties": {
                        "path": {
                            "type": "string",
                            "maxLength": MAX_PATH_CHARACTERS,
                            "description": "path relative to the given folder, using / between names",
                        }
                    },
                },
            ),
            _read_handler(vault),
        ),
        (
            McpToolDefinition(
                name="write_file",
                description=(
                    "Save text to a file in the one folder this worker has been given. Use it when "
                    "the user asked for a file to be produced or changed, not to keep working "
                    f"notes. It writes at most {write_limit} bytes. It will not replace a file that "
                    "already exists unless you pass overwrite as true, so try once without it "
                    "first. It cannot create folders and cannot write anywhere outside the given "
                    "folder."
                ),
                input_schema={
                    "type": "object",
                    "additionalProperties": False,
                    "required": ["path", "text"],
                    "properties": {
                        "path": {
                            "type": "string",
                            "maxLength": MAX_PATH_CHARACTERS,
                            "description": "path relative to the given folder, using / between names",
                        },
                        "text": {
                            "type": "string",
                            "maxLength": MAX_WRITE_BYTES,
                            "description": "the exact text to save",
                        },
                        "overwrite": {
                            "type": "boolean",
                            "description": "true to replace a file that already exists",
                        },
                    },
                },
            ),
            _write_handler(vault),
        ),
        (
            McpToolDefinition(
                name="list_files",
                description=(
                    "List what is in the one folder this worker has been given, or in a folder "
                    "inside it. Use it first when you have been asked about the files and do not "
                    "already know their exact names. It returns names and sizes only: call "
                    "read_file to see what a file says. Do not call it repeatedly for the same "
                    f"folder. It returns at most {MAX_ENTRIES} entries."
                ),
                input_schema={
                    "type": "object",
                    "additionalProperties": False,
                    "properties": {
                        "path": {
                            "type": "string",
                            "maxLength": MAX_PATH_CHARACTERS,
                            "description": "a folder inside the given one; leave it out for the top level",
                        }
                    },
                },
            ),
            _list_handler(vault),
        ),
    )


def file_mcp_server(
    root: object,
    *,
    max_read_bytes: int = MAX_READ_BYTES,
    max_write_bytes: int = MAX_WRITE_BYTES,
    max_files: int = MAX_FILES,
) -> InProcessMcpServer:
    """An in-process MCP server bounded to ``root`` and nothing else.

    ``root`` is the operator's decision and is required. Install the result
    with ``install_mcp_server(registry, file_mcp_server(path), FILE_MANIFEST)``,
    which produces ``files.read_file@1``, ``files.write_file@1`` and
    ``files.list_files@1``. Installing is still not approving: the operator's
    :class:`~distinct_tools.approval.OperatorPolicy` decides whether any of the
    three is ever advertised or callable.
    """

    vault = FileVault(
        root,
        max_read_bytes=max_read_bytes,
        max_write_bytes=max_write_bytes,
        max_files=max_files,
    )
    server = InProcessMcpServer()
    for definition, handler in _definitions(vault):
        server.add_tool(definition, handler)
    return server


__all__ = [
    "FILE_MANIFEST",
    "MAX_ENTRIES",
    "MAX_FILES",
    "MAX_PATH_CHARACTERS",
    "MAX_READ_BYTES",
    "MAX_WRITE_BYTES",
    "FileVault",
    "file_mcp_server",
]