15308555518/ MaxKB
1
0
Fork 0
mirror of https://github.com/1Panel-dev/MaxKB.git synced 2025-12-26 01:33:05 +00:00
Code Issues Actions 2 Packages Projects Releases Wiki Activity

fix: 白名单逻辑修改

This commit is contained in:
shaohuzhang1 2024-03-13 21:50:57 +08:00
parent 3e523903ba
commit 21a557ef43
9 changed files with 111 additions and 50 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
apps/application/serializers/application_serializers.py
View File

@ -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

5
apps/application/template/embed.js
View File

@ -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')

6
apps/application/views/application_views.py
View File

@ -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="创建应用",

14
apps/application/views/chat_views.py
View File

@ -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="获取对话列表",

35
apps/common/auth/authenticate.py
View File

@ -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=[

66
apps/common/middleware/chat_cookie_middleware.py Normal file
View File

@ -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)

23
apps/common/middleware/static_headers_middleware.py Normal file
View File

@ -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

1
apps/common/util/common.py
View File

@ -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

4
apps/smartdoc/settings/base.py
View File

@ -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'
]