demo-backend/ai-service/app/utils/jwt.py

import jwt
from fastapi import Depends, HTTPException, Request, Security, WebSocket, status
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer, SecurityScopes

from app.config.env import env

# Define permission sets
LOAN_OFFICER_PERMISSIONS = ["get:lead", "create:loan", "update:loan"]
LOAN_OFFICER_ADMIN_PERMISSIONS = ["get:reports"]

# Custom exception for unauthorized access


class UnauthorizedException(HTTPException):
    def __init__(self, detail: str, **kwargs):
        super().__init__(status.HTTP_403_FORBIDDEN, detail=detail)


# Custom exception for unauthenticated requests


class UnauthenticatedException(HTTPException):
    def __init__(self):
        super().__init__(status_code=status.HTTP_401_UNAUTHORIZED, detail="Requires authentication")


# Custom HTTPBearer for Auth0


class Auth0HTTPBearer(HTTPBearer):
    async def __call__(self, request: Request) -> HTTPAuthorizationCredentials:
        credentials: HTTPAuthorizationCredentials = await super().__call__(request)
        if credentials and credentials.scheme == "Bearer":
            return credentials
        else:
            raise HTTPException(
                status_code=status.HTTP_403_FORBIDDEN, detail="Invalid or missing authentication credentials"
            )


# Token verification class


class VerifyToken:
    USER_PERMISSIONS = []

    def __init__(self):
        jwks_url = f"https://{env.AUTH0_DOMAIN}/.well-known/jwks.json"
        self.jwks_client = jwt.PyJWKClient(jwks_url)

    async def verify(
        self, security_scopes: SecurityScopes, token: HTTPAuthorizationCredentials = Depends(Auth0HTTPBearer())
    ):
        if not token:
            raise UnauthenticatedException()

        # Extract and verify the JWT token from the credentials
        token_value = token.credentials

        try:
            signing_key = self.jwks_client.get_signing_key_from_jwt(token_value).key
            payload = jwt.decode(
                token_value,
                signing_key,
                algorithms=env.AUTH0_ALGORITHMS,
                audience=env.AUTH0_API_AUDIENCE,
                issuer=env.AUTH0_ISSUER,
            )
        except jwt.exceptions.PyJWKClientError as error:
            raise UnauthorizedException(str(error))
        except jwt.exceptions.DecodeError as error:
            raise UnauthorizedException(str(error))
        except Exception as error:
            raise UnauthorizedException(str(error))

        if security_scopes.scopes:
            self._check_claims(payload, "scope", security_scopes.scopes)

        self.USER_PERMISSIONS = payload.get("permissions", [])

        return payload

    async def verify_websocket(self, websocket: WebSocket, token: str):
        if not token:
            await websocket.close(code=1008)
            raise UnauthenticatedException()

        try:
            signing_key = self.jwks_client.get_signing_key_from_jwt(token).key
            payload = jwt.decode(
                token,
                signing_key,
                algorithms=env.AUTH0_ALGORITHMS,
                audience=env.AUTH0_API_AUDIENCE,
                issuer=env.AUTH0_ISSUER,
            )
            self.USER_PERMISSIONS = payload.get("permissions", [])
        except jwt.exceptions.PyJWKClientError as error:
            await websocket.close(code=1008)
            raise UnauthorizedException(str(error))
        except jwt.exceptions.DecodeError as error:
            await websocket.close(code=1008)
            raise UnauthorizedException(str(error))
        except Exception as error:
            await websocket.close(code=1008)
            raise UnauthorizedException(str(error))

        return payload

    def _check_claims(self, payload, claim_name, expected_value):
        if claim_name not in payload:
            raise UnauthorizedException(detail=f'No claim "{claim_name}" found in token')

        payload_claim = payload[claim_name].split(" ") if claim_name == "scope" else payload[claim_name]

        for value in expected_value:
            if value not in payload_claim:
                raise UnauthorizedException(detail=f'Missing required "{claim_name}" scope')

    def is_loan_officer(self):
        for permission in LOAN_OFFICER_PERMISSIONS:
            if permission not in self.USER_PERMISSIONS:
                raise UnauthorizedException(detail="Permission denied")

    def is_loan_officer_admin(self):
        for permission in LOAN_OFFICER_ADMIN_PERMISSIONS:
            if permission not in self.USER_PERMISSIONS:
                raise UnauthorizedException(detail="Permission denied")