diff --git a/apps/chat/tests.py b/apps/chat/tests.py deleted file mode 100644 index 7ce503c2dd9..00000000000 --- a/apps/chat/tests.py +++ /dev/null @@ -1,3 +0,0 @@ -from django.test import TestCase - -# Create your tests here. diff --git a/apps/chat/tests/__init__.py b/apps/chat/tests/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/apps/chat/tests/test_auth.py b/apps/chat/tests/test_auth.py new file mode 100644 index 00000000000..d28bb2599e9 --- /dev/null +++ b/apps/chat/tests/test_auth.py @@ -0,0 +1,103 @@ +from types import SimpleNamespace +from unittest.mock import Mock, patch + +from django.core import signing +from django.test import SimpleTestCase + +from chat.views.v3.chat import BaseAuthView +from common.auth.constants.operate_constants import Operate +from common.exception.app_exception import AppApiException + + +class BaseAuthViewTest(SimpleTestCase): + def setUp(self): + self.user = SimpleNamespace(id="chat-user-id", username="chat-user") + self.view = BaseAuthView() + self.view.request = Mock() + + def test_access_token_login_validates_authorization_and_keeps_operate_token(self): + access_token = "application-access-token" + application_access_token = SimpleNamespace( + application_id="application-id", + application=SimpleNamespace(workspace_id="application-workspace"), + ) + + with ( + patch( + "chat.views.v3.chat.ChatUserAccessTokenSerializer.create_token_and_cache", + ) as validate_authorization, + patch("chat.views.v3.chat.ApplicationAccessToken.objects.filter") as filter_access_token, + patch("chat.views.v3.chat.cache.set") as cache_set, + patch("chat.views.v3.chat.record_log") as record_log, + patch("common.auth.common.get_key", return_value="test-signing-key"), + ): + filter_access_token.return_value.select_related.return_value.first.return_value = application_access_token + filter_access_token.return_value.values_list.return_value.first.return_value = "application-id" + + token = self.view.create_token_and_cache(self.user, access_token, Operate.OIDC) + + validate_authorization.assert_called_once_with(access_token, self.user) + self.assertEqual( + signing.loads(token, key="test-signing-key"), + { + "id": self.user.id, + "type": "CHAT_USER", + "login_type": "OIDC", + "kwargs": {"application_id": "application-id"}, + }, + ) + cache_set.assert_called_once_with(token, self.user, timeout=7200, version="CHAT_USER_TOKEN") + record_log.assert_called_once_with( + menu="Chat User/login", + operate="Log in", + request=self.view.request, + user={"username": "chat-user"}, + status=200, + operation_object={"name": "chat-user"}, + workspace_id="application-workspace", + ) + + def test_login_without_access_token_records_audit_log(self): + with ( + patch("chat.views.v3.chat.record_log", create=True) as record_log, + patch("chat.views.v3.chat.ChatUserAccessTokenSerializer.create_token_and_cache") as validate_authorization, + patch("chat.views.v3.chat.ApplicationAccessToken.objects.filter") as filter_access_token, + patch("chat.views.v3.chat.cache.set"), + patch("common.auth.common.get_key", return_value="test-signing-key"), + ): + self.view.create_token_and_cache(self.user, None, Operate.LOCAL) + + validate_authorization.assert_not_called() + filter_access_token.assert_not_called() + record_log.assert_called_once_with( + menu="Chat User/login", + operate="Log in", + request=self.view.request, + user={"username": "chat-user"}, + status=200, + operation_object={"name": "chat-user"}, + workspace_id="default", + ) + + def test_failed_authorization_records_failed_login_log(self): + with ( + patch( + "chat.views.v3.chat.ChatUserAccessTokenSerializer.create_token_and_cache", + side_effect=AppApiException(1005, "access denied"), + ), + patch("chat.views.v3.chat.record_log") as record_log, + patch("chat.views.v3.chat.cache.set") as cache_set, + ): + with self.assertRaises(AppApiException): + self.view.create_token_and_cache(self.user, "application-access-token", Operate.OIDC) + + cache_set.assert_not_called() + record_log.assert_called_once_with( + menu="Chat User/login", + operate="Log in", + request=self.view.request, + user={"username": "chat-user"}, + status=500, + operation_object={"name": "chat-user"}, + workspace_id="default", + ) diff --git a/apps/chat/views/v2/chat.py b/apps/chat/views/v2/chat.py index fd76e575bb8..f19b3c25130 100644 --- a/apps/chat/views/v2/chat.py +++ b/apps/chat/views/v2/chat.py @@ -24,7 +24,6 @@ from application.models import ChatUserType, ChatSourceChoices from chat.api.chat_api import ChatAPI from chat.api.chat_authentication_api import ( - ChatAuthenticationAPI, ChatAuthenticationProfileAPIV2, ChatOpenAPI, OpenAIAPI, @@ -44,7 +43,6 @@ from common.auth import ChatTokenAuth from common.auth.authentication import has_permissions from common.auth.constants.chat_permission_constants import ChatPermissionConstants -from common.constants.authentication_type import AuthenticationType from common.constants.cache_version import Cache_Version from common.exception.app_exception import AppAuthenticationFailed, AppApiException from common.log.log import _get_ip_address, log @@ -55,8 +53,7 @@ from models_provider.api.model import DefaultModelResponse from oss.serializers.file import FileSerializer from system_manage.serializers.chat_user import RePasswordSerializer, ChatUserProfileSerializer -from system_manage.serializers.chat_user_serializer import ChatUserAccessTokenSerializer -from users.api import CaptchaAPI, LoginAPI +from users.api import CaptchaAPI from users.api.user import ResetPasswordAPI, UserProfileAPI from users.serializers.login import CaptchaSerializer from users.views import get_re_password_details @@ -409,7 +406,7 @@ def post(self, request: Request): decrypted_data = json.loads(decrypted_raw) if decrypted_raw else {} if isinstance(decrypted_data, dict): request_data = decrypted_data - except Exception as e: + except Exception: raise AppApiException(500, _("Invalid encrypted data")) serializer_obj = RePasswordSerializer(data=request_data) if serializer_obj.reset_password(request.user.id): @@ -435,48 +432,6 @@ def get(self, request: Request): return result.success(ChatUserProfileSerializer().profile(request.user)) -class BaseAuthView(APIView): - @staticmethod - def create_token_and_cache(access_token, user, request): - token = ChatUserAccessTokenSerializer.create_token_and_cache(access_token, user, request) - version, get_key = Cache_Version.CHAT_USER_TOKEN.value - cache.set(get_key(token), user, timeout=60 * 60 * 2, version=version) - return token - - @classmethod - def generate(self, request, f_token: str, response: HttpResponse, path: str = "/chat"): - secure = request.is_secure() - response.set_cookie( - "mk_file_auth", - value=f_token, - max_age=7 * 24 * 3600, - path=path, - domain=None, - secure=secure, - httponly=True, - samesite="Lax", - ) - return response - - -class LocalLoginView(BaseAuthView): - @extend_schema( - methods=["POST"], - description=_("Log in"), - summary=_("Log in"), - operation_id=_("Log in"), # type: ignore - tags=[_("Chat User/login")], # type: ignore - request=LoginAPI.get_request(), - responses=LoginAPI.get_response(), - ) - def post(self, request: Request, access_token: str = None): - user = ChatUserAccessTokenSerializer.local_login(request.data, access_token) - user.source = "LOCAL" - token = self.create_token_and_cache(access_token, user, request) - response = result.success({"token": token}) - return self.generate(request, token, response, path=f"/chat/{access_token}/") - - class Logout(APIView): authentication_classes = [ChatTokenAuth] diff --git a/apps/chat/views/v3/chat.py b/apps/chat/views/v3/chat.py index bc364f42c55..47f417c310d 100644 --- a/apps/chat/views/v3/chat.py +++ b/apps/chat/views/v3/chat.py @@ -45,7 +45,7 @@ from common.constants.authentication_type import AuthenticationType from common.constants.cache_version import Cache_Version from common.exception.app_exception import AppApiException -from common.log.log import _get_ip_address, log +from common.log.log import _get_ip_address, log, record_log from common.result import result from common.utils.rsa_util import decrypt from knowledge.models import FileSourceType @@ -54,6 +54,7 @@ from oss.serializers.file import FileSerializer from system_manage.serializers.chat_user import RePasswordSerializer, ChatUserProfileSerializer from chat.serializers.chat_user_serializer import ChatUserAccessTokenV3Serializer +from system_manage.serializers.chat_user_serializer import ChatUserAccessTokenSerializer from users.api import CaptchaAPI, LoginAPI from users.api.user import ResetPasswordAPI, UserProfileAPI from users.serializers.login import CaptchaSerializer @@ -475,21 +476,38 @@ def get(self, request: Request): class BaseAuthView(APIView): - @staticmethod - def create_token_and_cache(user, access_token, operate): + def create_token_and_cache(self, user, access_token, operate): application_id = None - if access_token: - application_id = ( - ApplicationAccessToken.objects.filter(access_token=access_token, is_active=True) - .values_list("application_id", flat=True) - .first() + workspace_id = "default" + status = 500 + try: + if access_token: + ChatUserAccessTokenSerializer.create_token_and_cache(access_token, user) + application_access_token = ( + ApplicationAccessToken.objects.filter(access_token=access_token, is_active=True) + .select_related("application") + .first() + ) + if application_access_token: + application_id = application_access_token.application_id + workspace_id = application_access_token.application.workspace_id + token = ChatToken( + str(user.id), AuthenticationType.CHAT_USER, str(operate), application_id=str(application_id) + ).to_token() + version, get_key = Cache_Version.CHAT_USER_TOKEN.value + cache.set(get_key(token), user, timeout=60 * 60 * 2, version=version) + status = 200 + return token + finally: + record_log( + menu="Chat User/login", + operate="Log in", + request=self.request, + user={"username": user.username}, + status=status, + operation_object={"name": user.username}, + workspace_id=workspace_id, ) - token = ChatToken( - str(user.id), AuthenticationType.CHAT_USER, str(operate), application_id=str(application_id) - ).to_token() - version, get_key = Cache_Version.CHAT_USER_TOKEN.value - cache.set(get_key(token), user, timeout=60 * 60 * 2, version=version) - return token @classmethod def generate(self, request, token: str, response: HttpResponse, path: str = "/chat"): diff --git a/apps/system_manage/serializers/chat_user_serializer.py b/apps/system_manage/serializers/chat_user_serializer.py index 4a650ff8345..c0ea53d290d 100644 --- a/apps/system_manage/serializers/chat_user_serializer.py +++ b/apps/system_manage/serializers/chat_user_serializer.py @@ -5,12 +5,8 @@ from rest_framework import serializers from application.models import ApplicationAccessToken -from common.auth.common import ChatToken -from common.auth.constants.operate_constants import Operate -from common.constants.authentication_type import AuthenticationType from common.constants.cache_version import Cache_Version from common.exception.app_exception import AppApiException -from common.log.log import record_log from common.utils.common import password_encrypt from common.utils.common import password_verify, needs_password_upgrade from common.utils.rsa_util import decrypt @@ -28,55 +24,33 @@ class ChatUserAccessTokenSerializer(serializers.Serializer): @staticmethod - def create_token_and_cache(access_token, user, request): - status = 500 # 默认失败状态 - workspace_id = "default" - try: - application_access_token = ApplicationAccessToken.objects.filter(access_token=access_token).first() - - if not application_access_token: - raise AppApiException(1005, _("Invalid access token")) - - application_id = application_access_token.application_id - workspace_id = application_access_token.application.workspace_id - # 检查用户是否有权限访问该应用 - is_authorized = ResourceChatUserAuthorize.objects.filter( - resource_id=application_id, resource_type=ResourceType.APPLICATION.value, is_auth=True, user_id=user.id - ).exists() - if not is_authorized: - # 获取资源组授权的用户组ID - resource_group_ids = ResourceChatUserGroupAuthorize.objects.filter( - resource_id=application_id, - resource_type=ResourceType.APPLICATION.value, - is_auth=True, - ).values_list("user_group_id", flat=True) - - # 如果有资源组授权，则检查用户是否属于这些用户组 - if resource_group_ids.exists(): - is_authorized = UserGroupRelation.objects.filter( - user_id=user.id, group_id__in=resource_group_ids - ).exists() - - if not is_authorized: - raise AppApiException(1005, _("The user does not have permission to access the application")) - token = ChatToken( - str(user.id), - AuthenticationType.CHAT_USER, - str(Operate.LOCAL).lower(), - application_id=str(application_id), - ).to_token() - status = 200 - return token - finally: - record_log( - menu="Chat User/login", - operate="Log in", - request=request, - user={"username": user.username}, - status=status, - operation_object={"name": user.username}, - workspace_id=workspace_id, - ) + def create_token_and_cache(access_token, user): + application_access_token = ApplicationAccessToken.objects.filter(access_token=access_token).first() + + if not application_access_token: + raise AppApiException(1005, _("Invalid access token")) + + application_id = application_access_token.application_id + # 检查用户是否有权限访问该应用 + is_authorized = ResourceChatUserAuthorize.objects.filter( + resource_id=application_id, resource_type=ResourceType.APPLICATION.value, is_auth=True, user_id=user.id + ).exists() + if not is_authorized: + # 获取资源组授权的用户组ID + resource_group_ids = ResourceChatUserGroupAuthorize.objects.filter( + resource_id=application_id, + resource_type=ResourceType.APPLICATION.value, + is_auth=True, + ).values_list("user_group_id", flat=True) + + # 如果有资源组授权，则检查用户是否属于这些用户组 + if resource_group_ids.exists(): + is_authorized = UserGroupRelation.objects.filter( + user_id=user.id, group_id__in=resource_group_ids + ).exists() + + if not is_authorized: + raise AppApiException(1005, _("The user does not have permission to access the application")) @staticmethod def get_auth_setting(access_token):