inbox/services/auth.py

from typing import Optional

from aiohttp.web import HTTPUnauthorized
from aiohttp.client import ClientSession
from pydantic import BaseModel
from functools import wraps
from starlette.authentication import AuthenticationBackend
from starlette.requests import HTTPConnection
from services.db import local_session
from settings import AUTH_URL
from orm.author import Author


class AuthUser(BaseModel):
    user_id: Optional[int]
    username: Optional[str]


class AuthCredentials(BaseModel):
    user_id: Optional[int] = None
    scopes: Optional[dict] = {}
    logged_in: bool = False
    error_message: str = ""


class JWTAuthenticate(AuthenticationBackend):
    async def authenticate(self, request: HTTPConnection):
        logged_in, user_id = await check_auth(request)
        return (
            AuthCredentials(user_id=user_id, logged_in=logged_in),
            AuthUser(user_id=user_id),
        )


async def check_auth(req):
    token = req.headers.get("Authorization")
    gql = (
        {"mutation": "{ getSession { user { id } } }"}
        if "v2" in AUTH_URL
        else {"query": "{ session { user { id } } }"}
    )
    headers = {"Authorization": token, "Content-Type": "application/json"}
    async with ClientSession(headers=headers) as session:
        async with session.post(AUTH_URL, data=gql) as response:
            if response.status != 200:
                return False, None
            r = await response.json()
            user_id = (
                r.get("data", {}).get("session", {}).get("user", {}).get("id", None)
            )
            is_authenticated = user_id is not None
            return is_authenticated, user_id


async def author_id_by_user_id(user_id):
    async with local_session() as session:
        author = session(Author).where(Author.user == user_id).first()
        return author.id


def login_required(f):
    @wraps(f)
    async def decorated_function(*args, **kwargs):
        info = args[1]
        context = info.context
        req = context.get("request")
        is_authenticated, user_id = await check_auth(req)
        if not is_authenticated:
            raise Exception("You are not logged in")
        else:
            # Добавляем author_id в контекст
            author_id = await author_id_by_user_id(user_id)
            context["author_id"] = author_id

        # Если пользователь аутентифицирован, выполняем резолвер
        return await f(*args, **kwargs)

    return decorated_function


def auth_request(f):
    @wraps(f)
    async def decorated_function(*args, **kwargs):
        req = args[0]
        is_authenticated, user_id = await check_auth(req)
        if not is_authenticated:
            raise HTTPUnauthorized()
        else:
            author_id = await author_id_by_user_id(user_id)
            req["author_id"] = author_id
        return await f(*args, **kwargs)

    return decorated_function