diff --git a/apps/application/serializers/application_serializers.py b/apps/application/serializers/application_serializers.py index 55114a3f9..153153ded 100644 --- a/apps/application/serializers/application_serializers.py +++ b/apps/application/serializers/application_serializers.py @@ -126,10 +126,15 @@ class ApplicationSerializer(serializers.Serializer): ApplicationSerializer.Authentication(data={'access_token': self.data.get('token')}).auth() except Exception as e: is_auth = 'false' + application_access_token = QuerySet(ApplicationAccessToken).filter( + access_token=self.data.get('token')).first() t = Template(content) s = t.render( Context( - {'is_auth': is_auth, 'protocol': 'http', 'host': 'localhost:8000', 'token': '0a8d892c755f1a75'})) + {'is_auth': is_auth, 'protocol': 'http', 'host': 'localhost:8000', 'token': '0a8d892c755f1a75', + 'white_list_str': ",".join( + application_access_token.white_list), + 'white_active': 'true' if application_access_token.white_active else 'false'})) response = HttpResponse(s, status=200, headers={'Content-Type': 'text/javascript'}) set_embed_identity_cookie(request, response) return response diff --git a/apps/application/template/embed.js b/apps/application/template/embed.js index 69f1ff4d7..a9a8a8e54 100644 --- a/apps/application/template/embed.js +++ b/apps/application/template/embed.js @@ -286,7 +286,10 @@ function initMaxkbStyle(){ } function embedChatbot() { - if ({{is_auth}}) { + white_list_str='{{white_list_str}}' + white_list=white_list_str.split(',') + + if ({{is_auth}}&&{{white_active}}?white_list.includes(window.location.origin):true) { // 初始化maxkb智能小助手 initMaxkb() } else console.error('invalid parameter') diff --git a/apps/application/views/application_views.py b/apps/application/views/application_views.py index d894f61fd..cf4885ea1 100644 --- a/apps/application/views/application_views.py +++ b/apps/application/views/application_views.py @@ -21,7 +21,7 @@ from common.constants.permission_constants import CompareConstants, PermissionCo from common.exception.app_exception import AppAuthenticationFailed from common.response import result from common.swagger_api.common_api import CommonApi -from common.util.common import query_params_to_single_dict, set_embed_identity_cookie +from common.util.common import query_params_to_single_dict from dataset.serializers.dataset_serializers import DataSetSerializers @@ -191,14 +191,12 @@ class Application(APIView): tags=["应用/认证"], security=[]) def post(self, request: Request): - response = result.success( + return result.success( ApplicationSerializer.Authentication(data={'access_token': request.data.get("access_token")}).auth(), headers={"Access-Control-Allow-Origin": "*", "Access-Control-Allow-Credentials": "true", "Access-Control-Allow-Methods": "POST", "Access-Control-Allow-Headers": "Origin,Content-Type,Cookie,Accept,Token"} ) - set_embed_identity_cookie(request, response) - return response @action(methods=['POST'], detail=False) @swagger_auto_schema(operation_summary="创建应用", diff --git a/apps/application/views/chat_views.py b/apps/application/views/chat_views.py index 936ed4b87..7f7bed080 100644 --- a/apps/application/views/chat_views.py +++ b/apps/application/views/chat_views.py @@ -18,7 +18,7 @@ from common.auth import TokenAuth, has_permissions from common.constants.permission_constants import Permission, Group, Operate, \ RoleConstants, ViewPermission, CompareConstants from common.response import result -from common.util.common import query_params_to_single_dict, set_embed_identity_cookie +from common.util.common import query_params_to_single_dict class ChatView(APIView): @@ -71,13 +71,11 @@ class ChatView(APIView): dynamic_tag=keywords.get('application_id'))]) ) def post(self, request: Request, chat_id: str): - response = ChatMessageSerializer(data={'chat_id': chat_id}).chat(request.data.get('message'), - request.data.get( - 're_chat') if 're_chat' in request.data else False, - request.data.get( - 'stream') if 'stream' in request.data else True) - set_embed_identity_cookie(request, response) - return response + return ChatMessageSerializer(data={'chat_id': chat_id}).chat(request.data.get('message'), + request.data.get( + 're_chat') if 're_chat' in request.data else False, + request.data.get( + 'stream') if 'stream' in request.data else True) @action(methods=['GET'], detail=False) @swagger_auto_schema(operation_summary="获取对话列表", diff --git a/apps/common/auth/authenticate.py b/apps/common/auth/authenticate.py index 9cc6cac8d..88d724459 100644 --- a/apps/common/auth/authenticate.py +++ b/apps/common/auth/authenticate.py @@ -6,14 +6,11 @@ @date：2023/9/4 11:16 @desc: 认证类 """ -import datetime import traceback -from urllib.parse import urlparse from django.core import cache from django.core import signing from django.db.models import QuerySet -from ipware import get_client_ip from rest_framework.authentication import TokenAuthentication from application.models.api_key_model import ApplicationAccessToken, ApplicationApiKey @@ -21,13 +18,10 @@ from common.constants.authentication_type import AuthenticationType from common.constants.permission_constants import Auth, get_permission_list_by_role, RoleConstants, Permission, Group, \ Operate from common.exception.app_exception import AppAuthenticationFailed, AppEmbedIdentityFailed, AppChatNumOutOfBoundsFailed -from common.util.common import getRestSeconds -from common.util.rsa_util import decrypt from smartdoc.settings import JWT_AUTH from users.models.user import User, get_user_dynamics_permission token_cache = cache.caches['token_cache'] -chat_cache = cache.caches['chat_cache'] class AnonymousAuthentication(TokenAuthentication): @@ -87,35 +81,6 @@ class TokenAuth(TokenAuthentication): raise AppAuthenticationFailed(1002, "身份验证信息不正确") if not application_access_token.access_token == auth_details.get('access_token'): raise AppAuthenticationFailed(1002, "身份验证信息不正确") - if application_access_token.white_active: - referer = request.META.get('HTTP_REFERER') - if referer is not None: - client_ip = urlparse(referer).hostname - else: - client_ip = get_client_ip(request) - if not application_access_token.white_list.__contains__(client_ip): - raise AppAuthenticationFailed(1002, "身份验证信息不正确") - if 'embed_identity' in request.COOKIES and request.path.__contains__('/api/application/chat_message/'): - embed_identity = request.COOKIES['embed_identity'] - try: - # 如果无法解密 说明embed_identity并非系统颁发 - value = decrypt(embed_identity) - except Exception as e: - raise AppEmbedIdentityFailed(1004, '嵌入cookie不正确') - embed_identity_number = chat_cache.get(value) - if embed_identity_number is not None: - if application_access_token.access_num <= embed_identity_number: - raise AppChatNumOutOfBoundsFailed(1003, '访问次数超过今日访问量') - # 对话次数+1 - try: - if not chat_cache.incr(value): - # 如果修改失败则设置为1 - chat_cache.set(value, 1, - timeout=getRestSeconds()) - except Exception as e: - # 如果修改失败则设置为1 证明 key不存在 - chat_cache.add(value, 1, - timeout=getRestSeconds()) return application_access_token.application.user, Auth( role_list=[RoleConstants.APPLICATION_ACCESS_TOKEN], permission_list=[ diff --git a/apps/common/middleware/chat_cookie_middleware.py b/apps/common/middleware/chat_cookie_middleware.py new file mode 100644 index 000000000..60a26fe5f --- /dev/null +++ b/apps/common/middleware/chat_cookie_middleware.py @@ -0,0 +1,66 @@ +# coding=utf-8 +""" + @project: maxkb + @Author：虎 + @file： chat_cookie_middleware.py + @date：2024/3/13 20:13 + @desc: +""" +from django.core import cache +from django.core import signing +from django.db.models import QuerySet +from django.utils.deprecation import MiddlewareMixin + +from application.models.api_key_model import ApplicationAccessToken +from common.exception.app_exception import AppEmbedIdentityFailed +from common.response import result +from common.util.common import set_embed_identity_cookie, getRestSeconds +from common.util.rsa_util import decrypt + +chat_cache = cache.caches['chat_cache'] + + +class ChatCookieMiddleware(MiddlewareMixin): + + def process_response(self, request, response): + if request.path.startswith('/api/application/chat_message') or request.path.startswith( + '/api/application/authentication') or request.path.startswith('/api/application/profile'): + set_embed_identity_cookie(request, response) + if 'embed_identity' in request.COOKIES and request.path.__contains__('/api/application/chat_message/'): + embed_identity = request.COOKIES['embed_identity'] + try: + # 如果无法解密 说明embed_identity并非系统颁发 + value = decrypt(embed_identity) + except Exception as e: + raise AppEmbedIdentityFailed(1004, '嵌入cookie不正确') + # 对话次数+1 + try: + if not chat_cache.incr(value): + # 如果修改失败则设置为1 + chat_cache.set(value, 1, + timeout=getRestSeconds()) + except Exception as e: + # 如果修改失败则设置为1 证明 key不存在 + chat_cache.set(value, 1, + timeout=getRestSeconds()) + return response + + def process_request(self, request): + if 'embed_identity' in request.COOKIES and request.path.__contains__('/api/application/chat_message/'): + auth = request.META.get('HTTP_AUTHORIZATION', None + ) + auth_details = signing.loads(auth) + application_access_token = QuerySet(ApplicationAccessToken).filter( + application_id=auth_details.get('application_id')).first() + embed_identity = request.COOKIES['embed_identity'] + try: + # 如果无法解密 说明embed_identity并非系统颁发 + value = decrypt(embed_identity) + except Exception as e: + return result.Result(1003, + message='访问次数超过今日访问量', response_status=460) + embed_identity_number = chat_cache.get(value) + if embed_identity_number is not None: + if application_access_token.access_num <= embed_identity_number: + return result.Result(1003, + message='访问次数超过今日访问量', response_status=461) diff --git a/apps/common/middleware/static_headers_middleware.py b/apps/common/middleware/static_headers_middleware.py new file mode 100644 index 000000000..8c1fa622f --- /dev/null +++ b/apps/common/middleware/static_headers_middleware.py @@ -0,0 +1,23 @@ +# coding=utf-8 +""" + @project: maxkb + @Author：虎 + @file： static_headers_middleware.py + @date：2024/3/13 18:26 + @desc: +""" +from django.db.models import QuerySet +from django.utils.deprecation import MiddlewareMixin + +from application.models.api_key_model import ApplicationAccessToken + + +class StaticHeadersMiddleware(MiddlewareMixin): + def process_response(self, request, response): + if request.path.startswith('/ui/chat/'): + access_token = request.path.replace('/ui/chat/', '') + application_access_token = QuerySet(ApplicationAccessToken).filter(access_token=access_token).first() + if application_access_token.white_active: + # 添加自定义的响应头 + response['Content-Security-Policy'] = f'frame-ancestors {" ".join(application_access_token.white_list)}' + return response diff --git a/apps/common/util/common.py b/apps/common/util/common.py index d6492e353..f2eb0120c 100644 --- a/apps/common/util/common.py +++ b/apps/common/util/common.py @@ -11,6 +11,7 @@ import importlib import uuid from functools import reduce from typing import Dict, List + from django.core import cache from .rsa_util import encrypt diff --git a/apps/smartdoc/settings/base.py b/apps/smartdoc/settings/base.py index 53f723a56..69e9a2bd3 100644 --- a/apps/smartdoc/settings/base.py +++ b/apps/smartdoc/settings/base.py @@ -1,9 +1,9 @@ import datetime +import mimetypes import os from pathlib import Path from ..const import CONFIG, PROJECT_DIR -import mimetypes mimetypes.add_type("text/css", ".css", True) mimetypes.add_type("text/javascript", ".js", True) @@ -46,6 +46,8 @@ MIDDLEWARE = [ 'django.contrib.sessions.middleware.SessionMiddleware', 'django.middleware.common.CommonMiddleware', 'django.contrib.messages.middleware.MessageMiddleware', + 'common.middleware.static_headers_middleware.StaticHeadersMiddleware', + 'common.middleware.chat_cookie_middleware.ChatCookieMiddleware' ]