beaver-agent/policy.py

"""§3.7: правила PreToolUse - граница без permission-промптов.

Модель везде в `bypassPermissions`; что ей можно, решают маунты,
`disallowed_tools` и правила отсюда. Правило получает `ToolCall` и
возвращает `Deny` с причиной, которую модель читает как результат тулзы,
или `None`. Маунт `/vault:ro` - первая линия, хук - вторая: он ловит то,
что маунт не различает (зоны внутри одного тома, порядок «скилл - тулза»).
"""

from __future__ import annotations

import fnmatch
import shlex
from dataclasses import dataclass
from typing import TYPE_CHECKING

from beaver_gateway.core.policy import Deny, PolicyRule, ToolCall

if TYPE_CHECKING:
    from collections.abc import Iterator
    from pathlib import Path

__all__ = [
    "DEEP_DISALLOWED",
    "DISPATCHER_DISALLOWED",
    "DISTILLER_DISALLOWED",
    "TRIAGE_DISALLOWED",
    "Zones",
    "bash_zones",
    "requires_skill",
    "skill_tracker",
    "vault_zones",
]

# §3.7: у диспетчера AskUserQuestion остаётся, планов и ноутбуков нет;
# у глубоких вопросы текстом, сабагентов нет; дистиллятор и триаж - без веба.
DISPATCHER_DISALLOWED = ("ExitPlanMode", "EnterPlanMode", "NotebookEdit")
DEEP_DISALLOWED = (
    "AskUserQuestion",
    "ExitPlanMode",
    "EnterPlanMode",
    "NotebookEdit",
    "Task",
)
DISTILLER_DISALLOWED = ("AskUserQuestion", "Task", "WebSearch", "WebFetch")
TRIAGE_DISALLOWED = DISTILLER_DISALLOWED

FILE_TOOLS = ("Write", "Edit", "MultiEdit", "NotebookEdit")
MUTATING = frozenset(
    {
        "rm",
        "rmdir",
        "unlink",
        "mv",
        "cp",
        "tee",
        "truncate",
        "touch",
        "mkdir",
        "ln",
        "install",
        "rsync",
        "dd",
        "chmod",
        "chown",
        "shred",
    }
)
COPYING = frozenset({"cp", "install", "rsync"})
INPLACE = frozenset({"sed", "perl"})
REDIRECTS = frozenset({">", ">>", ">|", "&>", "&>>", ">&"})
SEPARATORS = frozenset({";", "&&", "||", "|", "&", "(", ")"})


@dataclass(frozen=True, slots=True)
class Zones:
    """Зоны vault: ``write`` - полная запись, ``create`` - только новые файлы."""

    vault: Path
    write: tuple[Path, ...]
    create: tuple[Path, ...]

    def verdict(self, path: Path, *, creating: bool) -> Deny | None:
        if not _under(path, self.vault):
            return None
        if any(_under(path, zone) for zone in self.write):
            return None
        rel = path.relative_to(self.vault)
        for zone in self.create:
            if _under(path, zone):
                if creating and not path.exists():
                    return None
                return Deny(
                    reason=f"{rel}: в «{zone.relative_to(self.vault)}» можно только "
                    "создавать новые файлы, существующие не трогаем"
                )
        zones = ", ".join(f"«{z.relative_to(self.vault)}»" for z in self.write)
        return Deny(reason=f"{rel}: vault только на чтение; писать можно в {zones}")


def _under(path: Path, root: Path) -> bool:
    return path == root or root in path.parents


def vault_zones(zones: Zones) -> PolicyRule:
    """Write/Edit/NotebookEdit только в зонах; Write в create-зону - новый файл."""

    def rule(call: ToolCall) -> Deny | None:
        if call.tool not in FILE_TOOLS:
            return None
        path = call.path()
        if path is None:
            return None
        return zones.verdict(path, creating=call.tool == "Write")

    return rule


def bash_zones(zones: Zones) -> PolicyRule:
    """`rm`/`mv`/`cp`/`tee`/`sed -i`/редирект с путём в vault вне зон - отказ."""

    def rule(call: ToolCall) -> Deny | None:
        if call.tool != "Bash":
            return None
        command = call.input.get("command")
        if not isinstance(command, str):
            return None
        for target, creating in _bash_targets(command):
            deny = zones.verdict(call.resolve(target), creating=creating)
            if deny is not None:
                return Deny(reason=f"Bash: {deny.reason}")
        return None

    return rule


def _bash_targets(command: str) -> Iterator[tuple[str, bool]]:
    for segment in _segments(command):
        plain: list[str] = []
        i = 0
        while i < len(segment):
            word = segment[i]
            if word in REDIRECTS:
                if i + 1 < len(segment) and _pathlike(segment[i + 1]):
                    yield segment[i + 1], True
                i += 2
                continue
            plain.append(word)
            i += 1
        for j, word in enumerate(plain):
            rest = plain[j + 1 :]
            inplace = word in INPLACE and any(a.startswith("-i") for a in rest)
            if word not in MUTATING and not inplace:
                continue
            args = [a for a in rest if _pathlike(a)]
            if word in COPYING:
                args = args[-1:]
            for arg in args:
                yield arg, False
            break


def _segments(command: str) -> Iterator[list[str]]:
    lexer = shlex.shlex(command, posix=True, punctuation_chars=True)
    lexer.whitespace_split = True
    try:
        tokens = list(lexer)
    except ValueError:
        tokens = command.split()
    segment: list[str] = []
    for token in tokens:
        if token in SEPARATORS:
            if segment:
                yield segment
            segment = []
        else:
            segment.append(token)
    if segment:
        yield segment


def _pathlike(word: str) -> bool:
    return bool(word) and not (
        word.startswith(("-", "&", "$", "/dev/"))
        or word.isdigit()
        or ("=" in word and "/" not in word.split("=", 1)[0])
    )


def skill_tracker() -> PolicyRule:
    """Запоминает открытые скиллы в состоянии сессии (для `requires_skill`)."""

    def rule(call: ToolCall) -> None:
        if call.tool != "Skill":
            return
        name = str(call.input.get("skill") or call.input.get("name") or "")
        opened: set[str] = call.state.setdefault("skills", set())
        opened.add(name.rsplit(":", 1)[-1])

    return rule


def requires_skill(skill: str, tools: tuple[str, ...]) -> PolicyRule:
    """§4.3: жёсткий протокол - тулзы из ``tools`` только после `Skill(skill)`."""

    def rule(call: ToolCall) -> Deny | None:
        if not any(fnmatch.fnmatchcase(call.tool, pattern) for pattern in tools):
            return None
        if skill in call.state.get("skills", ()):
            return None
        return Deny(
            reason=f"сначала открой скилл «{skill}» (тулза Skill), потом {call.tool}"
        )

    return rule