honcho/src/security.py

import datetime
import logging
import os
from typing import Annotated, Optional

import jwt
from fastapi import Depends, Request
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from pydantic import BaseModel
from sqlalchemy.ext.asyncio import AsyncSession

from src.dependencies import get_db

from .exceptions import AuthenticationException

logger = logging.getLogger(__name__)

USE_AUTH = os.getenv("USE_AUTH", "False").lower() == "true"
AUTH_JWT_SECRET = os.getenv("AUTH_JWT_SECRET", "") if USE_AUTH else ""

if USE_AUTH and AUTH_JWT_SECRET == "":
    print(
        "\n    ERROR: No JWT secret provided. Set the AUTH_JWT_SECRET environment variable.\n"
    )
    exit(1)

security = HTTPBearer(
    auto_error=False,
)


#
#  jwt params
#  all optional, used to produce tokens valid for different routes
#  hierarchy: app > user > ( session / collection )
#  routes that involve a 'name' parameter require permissions for the parent object
#  name routes are considered 'queries' as names are mutable properties
#
#  note: add routes without parameters that assume the most immediately scoped key is providing
#
class JWTParams(BaseModel):
    """
    JWT parameters used to produce tokens valid for different routes.
    Hierarchy: app > user > (session / collection)

    All routers require at least the most tightly scoped parameter.
    Routes will accept a JWT with a scope higher in the hierarchy.

    Names shortened to minimize token size. Timestamp is included
    so that many unique tokens can be generated for the same resource.
    Note that the timestamp itself is not used for security, and can
    be omitted, such as when Honcho generates the initial admin JWT.

    Fields (all optional other than `t`):

    `t`: a string timestamp of when the JWT was created
    `exp`: a string timestamp of when the JWT expires (optional)
    `ad`: a boolean flag indicating if the JWT is an admin JWT
    `ap`: (string) app id
    `us`: (string) user id
    `se`: (string) session id
    `co`: (string) collection id
    """

    t: str = datetime.datetime.now().isoformat()
    exp: Optional[str] = None
    ad: Optional[bool] = None
    ap: Optional[str] = None
    us: Optional[str] = None
    se: Optional[str] = None
    co: Optional[str] = None


def create_admin_jwt() -> str:
    """Create a JWT for admin operations."""
    params = JWTParams(t="", ad=True)
    key = create_jwt(params)
    return key


def create_jwt(params: JWTParams) -> str:
    """Create a JWT token from the given parameters."""
    payload = {k: v for k, v in params.__dict__.items() if v is not None}
    return jwt.encode(payload, AUTH_JWT_SECRET.encode("utf-8"), algorithm="HS256")


async def verify_jwt(token: str) -> JWTParams:
    """Verify a JWT token and return the decoded parameters."""

    params = JWTParams()
    try:
        decoded = jwt.decode(
            token, AUTH_JWT_SECRET.encode("utf-8"), algorithms=["HS256"]
        )
        if "t" in decoded:
            params.t = decoded["t"]
        if "exp" in decoded:
            params.exp = decoded["exp"]
            if (
                params.exp
                and datetime.datetime.fromisoformat(params.exp)
                < datetime.datetime.now()
            ):
                raise AuthenticationException("JWT expired")
        if "ad" in decoded:
            params.ad = decoded["ad"]
        if "ap" in decoded:
            params.ap = decoded["ap"]
        if "us" in decoded:
            params.us = decoded["us"]
        if "se" in decoded:
            params.se = decoded["se"]
        if "co" in decoded:
            params.co = decoded["co"]
        return params
    except jwt.PyJWTError:
        raise AuthenticationException("Invalid JWT") from None


def require_auth(
    admin: Optional[bool] = None,
    app_id: Optional[str] = None,
    user_id: Optional[str] = None,
    session_id: Optional[str] = None,
    collection_id: Optional[str] = None,
):
    """
    Generate a dependency that requires authentication for the given parameters.
    """

    async def auth_dependency(
        request: Request,
        credentials: HTTPAuthorizationCredentials = Depends(security),
        db: AsyncSession = Depends(get_db),
    ):
        app_id_param = (
            request.path_params.get(app_id) or request.query_params.get(app_id)
            if app_id
            else None
        )
        user_id_param = (
            request.path_params.get(user_id) or request.query_params.get(user_id)
            if user_id
            else None
        )
        session_id_param = (
            request.path_params.get(session_id) or request.query_params.get(session_id)
            if session_id
            else None
        )
        collection_id_param = (
            request.path_params.get(collection_id)
            or request.query_params.get(collection_id)
            if collection_id
            else None
        )

        return await auth(
            credentials=credentials,
            admin=admin,
            app_id=app_id_param,
            user_id=user_id_param,
            session_id=session_id_param,
            collection_id=collection_id_param,
        )

    return auth_dependency


async def auth(
    credentials: Annotated[HTTPAuthorizationCredentials, Depends(security)],
    admin: Optional[bool] = None,
    app_id: Optional[str] = None,
    user_id: Optional[str] = None,
    session_id: Optional[str] = None,
    collection_id: Optional[str] = None,
) -> JWTParams:
    """Authenticate the given JWT and return the decoded parameters."""
    if not USE_AUTH:
        return JWTParams(t="", ad=True)
    if not credentials or not credentials.credentials:
        logger.warning("No access token provided")
        raise AuthenticationException("No access token provided")

    jwt_params = await verify_jwt(credentials.credentials)

    # based on api operation, verify api key based on that key's permissions
    if jwt_params.ad:
        return jwt_params
    if admin:
        raise AuthenticationException("Resource requires admin privileges")

    # Check if the JWT has direct access to the requested resource
    # For session or collection level access
    if session_id and jwt_params.se == session_id:
        return jwt_params
    if collection_id and jwt_params.co == collection_id:
        return jwt_params

    # For user level access - can access all sessions/collections under this user
    if user_id and jwt_params.us == user_id:
        return jwt_params

    # For app level access - can access all users/sessions/collections under this app
    if app_id and jwt_params.ap == app_id:
        return jwt_params

    if any([session_id, collection_id, user_id, app_id]):
        print([session_id, collection_id, user_id, app_id])
        print(jwt_params)
        raise AuthenticationException("JWT not permissioned for this resource")

    # Route did not specify any parameters, so it should parse parameters itself
    return jwt_params