mirror of
https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git
synced 2026-11-08 01:06:59 +08:00
382 lines
17 KiB
Python
382 lines
17 KiB
Python
import asyncio
|
||
import json
|
||
import os
|
||
from datetime import datetime, timedelta, timezone
|
||
from zoneinfo import ZoneInfo
|
||
from typing import List, Optional
|
||
|
||
from fastapi import APIRouter, Form, HTTPException, Path, Query, Request, Response, UploadFile
|
||
from fastapi.responses import StreamingResponse
|
||
from sqlbot_xpack.file_utils import SQLBotFileUtils
|
||
from sqlmodel import select
|
||
|
||
from apps.datasource.models.datasource import CoreDatasource
|
||
from apps.db.constant import DB
|
||
from apps.swagger.i18n import PLACEHOLDER_PREFIX
|
||
from apps.system.crud.assistant import AssistantOutDs, AssistantOutDsFactory, get_assistant_info
|
||
from apps.system.crud.assistant_manage import dynamic_upgrade_cors, get_ws_assistant, save
|
||
from apps.system.models.system_model import AssistantModel
|
||
from apps.system.schemas.auth import CacheName, CacheNamespace
|
||
from apps.system.schemas.permission import SqlbotPermission, require_permissions
|
||
from apps.system.schemas.system_schema import AssistantBase, AssistantDTO, AssistantUiSchema, AssistantValidator
|
||
from common.core.config import settings
|
||
from common.core.deps import CurrentAssistant, SessionDep, Trans, CurrentUser
|
||
from common.core.security import create_access_token
|
||
from common.core.sqlbot_cache import clear_cache
|
||
from common.utils.utils import get_origin_from_referer, origin_match_domain
|
||
router = APIRouter(tags=["system_assistant"], prefix="/system/assistant")
|
||
from common.audit.models.log_model import OperationType, OperationModules
|
||
from common.audit.schemas.logger_decorator import LogConfig, system_log
|
||
from sqlbot_xpack.core import decrypt_embedded_sign
|
||
|
||
@router.get("/info/{id}", include_in_schema=False)
|
||
async def info(request: Request, response: Response, session: SessionDep, trans: Trans, id: int, virtual: Optional[int] = Query(None)):
|
||
if not id:
|
||
raise Exception('miss assistant id')
|
||
db_model = await get_assistant_info(session=session, assistant_id=id)
|
||
if not db_model:
|
||
raise RuntimeError(f"assistant application not exist")
|
||
db_model = AssistantModel.model_validate(db_model)
|
||
|
||
# 校验 SQLBOT-EMBEDDED-SIGN 请求头
|
||
sign_header = request.headers.get("SQLBOT-EMBEDDED-SIGN")
|
||
if not sign_header:
|
||
raise RuntimeError(trans('i18n_embedded.missing_sign'))
|
||
|
||
try:
|
||
sign_data = await decrypt_embedded_sign(sign_header)
|
||
except Exception:
|
||
raise RuntimeError(trans('i18n_embedded.invalid_sign'))
|
||
|
||
# 校验 assistant_id 与 id 参数一致
|
||
if str(sign_data.get("assistant_id")) != str(id):
|
||
raise RuntimeError(trans('i18n_embedded.assistant_id_mismatch'))
|
||
|
||
# 校验 target（来源域名）是否合法
|
||
target = sign_data.get("target", "")
|
||
|
||
request_origin = request.headers.get("origin") or get_origin_from_referer(request)
|
||
if not request_origin:
|
||
raise RuntimeError(trans('i18n_embedded.missing_origin'))
|
||
request_origin = request_origin.rstrip('/')
|
||
if not target or target == "null":
|
||
target = request_origin
|
||
elif target != request_origin:
|
||
raise RuntimeError(trans('i18n_embedded.target_origin_mismatch'))
|
||
|
||
if not origin_match_domain(target, db_model.domain):
|
||
raise RuntimeError(trans('i18n_embedded.invalid_origin', origin=target or ''))
|
||
|
||
# 校验 sign_time 是否在 10 秒内
|
||
sign_time_str = sign_data.get("sign_time", "")
|
||
try:
|
||
sign_time = datetime.fromisoformat(sign_time_str)
|
||
except (ValueError, TypeError):
|
||
raise RuntimeError(trans('i18n_embedded.sign_time_invalid'))
|
||
now_utc = datetime.now(timezone.utc)
|
||
sign_time_utc = sign_time.astimezone(timezone.utc)
|
||
if abs((now_utc - sign_time_utc).total_seconds()) > 10:
|
||
raise RuntimeError(trans('i18n_embedded.sign_time_expired'))
|
||
|
||
# 校验是否为真实浏览器请求（非自动化工具）
|
||
if sign_data.get("webdriver", False):
|
||
raise RuntimeError(trans('i18n_embedded.automation_tool_detected'))
|
||
|
||
# 校验 User-Agent 一致性（签名中的 navigator.userAgent 与请求头一致）
|
||
sign_user_agent = sign_data.get("user_agent", "")
|
||
request_user_agent = request.headers.get("User-Agent", "")
|
||
if sign_user_agent != request_user_agent:
|
||
raise RuntimeError(trans('i18n_embedded.user_agent_mismatch'))
|
||
|
||
# 校验 timezone 与 sign_time 偏移一致性（防时区伪造）
|
||
tz_name = sign_data.get("timezone", "")
|
||
try:
|
||
tz = ZoneInfo(tz_name)
|
||
except Exception:
|
||
raise RuntimeError(trans('i18n_embedded.timezone_mismatch'))
|
||
sign_time_naive = sign_time.replace(tzinfo=None)
|
||
if tz.utcoffset(sign_time_naive) != sign_time.utcoffset():
|
||
raise RuntimeError(trans('i18n_embedded.timezone_mismatch'))
|
||
|
||
origin = target.rstrip('/')
|
||
|
||
response.headers["Access-Control-Allow-Origin"] = origin
|
||
|
||
|
||
assistant_oid = 1
|
||
if (db_model.type == 0):
|
||
configuration = db_model.configuration
|
||
config_obj = json.loads(configuration) if configuration else {}
|
||
assistant_oid = config_obj.get('oid', 1)
|
||
|
||
access_token_expires = timedelta(minutes=settings.ACCESS_TOKEN_EXPIRE_MINUTES)
|
||
assistantDict = {
|
||
"id": virtual, "account": 'sqlbot-inner-assistant', "oid": assistant_oid, "assistant_id": id
|
||
}
|
||
access_token = create_access_token(
|
||
assistantDict, expires_delta=access_token_expires
|
||
)
|
||
|
||
result = db_model.model_dump()
|
||
result["token"] = access_token
|
||
return result
|
||
|
||
|
||
@router.get("/app/{appId}", include_in_schema=False)
|
||
async def getApp(request: Request, response: Response, session: SessionDep, trans: Trans, appId: str) -> AssistantModel:
|
||
if not appId:
|
||
raise Exception('miss assistant appId')
|
||
db_model = session.exec(select(AssistantModel).where(AssistantModel.app_id == appId)).first()
|
||
if not db_model:
|
||
raise RuntimeError(f"assistant application not exist")
|
||
db_model = AssistantModel.model_validate(db_model)
|
||
origin = request.headers.get("origin") or get_origin_from_referer(request)
|
||
if not origin:
|
||
raise RuntimeError(trans('i18n_embedded.invalid_origin', origin=origin or ''))
|
||
origin = origin.rstrip('/')
|
||
if not origin_match_domain(origin, db_model.domain):
|
||
raise RuntimeError(trans('i18n_embedded.invalid_origin', origin=origin or ''))
|
||
|
||
response.headers["Access-Control-Allow-Origin"] = origin
|
||
return db_model
|
||
|
||
|
||
""" @router.get("/validator", response_model=AssistantValidator, include_in_schema=False)
|
||
async def validator(session: SessionDep, id: int, virtual: Optional[int] = Query(None)):
|
||
if not id:
|
||
raise Exception('miss assistant id')
|
||
|
||
db_model = await get_assistant_info(session=session, assistant_id=id)
|
||
if not db_model:
|
||
return AssistantValidator()
|
||
db_model = AssistantModel.model_validate(db_model)
|
||
assistant_oid = 1
|
||
if (db_model.type == 0):
|
||
configuration = db_model.configuration
|
||
config_obj = json.loads(configuration) if configuration else {}
|
||
assistant_oid = config_obj.get('oid', 1)
|
||
|
||
access_token_expires = timedelta(minutes=settings.ACCESS_TOKEN_EXPIRE_MINUTES)
|
||
assistantDict = {
|
||
"id": virtual, "account": 'sqlbot-inner-assistant', "oid": assistant_oid, "assistant_id": id
|
||
}
|
||
access_token = create_access_token(
|
||
assistantDict, expires_delta=access_token_expires
|
||
)
|
||
return AssistantValidator(True, True, True, access_token) """
|
||
|
||
|
||
@router.get('/picture/{file_id}', summary=f"{PLACEHOLDER_PREFIX}assistant_picture_api", description=f"{PLACEHOLDER_PREFIX}assistant_picture_api")
|
||
async def picture(file_id: str = Path(description="file_id")):
|
||
file_path = SQLBotFileUtils.get_file_path(file_id=file_id)
|
||
if not os.path.exists(file_path):
|
||
raise HTTPException(status_code=404, detail="File not found")
|
||
|
||
if file_id.lower().endswith(".svg"):
|
||
media_type = "image/svg+xml"
|
||
else:
|
||
media_type = "image/jpeg"
|
||
|
||
def iterfile():
|
||
with open(file_path, mode="rb") as f:
|
||
yield from f
|
||
|
||
return StreamingResponse(iterfile(), media_type=media_type)
|
||
|
||
|
||
@router.patch('/ui', summary=f"{PLACEHOLDER_PREFIX}assistant_ui_api", description=f"{PLACEHOLDER_PREFIX}assistant_ui_api")
|
||
@require_permissions(permission=SqlbotPermission(role=['ws_admin']))
|
||
@system_log(LogConfig(operation_type=OperationType.UPDATE, module=OperationModules.APPLICATION, result_id_expr="id"))
|
||
async def ui(session: SessionDep, current_user: CurrentUser, trans: Trans, data: str = Form(), files: List[UploadFile] = []):
|
||
json_data = json.loads(data)
|
||
uiSchema = AssistantUiSchema(**json_data)
|
||
id = uiSchema.id
|
||
db_model = get_ws_assistant(session, id, current_user, trans)
|
||
configuration = db_model.configuration
|
||
config_obj = json.loads(configuration) if configuration else {}
|
||
|
||
ui_schema_dict = uiSchema.model_dump(exclude_none=True, exclude_unset=True)
|
||
if files:
|
||
for file in files:
|
||
origin_file_name = file.filename
|
||
file_name, flag_name = SQLBotFileUtils.split_filename_and_flag(origin_file_name)
|
||
file.filename = file_name
|
||
if flag_name == 'logo' or flag_name == 'float_icon':
|
||
try:
|
||
SQLBotFileUtils.check_file(file=file, file_types=[".jpg", ".jpeg", ".png"],
|
||
limit_file_size=(10 * 1024 * 1024))
|
||
except ValueError as e:
|
||
error_msg = str(e)
|
||
if '文件大小超过限制' in error_msg:
|
||
raise ValueError(f"文件大小超过限制（最大 10 M）")
|
||
else:
|
||
raise e
|
||
if config_obj.get(flag_name):
|
||
SQLBotFileUtils.delete_file(config_obj.get(flag_name))
|
||
file_id = await SQLBotFileUtils.upload(file)
|
||
ui_schema_dict[flag_name] = file_id
|
||
else:
|
||
raise ValueError(f"Unsupported file flag: {flag_name}")
|
||
|
||
for flag_name in ['logo', 'float_icon']:
|
||
file_val = config_obj.get(flag_name)
|
||
if file_val and not ui_schema_dict.get(flag_name):
|
||
config_obj[flag_name] = None
|
||
SQLBotFileUtils.delete_file(file_val)
|
||
|
||
for attr, value in ui_schema_dict.items():
|
||
if attr != 'id' and not attr.startswith("__"):
|
||
config_obj[attr] = value
|
||
|
||
db_model.configuration = json.dumps(config_obj, ensure_ascii=False)
|
||
session.add(db_model)
|
||
session.commit()
|
||
await clear_ui_cache(db_model.id)
|
||
return db_model
|
||
|
||
|
||
@clear_cache(namespace=CacheNamespace.EMBEDDED_INFO, cacheName=CacheName.ASSISTANT_INFO, keyExpression="id")
|
||
async def clear_ui_cache(id: int):
|
||
pass
|
||
|
||
@router.get("/validate/{id}", include_in_schema=False)
|
||
async def validate(request: Request, response: Response, session: SessionDep, trans: Trans, id: int):
|
||
if not id:
|
||
raise Exception('miss assistant id')
|
||
db_model = session.get(AssistantModel, id)
|
||
if not db_model:
|
||
raise RuntimeError(f"assistant application not exist")
|
||
|
||
origin = request.headers.get("origin") or get_origin_from_referer(request)
|
||
if not origin:
|
||
raise RuntimeError(trans('i18n_embedded.invalid_origin', origin=origin or ''))
|
||
origin = origin.rstrip('/')
|
||
|
||
if not origin_match_domain(origin, db_model.domain):
|
||
raise RuntimeError(trans('i18n_embedded.invalid_origin', origin=origin or ''))
|
||
|
||
response.headers["Access-Control-Allow-Origin"] = origin
|
||
return {"valid": True, "origin": origin}
|
||
|
||
|
||
@router.get("/ds", include_in_schema=False, response_model=list[dict])
|
||
async def ds(session: SessionDep, current_assistant: CurrentAssistant):
|
||
if current_assistant.type == 0:
|
||
online = current_assistant.online
|
||
configuration = current_assistant.configuration
|
||
config: dict[any] = json.loads(configuration)
|
||
oid: int = int(config['oid'])
|
||
stmt = select(CoreDatasource.id, CoreDatasource.name, CoreDatasource.description, CoreDatasource.type, CoreDatasource.type_name, CoreDatasource.num).where(
|
||
CoreDatasource.oid == oid)
|
||
if not online:
|
||
public_list: list[int] = config.get('public_list') or None
|
||
if public_list:
|
||
stmt = stmt.where(CoreDatasource.id.in_(public_list))
|
||
else:
|
||
return []
|
||
db_ds_list = session.exec(stmt)
|
||
return [
|
||
{
|
||
"id": ds.id,
|
||
"name": ds.name,
|
||
"description": ds.description,
|
||
"type": ds.type,
|
||
"type_name": ds.type_name,
|
||
"num": ds.num,
|
||
}
|
||
for ds in db_ds_list]
|
||
if current_assistant.type == 1:
|
||
# 构造 AssistantOutDs 会同步请求宿主 API（requests.get，默认 10s 超时），
|
||
# 必须在工作线程执行，避免阻塞事件循环导致宿主回调死锁（Issue #1288）
|
||
def inner():
|
||
return AssistantOutDsFactory.get_instance(current_assistant)
|
||
|
||
out_ds_instance: AssistantOutDs = await asyncio.to_thread(inner)
|
||
return [
|
||
{
|
||
"id": str(ds.id),
|
||
"name": ds.name,
|
||
"description": ds.description or ds.comment,
|
||
"type": ds.type,
|
||
"type_name": get_db_type(ds.type),
|
||
"num": len(ds.tables) if ds.tables else 0,
|
||
}
|
||
for ds in out_ds_instance.ds_list
|
||
if get_db_type(ds.type)
|
||
]
|
||
|
||
return None
|
||
|
||
def get_db_type(type):
|
||
try:
|
||
db = DB.get_db(type)
|
||
return db.db_name
|
||
except Exception:
|
||
return None
|
||
|
||
|
||
@router.get("", response_model=list[AssistantModel], summary=f"{PLACEHOLDER_PREFIX}assistant_grid_api", description=f"{PLACEHOLDER_PREFIX}assistant_grid_api")
|
||
@require_permissions(permission=SqlbotPermission(role=['ws_admin']))
|
||
async def query(session: SessionDep, current_user: CurrentUser):
|
||
list_result = session.exec(select(AssistantModel).where(AssistantModel.oid == current_user.oid, AssistantModel.type != 4).order_by(AssistantModel.name,
|
||
AssistantModel.create_time)).all()
|
||
for model in list_result:
|
||
model.enable_custom_model = model.enable_custom_model or False
|
||
return list_result
|
||
|
||
|
||
@router.get("/advanced_application", response_model=list[AssistantModel], include_in_schema=False)
|
||
@require_permissions(permission=SqlbotPermission(role=['ws_admin']))
|
||
async def query_advanced_application(session: SessionDep, current_user: CurrentUser):
|
||
list_result = session.exec(select(AssistantModel).where(AssistantModel.type == 1, AssistantModel.oid == current_user.oid).order_by(AssistantModel.name,
|
||
AssistantModel.create_time)).all()
|
||
return list_result
|
||
|
||
|
||
@router.post("", summary=f"{PLACEHOLDER_PREFIX}assistant_create_api", description=f"{PLACEHOLDER_PREFIX}assistant_create_api")
|
||
@require_permissions(permission=SqlbotPermission(role=['ws_admin']))
|
||
@system_log(LogConfig(operation_type=OperationType.CREATE, module=OperationModules.APPLICATION, result_id_expr="id"))
|
||
async def add(request: Request, session: SessionDep, current_user: CurrentUser, creator: AssistantBase):
|
||
oid = current_user.oid if creator.type != 4 else 1
|
||
return await save(request, session, creator, oid)
|
||
|
||
|
||
@router.put("", summary=f"{PLACEHOLDER_PREFIX}assistant_update_api", description=f"{PLACEHOLDER_PREFIX}assistant_update_api")
|
||
@require_permissions(permission=SqlbotPermission(role=['ws_admin']))
|
||
@clear_cache(namespace=CacheNamespace.EMBEDDED_INFO, cacheName=CacheName.ASSISTANT_INFO, keyExpression="editor.id")
|
||
@system_log(LogConfig(operation_type=OperationType.UPDATE, module=OperationModules.APPLICATION, resource_id_expr="editor.id"))
|
||
async def update(request: Request, session: SessionDep, current_user: CurrentUser, trans: Trans, editor: AssistantDTO):
|
||
id = editor.id
|
||
db_model = get_ws_assistant(session, id, current_user, trans)
|
||
update_data = AssistantModel.model_validate(editor)
|
||
if not current_user.isAdmin:
|
||
# 忽略请求体 oid（含默认值 1）：非 admin 不允许跨工作空间移动/接管
|
||
update_data.oid = db_model.oid
|
||
# create_time 不在请求体中，model_validate 会填默认值 0，更新时保留原值
|
||
update_data.create_time = db_model.create_time
|
||
db_model.sqlmodel_update(update_data)
|
||
session.add(db_model)
|
||
session.commit()
|
||
dynamic_upgrade_cors(request=request, session=session)
|
||
|
||
|
||
""" @router.get("/{id}", response_model=AssistantModel, summary=f"{PLACEHOLDER_PREFIX}assistant_query_api", description=f"{PLACEHOLDER_PREFIX}assistant_query_api")
|
||
async def get_one(session: SessionDep, id: int = Path(description="ID")):
|
||
db_model = await get_assistant_info(session=session, assistant_id=id)
|
||
if not db_model:
|
||
raise ValueError(f"AssistantModel with id {id} not found")
|
||
db_model = AssistantModel.model_validate(db_model)
|
||
return db_model """
|
||
|
||
|
||
@router.delete("/{id}", summary=f"{PLACEHOLDER_PREFIX}assistant_del_api", description=f"{PLACEHOLDER_PREFIX}assistant_del_api")
|
||
@require_permissions(permission=SqlbotPermission(role=['ws_admin']))
|
||
@clear_cache(namespace=CacheNamespace.EMBEDDED_INFO, cacheName=CacheName.ASSISTANT_INFO, keyExpression="id")
|
||
@system_log(LogConfig(operation_type=OperationType.DELETE, module=OperationModules.APPLICATION, resource_id_expr="id"))
|
||
async def delete(request: Request, session: SessionDep, current_user: CurrentUser, trans: Trans, id: int = Path(description="ID")):
|
||
db_model = get_ws_assistant(session, id, current_user, trans)
|
||
session.delete(db_model)
|
||
session.commit()
|
||
dynamic_upgrade_cors(request=request, session=session)
|
||
|