agentauthlayer 0.1.0__py3-none-any.whl

agent_auth/core.py

@@ -0,0 +1,638 @@
+"""agent_auth.core — the main ``AgentAuth`` class."""
+
+from __future__ import annotations
+
+import json
+import secrets
+from datetime import datetime, timedelta, timezone
+from functools import wraps
+from typing import Any, Callable
+from uuid import uuid4
+
+from jose import ExpiredSignatureError, JWTError, jwt
+
+from agent_auth.audit import AuditBackend, AuditLogger
+from agent_auth.exceptions import (
+    AgentNotFoundError,
+    DuplicateAgentError,
+    DuplicateUserError,
+    InvalidCredentialsError,
+    ScopeError,
+    TokenExpiredError,
+    TokenInvalidError,
+    TokenMissingError,
+    TokenRevokedError,
+    TokenTypeMismatchError,
+    UserNotFoundError,
+)
+from agent_auth.models import Agent, TokenClaims, User, UserClaims
+from agent_auth.policy import check_scope
+from agent_auth.registry import ToolNotRegisteredError, ToolPolicy, ToolRegistry
+from agent_auth.storage import (
+    AgentStore,
+    AuditStore,
+    TokenStore,
+    ToolStore,
+    UserStore,
+    create_db_engine,
+)
+from agent_auth.tokens import issue_refresh_token, issue_token, verify_token
+
+
+class AgentAuth:
+    """Main entry-point for the agent_auth library.
+
+    Handles agents, human users, JWT tokens (access + refresh),
+    tool registry, scope enforcement, and audit logging.
+
+    Usage::
+
+        auth = AgentAuth(secret_key="your-secret")
+
+        # --- agents ---
+        auth.register_agent("email-bot", scopes=["send_email"], owner="team-a")
+        token = auth.issue_token("email-bot")
+
+        @auth.require_scope("send_email")
+        def send_email(to, body, **ctx): ...
+
+        # --- users ---
+        user = auth.register_user("alice@co.com", "s3cret")
+        user_token = auth.issue_user_token(user)
+        claims = auth.verify_user_token(user_token)
+
+        # --- refresh tokens ---
+        refresh = auth.issue_refresh_token("email-bot")
+        new_access, new_refresh = auth.rotate_refresh_token(refresh)
+    """
+
+    def __init__(
+        self,
+        secret_key: str,
+        algorithm: str = "HS256",
+        token_expire_minutes: int = 30,
+        refresh_token_expire_days: int = 7,
+        audit_backend: AuditBackend | None = None,
+        database_url: str | None = None,
+        store: AgentStore | None = None,
+    ) -> None:
+        self._secret = secret_key
+        self._algorithm = algorithm
+        self._expire_minutes = token_expire_minutes
+        self._refresh_expire_days = refresh_token_expire_days
+
+        self._store = store
+        if self._store is None:
+            engine = create_db_engine(database_url)
+            self._store = AgentStore(engine)
+        self._store.init()
+
+        _engine = self._store.engine
+        self._token_store = TokenStore(_engine)
+        self._tool_store = ToolStore(_engine)
+        self._audit_store = AuditStore(_engine)
+        self._user_store = UserStore(_engine)
+
+        self._agents: dict[str, Agent] = {}
+        self._revoked_jtis: set[str] = set()
+        self._tool_registry = ToolRegistry()
+        self._audit = AuditLogger(audit_backend)
+
+    # ------------------------------------------------------------------
+    # Agent registry
+    # ------------------------------------------------------------------
+
+    def register_agent(
+        self,
+        agent_id: str,
+        scopes: list[str] | None = None,
+        metadata: dict | None = None,
+        name: str | None = None,
+        owner: str | None = None,
+        role: str | None = None,
+        *,
+        allow_update: bool = False,
+    ) -> Agent:
+        """Register an agent identity.
+
+        Raises ``DuplicateAgentError`` if *agent_id* already exists
+        and *allow_update* is False.
+        """
+        existing = self._store.get_agent(agent_id)
+        if existing is not None and not allow_update:
+            raise DuplicateAgentError(agent_id)
+
+        agent = self._store.upsert_agent(
+            agent_id,
+            scopes or [],
+            name=name,
+            owner=owner,
+            role=role,
+            metadata=metadata,
+        )
+        self._agents[agent_id] = agent
+        self._audit.record("agent_registered", agent_id=agent_id, scopes=agent.scopes)
+        return agent
+
+    def get_agent(self, agent_id: str) -> Agent:
+        """Fetch a registered agent or raise ``AgentNotFoundError``."""
+        if agent_id in self._agents:
+            return self._agents[agent_id]
+
+        agent = self._store.get_agent(agent_id)
+        if agent is None:
+            raise AgentNotFoundError(agent_id)
+        self._agents[agent_id] = agent
+        return agent
+
+    def list_agents(self) -> list[Agent]:
+        """Return all registered agents."""
+        return self._store.list_agents()
+
+    def delete_agent(self, agent_id: str) -> bool:
+        """Delete an agent. Returns True if deleted, False if it didn't exist."""
+        deleted = self._store.delete_agent(agent_id)
+        if agent_id in self._agents:
+            del self._agents[agent_id]
+        if deleted:
+            self._audit.record("agent_deleted", agent_id=agent_id)
+            self._audit_store.log("agent_deleted", agent_id=agent_id)
+        return deleted
+
+    # ------------------------------------------------------------------
+    # Access token lifecycle
+    # ------------------------------------------------------------------
+
+    def issue_token(self, agent_id: str) -> str:
+        """Issue a signed access JWT for a registered agent.
+
+        Returns the encoded JWT string.
+        """
+        agent = self.get_agent(agent_id)
+        token, jti = issue_token(
+            agent_id=agent.agent_id,
+            scopes=agent.scopes,
+            secret_key=self._secret,
+            algorithm=self._algorithm,
+            expire_minutes=self._expire_minutes,
+        )
+        claims = verify_token(token, self._secret, self._algorithm)
+        self._token_store.save(
+            jti=jti,
+            agent_id=agent_id,
+            scopes=agent.scopes,
+            expires_at=claims.expires_at,
+            token_type="access",
+        )
+        self._audit.record("token_issued", agent_id=agent_id, jti=jti)
+        self._audit_store.log("token_issued", agent_id=agent_id, details={"jti": jti})
+        return token
+
+    def verify(self, token: str) -> TokenClaims:
+        """Verify an access JWT and return decoded claims.
+
+        Raises on expiry, bad signature, wrong type, or revoked token.
+        """
+        claims = verify_token(token, self._secret, self._algorithm)
+        if claims.token_type != "access":
+            raise TokenTypeMismatchError("access", claims.token_type)
+        if claims.jti in self._revoked_jtis or self._token_store.is_revoked(claims.jti):
+            raise TokenRevokedError(claims.jti)
+        return claims
+
+    def revoke_token(self, token: str) -> str:
+        """Revoke a token by its encoded JWT. Returns the revoked jti."""
+        claims = verify_token(token, self._secret, self._algorithm)
+        self._revoked_jtis.add(claims.jti)
+        self._token_store.revoke(claims.jti)
+        self._audit.record("token_revoked", agent_id=claims.agent_id, jti=claims.jti)
+        self._audit_store.log("token_revoked", agent_id=claims.agent_id,
+                              details={"jti": claims.jti})
+        return claims.jti
+
+    def revoke_jti(self, jti: str) -> None:
+        """Revoke a token directly by its jti."""
+        self._revoked_jtis.add(jti)
+        self._token_store.revoke(jti)
+        self._audit.record("token_revoked_by_jti", jti=jti)
+        self._audit_store.log("token_revoked_by_jti", details={"jti": jti})
+
+    def introspect_token(self, jti: str) -> TokenClaims | None:
+        """Return claims for an active non-expired token by its jti, or None.
+
+        Works for both access and refresh tokens.
+        """
+        row = self._token_store.get(jti)
+        if row is None or row.status != "active":
+            return None
+        expires = row.expires_at
+        if expires.tzinfo is None:
+            expires = expires.replace(tzinfo=timezone.utc)
+        if expires <= datetime.now(timezone.utc):
+            return None
+        scopes = json.loads(row.scopes_json) if row.scopes_json else []
+        return TokenClaims(
+            agent_id=row.agent_id,
+            jti=jti,
+            scopes=scopes,
+            issued_at=row.issued_at,
+            expires_at=expires,
+            token_type=row.token_type,
+        )
+
+    # ------------------------------------------------------------------
+    # Refresh token lifecycle
+    # ------------------------------------------------------------------
+
+    def issue_refresh_token(self, agent_id: str) -> str:
+        """Issue a signed refresh JWT for a registered agent.
+
+        Returns the encoded JWT string.
+        """
+        agent = self.get_agent(agent_id)
+        token, jti = issue_refresh_token(
+            agent_id=agent.agent_id,
+            scopes=agent.scopes,
+            secret_key=self._secret,
+            algorithm=self._algorithm,
+            expire_days=self._refresh_expire_days,
+        )
+        now = datetime.now(timezone.utc)
+        expires_at = now + timedelta(days=self._refresh_expire_days)
+        self._token_store.save(
+            jti=jti,
+            agent_id=agent_id,
+            scopes=agent.scopes,
+            expires_at=expires_at,
+            token_type="refresh",
+        )
+        self._audit.record("refresh_token_issued", agent_id=agent_id, jti=jti)
+        self._audit_store.log("refresh_token_issued", agent_id=agent_id,
+                              details={"jti": jti})
+        return token
+
+    def verify_refresh_token(self, token: str) -> TokenClaims:
+        """Verify a refresh JWT and return decoded claims.
+
+        Raises ``TokenTypeMismatchError`` if the token is not a refresh token.
+        Raises on expiry, bad signature, or revoked token.
+        """
+        claims = verify_token(token, self._secret, self._algorithm)
+        if claims.token_type != "refresh":
+            raise TokenTypeMismatchError("refresh", claims.token_type)
+        if claims.jti in self._revoked_jtis or self._token_store.is_revoked(claims.jti):
+            raise TokenRevokedError(claims.jti)
+        return claims
+
+    def rotate_refresh_token(self, refresh_token: str) -> tuple[str, str]:
+        """Validate the refresh token, revoke it, and issue a new access + refresh pair.
+
+        Returns:
+            (new_access_token, new_refresh_token)
+
+        The old refresh token is revoked (one-time use). Raises on invalid or
+        already-revoked refresh token.
+        """
+        claims = self.verify_refresh_token(refresh_token)
+        self._revoked_jtis.add(claims.jti)
+        self._token_store.revoke(claims.jti)
+        new_access = self.issue_token(claims.agent_id)
+        new_refresh = self.issue_refresh_token(claims.agent_id)
+        self._audit.record(
+            "refresh_token_rotated",
+            agent_id=claims.agent_id,
+            old_jti=claims.jti,
+        )
+        return new_access, new_refresh
+
+    # ------------------------------------------------------------------
+    # User management
+    # ------------------------------------------------------------------
+
+    def register_user(
+        self,
+        email: str,
+        password: str,
+        name: str | None = None,
+        role: str = "user",
+    ) -> User:
+        """Register a new human user.
+
+        Raises ``DuplicateUserError`` if the email is already registered.
+        """
+        if self._user_store.get_by_email(email) is not None:
+            raise DuplicateUserError(email)
+        user = self._user_store.create_user(email, password, name=name, role=role)
+        self._audit.record("user_registered", user_id=user.id, email=user.email)
+        self._audit_store.log("user_registered", details={"user_id": user.id, "email": user.email})
+        return user
+
+    def authenticate_user(self, email: str, password: str) -> User:
+        """Verify email + password.
+
+        Returns the ``User`` on success.
+        Raises ``InvalidCredentialsError`` on failure.
+        """
+        user = self._user_store.authenticate(email, password)
+        if user is None:
+            raise InvalidCredentialsError()
+        self._audit.record("user_authenticated", user_id=user.id)
+        return user
+
+    def get_user(self, user_id: str) -> User:
+        """Fetch a user by id or raise ``UserNotFoundError``."""
+        user = self._user_store.get_by_id(user_id)
+        if user is None:
+            raise UserNotFoundError(user_id)
+        return user
+
+    def get_user_by_email(self, email: str) -> User:
+        """Fetch a user by email or raise ``UserNotFoundError``."""
+        user = self._user_store.get_by_email(email)
+        if user is None:
+            raise UserNotFoundError(email)
+        return user
+
+    def issue_user_token(self, user: User) -> str:
+        """Issue a signed JWT for a human user.
+
+        The token carries ``type: "user"`` so it cannot be used as an agent
+        bearer token.
+
+        Returns the encoded JWT string.
+        """
+        now = datetime.now(timezone.utc)
+        exp = now + timedelta(minutes=self._expire_minutes)
+        jti = str(uuid4())
+        payload = {
+            "sub": user.id,
+            "email": user.email,
+            "role": user.role,
+            "type": "user",
+            "jti": jti,
+            "iat": now,
+            "exp": exp,
+        }
+        token = jwt.encode(payload, self._secret, algorithm=self._algorithm)
+        self._audit.record("user_token_issued", user_id=user.id, jti=jti)
+        return token
+
+    def verify_user_token(self, token: str) -> UserClaims:
+        """Verify a user JWT and return decoded claims.
+
+        Raises ``TokenTypeMismatchError`` if this is not a user token.
+        Raises on expiry or bad signature.
+        """
+        try:
+            payload = jwt.decode(token, self._secret, algorithms=[self._algorithm])
+        except ExpiredSignatureError:
+            raise TokenExpiredError()
+        except JWTError as exc:
+            raise TokenInvalidError(str(exc))
+
+        if payload.get("type") != "user":
+            raise TokenTypeMismatchError("user", payload.get("type", "unknown"))
+
+        return UserClaims(
+            user_id=payload["sub"],
+            email=payload["email"],
+            role=payload["role"],
+            jti=payload["jti"],
+            issued_at=datetime.fromtimestamp(payload["iat"], tz=timezone.utc),
+            expires_at=datetime.fromtimestamp(payload["exp"], tz=timezone.utc),
+            raw=payload,
+        )
+
+    def request_password_reset(self, email: str) -> str | None:
+        """Generate a password reset token for the given email.
+
+        Returns the reset token string, or None if the email is not registered.
+        The caller is responsible for delivering the token (email, SMS, etc.).
+
+        Returns None silently for unknown emails to prevent enumeration.
+        """
+        user = self._user_store.get_by_email(email)
+        if user is None:
+            return None
+        token = secrets.token_urlsafe(32)
+        expires_at = datetime.now(timezone.utc) + timedelta(hours=1)
+        self._user_store.save_reset_token(token, user.id, expires_at)
+        self._audit.record("password_reset_requested", user_id=user.id)
+        return token
+
+    def reset_password(self, reset_token: str, new_password: str) -> bool:
+        """Apply a new password using a reset token.
+
+        Returns True on success, False if the token is invalid/expired/used.
+        """
+        user_id = self._user_store.validate_reset_token(reset_token)
+        if user_id is None:
+            return False
+        self._user_store.update_password(user_id, new_password)
+        self._user_store.mark_reset_token_used(reset_token)
+        self._audit.record("password_reset_completed", user_id=user_id)
+        self._audit_store.log("password_reset_completed", details={"user_id": user_id})
+        return True
+
+    # ------------------------------------------------------------------
+    # Decorators
+    # ------------------------------------------------------------------
+
+    def require_scope(self, scope: str) -> Callable:
+        """Decorator that secures a tool function behind a scope check.
+
+        The caller **must** pass ``token=<jwt>`` as a keyword argument.
+        On success the decorator injects ``claims`` into kwargs.
+
+        Example::
+
+            @auth.require_scope("send_email")
+            def send_email(to, body, **ctx):
+                caller = ctx["claims"].agent_id
+                ...
+
+            send_email(to="a@b.com", body="hi", token=my_jwt)
+        """
+        def decorator(func: Callable) -> Callable:
+            @wraps(func)
+            def wrapper(*args: Any, **kwargs: Any) -> Any:
+                token = kwargs.pop("token", None)
+                if token is None:
+                    raise TokenMissingError()
+
+                claims = self.verify(token)
+                check_scope(claims.scopes, scope)
+
+                self._audit.record(
+                    "tool_call_authorized",
+                    agent_id=claims.agent_id,
+                    tool=func.__qualname__,
+                    scope=scope,
+                )
+
+                kwargs["claims"] = claims
+                return func(*args, **kwargs)
+
+            return wrapper
+
+        return decorator
+
+    def require_scopes(self, *scopes: str) -> Callable:
+        """Like ``require_scope`` but checks all listed scopes."""
+        from agent_auth.policy import check_scopes as _check_scopes
+
+        def decorator(func: Callable) -> Callable:
+            @wraps(func)
+            def wrapper(*args: Any, **kwargs: Any) -> Any:
+                token = kwargs.pop("token", None)
+                if token is None:
+                    raise TokenMissingError()
+
+                claims = self.verify(token)
+                _check_scopes(claims.scopes, list(scopes))
+
+                self._audit.record(
+                    "tool_call_authorized",
+                    agent_id=claims.agent_id,
+                    tool=func.__qualname__,
+                    scopes=list(scopes),
+                )
+
+                kwargs["claims"] = claims
+                return func(*args, **kwargs)
+
+            return wrapper
+
+        return decorator
+
+    # ------------------------------------------------------------------
+    # Tool registry
+    # ------------------------------------------------------------------
+
+    def register_tool(
+        self,
+        name: str,
+        scope: str,
+        description: str = "",
+        *,
+        allow_update: bool = False,
+    ) -> ToolPolicy:
+        """Register a tool and the scope it requires."""
+        policy = self._tool_registry.register(
+            name, scope, description, allow_update=allow_update,
+        )
+        self._tool_store.upsert(name, scope, description)
+        self._audit.record("tool_registered", tool=name, scope=scope)
+        self._audit_store.log("tool_registered", tool_name=name, details={"scope": scope})
+        return policy
+
+    def authorize_tool_call(self, token: str, tool_name: str) -> TokenClaims:
+        """Verify *token* and check that it carries the scope required by *tool_name*.
+
+        Returns ``TokenClaims`` on success. Raises on any failure.
+        """
+        claims = self.verify(token)
+        policy = self._tool_registry.get(tool_name)
+        check_scope(claims.scopes, policy.required_scope)
+        self._audit.record(
+            "tool_call_authorized",
+            agent_id=claims.agent_id,
+            tool=tool_name,
+            scope=policy.required_scope,
+        )
+        return claims
+
+    def list_tools(self) -> list[ToolPolicy]:
+        """Return all registered tool policies."""
+        return self._tool_registry.list_tools()
+
+    def list_agent_tools(self, agent_id: str) -> list[ToolPolicy]:
+        """Return tools that *agent_id* is allowed to call."""
+        agent = self.get_agent(agent_id)
+        return self._tool_registry.tools_for_scopes(agent.scopes)
+
+    def grant_tool_to_agent(self, agent_id: str, tool_name: str) -> Agent:
+        """Grant an agent the scope required by *tool_name*."""
+        policy = self._tool_registry.get(tool_name)
+        agent = self.get_agent(agent_id)
+        stored = self._store.grant_scope(agent_id, policy.required_scope)
+        self._agents[agent_id] = stored
+        self._audit.record(
+            "tool_granted",
+            agent_id=agent_id,
+            tool=tool_name,
+            scope=policy.required_scope,
+        )
+        return stored
+
+    def revoke_tool_from_agent(self, agent_id: str, tool_name: str) -> Agent:
+        """Remove the scope required by *tool_name* from an agent."""
+        policy = self._tool_registry.get(tool_name)
+        stored = self._store.revoke_scope(agent_id, policy.required_scope)
+        self._agents[agent_id] = stored
+        self._audit.record(
+            "tool_revoked",
+            agent_id=agent_id,
+            tool=tool_name,
+            scope=policy.required_scope,
+        )
+        return stored
+
+    def tool(self, name: str, scope: str, description: str = "") -> Callable:
+        """Decorator that registers a tool AND enforces its scope.
+
+        Combines ``register_tool`` + ``require_scope`` in one step::
+
+            @auth.tool("send_email", scope="email:send")
+            def send_email(to, body, **ctx): ...
+
+            send_email(to="a@b.com", body="hi", token=token)
+        """
+        self._tool_registry.register(name, scope, description, allow_update=True)
+        self._tool_store.upsert(name, scope, description)
+
+        def decorator(func: Callable) -> Callable:
+            @wraps(func)
+            def wrapper(*args: Any, **kwargs: Any) -> Any:
+                token = kwargs.pop("token", None)
+                if token is None:
+                    raise TokenMissingError()
+
+                claims = self.verify(token)
+                check_scope(claims.scopes, scope)
+
+                self._audit.record(
+                    "tool_call_authorized",
+                    agent_id=claims.agent_id,
+                    tool=name,
+                    scope=scope,
+                )
+
+                kwargs["claims"] = claims
+                return func(*args, **kwargs)
+
+            return wrapper
+
+        return decorator
+
+    def secure(self, func: Callable) -> Callable:
+        """Decorator that requires a valid token but no specific scope."""
+
+        @wraps(func)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            token = kwargs.pop("token", None)
+            if token is None:
+                raise TokenMissingError()
+
+            claims = self.verify(token)
+
+            self._audit.record(
+                "tool_call_authenticated",
+                agent_id=claims.agent_id,
+                tool=func.__qualname__,
+            )
+
+            kwargs["claims"] = claims
+            return func(*args, **kwargs)
+
+        return wrapper

agent_auth/delegation.py

@@ -0,0 +1,15 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from datetime import datetime
+
+
+@dataclass(frozen=True, slots=True)
+class DelegationGrant:
+    grant_id: str
+    delegator_id: str
+    delegatee_id: str
+    action: str
+    resource: str | None = None
+    expires_at: datetime | None = None
+    context: dict[str, str] = field(default_factory=dict)