@@ -1,8 +1,6 @@ from datetime import date, timedelta -from decimal import Decimal from uuid import UUID -from asgiref.sync import sync_to_async from django.db.models import Prefetch from django.utils import timezone from django.utils.translation import gettext_lazy as _ @@ -52,8 +50,6 @@ class UserSelector: 'business_admin', ): user.show_balance = user.business_account.show_balance - if not user.show_balance: - user.payment_plan.current_token_balance = Decimal('0') user.payment_plan.plan = user.business_account.parent_company.user.payment_plan.plan refresh_token = OutstandingToken.objects.filter(user=user).order_by('-created_at').first() if not refresh_token or refresh_token.expires_at < timezone.now() and provider == 'yandex': @@ -0,0 +1,54 @@ +# Authentication mapper +from django.db.models import Prefetch + +from ml_model.models import NeuronModel + + +def _gen_only(chain: str, *fields: str): + return [(f"{chain}__{f}") for f in fields] + + +PATH_PREFETCH_MAP = { + '/api/v1/auth/me': { + 'select': ( + 'host_account', + 'business_account__parent_company__user__payment_plan__plan', + 'payment_plan__plan', + ), + 'prefetch': ( + Prefetch('payment_plan__plan__accessed_models', queryset=NeuronModel.objects.only('slug')), + Prefetch( + 'business_account__parent_company__user__payment_plan__plan', + queryset=NeuronModel.objects.only('slug'), + ), + 'social_auth' + ), + 'only': ( + 'uid', 'first_name', 'last_name', 'username', 'created_at', 'email', 'active', + 'is_superuser', 'is_staff', 'is_confirmed', 'is_subscribed_to_emails', 'profile_picture_name', + *_gen_only('business_account', 'uid', 'parent_company__uid', 'show_balance', + 'account_privileges', 'parent_company__user__uid'), + *_gen_only('payment_plan', 'uid', 'last_payment_at', 'next_payment_at'), + *_gen_only('payment_plan__plan', 'uid', 'title', 'price', 'tokens_per_plan', 'duration', 'points'), + *_gen_only('business_account__parent_company__user__payment_plan', 'uid', 'last_payment_at', + 'next_payment_at', 'plan__uid', 'plan__title', 'plan__price', 'plan__tokens_per_plan', + 'plan__duration', 'plan__points') + ) + }, + '/api/v1/payments/user-balance': { + 'select': ( + 'host_account', + 'business_account__group', + 'payment_plan', + 'business_account__parent_company__user__payment_plan', + ), + 'prefetch': (), + 'only': ( + 'uid', 'is_staff', 'is_superuser', + *_gen_only('business_account', 'account_privileges', 'acceptance_status', 'token_limit'), + *_gen_only('business_account__group', 'uid', 'token_limit'), + *_gen_only('payment_plan', 'uid', 'current_token_balance'), + *_gen_only('business_account__parent_company__user__payment_plan', 'uid', 'current_token_balance') + ) + } +} @@ -1,36 +0,0 @@ -import logging -from typing import Any - -from django.contrib.auth import logout -from django.core.exceptions import PermissionDenied - -from authentication.models.business_host import BusinessUserHost -from authentication.selectors.user_selector import UserSelector -from authentication.utils import get_client_ip - -logger = logging.getLogger(__name__) - - -class CompanyIPMiddleware: - def __init__(self, get_response) -> None: - self.get_response = get_response - - def __call__(self, request) -> Any: - if request.user.is_authenticated: - acc_type = UserSelector(request.user).check_account_type() - if acc_type == 'regular' or request.user.is_staff or request.user.is_superuser: - pass - else: - company: BusinessUserHost = ( - request.user.host_account - if acc_type == 'business_host' - else request.user.business_account.parent_company - ) - if ( - company.ip_whitelist.is_enabled - and not company.ip_whitelist.ips.filter(ip=get_client_ip(request)).exists() - ): - logout(request) - raise PermissionDenied - response = self.get_response(request) - return response @@ -1,7 +1,7 @@ from typing import Any import jwt -from asgiref.sync import async_to_sync +from asgiref.sync import async_to_sync, sync_to_async from django.conf import settings from django.http import HttpRequest from django.utils.translation import gettext as _ @@ -15,8 +15,24 @@ from rest_framework.authentication import BaseAuthentication from rest_framework.exceptions import AuthenticationFailed from authentication.exceptions import InvalidToken -from authentication.models import CustomUserModel +from authentication.mapper import PATH_PREFETCH_MAP +from authentication.models import CustomUserModel, BusinessUserHost from authentication.services.token import TokenService +from authentication.utils import get_client_ip + + +def _check_ip_client(user: CustomUserModel, request: HttpRequest): + if (acc_type := user.account_type) != 'regular' and not user.is_staff and not user.is_superuser: + company: BusinessUserHost = ( + user.host_account + if acc_type == 'business_host' + else user.business_account.parent_company + ) + if ( + company.ip_whitelist.is_enabled + and not company.ip_whitelist.ips.filter(ip=get_client_ip(request)).exists() + ): + raise AuthenticationFailed(_('Forbidden'), code='ip_address_not_allowed') class JWTAuthentication(BaseAuthentication): @@ -40,6 +56,7 @@ class JWTAuthentication(BaseAuthentication): 'business_account__parent_company__user__payment_plan__plan', 'host_account' ).get(uid=payload['uid']) + _check_ip_client(user, request) except CustomUserModel.DoesNotExist as exc: raise AuthenticationFailed(_('User not found'), code='user_not_found') from exc return user, None @@ -55,62 +72,48 @@ class SimpleJWTScheme(BaseSimpleJWTScheme): class SyncAuthBearer(HttpBearer): def authenticate(self, request: HttpRequest, token: str) -> Any | None: + mapper = PATH_PREFETCH_MAP.get(request.path, {}) try: user_payload = async_to_sync(TokenService.decode)(token=token) request.provider = 'air' - return CustomUserModel.objects.select_related( - 'payment_plan', - 'payment_plan__plan', - 'business_account', - 'business_account__group', - 'business_account__parent_company__user__payment_plan', - 'business_account__parent_company__user__payment_plan__plan', - 'host_account' + user = CustomUserModel.objects.select_related( + *mapper.get('select', []) ).prefetch_related( - 'payment_plan__plan__accessed_models', - 'business_account__parent_company__user__payment_plan__plan__accessed_models', - 'social_auth' + *mapper.get('prefetch', []) + ).only( + *mapper.get('only', []) ).get( **{key: user_payload[f'{key}'] for key in settings.JWT_SETTINGS['encode_attributes']} ) except InvalidToken: access = AccessToken.objects.select_related( 'user', - 'user__payment_plan', - 'user__payment_plan__plan', - 'user__business_account', - 'user__business_account__group', - 'user__business_account__parent_company__user__payment_plan', - 'user__business_account__parent_company__user__payment_plan__plan', - 'user__host_account' + *(f'user__{el}' for el in mapper.get('select', [])) ).prefetch_related( - 'user__payment_plan__plan__accessed_models', - 'user__business_account__parent_company__user__payment_plan__plan__accessed_models', - 'user__social_auth' + *(f'user__{el}' for el in mapper.get('prefetch', [])) + ).only( + *(f'user__{el}' for el in mapper.get('only', [])) ).filter(token=token).order_by('-created') request.provider = 'yandex' if a := access.first(): - return a.user + user = a.user raise HttpError(401, _('Access token expired or does not exist')) + _check_ip_client(user, request) + return user class AsyncAuthBearer(HttpBearer): async def authenticate(self, request: HttpRequest, token: str) -> Any | None: + mapper = PATH_PREFETCH_MAP.get(request.path, {}) try: user_payload = await TokenService.decode(token=token) request.provider = 'air' - return await CustomUserModel.objects.select_related( - 'payment_plan', - 'payment_plan__plan', - 'business_account', - 'business_account__group', - 'business_account__parent_company__user__payment_plan', - 'business_account__parent_company__user__payment_plan__plan', - 'host_account' + user = await CustomUserModel.objects.select_related( + *mapper.get('select', []) ).prefetch_related( - 'payment_plan__plan__accessed_models', - 'business_account__parent_company__user__payment_plan__plan__accessed_models', - 'social_auth' + *mapper.get('prefetch', []) + ).only( + *mapper.get('only', []) ).aget( **{key: user_payload[f'{key}'] for key in settings.JWT_SETTINGS['encode_attributes']} ) @@ -118,19 +121,15 @@ class AsyncAuthBearer(HttpBearer): try: access = await AccessToken.objects.select_related( 'user', - 'user__payment_plan', - 'user__payment_plan__plan', - 'user__business_account', - 'user__business_account__group', - 'user__business_account__parent_company__user__payment_plan', - 'user__business_account__parent_company__user__payment_plan__plan', - 'user__host_account' + *(f'user__{el}' for el in mapper.get('select', [])) ).prefetch_related( - 'user__payment_plan__plan__accessed_models', - 'user__business_account__parent_company__user__payment_plan__plan__accessed_models', - 'user__social_auth' + *(f'user__{el}' for el in mapper.get('prefetch', [])) + ).only( + *(f'user__{el}' for el in mapper.get('only', [])) ).aget(token=token) request.provider = 'yandex' - return access.user + user = access.user except AccessToken.DoesNotExist: raise HttpError(401, _('Access token expired or does not exist')) + await sync_to_async(_check_ip_client)(user, request) + return user @@ -31,7 +31,6 @@ MIDDLEWARE = [ 'core.middleware.PyroscopeWrapper', 'django.middleware.csrf.CsrfViewMiddleware', 'django.contrib.auth.middleware.AuthenticationMiddleware', - 'authentication.middleware.CompanyIPMiddleware', 'django.contrib.messages.middleware.MessageMiddleware', 'django.middleware.clickjacking.XFrameOptionsMiddleware', 'social_django.middleware.SocialAuthExceptionMiddleware', @@ -436,6 +435,7 @@ CACHEOPS_DEGRADE_ON_FAILURE = True if CACHEOPS_REDIS: CACHEOPS = { # 'authentication.*': {'ops': 'all', 'timeout': 60 * 60}, + 'authentication.companyipwhitelist': {'ops': 'all', 'timeout': 60 * 60}, 'ml_model.*': {'ops': 'all', 'timeout': 60 * 60}, 'tools.chats.*': {'ops': 'all', 'timeout': 60 * 60}, 'tools.media.*': {'ops': 'all', 'timeout': 60 * 60}, @@ -59,11 +59,8 @@ class PaymentPlanSelector: return PaymentPlanSerializer(plan) def get_current_balance(self) -> Decimal: - user_type = self.user.account_type if ( - user_type == 'business_account' - or user_type == 'business_admin' - or user_type == 'business_security' + self.user.account_type in ('business_account', 'business_admin', 'business_security') ) and self.user.business_account.acceptance_status == InvitationStatus.ACCEPTED: balance = ( self.user.business_account.group.token_limit @@ -1,5 +1,4 @@ from datetime import date -from decimal import Decimal from typing import List from uuid import UUID @@ -29,12 +28,6 @@ class UserPlanDetailSchema(Schema): plan: PaymentPlanSchema last_payment_at: date next_payment_at: date - current_token_balance: Decimal - - @field_validator('current_token_balance', mode='before') - @classmethod - def get_current_token_balance(cls, value: Decimal) -> Decimal: - return Decimal(f"{value:.2f}" if value == value.to_integral() else value.normalize().to_eng_string()) class PromoCodeSchema(ModelSchema):