mirror of
https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git
synced 2026-11-08 01:06:59 +08:00
fit2cloud-chenywandClaude
bbfc4cc8f5
- TokenMiddleware/ResponseMiddleware/audit logger use scope['path'] instead of request.url.path for whitelist checks: a crafted Host header (e.g. x/api/v1/mcp) was concatenated into Starlette's URL and made protected routes match the /mcp* whitelist entry - validateEmbedded rejects admin accounts and non-type-4 apps, so a self-signed token with account=admin can no longer impersonate the administrator - Add HostValidationMiddleware rejecting malformed Host headers at the outermost layer (defense in depth) - Add regression tests for host validation, whitelist matching and source-level guards Co-Authored-By: Claude <[email protected]>
266 lines
14 KiB
Python
266 lines
14 KiB
Python
|
|
import base64
|
|
import json
|
|
import re
|
|
from typing import Optional
|
|
from fastapi import Request
|
|
from fastapi.responses import JSONResponse
|
|
from starlette.responses import Response
|
|
import jwt
|
|
from sqlmodel import Session
|
|
from starlette.middleware.base import BaseHTTPMiddleware
|
|
from apps.system.crud.apikey_manage import get_api_key
|
|
from apps.system.models.system_model import ApiKeyModel, AssistantModel
|
|
from common.core.db import engine
|
|
from apps.system.crud.assistant import get_assistant_info, get_assistant_user
|
|
from apps.system.crud.user import get_user_by_account, get_user_info
|
|
from apps.system.schemas.system_schema import AssistantHeader, UserInfoDTO
|
|
from common.core import security
|
|
from common.core.config import settings
|
|
from common.core.schemas import TokenPayload
|
|
from common.utils.locale import I18n
|
|
from common.utils.utils import SQLBotLogUtil, get_origin_from_referer, origin_match_domain
|
|
from common.utils.whitelist import whiteUtils
|
|
from fastapi.security.utils import get_authorization_scheme_param
|
|
from common.core.deps import get_i18n
|
|
class TokenMiddleware(BaseHTTPMiddleware):
|
|
|
|
|
|
|
|
def __init__(self, app):
|
|
super().__init__(app)
|
|
|
|
async def dispatch(self, request, call_next):
|
|
# 使用 scope["path"]（请求行真实路径）做白名单判断：
|
|
# request.url.path 会拼接未校验的 Host 头（Starlette URL(scope) 行为），
|
|
# 攻击者可通过伪造 Host: x/api/v1/mcp 将受保护接口伪装成白名单路径绕过认证
|
|
request_path = request.scope.get("path") or request.url.path
|
|
|
|
if self.is_options(request) or whiteUtils.is_whitelisted(request_path):
|
|
# 动态处理 /system/assistant/info/{id} 的 CORS 预检
|
|
if request.method == "OPTIONS":
|
|
origin = request.headers.get("origin", "")
|
|
if origin:
|
|
match = re.search(r'/system/assistant/info/(\d+)', request_path)
|
|
if match:
|
|
assistant_id = int(match.group(1))
|
|
with Session(engine) as session:
|
|
db_model = session.get(AssistantModel, assistant_id)
|
|
if db_model and origin_match_domain(origin, db_model.domain):
|
|
return Response(
|
|
status_code=200,
|
|
headers={
|
|
"Access-Control-Allow-Origin": origin,
|
|
"Access-Control-Allow-Methods": "GET, OPTIONS",
|
|
"Access-Control-Allow-Headers": "*",
|
|
"Access-Control-Allow-Credentials": "true",
|
|
"Access-Control-Max-Age": "600",
|
|
},
|
|
)
|
|
return await call_next(request)
|
|
assistantTokenKey = settings.ASSISTANT_TOKEN_KEY
|
|
assistantToken = request.headers.get(assistantTokenKey)
|
|
askToken = request.headers.get("X-SQLBOT-ASK-TOKEN")
|
|
trans = await get_i18n(request)
|
|
if askToken:
|
|
validate_pass, data = await self.validateAskToken(askToken, trans)
|
|
if validate_pass:
|
|
request.state.current_user = data
|
|
return await call_next(request)
|
|
message = trans('i18n_permission.authenticate_invalid', msg = data)
|
|
return JSONResponse(message, status_code=401, headers={"Access-Control-Allow-Origin": "*"})
|
|
#if assistantToken and assistantToken.lower().startswith("assistant "):
|
|
if assistantToken:
|
|
validator: tuple[any] = await self.validateAssistant(assistantToken, trans)
|
|
if validator[0]:
|
|
request.state.current_user = validator[1]
|
|
if request.state.current_user and trans.lang:
|
|
request.state.current_user.language = trans.lang
|
|
request.state.assistant = validator[2]
|
|
origin = request.headers.get("X-SQLBOT-HOST-ORIGIN") or get_origin_from_referer(request)
|
|
if origin and validator[2]:
|
|
request.state.assistant.request_origin = origin
|
|
return await call_next(request)
|
|
message = trans('i18n_permission.authenticate_invalid', msg = validator[1])
|
|
return JSONResponse(message, status_code=401, headers={"Access-Control-Allow-Origin": "*"})
|
|
#validate pass
|
|
tokenkey = settings.TOKEN_KEY
|
|
token = request.headers.get(tokenkey)
|
|
validate_pass, data = await self.validateToken(token, trans)
|
|
if validate_pass:
|
|
request.state.current_user = data
|
|
return await call_next(request)
|
|
|
|
message = trans('i18n_permission.authenticate_invalid', msg = data)
|
|
return JSONResponse(message, status_code=401, headers={"Access-Control-Allow-Origin": "*"})
|
|
|
|
def is_options(self, request: Request):
|
|
return request.method == "OPTIONS"
|
|
|
|
async def validateAskToken(self, askToken: Optional[str], trans: I18n):
|
|
if not askToken:
|
|
return False, f"Miss Token[X-SQLBOT-ASK-TOKEN]!"
|
|
schema, param = get_authorization_scheme_param(askToken)
|
|
if schema.lower() != "sk":
|
|
return False, f"Token schema error!"
|
|
try:
|
|
payload = jwt.decode(
|
|
param, options={"verify_signature": False, "verify_exp": False}, algorithms=[security.ALGORITHM]
|
|
)
|
|
access_key = payload.get('access_key', None)
|
|
|
|
if not access_key:
|
|
return False, f"Miss access_key payload error!"
|
|
with Session(engine) as session:
|
|
api_key_model = await get_api_key(session, access_key)
|
|
api_key_model = ApiKeyModel.model_validate(api_key_model) if api_key_model else None
|
|
if not api_key_model:
|
|
return False, f"Invalid access_key!"
|
|
if not api_key_model.status:
|
|
return False, f"Disabled access_key!"
|
|
payload = jwt.decode(
|
|
param, api_key_model.secret_key, algorithms=[security.ALGORITHM]
|
|
)
|
|
uid = api_key_model.uid
|
|
session_user = await get_user_info(session = session, user_id = uid)
|
|
if not session_user:
|
|
message = trans('i18n_not_exist', msg = trans('i18n_user.account'))
|
|
raise Exception(message)
|
|
session_user = UserInfoDTO.model_validate(session_user)
|
|
if session_user.status != 1:
|
|
message = trans('i18n_login.user_disable', msg = trans('i18n_concat_admin'))
|
|
raise Exception(message)
|
|
if not session_user.oid or session_user.oid == 0:
|
|
message = trans('i18n_login.no_associated_ws', msg = trans('i18n_concat_admin'))
|
|
raise Exception(message)
|
|
return True, session_user
|
|
except Exception as e:
|
|
msg = str(e)
|
|
SQLBotLogUtil.exception(f"Token validation error: {msg}")
|
|
if 'expired' in msg:
|
|
return False, jwt.ExpiredSignatureError(trans('i18n_permission.token_expired'))
|
|
return False, e
|
|
|
|
async def validateToken(self, token: Optional[str], trans: I18n):
|
|
if not token:
|
|
return False, f"Miss Token[{settings.TOKEN_KEY}]!"
|
|
schema, param = get_authorization_scheme_param(token)
|
|
if schema.lower() != "bearer":
|
|
return False, f"Token schema error!"
|
|
try:
|
|
payload = jwt.decode(
|
|
param, settings.SECRET_KEY, algorithms=[security.ALGORITHM]
|
|
)
|
|
token_data = TokenPayload(**payload)
|
|
with Session(engine) as session:
|
|
session_user = await get_user_info(session = session, user_id = token_data.id)
|
|
if not session_user:
|
|
message = trans('i18n_not_exist', msg = trans('i18n_user.account'))
|
|
raise Exception(message)
|
|
session_user = UserInfoDTO.model_validate(session_user)
|
|
if session_user.status != 1:
|
|
message = trans('i18n_login.user_disable', msg = trans('i18n_concat_admin'))
|
|
raise Exception(message)
|
|
if not session_user.oid or session_user.oid == 0:
|
|
message = trans('i18n_login.no_associated_ws', msg = trans('i18n_concat_admin'))
|
|
raise Exception(message)
|
|
return True, session_user
|
|
except Exception as e:
|
|
msg = str(e)
|
|
SQLBotLogUtil.exception(f"Token validation error: {msg}")
|
|
if 'expired' in msg:
|
|
return False, jwt.ExpiredSignatureError(trans('i18n_permission.token_expired'))
|
|
return False, e
|
|
|
|
|
|
async def validateAssistant(self, assistantToken: Optional[str], trans: I18n) -> tuple[any]:
|
|
if not assistantToken:
|
|
return False, f"Miss Token[{settings.TOKEN_KEY}]!"
|
|
schema, param = get_authorization_scheme_param(assistantToken)
|
|
|
|
|
|
try:
|
|
if schema.lower() == 'embedded':
|
|
return await self.validateEmbedded(param, trans)
|
|
if schema.lower() != "assistant":
|
|
return False, f"Token schema error!"
|
|
payload = jwt.decode(
|
|
param, settings.SECRET_KEY, algorithms=[security.ALGORITHM]
|
|
)
|
|
token_data = TokenPayload(**payload)
|
|
if not payload['assistant_id']:
|
|
return False, f"Miss assistant payload error!"
|
|
with Session(engine) as session:
|
|
""" session_user = await get_user_info(session = session, user_id = token_data.id)
|
|
session_user = UserInfoDTO.model_validate(session_user) """
|
|
session_user = get_assistant_user(id = token_data.id)
|
|
assistant_info = await get_assistant_info(session=session, assistant_id=payload['assistant_id'])
|
|
assistant_info = AssistantModel.model_validate(assistant_info)
|
|
assistant_info = AssistantHeader.model_validate(assistant_info.model_dump(exclude_unset=True))
|
|
session_user.oid = int(assistant_info.oid)
|
|
|
|
return True, session_user, assistant_info
|
|
except Exception as e:
|
|
SQLBotLogUtil.exception(f"Assistant validation error: {str(e)}")
|
|
# Return False and the exception message
|
|
return False, e
|
|
|
|
async def validateEmbedded(self, param: str, trans: I18n) -> tuple[any]:
|
|
try:
|
|
# WARNING: Signature verification is disabled for embedded tokens
|
|
# This is a security risk and should only be used if absolutely necessary
|
|
# Consider implementing proper signature verification with a shared secret
|
|
payload: dict = jwt.decode(
|
|
param,
|
|
options={"verify_signature": False, "verify_exp": False},
|
|
algorithms=[security.ALGORITHM]
|
|
)
|
|
app_key = payload.get('appId', '')
|
|
embeddedId = payload.get('embeddedId', None)
|
|
if not embeddedId:
|
|
embeddedId = xor_decrypt(app_key)
|
|
if not payload['account']:
|
|
return False, f"Miss account payload error!"
|
|
account = payload['account']
|
|
with Session(engine) as session:
|
|
assistant_info = await get_assistant_info(session=session, assistant_id=embeddedId)
|
|
assistant_info = AssistantModel.model_validate(assistant_info)
|
|
# embedded 协议（app_secret + account）仅适用于页面嵌入（type=4）应用
|
|
if assistant_info.type != 4:
|
|
return False, f"Invalid embedded app type!"
|
|
payload = jwt.decode(
|
|
param, assistant_info.app_secret, algorithms=[security.ALGORITHM]
|
|
)
|
|
assistant_info = AssistantHeader.model_validate(assistant_info.model_dump(exclude_unset=True))
|
|
""" session_user = await get_user_info(session = session, user_id = token_data.id)
|
|
session_user = UserInfoDTO.model_validate(session_user) """
|
|
session_user = get_user_by_account(session = session, account=account)
|
|
if not session_user:
|
|
message = trans('i18n_not_exist', msg = trans('i18n_user.account'))
|
|
raise Exception(message)
|
|
session_user = await get_user_info(session = session, user_id = session_user.id)
|
|
|
|
session_user = UserInfoDTO.model_validate(session_user)
|
|
if session_user.status != 1:
|
|
message = trans('i18n_login.user_disable', msg = trans('i18n_concat_admin'))
|
|
raise Exception(message)
|
|
if not session_user.oid or session_user.oid == 0:
|
|
message = trans('i18n_login.no_associated_ws', msg = trans('i18n_concat_admin'))
|
|
raise Exception(message)
|
|
# 管理员账号不允许通过 embedded token 使用：app_secret 由集成方持有，
|
|
# 攻击者若取得任意应用 app_secret 即可伪造 account=admin 的管理员身份
|
|
if session_user.isAdmin:
|
|
return False, f"Admin account is not allowed for embedded token!"
|
|
if session_user.oid:
|
|
assistant_info.oid = int(session_user.oid)
|
|
return True, session_user, assistant_info
|
|
except Exception as e:
|
|
SQLBotLogUtil.exception(f"Embedded validation error: {str(e)}")
|
|
# Return False and the exception message
|
|
return False, e
|
|
|
|
def xor_decrypt(encrypted_str: str, key: int = 0xABCD1234) -> int:
|
|
encrypted_bytes = base64.urlsafe_b64decode(encrypted_str)
|
|
hex_str = encrypted_bytes.hex()
|
|
encrypted_num = int(hex_str, 16)
|
|
return encrypted_num ^ key |