agentauthlayer 0.1.0__py3-none-any.whl

agent_auth/registry.py

@@ -0,0 +1,90 @@
+"""agent_auth.registry — central tool registry mapping tools → required scopes."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from agent_auth.exceptions import AgentAuthError
+
+
+class ToolNotRegisteredError(AgentAuthError):
+    """Raised when a tool name is not found in the registry."""
+
+    def __init__(self, tool_name: str) -> None:
+        self.tool_name = tool_name
+        super().__init__(f"Tool not registered: {tool_name}")
+
+
+class DuplicateToolError(AgentAuthError):
+    """Raised when registering a tool name that already exists."""
+
+    def __init__(self, tool_name: str) -> None:
+        self.tool_name = tool_name
+        super().__init__(f"Tool already registered: {tool_name}")
+
+
+@dataclass(frozen=True, slots=True)
+class ToolPolicy:
+    """Policy for a single tool."""
+
+    name: str
+    required_scope: str
+    description: str = ""
+
+
+class ToolRegistry:
+    """Central registry: tool_name → required scope.
+
+    Usage::
+
+        registry = ToolRegistry()
+        registry.register("send_email", scope="email:send")
+        registry.register("create_ticket", scope="ticket:create")
+
+        policy = registry.get("send_email")
+        # policy.required_scope == "email:send"
+    """
+
+    def __init__(self) -> None:
+        self._tools: dict[str, ToolPolicy] = {}
+
+    def register(
+        self,
+        name: str,
+        scope: str,
+        description: str = "",
+        *,
+        allow_update: bool = False,
+    ) -> ToolPolicy:
+        """Register a tool with its required scope."""
+        if name in self._tools and not allow_update:
+            raise DuplicateToolError(name)
+        policy = ToolPolicy(name=name, required_scope=scope, description=description)
+        self._tools[name] = policy
+        return policy
+
+    def get(self, name: str) -> ToolPolicy:
+        """Get tool policy or raise ``ToolNotRegisteredError``."""
+        try:
+            return self._tools[name]
+        except KeyError:
+            raise ToolNotRegisteredError(name)
+
+    def list_tools(self) -> list[ToolPolicy]:
+        """Return all registered tool policies."""
+        return list(self._tools.values())
+
+    def tools_for_scopes(self, scopes: list[str]) -> list[ToolPolicy]:
+        """Return tools that the given scopes grant access to.
+
+        Supports wildcard ``"*"`` scope (grants all tools).
+        """
+        if "*" in scopes:
+            return self.list_tools()
+        return [t for t in self._tools.values() if t.required_scope in scopes]
+
+    def __contains__(self, name: str) -> bool:
+        return name in self._tools
+
+    def __len__(self) -> int:
+        return len(self._tools)

agent_auth/session.py

@@ -0,0 +1,135 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from datetime import datetime, timedelta, timezone
+from typing import Protocol
+from uuid import uuid4
+
+from jose import JWTError, jwt
+
+from agent_auth.models import TokenRecord
+from agent_auth.principals import Principal
+
+
+class TokenRecordRepository(Protocol):
+    def save(self, record: TokenRecord) -> TokenRecord:
+        ...
+
+    def get(self, jti: str) -> TokenRecord | None:
+        ...
+
+    def revoke_all_for_subject(self, subject_id: str) -> int:
+        ...
+
+
+@dataclass(frozen=True, slots=True)
+class TokenSettings:
+    secret_key: str
+    algorithm: str = "HS256"
+    access_token_expire_minutes: int = 60
+    refresh_token_expire_days: int = 7
+
+
+class PrincipalTokenService:
+    """Core token/session lifecycle for user, agent, and system principals."""
+
+    def __init__(self, settings: TokenSettings, repo: TokenRecordRepository):
+        self.settings = settings
+        self.repo = repo
+
+    def issue_access_token(self, principal: Principal) -> tuple[str, TokenRecord]:
+        return self._issue_token(principal, expires_delta=timedelta(minutes=self.settings.access_token_expire_minutes))
+
+    def issue_refresh_token(self, principal: Principal) -> tuple[str, TokenRecord]:
+        return self._issue_token(
+            principal,
+            expires_delta=timedelta(days=self.settings.refresh_token_expire_days),
+            token_type="refresh",
+        )
+
+    def rotate_refresh_token(self, raw_refresh: str, principal: Principal) -> tuple[str, TokenRecord, str, TokenRecord] | None:
+        payload = self.decode_token(raw_refresh)
+        if payload is None or payload.get("type") != "refresh":
+            return None
+        old_jti = payload.get("jti")
+        if not old_jti or self.introspect(old_jti) is None:
+            return None
+        self.revoke_token(old_jti)
+        access_token, access_record = self.issue_access_token(principal)
+        refresh_token, refresh_record = self.issue_refresh_token(principal)
+        return access_token, access_record, refresh_token, refresh_record
+
+    def revoke_token(self, jti: str) -> TokenRecord | None:
+        record = self.repo.get(jti)
+        if record:
+            record.status = "revoked"
+            self.repo.save(record)
+        return record
+
+    def revoke_all_for_subject(self, subject_id: str) -> int:
+        return self.repo.revoke_all_for_subject(subject_id)
+
+    def introspect(self, jti: str) -> TokenRecord | None:
+        record = self.repo.get(jti)
+        if not record or record.status != "active":
+            return None
+        expires_at = record.expires_at
+        if expires_at.tzinfo is None:
+            expires_at = expires_at.replace(tzinfo=timezone.utc)
+        if expires_at <= datetime.now(timezone.utc):
+            return None
+        return record
+
+    def introspect_by_token(self, raw_token: str) -> TokenRecord | None:
+        payload = self.decode_token(raw_token)
+        if payload is None:
+            return None
+        jti = payload.get("jti")
+        if not jti:
+            return None
+        return self.introspect(jti)
+
+    def authenticate_bearer(self, raw_token: str) -> tuple[dict, TokenRecord] | None:
+        payload = self.decode_token(raw_token)
+        if payload is None:
+            return None
+        principal_id = payload.get("sub")
+        jti = payload.get("jti")
+        if not principal_id or not jti:
+            return None
+        record = self.introspect(jti)
+        if record is None:
+            return None
+        return payload, record
+
+    def decode_token(self, raw_token: str) -> dict | None:
+        try:
+            return jwt.decode(raw_token, self.settings.secret_key, algorithms=[self.settings.algorithm])
+        except JWTError:
+            return None
+
+    def _issue_token(self, principal: Principal, *, expires_delta: timedelta, token_type: str | None = None) -> tuple[str, TokenRecord]:
+        now = datetime.now(timezone.utc)
+        expires_at = now + expires_delta
+        jti = str(uuid4())
+        payload = {
+            "sub": principal.principal_id,
+            "sub_type": principal.principal_type,
+            "jti": jti,
+            "scopes": list(principal.scopes),
+            "role": principal.role,
+            "email": principal.email,
+            "exp": expires_at,
+        }
+        if token_type:
+            payload["type"] = token_type
+        token = jwt.encode(payload, self.settings.secret_key, algorithm=self.settings.algorithm)
+        record = TokenRecord(
+            jti=jti,
+            principal_id=principal.principal_id,
+            principal_type=principal.principal_type,
+            scopes=list(principal.scopes),
+            expires_at=expires_at,
+        )
+        self.repo.save(record)
+        return token, record