agentauthlayer 0.1.0__py3-none-any.whl

agent_auth/storage.py

@@ -0,0 +1,536 @@
+"""agent_auth.storage — SQLAlchemy-backed persistence layer.
+
+Tables:
+- agents               — agent identities (id, name, owner, role, scopes, status)
+- agent_scopes         — granted scopes per agent
+- tools                — registered tool catalog
+- token_records        — issued tokens (access + refresh) for revocation / introspection
+- audit_events         — immutable audit trail
+- users                — human user accounts
+- password_reset_tokens — one-time password reset tokens
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import uuid
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from typing import Any
+
+import bcrypt as _bcrypt
+from sqlalchemy import (
+    Boolean,
+    DateTime,
+    ForeignKey,
+    Integer,
+    String,
+    Text,
+    UniqueConstraint,
+    create_engine,
+    delete,
+    select,
+)
+from sqlalchemy.engine import Engine
+from sqlalchemy.orm import DeclarativeBase, Mapped, Session, mapped_column, relationship
+from sqlalchemy.pool import StaticPool
+
+from agent_auth.models import Agent, User
+
+
+# ---------------------------------------------------------------------------
+# ORM base + table definitions
+# ---------------------------------------------------------------------------
+
+class Base(DeclarativeBase):
+    pass
+
+
+class AgentRow(Base):
+    __tablename__ = "agents"
+
+    agent_id: Mapped[str] = mapped_column(String(255), primary_key=True)
+    name: Mapped[str | None] = mapped_column(String(255), nullable=True)
+    owner: Mapped[str | None] = mapped_column(String(255), nullable=True)
+    role: Mapped[str | None] = mapped_column(String(255), nullable=True)
+    status: Mapped[str] = mapped_column(String(32), default="active")
+    metadata_json: Mapped[str | None] = mapped_column(Text, nullable=True)
+    created_at: Mapped[datetime] = mapped_column(
+        DateTime(timezone=True), default=lambda: datetime.now(timezone.utc),
+    )
+    updated_at: Mapped[datetime] = mapped_column(
+        DateTime(timezone=True),
+        default=lambda: datetime.now(timezone.utc),
+        onupdate=lambda: datetime.now(timezone.utc),
+    )
+
+    scopes: Mapped[list["AgentScopeRow"]] = relationship(
+        back_populates="agent",
+        cascade="all, delete-orphan",
+        lazy="selectin",
+    )
+
+
+class AgentScopeRow(Base):
+    __tablename__ = "agent_scopes"
+    __table_args__ = (
+        UniqueConstraint("agent_id", "scope", name="uq_agent_scope"),
+    )
+
+    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
+    agent_id: Mapped[str] = mapped_column(ForeignKey("agents.agent_id", ondelete="CASCADE"))
+    scope: Mapped[str] = mapped_column(String(255), nullable=False)
+
+    agent: Mapped[AgentRow] = relationship(back_populates="scopes")
+
+
+class ToolRow(Base):
+    __tablename__ = "tools"
+
+    name: Mapped[str] = mapped_column(String(255), primary_key=True)
+    required_scope: Mapped[str] = mapped_column(String(255), nullable=False)
+    description: Mapped[str] = mapped_column(Text, default="")
+    created_at: Mapped[datetime] = mapped_column(
+        DateTime(timezone=True), default=lambda: datetime.now(timezone.utc),
+    )
+
+
+class TokenRecordRow(Base):
+    __tablename__ = "token_records"
+
+    jti: Mapped[str] = mapped_column(String(255), primary_key=True)
+    agent_id: Mapped[str] = mapped_column(String(255), nullable=False)
+    scopes_json: Mapped[str] = mapped_column(Text, default="[]")
+    token_type: Mapped[str] = mapped_column(String(32), default="access")
+    status: Mapped[str] = mapped_column(String(32), default="active")
+    issued_at: Mapped[datetime] = mapped_column(
+        DateTime(timezone=True), default=lambda: datetime.now(timezone.utc),
+    )
+    expires_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False)
+
+
+class AuditEventRow(Base):
+    __tablename__ = "audit_events"
+
+    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
+    event_type: Mapped[str] = mapped_column(String(128), nullable=False)
+    agent_id: Mapped[str | None] = mapped_column(String(255), nullable=True)
+    tool_name: Mapped[str | None] = mapped_column(String(255), nullable=True)
+    details_json: Mapped[str] = mapped_column(Text, default="{}")
+    timestamp: Mapped[datetime] = mapped_column(
+        DateTime(timezone=True), default=lambda: datetime.now(timezone.utc),
+    )
+
+
+class UserRow(Base):
+    __tablename__ = "users"
+
+    id: Mapped[str] = mapped_column(String(255), primary_key=True,
+                                    default=lambda: str(uuid.uuid4()))
+    email: Mapped[str] = mapped_column(String(255), unique=True, nullable=False, index=True)
+    name: Mapped[str | None] = mapped_column(String(255), nullable=True)
+    hashed_password: Mapped[str] = mapped_column(Text, nullable=False)
+    role: Mapped[str] = mapped_column(String(64), nullable=False, default="user")
+    created_at: Mapped[datetime] = mapped_column(
+        DateTime(timezone=True), default=lambda: datetime.now(timezone.utc),
+    )
+    updated_at: Mapped[datetime] = mapped_column(
+        DateTime(timezone=True), default=lambda: datetime.now(timezone.utc),
+    )
+
+
+class PasswordResetRow(Base):
+    __tablename__ = "password_reset_tokens"
+
+    token: Mapped[str] = mapped_column(String(255), primary_key=True)
+    user_id: Mapped[str] = mapped_column(String(255), nullable=False)
+    expires_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False)
+    used: Mapped[bool] = mapped_column(Boolean, default=False)
+
+
+# ---------------------------------------------------------------------------
+# Engine helpers
+# ---------------------------------------------------------------------------
+
+def default_database_url() -> str:
+    url = os.getenv("DATABASE_URL")
+    if url:
+        return url
+    return "sqlite:///agent_auth.db"
+
+
+def create_db_engine(database_url: str | None = None) -> Engine:
+    url = database_url or default_database_url()
+
+    connect_args: dict = {}
+    engine_kwargs: dict = {"future": True}
+
+    if url.startswith("sqlite"):
+        connect_args["check_same_thread"] = False
+
+        if ":memory:" in url:
+            engine_kwargs["poolclass"] = StaticPool
+
+        if "file:" in url:
+            connect_args["uri"] = True
+
+    return create_engine(url, connect_args=connect_args, **engine_kwargs)
+
+
+def init_db(engine: Engine) -> None:
+    Base.metadata.create_all(engine)
+
+
+# ---------------------------------------------------------------------------
+# AgentStore
+# ---------------------------------------------------------------------------
+
+@dataclass(slots=True)
+class AgentStore:
+    """Persistent storage for agents + scopes."""
+
+    engine: Engine
+
+    def init(self) -> None:
+        init_db(self.engine)
+
+    def _row_to_agent(self, row: AgentRow) -> Agent:
+        meta = json.loads(row.metadata_json) if row.metadata_json else {}
+        return Agent(
+            agent_id=row.agent_id,
+            scopes=[s.scope for s in row.scopes],
+            metadata=meta,
+            name=row.name,
+            owner=row.owner,
+            role=row.role,
+            status=row.status,
+            created_at=row.created_at,
+        )
+
+    def upsert_agent(
+        self,
+        agent_id: str,
+        scopes: list[str],
+        name: str | None = None,
+        owner: str | None = None,
+        role: str | None = None,
+        metadata: dict | None = None,
+    ) -> Agent:
+        with Session(self.engine) as session:
+            row = session.get(AgentRow, agent_id)
+            if row is None:
+                row = AgentRow(agent_id=agent_id)
+                session.add(row)
+
+            if name is not None:
+                row.name = name
+            if owner is not None:
+                row.owner = owner
+            if role is not None:
+                row.role = role
+            if metadata is not None:
+                row.metadata_json = json.dumps(metadata, default=str)
+            row.updated_at = datetime.now(timezone.utc)
+            session.flush()
+
+            session.execute(delete(AgentScopeRow).where(AgentScopeRow.agent_id == agent_id))
+            session.flush()
+            for scope in scopes:
+                session.add(AgentScopeRow(agent_id=agent_id, scope=scope))
+
+            session.commit()
+            session.refresh(row)
+            return self._row_to_agent(row)
+
+    def get_agent(self, agent_id: str) -> Agent | None:
+        with Session(self.engine) as session:
+            row = session.get(AgentRow, agent_id)
+            if row is None:
+                return None
+            return self._row_to_agent(row)
+
+    def list_agents(self) -> list[Agent]:
+        with Session(self.engine) as session:
+            rows = session.scalars(select(AgentRow)).all()
+            return [self._row_to_agent(r) for r in rows]
+
+    def grant_scope(self, agent_id: str, scope: str) -> Agent:
+        with Session(self.engine) as session:
+            row = session.get(AgentRow, agent_id)
+            if row is None:
+                row = AgentRow(agent_id=agent_id)
+                session.add(row)
+                session.flush()
+
+            existing = session.scalar(
+                select(AgentScopeRow).where(
+                    AgentScopeRow.agent_id == agent_id,
+                    AgentScopeRow.scope == scope,
+                )
+            )
+            if existing is None:
+                session.add(AgentScopeRow(agent_id=agent_id, scope=scope))
+            session.commit()
+            session.refresh(row)
+            return self._row_to_agent(row)
+
+    def revoke_scope(self, agent_id: str, scope: str) -> Agent:
+        with Session(self.engine) as session:
+            session.execute(
+                delete(AgentScopeRow).where(
+                    AgentScopeRow.agent_id == agent_id,
+                    AgentScopeRow.scope == scope,
+                )
+            )
+            session.commit()
+
+            row = session.get(AgentRow, agent_id)
+            if row is None:
+                return Agent(agent_id=agent_id, scopes=[], metadata={})
+            return self._row_to_agent(row)
+
+    def delete_agent(self, agent_id: str) -> bool:
+        with Session(self.engine) as session:
+            row = session.get(AgentRow, agent_id)
+            if row is None:
+                return False
+            session.delete(row)
+            session.commit()
+            return True
+
+
+# ---------------------------------------------------------------------------
+# ToolStore
+# ---------------------------------------------------------------------------
+
+@dataclass(slots=True)
+class ToolStore:
+    """Persistent storage for the tool catalog."""
+
+    engine: Engine
+
+    def upsert(self, name: str, required_scope: str, description: str = "") -> None:
+        with Session(self.engine) as session:
+            row = session.get(ToolRow, name)
+            if row is None:
+                row = ToolRow(name=name, required_scope=required_scope, description=description)
+                session.add(row)
+            else:
+                row.required_scope = required_scope
+                row.description = description
+            session.commit()
+
+    def get(self, name: str) -> ToolRow | None:
+        with Session(self.engine) as session:
+            return session.get(ToolRow, name)
+
+    def list_all(self) -> list[ToolRow]:
+        with Session(self.engine) as session:
+            return list(session.scalars(select(ToolRow)).all())
+
+
+# ---------------------------------------------------------------------------
+# TokenStore
+# ---------------------------------------------------------------------------
+
+@dataclass(slots=True)
+class TokenStore:
+    """Persistent storage for token records (revocation + introspection)."""
+
+    engine: Engine
+
+    def save(
+        self,
+        jti: str,
+        agent_id: str,
+        scopes: list[str],
+        expires_at: datetime,
+        token_type: str = "access",
+    ) -> None:
+        with Session(self.engine) as session:
+            row = TokenRecordRow(
+                jti=jti,
+                agent_id=agent_id,
+                scopes_json=json.dumps(scopes),
+                token_type=token_type,
+                status="active",
+                expires_at=expires_at,
+            )
+            session.add(row)
+            session.commit()
+
+    def revoke(self, jti: str) -> bool:
+        with Session(self.engine) as session:
+            row = session.get(TokenRecordRow, jti)
+            if row is None:
+                return False
+            row.status = "revoked"
+            session.commit()
+            return True
+
+    def get(self, jti: str) -> TokenRecordRow | None:
+        with Session(self.engine) as session:
+            return session.get(TokenRecordRow, jti)
+
+    def is_revoked(self, jti: str) -> bool:
+        with Session(self.engine) as session:
+            row = session.get(TokenRecordRow, jti)
+            if row is None:
+                return False
+            return row.status == "revoked"
+
+
+# ---------------------------------------------------------------------------
+# AuditStore
+# ---------------------------------------------------------------------------
+
+@dataclass(slots=True)
+class AuditStore:
+    """Persistent audit event log."""
+
+    engine: Engine
+
+    def log(self, event_type: str, agent_id: str | None = None,
+            tool_name: str | None = None, details: dict[str, Any] | None = None) -> None:
+        with Session(self.engine) as session:
+            row = AuditEventRow(
+                event_type=event_type,
+                agent_id=agent_id,
+                tool_name=tool_name,
+                details_json=json.dumps(details or {}, default=str),
+            )
+            session.add(row)
+            session.commit()
+
+    def list_events(self, limit: int = 100) -> list[dict[str, Any]]:
+        with Session(self.engine) as session:
+            rows = session.scalars(
+                select(AuditEventRow).order_by(AuditEventRow.id.desc()).limit(limit)
+            ).all()
+            return [
+                {
+                    "id": r.id,
+                    "event_type": r.event_type,
+                    "agent_id": r.agent_id,
+                    "tool_name": r.tool_name,
+                    "details": json.loads(r.details_json) if r.details_json else {},
+                    "timestamp": r.timestamp.isoformat() if r.timestamp else None,
+                }
+                for r in reversed(rows)
+            ]
+
+
+# ---------------------------------------------------------------------------
+# UserStore
+# ---------------------------------------------------------------------------
+
+def _hash_password(password: str) -> str:
+    return _bcrypt.hashpw(password.encode(), _bcrypt.gensalt()).decode()
+
+
+def _verify_password(password: str, hashed: str) -> bool:
+    return _bcrypt.checkpw(password.encode(), hashed.encode())
+
+
+@dataclass(slots=True)
+class UserStore:
+    """Persistent storage for human user accounts and password reset tokens.
+
+    Password hashing is handled internally — callers pass plaintext passwords.
+    """
+
+    engine: Engine
+
+    def _row_to_user(self, row: UserRow) -> User:
+        return User(
+            id=row.id,
+            email=row.email,
+            name=row.name,
+            role=row.role,
+            created_at=row.created_at,
+        )
+
+    def create_user(
+        self,
+        email: str,
+        password: str,
+        name: str | None = None,
+        role: str = "user",
+    ) -> User:
+        hashed = _hash_password(password)
+        with Session(self.engine) as session:
+            row = UserRow(
+                id=str(uuid.uuid4()),
+                email=email.lower(),
+                name=name,
+                hashed_password=hashed,
+                role=role,
+            )
+            session.add(row)
+            session.commit()
+            session.refresh(row)
+            return self._row_to_user(row)
+
+    def authenticate(self, email: str, password: str) -> User | None:
+        """Verify email + password. Returns User on success, None on failure."""
+        with Session(self.engine) as session:
+            row = session.scalar(
+                select(UserRow).where(UserRow.email == email.lower())
+            )
+            if row is None:
+                return None
+            if not _verify_password(password, row.hashed_password):
+                return None
+            return self._row_to_user(row)
+
+    def get_by_id(self, user_id: str) -> User | None:
+        with Session(self.engine) as session:
+            row = session.get(UserRow, user_id)
+            if row is None:
+                return None
+            return self._row_to_user(row)
+
+    def get_by_email(self, email: str) -> User | None:
+        with Session(self.engine) as session:
+            row = session.scalar(
+                select(UserRow).where(UserRow.email == email.lower())
+            )
+            if row is None:
+                return None
+            return self._row_to_user(row)
+
+    def update_password(self, user_id: str, new_password: str) -> None:
+        hashed = _hash_password(new_password)
+        with Session(self.engine) as session:
+            row = session.get(UserRow, user_id)
+            if row:
+                row.hashed_password = hashed
+                row.updated_at = datetime.now(timezone.utc)
+                session.commit()
+
+    def save_reset_token(self, token: str, user_id: str, expires_at: datetime) -> None:
+        with Session(self.engine) as session:
+            row = PasswordResetRow(token=token, user_id=user_id, expires_at=expires_at)
+            session.add(row)
+            session.commit()
+
+    def validate_reset_token(self, token: str) -> str | None:
+        """Return user_id if the token is valid and unused, else None."""
+        with Session(self.engine) as session:
+            row = session.get(PasswordResetRow, token)
+            if row is None or row.used:
+                return None
+            expires = row.expires_at
+            if expires.tzinfo is None:
+                expires = expires.replace(tzinfo=timezone.utc)
+            if datetime.now(timezone.utc) > expires:
+                return None
+            return row.user_id
+
+    def mark_reset_token_used(self, token: str) -> None:
+        with Session(self.engine) as session:
+            row = session.get(PasswordResetRow, token)
+            if row:
+                row.used = True
+                session.commit()

agent_auth/tokens.py

@@ -0,0 +1,92 @@
+"""agent_auth JWT token operations — issue & verify for agents and users."""
+
+from __future__ import annotations
+
+from datetime import datetime, timedelta, timezone
+from uuid import uuid4
+
+from jose import ExpiredSignatureError, JWTError, jwt
+
+from agent_auth.exceptions import TokenExpiredError, TokenInvalidError
+from agent_auth.models import TokenClaims
+
+
+def issue_token(
+    agent_id: str,
+    scopes: list[str],
+    secret_key: str,
+    algorithm: str = "HS256",
+    expire_minutes: int = 30,
+) -> tuple[str, str]:
+    """Create a signed access JWT.
+
+    Returns:
+        (encoded_jwt, jti)
+    """
+    now = datetime.now(timezone.utc)
+    jti = str(uuid4())
+    payload = {
+        "sub": agent_id,
+        "jti": jti,
+        "scopes": scopes,
+        "type": "access",
+        "iat": now,
+        "exp": now + timedelta(minutes=expire_minutes),
+    }
+    encoded = jwt.encode(payload, secret_key, algorithm=algorithm)
+    return encoded, jti
+
+
+def issue_refresh_token(
+    agent_id: str,
+    scopes: list[str],
+    secret_key: str,
+    algorithm: str = "HS256",
+    expire_days: int = 7,
+) -> tuple[str, str]:
+    """Create a signed refresh JWT.
+
+    Returns:
+        (encoded_jwt, jti)
+    """
+    now = datetime.now(timezone.utc)
+    jti = str(uuid4())
+    payload = {
+        "sub": agent_id,
+        "jti": jti,
+        "scopes": scopes,
+        "type": "refresh",
+        "iat": now,
+        "exp": now + timedelta(days=expire_days),
+    }
+    encoded = jwt.encode(payload, secret_key, algorithm=algorithm)
+    return encoded, jti
+
+
+def verify_token(
+    token: str,
+    secret_key: str,
+    algorithm: str = "HS256",
+) -> TokenClaims:
+    """Decode and validate a JWT (access or refresh).
+
+    Raises:
+        TokenExpiredError: if the token has passed its ``exp``.
+        TokenInvalidError: if the signature or structure is invalid.
+    """
+    try:
+        payload = jwt.decode(token, secret_key, algorithms=[algorithm])
+    except ExpiredSignatureError:
+        raise TokenExpiredError()
+    except JWTError as exc:
+        raise TokenInvalidError(str(exc))
+
+    return TokenClaims(
+        agent_id=payload["sub"],
+        jti=payload["jti"],
+        scopes=payload.get("scopes", []),
+        issued_at=datetime.fromtimestamp(payload["iat"], tz=timezone.utc),
+        expires_at=datetime.fromtimestamp(payload["exp"], tz=timezone.utc),
+        raw=payload,
+        token_type=payload.get("type", "access"),
+    )