felicity-lims/backend/felicity_lims/felicity/main.py

from felicity.init import initialize_felicity
if initialize_felicity():
    import logging
    from typing import List
    from fastapi import FastAPI, WebSocket
    from starlette.middleware.cors import CORSMiddleware
    from graphql.execution.executors.asyncio import AsyncioExecutor

    from starlette.middleware.authentication import AuthenticationMiddleware
    from starlette.authentication import (
        AuthenticationBackend,
        AuthenticationError,
        SimpleUser,
        # UnauthenticatedUser,
        AuthCredentials
    )
    import base64
    import binascii

    from felicity.database.session import database # noqa

    from felicity.api.api_v1.api import api_router # noqa
    from felicity.core.config import settings # noqa

    from starlette.graphql import GraphQLApp
    from felicity.gql.schema import gql_schema # noqa
    from felicity.gql.deps import get_current_active_user

    from felicity.apps.job.sched import felicity_workforce_init, felicity_halt_workforce

    logging.basicConfig(level=logging.INFO)
    logger = logging.getLogger(__name__)

    class FelicityAuthBackend(AuthenticationBackend):
        async def authenticate(self, request):
            # logger.info(f"trial middle ware: {request}")
            # logger.info(f"trial middle ware: {dir(request)}")
            if "Authorization" not in request.headers:
                logger.info(f"No Authorization Data")
                return

            # logger.info(f"Authorization checking: {request.headers}")

            auth = request.headers["Authorization"]
            try:
                scheme, credentials = auth.split()
                if scheme.lower() == 'basic':
                    decoded = base64.b64decode(credentials).decode("ascii")
                    username, _, password = decoded.partition(":")
                    # TODO: You'd want to verify the username and password here if needed
                elif scheme.lower() == 'bearer':
                    """"get is active user from token"""
                    logger.info(f"credentials bearer: {credentials}")
                    user = get_current_active_user(credentials)
                    logger.info(f"User from token: {user.auth.user_name}")
                    username, _, password = user.auth.user_name, None, None
                else:
                    raise AuthenticationError(f'UnKnown Authentication Backend: {scheme.lower()}')

                return AuthCredentials(["authenticated"]), SimpleUser(username)

            except (ValueError, UnicodeDecodeError, binascii.Error) as exc:
                raise AuthenticationError(f'Invalid auth credentials: {exc}')


    flims = FastAPI(
        title=settings.PROJECT_NAME, openapi_url=f"{settings.API_V1_STR}/openapi.json"
    )

    @flims.on_event("startup")
    async def startup():
        await database.connect()
        felicity_workforce_init()

    @flims.on_event("shutdown")
    async def shutdown():
        felicity_halt_workforce()
        await database.disconnect()
        
    # Set all CORS enabled origins
    if settings.BACKEND_CORS_ORIGINS:
        flims.add_middleware(
            CORSMiddleware,
            allow_origins=[str(origin) for origin in settings.BACKEND_CORS_ORIGINS],
            allow_credentials=True,
            allow_methods=["*"],
            allow_headers=["*"],
        )
        flims.add_middleware(
            AuthenticationMiddleware,
            backend=FelicityAuthBackend()
        )

    flims.include_router(api_router, prefix=settings.API_V1_STR)
    flims.add_route("/felicity-gql", GraphQLApp(schema=gql_schema, executor_class=AsyncioExecutor))

    class ConnectionManager:
        def __init__(self):
            self.connections: List[WebSocket] = []

        async def connect(self, websocket: WebSocket):
            await websocket.accept()
            self.connections.append(websocket)

        async def broadcast(self, data: str):
            for connection in self.connections:
                await connection.send_text(data)


    manager = ConnectionManager()

    @flims.websocket("/ws")
    async def websocket_endpoint(websocket: WebSocket):
        await manager.connect(websocket)
        while True:
            data = await websocket.receive_text()
            await manager.broadcast(f"Client {data}")