diff --git a/apps/application/views/application_views.py b/apps/application/views/application_views.py index 8c3e8059bcb..c27d820a9fb 100644 --- a/apps/application/views/application_views.py +++ b/apps/application/views/application_views.py @@ -12,7 +12,7 @@ from application.swagger_api.application_api import ApplicationApi from application.swagger_api.application_statistics_api import ApplicationStatisticsApi from application.views.common import get_application_operation_object -from common.auth import TokenAuth, has_permissions +from common.auth import TokenAuth, has_permissions, AllTokenAuth from common.constants.permission_constants import CompareConstants, PermissionConstants, Permission, Group, Operate, \ ViewPermission, RoleConstants from common.exception.app_exception import AppAuthenticationFailed @@ -315,7 +315,7 @@ def get(self, request: Request, application_id: str, app_id: str): 'user_id': request.user.id}).get_application(app_id)) class Profile(APIView): - authentication_classes = [TokenAuth] + authentication_classes = [AllTokenAuth] @action(methods=['GET'], detail=False) @swagger_auto_schema(operation_summary=_("Get application related information"), diff --git a/apps/application/views/chat_views.py b/apps/application/views/chat_views.py index 30d54fa65a4..78b28217243 100644 --- a/apps/application/views/chat_views.py +++ b/apps/application/views/chat_views.py @@ -20,7 +20,7 @@ from application.swagger_api.chat_api import ChatApi, VoteApi, ChatRecordApi, ImproveApi, ChatRecordImproveApi, \ ChatClientHistoryApi, OpenAIChatApi from application.views import get_application_operation_object -from common.auth import TokenAuth, has_permissions, OpenAIKeyAuth +from common.auth import TokenAuth, has_permissions, OpenAIKeyAuth, AllTokenAuth from common.constants.authentication_type import AuthenticationType from common.constants.permission_constants import Permission, Group, Operate, \ RoleConstants, ViewPermission, CompareConstants @@ -70,7 +70,7 @@ def post(self, request: Request, application_id: str): 'user_id': request.user.id}).export(request.data) class Open(APIView): - authentication_classes = [TokenAuth] + authentication_classes = [AllTokenAuth] @action(methods=['GET'], detail=False) @swagger_auto_schema(operation_summary=_("Get the session id according to the application id"), @@ -114,7 +114,7 @@ def post(self, request: Request): data={**request.data, 'user_id': request.user.id}).open()) class Message(APIView): - authentication_classes = [TokenAuth] + authentication_classes = [AllTokenAuth] @action(methods=['POST'], detail=False) @swagger_auto_schema(operation_summary=_("dialogue"), diff --git a/apps/common/auth/authenticate.py b/apps/common/auth/authenticate.py index 6eddb76b4c3..4deaddb5663 100644 --- a/apps/common/auth/authenticate.py +++ b/apps/common/auth/authenticate.py @@ -12,11 +12,12 @@ from django.conf import settings from django.core import cache from django.core import signing +from django.utils.translation import gettext_lazy as _ from rest_framework.authentication import TokenAuthentication from common.exception.app_exception import AppAuthenticationFailed, AppEmbedIdentityFailed, AppChatNumOutOfBoundsFailed, \ - ChatException, AppApiException -from django.utils.translation import gettext_lazy as _ + AppApiException + token_cache = cache.caches['token_cache'] @@ -93,3 +94,60 @@ def authenticate(self, request): AppApiException): raise e raise AppAuthenticationFailed(1002, _('Authentication information is incorrect! illegal user')) + + +agent_handle = [new_instance_by_class_path(class_path) for class_path in settings.AGENT_HANDLES] + +all_handles = handles + agent_handle + +class AgentTokenAuth(TokenAuthentication): + keyword = "Bearer" + + # 重新 authenticate 方法,自定义认证规则 + def authenticate(self, request): + auth = request.META.get('HTTP_AUTHORIZATION') + # 未认证 + if auth is None: + raise AppAuthenticationFailed(1003, _('Not logged in, please log in first')) + if not auth.startswith("Bearer "): + raise AppAuthenticationFailed(1002, _('Authentication information is incorrect! illegal user')) + try: + token = auth[7:] + token_details = TokenDetails(token) + for handle in agent_handle: + if handle.support(request, token, token_details.get_token_details): + return handle.handle(request, token, token_details.get_token_details) + raise AppAuthenticationFailed(1002, _('Authentication information is incorrect! illegal user')) + except Exception as e: + traceback.format_exc() + if isinstance(e, AppEmbedIdentityFailed) or isinstance(e, AppChatNumOutOfBoundsFailed) or isinstance(e, + AppApiException): + raise e + raise AppAuthenticationFailed(1002, _('Authentication information is incorrect! illegal user')) + + +class AllTokenAuth(TokenAuthentication): + keyword = "Bearer" + + # 重新 authenticate 方法,自定义认证规则 + def authenticate(self, request): + auth = request.META.get('HTTP_AUTHORIZATION') + # 未认证 + if auth is None: + raise AppAuthenticationFailed(1003, _('Not logged in, please log in first')) + if not auth.startswith("Bearer "): + raise AppAuthenticationFailed(1002, _('Authentication information is incorrect! illegal user')) + try: + token = auth[7:] + token_details = TokenDetails(token) + for handle in all_handles: + if handle.support(request, token, token_details.get_token_details): + return handle.handle(request, token, token_details.get_token_details) + raise AppAuthenticationFailed(1002, _('Authentication information is incorrect! illegal user')) + except Exception as e: + traceback.format_exc() + if isinstance(e, AppEmbedIdentityFailed) or isinstance(e, AppChatNumOutOfBoundsFailed) or isinstance(e, + AppApiException): + raise e + raise AppAuthenticationFailed(1002, _('Authentication information is incorrect! illegal user')) +