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.
This commit is contained in:
RRRosmontis
@@ -0,0 +1,361 @@
|
||||
"""Randomly addressed, path-confined artifacts captured from OneBot media."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import binascii
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import secrets
|
||||
import shutil
|
||||
import threading
|
||||
import urllib.parse
|
||||
import urllib.request
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Optional, Union
|
||||
|
||||
|
||||
DEFAULT_ARTIFACT_ROOT = Path(__file__).resolve().parents[1] / "data" / "artifacts"
|
||||
_ID_RE = re.compile(r"^[A-Za-z0-9_-]{16,128}$")
|
||||
_SAFE_NAME_RE = re.compile(r"[^A-Za-z0-9._-]+")
|
||||
|
||||
|
||||
class ArtifactError(ValueError):
|
||||
"""An artifact identifier, source, or confined path is invalid."""
|
||||
|
||||
|
||||
class ArtifactStore:
|
||||
"""Register captured media and materialize it for mail or explicit delivery.
|
||||
|
||||
Local files and base64 payloads are copied beneath ``root`` immediately.
|
||||
HTTP(S) sources are stored as references and downloaded on first
|
||||
``materialize`` call. Public artifact IDs are random URL-safe tokens; the
|
||||
caller can never provide a destination path. Every materialized path is
|
||||
resolved and checked against this store's root before it is returned.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
root: Optional[Union[str, os.PathLike]] = None,
|
||||
*,
|
||||
max_bytes: int = 50 * 1024 * 1024,
|
||||
download_timeout: float = 30.0,
|
||||
) -> None:
|
||||
self.root = Path(root) if root is not None else DEFAULT_ARTIFACT_ROOT
|
||||
self.root.mkdir(parents=True, exist_ok=True)
|
||||
self.root = self.root.resolve()
|
||||
self.max_bytes = int(max_bytes)
|
||||
self.download_timeout = float(download_timeout)
|
||||
if self.max_bytes < 1:
|
||||
raise ValueError("max_bytes must be positive")
|
||||
self._lock = threading.RLock()
|
||||
|
||||
@staticmethod
|
||||
def _filename(value: Optional[str], fallback: str = "artifact.bin") -> str:
|
||||
name = Path(value or fallback).name
|
||||
name = _SAFE_NAME_RE.sub("_", name).strip("._")
|
||||
return (name or fallback)[:160]
|
||||
|
||||
def _new_id(self) -> str:
|
||||
for _ in range(20):
|
||||
artifact_id = secrets.token_urlsafe(24)
|
||||
if not (self.root / artifact_id).exists():
|
||||
return artifact_id
|
||||
raise ArtifactError("could not allocate an artifact id")
|
||||
|
||||
@staticmethod
|
||||
def _public(record: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""Exclude local paths and remote source URLs from model-facing metadata."""
|
||||
return {
|
||||
"artifact_id": record["artifact_id"],
|
||||
"media_type": record["media_type"],
|
||||
"filename": record["filename"],
|
||||
"size": record.get("size"),
|
||||
"remote": record.get("source_kind") == "url",
|
||||
}
|
||||
|
||||
def _directory(self, artifact_id: str) -> Path:
|
||||
if not isinstance(artifact_id, str) or not _ID_RE.fullmatch(artifact_id):
|
||||
raise ArtifactError("invalid artifact id")
|
||||
path = (self.root / artifact_id).resolve()
|
||||
if path.parent != self.root:
|
||||
raise ArtifactError("artifact path escapes store")
|
||||
return path
|
||||
|
||||
def _write_record(self, record: Dict[str, Any]) -> None:
|
||||
directory = self._directory(record["artifact_id"])
|
||||
directory.mkdir(parents=True, exist_ok=True)
|
||||
target = directory / "artifact.json"
|
||||
temporary = directory / ".artifact.json.tmp"
|
||||
temporary.write_text(json.dumps(record, ensure_ascii=False), encoding="utf-8")
|
||||
os.replace(temporary, target)
|
||||
|
||||
def _record(self, artifact_id: str) -> Dict[str, Any]:
|
||||
directory = self._directory(artifact_id)
|
||||
try:
|
||||
record = json.loads((directory / "artifact.json").read_text(encoding="utf-8"))
|
||||
except (OSError, ValueError) as exc:
|
||||
raise ArtifactError("unknown artifact id") from exc
|
||||
if record.get("artifact_id") != artifact_id:
|
||||
raise ArtifactError("invalid artifact metadata")
|
||||
return record
|
||||
|
||||
@staticmethod
|
||||
def _context_metadata(ctx: Any) -> Dict[str, Any]:
|
||||
"""Return the stable owner and conversation scope for an artifact."""
|
||||
channel = getattr(getattr(ctx, "channel", None), "name", None)
|
||||
msg_type = getattr(ctx, "msg_type", None)
|
||||
owner_id = getattr(ctx, "user_id", None)
|
||||
target_id = getattr(ctx, "target_id", None)
|
||||
if (not isinstance(channel, str) or not channel
|
||||
or msg_type not in ("private", "group")
|
||||
or owner_id is None or target_id is None):
|
||||
raise ArtifactError("invalid artifact context")
|
||||
return {
|
||||
"owner_id": str(owner_id),
|
||||
"session": {
|
||||
"channel": channel,
|
||||
"msg_type": msg_type,
|
||||
"target_id": str(target_id),
|
||||
},
|
||||
}
|
||||
|
||||
def _authorize(self, record: Dict[str, Any], ctx: Any = None) -> None:
|
||||
"""Require a matching context for scoped artifacts.
|
||||
|
||||
Legacy records have no trustworthy owner/scope. They remain usable by
|
||||
old trusted local callers that omit ``ctx``, but are never exposed to a
|
||||
context-aware agent request. Scoped records always require ``ctx`` so
|
||||
dropping the context cannot bypass the check.
|
||||
"""
|
||||
has_owner = "owner_id" in record
|
||||
has_session = "session" in record
|
||||
if not has_owner and not has_session:
|
||||
if ctx is not None:
|
||||
raise ArtifactError("artifact is not available in this context")
|
||||
return
|
||||
if ctx is None:
|
||||
raise ArtifactError("artifact context required")
|
||||
expected = self._context_metadata(ctx)
|
||||
if (not has_owner or not has_session
|
||||
or record.get("owner_id") != expected["owner_id"]
|
||||
or record.get("session") != expected["session"]):
|
||||
raise ArtifactError("artifact is not available in this context")
|
||||
|
||||
def get(self, artifact_id: str, ctx: Any = None) -> Dict[str, Any]:
|
||||
"""Return public metadata after enforcing artifact context."""
|
||||
record = self._record(artifact_id)
|
||||
self._authorize(record, ctx)
|
||||
return self._public(record)
|
||||
|
||||
def _register_bytes(
|
||||
self,
|
||||
payload: bytes,
|
||||
*,
|
||||
media_type: str,
|
||||
filename: Optional[str],
|
||||
source_kind: str,
|
||||
source_url: Optional[str] = None,
|
||||
ctx: Any = None,
|
||||
) -> Dict[str, Any]:
|
||||
if len(payload) > self.max_bytes:
|
||||
raise ArtifactError("artifact exceeds size limit")
|
||||
artifact_id = self._new_id()
|
||||
directory = self._directory(artifact_id)
|
||||
directory.mkdir(parents=True, exist_ok=False)
|
||||
name = self._filename(filename)
|
||||
relative = str(Path(artifact_id) / name)
|
||||
record: Dict[str, Any] = {
|
||||
"artifact_id": artifact_id,
|
||||
"media_type": media_type,
|
||||
"filename": name,
|
||||
"size": len(payload),
|
||||
"source_kind": source_kind,
|
||||
"path": relative,
|
||||
}
|
||||
if source_url is not None:
|
||||
record["source_url"] = source_url
|
||||
if ctx is not None:
|
||||
record.update(self._context_metadata(ctx))
|
||||
try:
|
||||
(self.root / relative).write_bytes(payload)
|
||||
self._write_record(record)
|
||||
except Exception:
|
||||
shutil.rmtree(directory, ignore_errors=True)
|
||||
raise
|
||||
return self._public(record)
|
||||
|
||||
def register(
|
||||
self,
|
||||
source: Union[str, os.PathLike],
|
||||
media_type: str = "file",
|
||||
filename: Optional[str] = None,
|
||||
*,
|
||||
ctx: Any = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""Register a local path/file URI, HTTP(S) URL, or base64 media payload.
|
||||
|
||||
Raw base64 and ``base64://`` values are accepted because OneBot uses
|
||||
both forms. Data URLs retain their optional MIME-derived filename only
|
||||
when a filename is explicitly supplied; no payload is logged.
|
||||
"""
|
||||
if not isinstance(media_type, str) or not media_type:
|
||||
raise ValueError("media_type must be a non-empty string")
|
||||
raw = str(source)
|
||||
parsed = urllib.parse.urlparse(raw)
|
||||
if parsed.scheme in ("http", "https"):
|
||||
if not parsed.hostname:
|
||||
raise ArtifactError("invalid media URL")
|
||||
artifact_id = self._new_id()
|
||||
inferred = Path(urllib.parse.unquote(parsed.path)).name
|
||||
record = {
|
||||
"artifact_id": artifact_id,
|
||||
"media_type": media_type,
|
||||
"filename": self._filename(filename or inferred),
|
||||
"size": None,
|
||||
"source_kind": "url",
|
||||
"source_url": raw,
|
||||
"path": None,
|
||||
}
|
||||
if ctx is not None:
|
||||
record.update(self._context_metadata(ctx))
|
||||
self._write_record(record)
|
||||
return self._public(record)
|
||||
if parsed.scheme == "file":
|
||||
if parsed.netloc not in ("", "localhost"):
|
||||
raise ArtifactError("remote file URI is not allowed")
|
||||
local = Path(urllib.request.url2pathname(urllib.parse.unquote(parsed.path)))
|
||||
return self._register_local(local, media_type, filename, ctx)
|
||||
if parsed.scheme and len(parsed.scheme) > 1 and not isinstance(source, os.PathLike):
|
||||
if parsed.scheme == "base64":
|
||||
return self._register_base64(raw[len("base64://"):], media_type, filename, ctx)
|
||||
if parsed.scheme == "data":
|
||||
return self._register_data_url(raw, media_type, filename, ctx)
|
||||
raise ArtifactError("unsupported artifact source")
|
||||
path = Path(source)
|
||||
try:
|
||||
exists = path.exists()
|
||||
except OSError:
|
||||
# Long raw base64 strings are not valid filesystem names.
|
||||
exists = False
|
||||
if exists:
|
||||
return self._register_local(path, media_type, filename, ctx)
|
||||
if isinstance(source, os.PathLike):
|
||||
raise ArtifactError("local media file does not exist")
|
||||
return self._register_base64(raw, media_type, filename, ctx)
|
||||
|
||||
def _register_local(
|
||||
self,
|
||||
path: Path,
|
||||
media_type: str,
|
||||
filename: Optional[str],
|
||||
ctx: Any = None,
|
||||
) -> Dict[str, Any]:
|
||||
try:
|
||||
if not path.is_file():
|
||||
raise ArtifactError("local media is not a regular file")
|
||||
size = path.stat().st_size
|
||||
if size > self.max_bytes:
|
||||
raise ArtifactError("artifact exceeds size limit")
|
||||
with path.open("rb") as handle:
|
||||
payload = handle.read(self.max_bytes + 1)
|
||||
except ArtifactError:
|
||||
raise
|
||||
except OSError as exc:
|
||||
raise ArtifactError("cannot read local media") from exc
|
||||
return self._register_bytes(
|
||||
payload,
|
||||
media_type=media_type,
|
||||
filename=filename or path.name,
|
||||
source_kind="local",
|
||||
ctx=ctx,
|
||||
)
|
||||
|
||||
def _register_base64(
|
||||
self,
|
||||
raw: str,
|
||||
media_type: str,
|
||||
filename: Optional[str],
|
||||
ctx: Any = None,
|
||||
) -> Dict[str, Any]:
|
||||
compact = "".join(raw.split())
|
||||
if len(compact) > (self.max_bytes * 4 // 3 + 16):
|
||||
raise ArtifactError("artifact exceeds size limit")
|
||||
try:
|
||||
payload = base64.b64decode(compact, validate=True)
|
||||
except (binascii.Error, ValueError) as exc:
|
||||
raise ArtifactError("invalid or nonexistent media source") from exc
|
||||
return self._register_bytes(
|
||||
payload,
|
||||
media_type=media_type,
|
||||
filename=filename,
|
||||
source_kind="base64",
|
||||
ctx=ctx,
|
||||
)
|
||||
|
||||
def _register_data_url(
|
||||
self,
|
||||
raw: str,
|
||||
media_type: str,
|
||||
filename: Optional[str],
|
||||
ctx: Any = None,
|
||||
) -> Dict[str, Any]:
|
||||
header, separator, payload = raw.partition(",")
|
||||
if not separator or ";base64" not in header.lower():
|
||||
raise ArtifactError("only base64 data URLs are supported")
|
||||
return self._register_base64(payload, media_type, filename, ctx)
|
||||
|
||||
def _confined_file(self, record: Dict[str, Any]) -> Path:
|
||||
relative = record.get("path")
|
||||
if not isinstance(relative, str) or not relative:
|
||||
raise ArtifactError("artifact has not been materialized")
|
||||
candidate = (self.root / relative).resolve()
|
||||
try:
|
||||
candidate.relative_to(self.root)
|
||||
except ValueError as exc:
|
||||
raise ArtifactError("artifact path escapes store") from exc
|
||||
if candidate.parent != self._directory(record["artifact_id"]):
|
||||
raise ArtifactError("invalid artifact path")
|
||||
if not candidate.is_file():
|
||||
raise ArtifactError("artifact file is missing")
|
||||
return candidate
|
||||
|
||||
def materialize(self, artifact_id: str, ctx: Any = None) -> Path:
|
||||
"""Return a readable local path confined to this store and context.
|
||||
|
||||
URL artifacts are downloaded with a byte limit. Existing local/base64
|
||||
artifacts return their stored copy, never the original arbitrary path.
|
||||
"""
|
||||
with self._lock:
|
||||
record = self._record(artifact_id)
|
||||
self._authorize(record, ctx)
|
||||
if record.get("path"):
|
||||
return self._confined_file(record)
|
||||
if record.get("source_kind") != "url":
|
||||
raise ArtifactError("artifact has no local representation")
|
||||
url = record.get("source_url", "")
|
||||
if urllib.parse.urlparse(url).scheme not in ("http", "https"):
|
||||
raise ArtifactError("invalid artifact URL")
|
||||
directory = self._directory(artifact_id)
|
||||
destination = directory / self._filename(record.get("filename"))
|
||||
try:
|
||||
with urllib.request.urlopen(url, timeout=self.download_timeout) as response:
|
||||
final_scheme = urllib.parse.urlparse(response.geturl()).scheme
|
||||
if final_scheme not in ("http", "https"):
|
||||
raise ArtifactError("artifact redirected to a non-HTTP source")
|
||||
length = response.headers.get("Content-Length")
|
||||
if length and int(length) > self.max_bytes:
|
||||
raise ArtifactError("artifact exceeds size limit")
|
||||
payload = response.read(self.max_bytes + 1)
|
||||
if len(payload) > self.max_bytes:
|
||||
raise ArtifactError("artifact exceeds size limit")
|
||||
destination.write_bytes(payload)
|
||||
record["path"] = str(Path(artifact_id) / destination.name)
|
||||
record["size"] = len(payload)
|
||||
self._write_record(record)
|
||||
return self._confined_file(record)
|
||||
except ArtifactError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
raise ArtifactError("failed to materialize remote artifact") from exc
|
||||
Reference in New Issue
Block a user