google-adk-extras 0.2.6__py3-none-any.whl → 0.3.0__py3-none-any.whl

google_adk_extras/auth/attach.py

@@ -0,0 +1,227 @@
+from __future__ import annotations
+
+from typing import Optional
+
+from fastapi import APIRouter, Depends, FastAPI, HTTPException, Request
+import base64
+from fastapi.security import APIKeyHeader
+from starlette.middleware.base import BaseHTTPMiddleware
+
+from .config import AuthConfig, JwtIssuerConfig, JwtValidatorConfig
+from .jwt_utils import decode_jwt, encode_jwt, now_ts
+from .sql_store import AuthStore
+
+
+def attach_auth(app: FastAPI, cfg: Optional[AuthConfig]) -> None:
+    """Attach optional auth to the provided FastAPI app.
+
+    - Adds middleware that enforces auth on sensitive routes.
+    - Optionally registers token issuance endpoints if configured.
+    """
+    if not cfg or not cfg.enabled or cfg.allow_no_auth:
+        return
+
+    validator = cfg.jwt_validator
+    issuer_cfg = cfg.jwt_issuer
+    api_keys = set(cfg.api_keys or [])
+    basic_users = cfg.basic_users or {}
+    auth_store: Optional[AuthStore] = None
+    if issuer_cfg and issuer_cfg.database_url:
+        auth_store = AuthStore(issuer_cfg.database_url)
+
+    # Security helpers
+    api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)
+
+    async def _authenticate(request: Request) -> dict:
+        # API Key
+        api_key = request.query_params.get("api_key") or request.headers.get("x-api-key") or request.headers.get("X-API-Key")
+        if not api_key:
+            api_key = await api_key_header.__call__(request)
+        if api_key and api_key in api_keys:
+            return {"method": "api_key", "sub": "api_key_client"}
+        if api_key and auth_store and auth_store.verify_api_key(api_key):
+            return {"method": "api_key", "sub": "api_key_client"}
+
+        # Basic
+        authz = request.headers.get("authorization") or request.headers.get("Authorization")
+        if authz and authz.lower().startswith("basic "):
+            try:
+                b64 = authz.split(" ", 1)[1]
+                raw = base64.b64decode(b64).decode("utf-8")
+                username, _, password = raw.partition(":")
+            except Exception:
+                username, password = "", ""
+            # If SQL store present, try it first; else fall back to configured map
+            if auth_store:
+                uid = auth_store.authenticate_basic(username, password)
+                if uid:
+                    return {"method": "basic", "sub": uid, "username": username}
+            stored = basic_users.get(username)
+            if stored and (stored == password):
+                return {"method": "basic", "sub": username, "username": username}
+
+        # Bearer JWT
+        if authz and authz.lower().startswith("bearer "):
+            token = authz.split(" ", 1)[1]
+            if validator and (validator.jwks_url or validator.hs256_secret):
+                try:
+                    claims = decode_jwt(
+                        token,
+                        issuer=validator.issuer,
+                        audience=validator.audience,
+                        jwks_url=validator.jwks_url,
+                        hs256_secret=validator.hs256_secret,
+                    )
+                    sub = str(claims.get("sub"))
+                    if not sub:
+                        raise HTTPException(status_code=401, detail="Invalid token: no subject")
+                    return {"method": "jwt", "sub": sub, "claims": claims}
+                except Exception as e:
+                    raise HTTPException(status_code=401, detail=f"Invalid token: {e}")
+
+        raise HTTPException(status_code=401, detail="Unauthorized")
+
+    def _path_requires_auth(path: str, method: str) -> bool:
+        method = method.upper()
+        # Always protect core run endpoints
+        if path == "/run" and method == "POST":
+            return True
+        if path == "/run_sse" and method == "POST":
+            return True
+        # Sessions and artifacts under /apps
+        if path.startswith("/apps/"):
+            # Allow metrics to be toggled
+            if path.endswith("/metrics-info") and method == "GET":
+                return cfg.protect_metrics
+            return True
+        # Debug and builder are privileged
+        if path.startswith("/debug/") or path.startswith("/builder/"):
+            return True
+        # API key management endpoints
+        if path.startswith("/auth/api-keys"):
+            return True
+        # Optionally protect list-apps
+        if path == "/list-apps" and method == "GET":
+            return cfg.protect_list_apps
+        return False
+
+    class _AuthMiddleware(BaseHTTPMiddleware):
+        async def dispatch(self, request: Request, call_next):
+            path = request.url.path
+            if not _path_requires_auth(path, request.method):
+                return await call_next(request)
+            # Authenticate
+            try:
+                request.state.identity = await _authenticate(request)
+            except HTTPException as e:
+                from fastapi.responses import JSONResponse
+                return JSONResponse({"detail": e.detail}, status_code=e.status_code)
+            # Optional: Enforce user ownership when path has /users/{user_id}/
+            try:
+                parts = path.strip("/").split("/")
+                if "users" in parts:
+                    idx = parts.index("users")
+                    claimed = parts[idx + 1]
+                    sub = str(request.state.identity.get("sub"))
+                    # Allow api_key method to bypass ownership
+                    if request.state.identity.get("method") != "api_key" and sub != claimed:
+                        from fastapi.responses import JSONResponse
+                        return JSONResponse({"detail": "Forbidden: user mismatch"}, status_code=403)
+            except HTTPException:
+                raise
+            except Exception:
+                pass
+            return await call_next(request)
+
+    app.add_middleware(_AuthMiddleware)
+
+    # Token issuance endpoints (optional)
+    if issuer_cfg and issuer_cfg.enabled:
+        if issuer_cfg.algorithm == "HS256" and not issuer_cfg.hs256_secret:
+            raise RuntimeError("HS256 issuer requires hs256_secret")
+        router = APIRouter()
+
+        @router.post("/auth/register")
+        async def register(username: str, password: str):
+            if not auth_store:
+                raise HTTPException(status_code=400, detail="SQL store not configured")
+            uid = auth_store.create_user(username, password)
+            return {"user_id": uid}
+
+        @router.post("/auth/token")
+        async def token_grant(grant_type: str = "password", username: Optional[str] = None, password: Optional[str] = None,
+                              user_id: Optional[str] = None, fingerprint: Optional[str] = None):
+            sub: Optional[str] = None
+            if grant_type == "password":
+                if not auth_store or not username or password is None:
+                    raise HTTPException(status_code=400, detail="invalid_request")
+                uid = auth_store.authenticate_basic(username, password)
+                if not uid:
+                    raise HTTPException(status_code=401, detail="invalid_grant")
+                sub = uid
+            elif grant_type == "client_credentials":
+                # For simplicity map to provided user_id
+                if not user_id:
+                    raise HTTPException(status_code=400, detail="invalid_request")
+                sub = user_id
+            else:
+                raise HTTPException(status_code=400, detail="unsupported_grant_type")
+
+            now = now_ts()
+            access = {
+                "iss": issuer_cfg.issuer,
+                "aud": issuer_cfg.audience,
+                "sub": sub,
+                "iat": now,
+                "nbf": now,
+                "exp": now + issuer_cfg.access_ttl_seconds,
+            }
+            key = issuer_cfg.hs256_secret if issuer_cfg.algorithm == "HS256" else ""
+            access_token = encode_jwt(access, algorithm=issuer_cfg.algorithm, key=key)
+
+            refresh_token = None
+            if auth_store:
+                jti = auth_store.issue_refresh(sub, issuer_cfg.refresh_ttl_seconds, fingerprint=fingerprint)
+                refresh_token = jti
+            return {"access_token": access_token, "token_type": "bearer", "refresh_token": refresh_token}
+
+        @router.post("/auth/refresh")
+        async def refresh(user_id: str, refresh_token: str, fingerprint: Optional[str] = None):
+            if not auth_store:
+                raise HTTPException(status_code=400, detail="invalid_request")
+            if not auth_store.verify_refresh(refresh_token, user_id, fingerprint=fingerprint):
+                raise HTTPException(status_code=401, detail="invalid_grant")
+            now = now_ts()
+            access = {
+                "iss": issuer_cfg.issuer,
+                "aud": issuer_cfg.audience,
+                "sub": user_id,
+                "iat": now,
+                "nbf": now,
+                "exp": now + issuer_cfg.access_ttl_seconds,
+            }
+            key = issuer_cfg.hs256_secret if issuer_cfg.algorithm == "HS256" else ""
+            access_token = encode_jwt(access, algorithm=issuer_cfg.algorithm, key=key)
+            return {"access_token": access_token, "token_type": "bearer"}
+
+        app.include_router(router)
+
+    # API key management endpoints (require SQL store)
+    if auth_store:
+        api_router = APIRouter()
+
+        @api_router.post("/auth/api-keys")
+        async def create_api_key(user_id: Optional[str] = None, name: Optional[str] = None):
+            key_id, key_plain = auth_store.create_api_key(user_id=user_id, name=name)
+            return {"id": key_id, "api_key": key_plain}
+
+        @api_router.get("/auth/api-keys")
+        async def list_api_keys():
+            return auth_store.list_api_keys()
+
+        @api_router.delete("/auth/api-keys/{key_id}")
+        async def delete_api_key(key_id: str):
+            auth_store.revoke_api_key(key_id)
+            return {"ok": True}
+
+        app.include_router(api_router)

google_adk_extras/auth/config.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import List, Optional
+
+
+@dataclass
+class JwtValidatorConfig:
+    # Accept JWTs from external issuers (e.g., Google/Auth0/Okta) or our own issuer
+    jwks_url: Optional[str] = None
+    issuer: Optional[str] = None
+    audience: Optional[str] = None
+    # If you want to validate with an HS256 shared secret (tests/dev)
+    hs256_secret: Optional[str] = None
+
+
+@dataclass
+class JwtIssuerConfig:
+    # Configure our own issuer if we issue tokens
+    enabled: bool = False
+    issuer: str = "https://example-issuer"
+    audience: str = "adk-api"
+    algorithm: str = "HS256"  # HS256 or RS256/ES256 later
+    hs256_secret: Optional[str] = None
+    access_ttl_seconds: int = 3600
+    refresh_ttl_seconds: int = 60 * 60 * 24 * 14
+    # SQL store for users/refresh tokens
+    database_url: Optional[str] = None  # e.g. sqlite:///auth.db
+
+
+@dataclass
+class AuthConfig:
+    # Global toggle
+    enabled: bool = False
+    # Modes
+    allow_no_auth: bool = False  # if True, bypass checks entirely
+    api_keys: List[str] = field(default_factory=list)  # accepted API keys
+    basic_users: dict[str, str] = field(default_factory=dict)  # username -> password (PBKDF2 hash or plaintext for tests)
+    jwt_validator: Optional[JwtValidatorConfig] = None
+    jwt_issuer: Optional[JwtIssuerConfig] = None
+    # Route policy toggles
+    protect_list_apps: bool = True
+    protect_metrics: bool = True
+    # Scopes are advisory; we currently validate presence of a token and subject. Extend as needed.
+

google_adk_extras/auth/jwt_utils.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+import base64
+import json
+import time
+from typing import Any, Dict, Optional
+
+import jwt
+from jwt import PyJWKClient
+
+
+def _b64url(data: bytes) -> str:
+    return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii")
+
+
+def encode_jwt(payload: Dict[str, Any], *, algorithm: str, key: str, headers: Optional[Dict[str, Any]] = None) -> str:
+    return jwt.encode(payload, key, algorithm=algorithm, headers=headers)
+
+
+def decode_jwt(token: str, *, issuer: Optional[str] = None, audience: Optional[str] = None,
+               jwks_url: Optional[str] = None, hs256_secret: Optional[str] = None) -> Dict[str, Any]:
+    options = {"verify_signature": True, "verify_exp": True, "verify_nbf": True}
+    if jwks_url:
+        jwk_client = PyJWKClient(jwks_url)
+        signing_key = jwk_client.get_signing_key_from_jwt(token).key
+        return jwt.decode(token, signing_key, algorithms=["RS256", "ES256"], audience=audience, issuer=issuer, options=options)
+    elif hs256_secret:
+        return jwt.decode(token, hs256_secret, algorithms=["HS256"], audience=audience, issuer=issuer, options=options)
+    else:
+        raise ValueError("No validation method configured (jwks_url or hs256_secret required)")
+
+
+def now_ts() -> int:
+    return int(time.time())
+
+

google_adk_extras/auth/sql_store.py

@@ -0,0 +1,183 @@
+from __future__ import annotations
+
+import hashlib
+import os
+import secrets
+from datetime import datetime, timedelta, timezone
+from typing import Optional
+
+try:
+    from sqlalchemy import Column, String, DateTime, create_engine, Text, Boolean
+    from sqlalchemy.orm import declarative_base, sessionmaker
+except ImportError as e:
+    raise ImportError(
+        "SQLAlchemy is required for the auth SQL store. Install with: pip install sqlalchemy"
+    ) from e
+
+
+Base = declarative_base()
+
+
+def _pbkdf2(password: str, salt: str) -> str:
+    dk = hashlib.pbkdf2_hmac("sha256", password.encode(), salt.encode(), 200_000)
+    return dk.hex()
+
+
+def hash_password(password: str) -> str:
+    salt = secrets.token_hex(16)
+    return f"pbkdf2_sha256${salt}${_pbkdf2(password, salt)}"
+
+
+def verify_password(password: str, stored: str) -> bool:
+    try:
+        algo, salt, digest = stored.split("$", 2)
+        if algo != "pbkdf2_sha256":
+            # fallback for plaintext in tests
+            return secrets.compare_digest(password, stored)
+        return secrets.compare_digest(_pbkdf2(password, salt), digest)
+    except Exception:
+        return secrets.compare_digest(password, stored)
+
+
+class User(Base):
+    __tablename__ = "auth_users"
+    id = Column(String, primary_key=True)
+    username = Column(String, unique=True, index=True, nullable=False)
+    password_hash = Column(String, nullable=False)
+    roles = Column(String, default="")  # comma-separated
+    created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc))
+    disabled = Column(Boolean, default=False)
+
+
+class RefreshToken(Base):
+    __tablename__ = "auth_refresh_tokens"
+    jti = Column(String, primary_key=True)
+    user_id = Column(String, index=True, nullable=False)
+    expires_at = Column(DateTime(timezone=True), nullable=False)
+    revoked_at = Column(DateTime(timezone=True), nullable=True)
+    fingerprint = Column(String, nullable=True)
+
+
+class ApiKey(Base):
+    __tablename__ = "auth_api_keys"
+    id = Column(String, primary_key=True)
+    user_id = Column(String, index=True, nullable=True)
+    key_hash = Column(String, nullable=False)
+    name = Column(String, nullable=True)
+    created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc))
+    revoked_at = Column(DateTime(timezone=True), nullable=True)
+
+
+class AuthStore:
+    def __init__(self, database_url: str):
+        self.engine = create_engine(database_url)
+        Base.metadata.create_all(self.engine)
+        self.Session = sessionmaker(bind=self.engine, autoflush=False, autocommit=False)
+
+    def create_user(self, username: str, password: str, user_id: Optional[str] = None) -> str:
+        import uuid
+        uid = user_id or str(uuid.uuid4())
+        with self.Session() as s:
+            u = User(id=uid, username=username, password_hash=hash_password(password))
+            s.add(u)
+            s.commit()
+        return uid
+
+    def authenticate_basic(self, username: str, password: str) -> Optional[str]:
+        with self.Session() as s:
+            u: Optional[User] = s.query(User).filter_by(username=username).first()
+            if not u or u.disabled:
+                return None
+            if verify_password(password, u.password_hash):
+                return u.id
+        return None
+
+    def issue_refresh(self, user_id: str, ttl_seconds: int, fingerprint: Optional[str] = None) -> str:
+        import uuid
+        jti = str(uuid.uuid4())
+        with self.Session() as s:
+            rt = RefreshToken(
+                jti=jti,
+                user_id=user_id,
+                expires_at=datetime.now(timezone.utc) + timedelta(seconds=ttl_seconds),
+                fingerprint=fingerprint,
+            )
+            s.add(rt)
+            s.commit()
+        return jti
+
+    def verify_refresh(self, jti: str, user_id: str, fingerprint: Optional[str] = None) -> bool:
+        with self.Session() as s:
+            rt: Optional[RefreshToken] = s.query(RefreshToken).filter_by(jti=jti, user_id=user_id).first()
+            if not rt or rt.revoked_at is not None:
+                return False
+            if rt.expires_at <= datetime.now(timezone.utc):
+                return False
+            if fingerprint and rt.fingerprint and rt.fingerprint != fingerprint:
+                return False
+            return True
+
+    def revoke_refresh(self, jti: str) -> None:
+        with self.Session() as s:
+            rt: Optional[RefreshToken] = s.query(RefreshToken).filter_by(jti=jti).first()
+            if not rt:
+                return
+            rt.revoked_at = datetime.now(timezone.utc)
+            s.add(rt)
+            s.commit()
+
+    # API Keys
+    def _hash_api_key(self, key: str) -> str:
+        # Reuse PBKDF2; different prefix
+        salt = secrets.token_hex(16)
+        return f"api_pbkdf2_sha256${salt}${_pbkdf2(key, salt)}"
+
+    def _verify_api_key(self, key: str, stored: str) -> bool:
+        try:
+            algo, salt, digest = stored.split("$", 3)
+            if algo != "api_pbkdf2_sha256":
+                return secrets.compare_digest(key, stored)
+            return secrets.compare_digest(_pbkdf2(key, salt), digest)
+        except Exception:
+            return False
+
+    def create_api_key(self, user_id: Optional[str] = None, name: Optional[str] = None) -> tuple[str, str]:
+        import uuid
+        key_plain = secrets.token_urlsafe(32)
+        key_id = str(uuid.uuid4())
+        with self.Session() as s:
+            rec = ApiKey(id=key_id, user_id=user_id, key_hash=self._hash_api_key(key_plain), name=name)
+            s.add(rec)
+            s.commit()
+        return key_id, key_plain
+
+    def list_api_keys(self):
+        with self.Session() as s:
+            rows = s.query(ApiKey).all()
+            return [
+                {
+                    "id": r.id,
+                    "user_id": r.user_id,
+                    "name": r.name,
+                    "created_at": r.created_at.isoformat() if r.created_at else None,
+                    "revoked": r.revoked_at is not None,
+                }
+                for r in rows
+            ]
+
+    def revoke_api_key(self, key_id: str) -> None:
+        with self.Session() as s:
+            rec = s.query(ApiKey).filter_by(id=key_id).first()
+            if not rec:
+                return
+            rec.revoked_at = datetime.now(timezone.utc)
+            s.add(rec)
+            s.commit()
+
+    def verify_api_key(self, key: str) -> bool:
+        with self.Session() as s:
+            rows = s.query(ApiKey).filter(ApiKey.revoked_at.is_(None)).all()
+            for r in rows:
+                if self._verify_api_key(key, r.key_hash):
+                    return True
+        return False

google_adk_extras/custom_agent_loader.py

@@ -7,7 +7,7 @@ thread-safe registry management.
 
 import logging
 import threading
-from typing import Dict, List, Optional
+from typing import Dict, List
 
 from google.adk.agents.base_agent import BaseAgent
 from google.adk.cli.utils.base_agent_loader import BaseAgentLoader

google_adk_extras/enhanced_adk_web_server.py

@@ -5,13 +5,11 @@ AdkWebServer to use our EnhancedRunner with advanced features.
 """
 
 import os
-from typing import Optional
 
 from google.adk.cli.adk_web_server import AdkWebServer
 from google.adk.auth.credential_service.in_memory_credential_service import InMemoryCredentialService
 from google.adk.cli.utils import cleanup
 from google.adk.cli.utils import envs
-from google.adk.runners import Runner
 
 from .enhanced_runner import EnhancedRunner
 

google_adk_extras/enhanced_fastapi.py

@@ -22,7 +22,6 @@ from watchdog.observers import Observer
 
 from google.adk.artifacts.gcs_artifact_service import GcsArtifactService
 from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService
-from google.adk.auth.credential_service.in_memory_credential_service import InMemoryCredentialService
 from google.adk.auth.credential_service.base_credential_service import BaseCredentialService
 from google.adk.evaluation.local_eval_set_results_manager import LocalEvalSetResultsManager
 from google.adk.evaluation.local_eval_sets_manager import LocalEvalSetsManager
@@ -35,6 +34,7 @@ from google.adk.sessions.database_session_service import DatabaseSessionService
 from google.adk.utils.feature_decorator import working_in_progress
 from google.adk.cli.adk_web_server import AdkWebServer
 from .enhanced_adk_web_server import EnhancedAdkWebServer
+from .auth import attach_auth, AuthConfig, JwtIssuerConfig, JwtValidatorConfig
 from .streaming import StreamingController, StreamingConfig
 from google.adk.cli.utils import envs
 from google.adk.cli.utils import evals
@@ -69,6 +69,8 @@ def get_enhanced_fast_api_app(
     # Streaming layer (optional)
     enable_streaming: bool = False,
     streaming_config: Optional[StreamingConfig] = None,
+    # Auth layer (optional)
+    auth_config: Optional[AuthConfig] = None,
 ) -> FastAPI:
     """Enhanced version of Google ADK's get_fast_api_app with EnhancedRunner integration.
     
@@ -601,4 +603,7 @@ def get_enhanced_fast_api_app(
 
         app.include_router(router)
 
+    # Attach optional auth layer last so all routes are covered
+    attach_auth(app, auth_config)
+
     return app

google_adk_extras/memory/mongo_memory_service.py

@@ -1,6 +1,5 @@
 """MongoDB-based memory service implementation using PyMongo."""
 
-import json
 import logging
 from typing import Optional, List
 import re

google_adk_extras/memory/sql_memory_service.py

@@ -2,7 +2,7 @@
 
 import json
 import logging
-from typing import Optional, List
+from typing import Optional
 import re
 from datetime import datetime, timezone
 

google_adk_extras/memory/yaml_file_memory_service.py

@@ -1,9 +1,7 @@
 """YAML file-based memory service implementation."""
 
-import os
-import json
 import logging
-from typing import Optional, List
+from typing import List
 from pathlib import Path
 import re
 from datetime import datetime

google_adk_extras/sessions/mongo_session_service.py

@@ -1,6 +1,5 @@
 """MongoDB-based session service implementation."""
 
-import json
 import time
 import uuid
 from typing import Any, Optional, Dict

google_adk_extras/sessions/redis_session_service.py

@@ -3,7 +3,7 @@
 import json
 import time
 import uuid
-from typing import Any, Optional, Dict
+from typing import Any, Optional
 
 try:
     import redis

google_adk_extras/sessions/yaml_file_session_service.py

@@ -1,9 +1,7 @@
 """YAML file-based session service implementation."""
 
-import json
 import time
 import uuid
-import os
 from typing import Any, Optional
 from pathlib import Path
 

google_adk_extras/streaming/streaming_controller.py

@@ -1,9 +1,9 @@
 import asyncio
 import time
 from dataclasses import dataclass, field
-from typing import Any, Dict, Optional, Set, Callable, Awaitable
+from typing import Any, Dict, Optional, Callable, Awaitable
 
-from fastapi import WebSocket, HTTPException
+from fastapi import HTTPException
 from pydantic import BaseModel
 
 from google.adk.events.event import Event