Files
shirobana-bot/bot-utility/skill_registry.py
T
RRRosmontis 34089d0042 Initial import of Shirobana bot backend
Include utility modules, shared channels, prompts, read-only skills, wardrobe memory and regression tests. Externalize legacy credentials and exclude private runtime state.
2026-11-11 00:25:30 +08:00

374 lines
15 KiB
Python

"""Read-only Agent Skills, with pinned, immutable per-turn snapshots.
No imports of the bot, script execution, network access, or writes. Only explicitly
allowlisted directories are read. ``allowed-tools`` is descriptive metadata, not
authorization. Reference keys are POSIX paths such as ``references/notes.md``.
All symlinks are rejected, including config symlinks. A custom config must be
inside root.parent. Missing/empty configuration is an enabled, empty registry.
"""
from __future__ import annotations
import hashlib
import json
import logging
import os
import re
import stat
import threading
from dataclasses import dataclass, field
from pathlib import Path, PurePosixPath
from types import MappingProxyType
from typing import Mapping
import yaml
_LOG = logging.getLogger(__name__)
_BASE = Path(__file__).absolute().parent
_SCENES = ("private_chat", "group_chat", "proactive", "help", "notification",
"reminder", "tool_result")
_DEFAULT_LIMITS = {"body_bytes": 32768, "reference_bytes": 16384,
"max_skills": 3, "max_rounds": 2, "max_tokens": 12000}
_LIMIT_RANGES = {"body_bytes": (1, 131072), "reference_bytes": (1, 65536),
"max_skills": (1, 10), "max_rounds": (1, 8),
"max_tokens": (1, 64000)}
_NAME = re.compile(r"[a-z0-9]+(?:-[a-z0-9]+)*\Z")
_MAX_CONFIG = 65536
_MAX_HEADER = 16384
_MAX_REFERENCES = 256
_MAX_REFERENCE_TOTAL = 8 * 1024 * 1024
_MAX_ALLOWLIST = 128
class SkillError(ValueError):
"""A safe public error; code/message never contain filesystem paths/content."""
def __init__(self, code: str):
self.code = code
super().__init__({
"disabled": "Skills are disabled.",
"unknown_skill": "Skill is unavailable.",
"unknown_reference": "Reference is unavailable.",
"invalid_path": "Invalid skill reference path.",
"invalid_config": "Invalid skills configuration.",
"invalid_skill": "Invalid skill document.",
"unsafe_path": "Unsafe skill filesystem path.",
"invalid_document": "Invalid or oversized UTF-8 document.",
}.get(code, "Skill operation failed."))
class _UniqueLoader(yaml.SafeLoader):
pass
def _mapping(loader, node, deep=False):
result = {}
for key_node, value_node in node.value:
key = loader.construct_object(key_node, deep=deep)
if not isinstance(key, str) or key in result:
raise SkillError("invalid_config")
result[key] = loader.construct_object(value_node, deep=deep)
return result
_UniqueLoader.add_constructor(yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG, _mapping)
def _yaml(text):
try:
# Aliases are unnecessary for these small schemas; forbid cycles/bombs.
if any(isinstance(token, (yaml.tokens.AliasToken, yaml.tokens.AnchorToken))
for token in yaml.scan(text)):
raise SkillError("invalid_config")
return yaml.load(text, Loader=_UniqueLoader)
except (yaml.YAMLError, RecursionError, ValueError):
raise SkillError("invalid_config") from None
def _valid_name(name):
return isinstance(name, str) and 1 <= len(name) <= 64 and _NAME.fullmatch(name)
def _names(value):
if not isinstance(value, list) or any(not _valid_name(n) for n in value):
raise SkillError("invalid_config")
if len(value) != len(set(value)) or len(value) > _MAX_ALLOWLIST:
raise SkillError("invalid_config")
return tuple(value)
def _reference_path(path):
if not isinstance(path, str) or not path or "\\" in path or "\x00" in path:
raise SkillError("invalid_path")
parts = path.split("/")
if (len(parts) < 2 or parts[0] != "references"
or any(p in ("", ".", "..") for p in parts)
or PurePosixPath(path).suffix not in (".md", ".txt")):
raise SkillError("invalid_path")
return path
def _open_nofollow(path, directory=False):
"""Open each absolute component via dir_fd to prevent symlink TOCTOU escapes."""
path = Path(os.path.abspath(path))
fd = os.open(path.anchor, os.O_RDONLY | os.O_DIRECTORY)
try:
for index, part in enumerate(path.parts[1:]):
flags = os.O_RDONLY | os.O_NOFOLLOW | os.O_NONBLOCK
if index < len(path.parts) - 2 or directory:
flags |= os.O_DIRECTORY
new_fd = os.open(part, flags, dir_fd=fd)
os.close(fd)
fd = new_fd
result, fd = fd, None
return result
finally:
if fd is not None:
os.close(fd)
def _text(path, limit):
try:
fd = _open_nofollow(path)
with os.fdopen(fd, "rb") as stream:
if not stat.S_ISREG(os.fstat(stream.fileno()).st_mode):
raise SkillError("unsafe_path")
data = stream.read(limit + 1)
if len(data) > limit:
raise SkillError("invalid_document")
text = data.decode("utf-8")
if any((ord(c) < 32 and c not in "\t\n\r") or ord(c) == 127 for c in text):
raise SkillError("invalid_document")
return text
except UnicodeError:
raise SkillError("invalid_document") from None
except OSError:
raise SkillError("unsafe_path") from None
@dataclass(frozen=True)
class SkillRecord:
name: str
description: str
body: str
references: Mapping[str, str] = field(default_factory=dict)
metadata: Mapping[str, str] = field(default_factory=dict)
allowed_tools: tuple[str, ...] = ()
def __post_init__(self):
object.__setattr__(self, "references", MappingProxyType(dict(self.references)))
object.__setattr__(self, "metadata", MappingProxyType(dict(self.metadata)))
object.__setattr__(self, "allowed_tools", tuple(self.allowed_tools))
@dataclass(frozen=True)
class SkillSnapshot:
enabled: bool
records: Mapping[str, SkillRecord]
limits: Mapping[str, int]
always: tuple[str, ...]
scenes: Mapping[str, tuple[str, ...]]
fingerprint: str = field(init=False)
def __post_init__(self):
object.__setattr__(self, "records", MappingProxyType(dict(self.records)))
object.__setattr__(self, "limits", MappingProxyType(dict(self.limits)))
object.__setattr__(self, "always", tuple(self.always))
object.__setattr__(self, "scenes", MappingProxyType({k: tuple(v) for k, v in self.scenes.items()}))
data = {"enabled": self.enabled, "limits": dict(self.limits),
"always": self.always, "scenes": dict(self.scenes),
"records": {n: {"description": r.description, "body": r.body,
"references": dict(r.references),
"metadata": dict(r.metadata), "allowed_tools": r.allowed_tools}
for n, r in self.records.items()}}
canonical = json.dumps(data, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
object.__setattr__(self, "fingerprint", hashlib.sha256(canonical.encode("utf-8")).hexdigest())
def catalog_text(self) -> str:
if not self.enabled:
return ""
return "\n".join(f"{name}: {' '.join(self.records[name].description.split())}"
for name in sorted(self.records))
def bound_names(self, scene: str) -> tuple[str, ...]:
if not self.enabled:
return ()
return tuple(dict.fromkeys(self.always + self.scenes.get(scene, ())))
def _record(self, name):
if not self.enabled:
raise SkillError("disabled")
if not isinstance(name, str) or name not in self.records:
raise SkillError("unknown_skill")
return self.records[name]
def load(self, name: str) -> str:
return self._record(name).body
def read(self, name: str, path: str) -> str:
record = self._record(name)
path = _reference_path(path)
if path not in record.references:
raise SkillError("unknown_reference")
return record.references[path]
def _references(directory, limit):
references = {}
total = 0
base = directory / "references"
# lexists deliberately detects broken symlinks as well.
if not os.path.lexists(base):
return references
def visit(path, relative, depth=0):
nonlocal total
if depth > 16:
raise SkillError("unsafe_path")
try:
fd = _open_nofollow(path, directory=True)
try:
with os.scandir(fd) as entries:
entries = sorted(entries, key=lambda item: item.name)
for entry in entries:
info = os.stat(entry.name, dir_fd=fd, follow_symlinks=False)
if stat.S_ISLNK(info.st_mode):
raise SkillError("unsafe_path")
key = relative + "/" + entry.name
if stat.S_ISDIR(info.st_mode):
visit(path / entry.name, key, depth + 1)
elif PurePosixPath(entry.name).suffix in (".md", ".txt"):
_reference_path(key)
if len(references) >= _MAX_REFERENCES:
raise SkillError("invalid_document")
content = _text(path / entry.name, limit)
total += len(content.encode("utf-8"))
if total > _MAX_REFERENCE_TOTAL:
raise SkillError("invalid_document")
references[key] = content
finally:
os.close(fd)
except OSError:
raise SkillError("unsafe_path") from None
visit(base, "references")
return references
def _skill(root, name, limits):
directory = root / name
text = _text(directory / "SKILL.md", limits["body_bytes"] + _MAX_HEADER)
lines = text.splitlines(keepends=True)
if not lines or lines[0].strip() != "---":
raise SkillError("invalid_skill")
closing = next((i for i in range(1, len(lines)) if lines[i].strip() == "---"), None)
if closing is None:
raise SkillError("invalid_skill")
header = "".join(lines[1:closing])
body = "".join(lines[closing + 1:])
if len(header.encode("utf-8")) > _MAX_HEADER or len(body.encode("utf-8")) > limits["body_bytes"]:
raise SkillError("invalid_document")
meta = _yaml(header)
if not isinstance(meta, dict) or meta.get("name") != name:
raise SkillError("invalid_skill")
description = meta.get("description")
if not isinstance(description, str) or not description.strip() or len(description) > 1024:
raise SkillError("invalid_skill")
metadata = meta.get("metadata", {})
if not isinstance(metadata, dict) or any(not isinstance(v, str) for v in metadata.values()):
raise SkillError("invalid_skill")
allowed = meta.get("allowed-tools", ())
if isinstance(allowed, str):
allowed = tuple(allowed.split())
if not isinstance(allowed, (list, tuple)) or any(not isinstance(v, str) for v in allowed):
raise SkillError("invalid_skill")
return SkillRecord(name, description, body,
_references(directory, limits["reference_bytes"]), metadata, tuple(allowed))
class SkillRegistry:
"""Re-read on snapshot(); errors retain the last valid config/content revision.
Invalid unbound skills are quarantined (omitted and logged). Bindings to a
quarantined/missing skill invalidate the candidate snapshot. Disabled configs
intentionally bypass all other fields, so an emergency disable always works.
Runtime limits govern model use, not catalog size. No mtime is fingerprinted.
"""
def __init__(self, root=None, config_path=None):
self.root = Path(os.path.abspath(root if root is not None else _BASE / "skills"))
self.config_path = Path(os.path.abspath(config_path if config_path is not None else _BASE / "skills_config.yml"))
self._lock = threading.RLock()
self._last = SkillSnapshot(False, {}, _DEFAULT_LIMITS, (), {})
def _config(self):
try:
self.config_path.relative_to(self.root.parent)
except ValueError:
raise SkillError("unsafe_path") from None
# Missing config is equivalent to an empty config, not a load failure.
# Still validate every parent to reject symlinked/missing ancestors.
try:
fd = _open_nofollow(self.config_path.parent, directory=True)
os.close(fd)
except OSError:
raise SkillError("unsafe_path") from None
if not os.path.lexists(self.config_path):
return {}
config = _yaml(_text(self.config_path, _MAX_CONFIG))
if config is None:
return {}
if not isinstance(config, dict):
raise SkillError("invalid_config")
return config
def _build(self):
config = self._config()
enabled = config.get("enabled", True)
if not isinstance(enabled, bool):
raise SkillError("invalid_config")
if not enabled:
return SkillSnapshot(False, {}, _DEFAULT_LIMITS, (), {})
if set(config) - {"enabled", "allowlist", "always", "scenes", "limits"}:
raise SkillError("invalid_config")
allowlist = _names(config.get("allowlist", []))
always = _names(config.get("always", []))
limits = dict(_DEFAULT_LIMITS)
overrides = config.get("limits", {})
if not isinstance(overrides, dict) or set(overrides) - set(limits):
raise SkillError("invalid_config")
for key, value in overrides.items():
if type(value) is not int:
raise SkillError("invalid_config")
low, high = _LIMIT_RANGES[key]
limits[key] = max(low, min(value, high))
scenes = config.get("scenes", {})
if not isinstance(scenes, dict) or set(scenes) - set(_SCENES):
raise SkillError("invalid_config")
scenes = {scene: _names(scenes.get(scene, [])) for scene in _SCENES}
scenes["reminder"] = tuple(dict.fromkeys(scenes["notification"] + scenes["reminder"]))
if any(n not in allowlist for n in always + tuple(n for names in scenes.values() for n in names)):
raise SkillError("invalid_config")
records = {}
for name in sorted(allowlist):
try:
records[name] = _skill(self.root, name, limits)
except SkillError as exc:
_LOG.warning("Skill quarantined: %s (%s)", name, exc.code)
if any(n not in records for n in always + tuple(n for names in scenes.values() for n in names)):
raise SkillError("invalid_config")
return SkillSnapshot(True, records, limits, always, scenes)
def snapshot(self) -> SkillSnapshot:
with self._lock:
try:
candidate = self._build()
except (SkillError, OSError, RecursionError):
_LOG.warning("Skills configuration rejected; retaining last valid snapshot")
return self._last
if candidate.fingerprint != self._last.fingerprint:
self._last = candidate
return self._last
registry = SkillRegistry()