317 lines
11 KiB
Python
317 lines
11 KiB
Python
"""Signed tenant context and canonical message contract for AIOS V2R1."""
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
from dataclasses import asdict, dataclass
|
|
from datetime import datetime, timedelta, timezone
|
|
import hashlib
|
|
import hmac
|
|
import json
|
|
import re
|
|
import secrets
|
|
from typing import Any, Mapping
|
|
|
|
VERSION = "V2R1C0"
|
|
MAX_SAFE_INTEGER = 9_007_199_254_740_991
|
|
TENANT_RE = re.compile(r"^[a-z0-9](?:[a-z0-9-]{0,38}[a-z0-9])?$")
|
|
IDENT_RE = re.compile(r"^[A-Za-z0-9](?:[A-Za-z0-9._:@-]{0,126}[A-Za-z0-9])?$")
|
|
|
|
|
|
class V2R1Error(RuntimeError):
|
|
pass
|
|
|
|
|
|
class ValidationError(V2R1Error):
|
|
pass
|
|
|
|
|
|
class ScopeError(V2R1Error):
|
|
pass
|
|
|
|
|
|
class ConflictError(V2R1Error):
|
|
pass
|
|
|
|
|
|
def utc_now() -> datetime:
|
|
return datetime.now(timezone.utc)
|
|
|
|
|
|
def isoformat(value: datetime | None = None) -> str:
|
|
return (value or utc_now()).astimezone(timezone.utc).isoformat()
|
|
|
|
|
|
def parse_time(value: str) -> datetime:
|
|
if not isinstance(value, str) or not value.strip():
|
|
raise ValidationError("timestamp must be a non-empty string")
|
|
text = value.strip()
|
|
if text.endswith("Z"):
|
|
text = text[:-1] + "+00:00"
|
|
try:
|
|
parsed = datetime.fromisoformat(text)
|
|
except ValueError as exc:
|
|
raise ValidationError(f"invalid timestamp: {value!r}") from exc
|
|
if parsed.tzinfo is None:
|
|
raise ValidationError("timestamp must include timezone")
|
|
return parsed.astimezone(timezone.utc)
|
|
|
|
|
|
def require_ident(name: str, value: Any, *, tenant: bool = False) -> str:
|
|
if not isinstance(value, str):
|
|
raise ValidationError(f"{name} must be a string")
|
|
value = value.strip()
|
|
pattern = TENANT_RE if tenant else IDENT_RE
|
|
if not pattern.fullmatch(value):
|
|
raise ValidationError(f"invalid {name}: {value!r}")
|
|
return value
|
|
|
|
|
|
def normalize_message_id(value: Any, *, field_name: str = "message_id") -> str:
|
|
if isinstance(value, bool):
|
|
raise ValidationError(f"{field_name} cannot be boolean")
|
|
if isinstance(value, str):
|
|
value = value.strip()
|
|
if not value or len(value) > 256 or any(ord(ch) < 32 for ch in value):
|
|
raise ValidationError(f"invalid {field_name}")
|
|
return value
|
|
if isinstance(value, int):
|
|
if abs(value) > MAX_SAFE_INTEGER:
|
|
raise ValidationError(f"{field_name} is an unsafe numeric ID")
|
|
return str(value)
|
|
raise ValidationError(f"{field_name} must be a lossless string or safe integer")
|
|
|
|
|
|
def canonical_json(value: Any) -> bytes:
|
|
return json.dumps(
|
|
value, ensure_ascii=False, separators=(",", ":"), sort_keys=True
|
|
).encode("utf-8")
|
|
|
|
|
|
def _b64(data: bytes) -> str:
|
|
return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii")
|
|
|
|
|
|
def _unb64(value: str) -> bytes:
|
|
return base64.urlsafe_b64decode(value + "=" * (-len(value) % 4))
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class TenantContext:
|
|
tenant_id: str
|
|
cell_id: str
|
|
principal_id: str
|
|
principal_type: str
|
|
session_id: str
|
|
task_id: str | None
|
|
channel: str
|
|
channel_account_id: str
|
|
issued_at: str
|
|
expires_at: str
|
|
nonce: str
|
|
|
|
def __post_init__(self) -> None:
|
|
require_ident("tenant_id", self.tenant_id, tenant=True)
|
|
for name in (
|
|
"cell_id", "principal_id", "principal_type", "session_id",
|
|
"channel", "channel_account_id", "nonce",
|
|
):
|
|
require_ident(name, getattr(self, name))
|
|
if self.task_id is not None:
|
|
require_ident("task_id", self.task_id)
|
|
issued = parse_time(self.issued_at)
|
|
expires = parse_time(self.expires_at)
|
|
if expires <= issued or expires - issued > timedelta(hours=1):
|
|
raise ValidationError("invalid tenant context lifetime")
|
|
|
|
@classmethod
|
|
def issue(cls, *, lifetime_seconds: int = 300, now: datetime | None = None, **scope: Any) -> "TenantContext":
|
|
if not 1 <= lifetime_seconds <= 3600:
|
|
raise ValidationError("lifetime_seconds must be between 1 and 3600")
|
|
current = (now or utc_now()).astimezone(timezone.utc)
|
|
return cls(
|
|
**scope,
|
|
issued_at=isoformat(current),
|
|
expires_at=isoformat(current + timedelta(seconds=lifetime_seconds)),
|
|
nonce=secrets.token_hex(16),
|
|
)
|
|
|
|
def sign(self, key: bytes) -> str:
|
|
if len(key) < 32:
|
|
raise ValidationError("signing key must be at least 32 bytes")
|
|
payload = _b64(canonical_json(asdict(self)))
|
|
signature = _b64(hmac.new(key, payload.encode("ascii"), hashlib.sha256).digest())
|
|
return f"{payload}.{signature}"
|
|
|
|
@classmethod
|
|
def verify(
|
|
cls,
|
|
token: str,
|
|
key: bytes,
|
|
*,
|
|
now: datetime | None = None,
|
|
expected_tenant_id: str | None = None,
|
|
expected_cell_id: str | None = None,
|
|
) -> "TenantContext":
|
|
if len(key) < 32:
|
|
raise ValidationError("signing key must be at least 32 bytes")
|
|
try:
|
|
payload, encoded_signature = token.split(".", 1)
|
|
actual = _unb64(encoded_signature)
|
|
if encoded_signature != _b64(actual):
|
|
raise ValidationError("non-canonical tenant context signature")
|
|
data = json.loads(_unb64(payload))
|
|
except Exception as exc:
|
|
raise ValidationError("malformed tenant context token") from exc
|
|
expected = hmac.new(key, payload.encode("ascii"), hashlib.sha256).digest()
|
|
if not hmac.compare_digest(actual, expected):
|
|
raise ValidationError("tenant context signature mismatch")
|
|
context = cls(**data)
|
|
current = (now or utc_now()).astimezone(timezone.utc)
|
|
if current < parse_time(context.issued_at) - timedelta(seconds=60):
|
|
raise ValidationError("tenant context issued in future")
|
|
if current >= parse_time(context.expires_at):
|
|
raise ValidationError("tenant context expired")
|
|
if expected_tenant_id and context.tenant_id != expected_tenant_id:
|
|
raise ScopeError("tenant mismatch")
|
|
if expected_cell_id and context.cell_id != expected_cell_id:
|
|
raise ScopeError("cell mismatch")
|
|
return context
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class CanonicalMessage:
|
|
tenant_id: str
|
|
cell_id: str
|
|
principal_id: str
|
|
principal_type: str
|
|
origin_type: str
|
|
channel: str
|
|
channel_account_id: str
|
|
peer_id: str
|
|
conversation_id: str
|
|
session_id: str
|
|
task_id: str | None
|
|
event_id: str
|
|
message_id: str
|
|
reply_to_message_id: str | None
|
|
direction: str
|
|
command_capable: bool
|
|
occurred_at: str
|
|
received_at: str
|
|
body: str
|
|
quoted_body: str | None
|
|
adapter_name: str
|
|
adapter_version: str
|
|
schema_version: str = "aios.message.v1"
|
|
|
|
def __post_init__(self) -> None:
|
|
if self.schema_version != "aios.message.v1":
|
|
raise ValidationError("unsupported message schema version")
|
|
require_ident("tenant_id", self.tenant_id, tenant=True)
|
|
for name in (
|
|
"cell_id", "principal_id", "principal_type", "origin_type", "channel",
|
|
"channel_account_id", "peer_id", "conversation_id", "session_id",
|
|
"event_id", "adapter_name", "adapter_version",
|
|
):
|
|
require_ident(name, getattr(self, name))
|
|
if self.task_id is not None:
|
|
require_ident("task_id", self.task_id)
|
|
normalize_message_id(self.message_id)
|
|
if self.reply_to_message_id is not None:
|
|
normalize_message_id(self.reply_to_message_id, field_name="reply_to_message_id")
|
|
if self.direction not in {"inbound", "outbound", "internal"}:
|
|
raise ValidationError("invalid direction")
|
|
if not isinstance(self.command_capable, bool):
|
|
raise ValidationError("command_capable must be boolean")
|
|
parse_time(self.occurred_at)
|
|
parse_time(self.received_at)
|
|
if not isinstance(self.body, str):
|
|
raise ValidationError("body must be a string")
|
|
if self.quoted_body is not None and not isinstance(self.quoted_body, str):
|
|
raise ValidationError("quoted_body must be a string or null")
|
|
|
|
@classmethod
|
|
def from_mapping(cls, data: Mapping[str, Any]) -> "CanonicalMessage":
|
|
content = data.get("content") or {}
|
|
adapter = data.get("adapter") or {}
|
|
if not isinstance(content, Mapping) or not isinstance(adapter, Mapping):
|
|
raise ValidationError("content and adapter must be objects")
|
|
values = dict(data)
|
|
return cls(
|
|
schema_version=values.get("schema_version", "aios.message.v1"),
|
|
tenant_id=values["tenant_id"],
|
|
cell_id=values["cell_id"],
|
|
principal_id=values["principal_id"],
|
|
principal_type=values["principal_type"],
|
|
origin_type=values["origin_type"],
|
|
channel=values["channel"],
|
|
channel_account_id=values["channel_account_id"],
|
|
peer_id=values["peer_id"],
|
|
conversation_id=values["conversation_id"],
|
|
session_id=values["session_id"],
|
|
task_id=values.get("task_id"),
|
|
event_id=values["event_id"],
|
|
message_id=normalize_message_id(values["message_id"]),
|
|
reply_to_message_id=(
|
|
normalize_message_id(values["reply_to_message_id"], field_name="reply_to_message_id")
|
|
if values.get("reply_to_message_id") is not None else None
|
|
),
|
|
direction=values["direction"],
|
|
command_capable=values["command_capable"],
|
|
occurred_at=values["occurred_at"],
|
|
received_at=values["received_at"],
|
|
body=content.get("body", ""),
|
|
quoted_body=content.get("quoted_body"),
|
|
adapter_name=adapter["name"],
|
|
adapter_version=adapter["version"],
|
|
)
|
|
|
|
def assert_context(self, context: TenantContext) -> None:
|
|
fields = (
|
|
"tenant_id", "cell_id", "principal_id", "session_id", "task_id",
|
|
"channel", "channel_account_id",
|
|
)
|
|
mismatches = [name for name in fields if getattr(self, name) != getattr(context, name)]
|
|
if mismatches:
|
|
raise ScopeError("message/context mismatch: " + ", ".join(mismatches))
|
|
|
|
@property
|
|
def is_authenticated_human_command(self) -> bool:
|
|
return (
|
|
self.direction == "inbound"
|
|
and self.command_capable
|
|
and self.principal_type == "authenticated_user"
|
|
and self.origin_type == "channel_user"
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class PolicyDecision:
|
|
allowed: bool
|
|
approval_required: bool
|
|
code: str
|
|
reason: str
|
|
|
|
|
|
READ_ONLY_TOOLS = frozenset({
|
|
"read", "grep", "glob", "find", "ls", "web_search", "web_fetch",
|
|
"memory_search", "sessions_list", "sessions_history", "session_status",
|
|
})
|
|
|
|
|
|
def evaluate_tool_call(tool_name: str, message: CanonicalMessage) -> PolicyDecision:
|
|
normalized = tool_name.strip().lower() if isinstance(tool_name, str) else ""
|
|
if not normalized:
|
|
return PolicyDecision(False, False, "invalid_tool", "tool name missing")
|
|
if normalized in READ_ONLY_TOOLS:
|
|
return PolicyDecision(True, False, "read_only", "read-only operation")
|
|
if not message.is_authenticated_human_command:
|
|
return PolicyDecision(
|
|
False, False, "origin_denied",
|
|
"side effects require an authenticated current human command",
|
|
)
|
|
return PolicyDecision(
|
|
True, True, "approval_required",
|
|
"authenticated human side effect requires explicit approval",
|
|
)
|