Auth revamp

fai-rag-app/fai-backend/fai_backend/auth_v2/api_key/dependencies.py

@@ -0,0 +1,6 @@
+from fai_backend.auth_v2.api_key.service import ApiKeyService
+from fai_backend.repositories import api_key_repo
+
+
+async def get_api_key_service() -> ApiKeyService:
+    return ApiKeyService(repo=api_key_repo)

fai-rag-app/fai-backend/fai_backend/auth_v2/api_key/models.py

@@ -0,0 +1,21 @@
+from beanie import Document
+from pydantic import BaseModel
+
+
+class ApiKeyModel(BaseModel):
+    key_hash: str
+    key_hint: str
+    scopes: list[str]
+
+
+class ReadOnlyApiKeyModel(BaseModel):
+    revoke_id: str
+    key_hint: str
+    scopes: list[str]
+
+
+class ApiKeyDocumentModel(Document):
+    api_key: ApiKeyModel
+
+    class Settings:
+        name = 'api_key'

fai-rag-app/fai-backend/fai_backend/auth_v2/api_key/routes.py

@@ -0,0 +1,91 @@
+from fastapi import APIRouter, status, HTTPException
+from pydantic import BaseModel
+
+from fai_backend.auth_v2.api_key.dependencies import get_api_key_service
+from fai_backend.auth_v2.api_key.models import ReadOnlyApiKeyModel
+from fai_backend.auth_v2.authentication.models import AuthenticatedIdentity
+from fai_backend.auth_v2.authorization.factory import AuthorizationFactory
+from fai_backend.auth_v2.fastapi_auth import AuthRouterDecorator
+
+router = APIRouter(
+    prefix='/api/auth',
+    tags=['Auth', 'API Key']
+)
+auth = AuthRouterDecorator(router)
+
+
+class CreateApiKeyRequest(BaseModel):
+    scopes: list[str]
+
+
+class CreateApiKeyResponse(BaseModel):
+    key: str
+    revoke_id: str
+
+
+@auth.post(
+    '/apikey',
+    ['can_manage_api_keys'],
+    response_model=CreateApiKeyResponse,
+    summary='Create API Key',
+    description='''
+    Create a new API key with the given scopes.
+    ''',
+    status_code=status.HTTP_201_CREATED,
+    response_400_description='No scopes provided.'
+)
+async def create_api_key(body: CreateApiKeyRequest, auth_identity: AuthenticatedIdentity):
+    desired_scopes = [scope for scope in body.scopes if len(scope) > 0]
+
+    if len(desired_scopes) == 0:
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail='No scopes provided.'
+        )
+
+    service = await get_api_key_service()
+
+    authorization_provider = await AuthorizationFactory.get()
+    if not await authorization_provider.has_scopes(auth_identity, desired_scopes):
+        raise HTTPException(
+            status_code=status.HTTP_403_FORBIDDEN,
+            detail=f"Can not assign scope(s) outside of creator's scopes to an API key."
+        )
+
+    revoke_id, key = await service.create(scopes=desired_scopes)
+    return CreateApiKeyResponse(key=key, revoke_id=str(revoke_id))
+
+
+class ListApiKeyResponse(BaseModel):
+    api_keys: list[ReadOnlyApiKeyModel]
+
+
+@auth.get(
+    '/apikey',
+    ['can_manage_api_keys'],
+    response_model=ListApiKeyResponse,
+    description='''
+List information about all API keys.
+
+Keys themselves are redacted for security purposes.
+    ''',
+)
+async def list_api_keys():
+    service = await get_api_key_service()
+    keys = await service.list()
+    return ListApiKeyResponse(api_keys=keys)
+
+
+@auth.delete(
+    '/apikey/{revoke_id}',
+    ['can_manage_api_keys'],
+    description='''
+Revoke a specific API key (permanently delete it).
+
+The `revoke_id` can be found through the `create` and `list` operations.
+    ''',
+    status_code=status.HTTP_204_NO_CONTENT,
+)
+async def revoke_api_key(revoke_id: str):
+    service = await get_api_key_service()
+    await service.revoke(revoke_id)

fai-rag-app/fai-backend/fai_backend/auth_v2/api_key/service.py

@@ -0,0 +1,59 @@
+import crypt
+import hashlib
+import hmac
+import uuid
+
+from fai_backend.auth_v2.api_key.models import ApiKeyDocumentModel, ApiKeyModel, ReadOnlyApiKeyModel
+from fai_backend.config import settings
+from fai_backend.repository.interface import IAsyncRepo
+
+
+class ApiKeyService:
+    def __init__(self, repo: IAsyncRepo[ApiKeyDocumentModel]):
+        self._repo = repo
+
+    async def create(self, scopes: list[str]) -> (str, str):
+        key = f'fai-{uuid.uuid4().hex}'
+        key_hash = self._hash_api_key(key)
+        key_hint = self._create_key_hint(key)
+        api_key = ApiKeyModel(key_hash=key_hash, key_hint=key_hint, scopes=scopes)
+        new_entry = await self._repo.create(ApiKeyDocumentModel(api_key=api_key))
+        return new_entry.id, key
+
+    async def revoke(self, key_repo_id: str):
+        await self._repo.delete(key_repo_id)
+
+    async def list(self) -> list[ReadOnlyApiKeyModel]:
+        all_keys = await self._repo.list()
+        return [ApiKeyService._to_read_only_api_key(key_data) for key_data in all_keys]
+
+    async def find_by_key(self, key: str) -> ReadOnlyApiKeyModel | None:
+        all_keys = await self._repo.list()
+        key_hash = self._hash_api_key(key)
+        return next(
+            (ApiKeyService._to_read_only_api_key(key_data) for key_data in all_keys if
+             key_data.api_key.key_hash == key_hash),
+            None)
+
+    async def find_by_revoke_id(self, revoke_id: str) -> ReadOnlyApiKeyModel | None:
+        all_keys = await self._repo.list()
+        return next(
+            (ApiKeyService._to_read_only_api_key(key_data) for key_data in all_keys if
+             str(key_data.id) == revoke_id),
+            None)
+
+    @staticmethod
+    def _hash_api_key(key: str) -> str:
+        return hmac.new(settings.SECRET_KEY.encode(), key.encode(), hashlib.sha256).hexdigest()
+
+    @staticmethod
+    def _to_read_only_api_key(key_data: ApiKeyDocumentModel):
+        return ReadOnlyApiKeyModel(
+            revoke_id=str(key_data.id),
+            key_hint=key_data.api_key.key_hint,
+            scopes=key_data.api_key.scopes,
+        )
+
+    @staticmethod
+    def _create_key_hint(key: str) -> str:
+        return key[:8] + "..." + key[-4:]

fai-rag-app/fai-backend/fai_backend/auth_v2/authentication/factory.py

@@ -0,0 +1,17 @@
+from typing import Callable
+
+from fai_backend.auth_v2.authentication.impl.api_key import ApiKeyProvider
+from fai_backend.auth_v2.authentication.impl.bearer_token import BearerTokenProvider
+from fai_backend.auth_v2.authentication.models import IAuthenticationProvider, AuthenticationType
+
+
+class AuthenticationFactory:
+    @staticmethod
+    async def get(auth_type: str) -> IAuthenticationProvider:
+        mapping: dict[str, Callable[[], IAuthenticationProvider]] = {
+            AuthenticationType.API_KEY: ApiKeyProvider,
+            AuthenticationType.BEARER_TOKEN: BearerTokenProvider,
+        }
+        if auth_type in mapping:
+            return mapping[auth_type]()
+        raise ValueError(f"Authentication type '{auth_type}' not supported")

fai-rag-app/fai-backend/fai_backend/auth_v2/authentication/fastapi_sources.py

@@ -0,0 +1,18 @@
+from fastapi.security import APIKeyHeader
+
+from fai_backend.auth.security import access_security
+
+api_key_source = APIKeyHeader(
+    name='X-Api-Key',
+    description='API Key. Can be created/revoked by an administrator.',
+    auto_error=False
+)
+
+# bearer_token_source = HTTPBearer(
+#     scheme_name='User Token',
+#     description='JWT Token associated with an user. Generated by logging in.',
+#     auto_error=False
+# )
+
+# TODO: don't use fastapi_jwt
+bearer_token_source = access_security

fai-rag-app/fai-backend/fai_backend/auth_v2/authentication/impl/api_key.py

@@ -0,0 +1,15 @@
+from fai_backend.auth_v2.api_key.dependencies import get_api_key_service
+from fai_backend.auth_v2.authentication.models import AuthenticatedIdentity, IAuthenticationProvider, AuthenticationType
+
+
+class ApiKeyProvider(IAuthenticationProvider):
+    async def validate(self, data: str | None) -> AuthenticatedIdentity | None:
+        service = await get_api_key_service()
+        key = await service.find_by_key(data)
+
+        if not key:
+            return None
+
+        return AuthenticatedIdentity(
+            uid=key.revoke_id
+        )

fai-rag-app/fai-backend/fai_backend/auth_v2/authentication/impl/bearer_token.py

@@ -0,0 +1,10 @@
+import json
+
+from fai_backend.auth_v2.authentication.models import IAuthenticationProvider, AuthenticatedIdentity, AuthenticationType
+
+
+class BearerTokenProvider(IAuthenticationProvider):
+    async def validate(self, data: str | None) -> AuthenticatedIdentity | None:
+        return AuthenticatedIdentity(
+            uid=json.loads(data)['email']
+        )

fai-rag-app/fai-backend/fai_backend/auth_v2/authentication/models.py

@@ -0,0 +1,18 @@
+from typing import Protocol
+
+from pydantic import BaseModel
+
+
+class AuthenticationType:
+    NONE = 'NONE'
+    API_KEY = 'api_key'
+    BEARER_TOKEN = 'bearer_token'
+
+
+class AuthenticatedIdentity(BaseModel):
+    uid: str
+
+
+class IAuthenticationProvider(Protocol):
+    async def validate(self, data: str | None) -> AuthenticatedIdentity | None:
+        ...

fai-rag-app/fai-backend/fai_backend/auth_v2/authorization/factory.py

@@ -0,0 +1,8 @@
+from fai_backend.auth_v2.authorization.impl.repo import RepoAuthorizationProvider
+from fai_backend.auth_v2.authorization.models import IAuthorizationProvider
+
+
+class AuthorizationFactory:
+    @staticmethod
+    async def get() -> IAuthorizationProvider:
+        return RepoAuthorizationProvider()

fai-rag-app/fai-backend/fai_backend/auth_v2/authorization/impl/repo.py

@@ -0,0 +1,21 @@
+from fai_backend.auth_v2.api_key.dependencies import get_api_key_service
+from fai_backend.auth_v2.authentication.models import AuthenticatedIdentity, AuthenticationType
+from fai_backend.auth_v2.authorization.models import IAuthorizationProvider
+from fai_backend.repositories import users_repo
+
+
+class RepoAuthorizationProvider(IAuthorizationProvider):
+    async def has_scopes(self, identity: AuthenticatedIdentity, scopes: list[str]) -> bool:
+        user = await users_repo.get_user_by_email(identity.uid)
+
+        if user is not None:
+            permissions = user.projects[0].permissions
+            return all(key in permissions and permissions[key] is True for key in scopes)
+
+        api_key_service = await get_api_key_service()
+        api_key = await api_key_service.find_by_revoke_id(identity.uid)
+
+        if api_key is not None:
+            return all(scope in api_key.scopes for scope in scopes)
+
+        return False

fai-rag-app/fai-backend/fai_backend/auth_v2/authorization/models.py

@@ -0,0 +1,8 @@
+from typing import Protocol
+
+from fai_backend.auth_v2.authentication.models import AuthenticatedIdentity
+
+
+class IAuthorizationProvider(Protocol):
+    async def has_scopes(self, identity: AuthenticatedIdentity, scopes: list[str]) -> bool:
+        ...