mirror of
https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git
synced 2026-11-08 01:06:59 +08:00
338 lines
15 KiB
Python
338 lines
15 KiB
Python
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, 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.invalid_origin', origin=''))
|
|
|
|
sign_data = await decrypt_embedded_sign(sign_header)
|
|
|
|
# 校验 assistant_id 与 id 参数一致
|
|
if str(sign_data.get("assistant_id")) != str(id):
|
|
raise RuntimeError(trans('i18n_embedded.invalid_origin', origin=''))
|
|
|
|
# 校验 target（来源域名）是否合法
|
|
target = sign_data.get("target", "")
|
|
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", "")
|
|
sign_time = datetime.fromisoformat(sign_time_str)
|
|
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.invalid_origin', origin=target or ''))
|
|
|
|
# 校验是否为真实浏览器请求（非自动化工具）
|
|
if sign_data.get("webdriver", False):
|
|
raise RuntimeError(trans('i18n_embedded.invalid_origin', origin=target or ''))
|
|
|
|
# 校验 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.invalid_origin', origin=target or ''))
|
|
|
|
# 校验 timezone 与 sign_time 偏移一致性（防时区伪造）
|
|
tz_name = sign_data.get("timezone", "")
|
|
tz = ZoneInfo(tz_name)
|
|
sign_time_naive = sign_time.replace(tzinfo=None)
|
|
if tz.utcoffset(sign_time_naive) != sign_time.utcoffset():
|
|
raise RuntimeError(trans('i18n_embedded.invalid_origin', origin=target or ''))
|
|
|
|
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, data: str = Form(), files: List[UploadFile] = []):
|
|
json_data = json.loads(data)
|
|
uiSchema = AssistantUiSchema(**json_data)
|
|
id = uiSchema.id
|
|
db_model = session.get(AssistantModel, id)
|
|
if not db_model:
|
|
raise ValueError(f"AssistantModel with id {id} not found")
|
|
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("/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:
|
|
out_ds_instance: AssistantOutDs = AssistantOutDsFactory.get_instance(current_assistant)
|
|
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, editor: AssistantDTO):
|
|
id = editor.id
|
|
db_model = session.get(AssistantModel, id)
|
|
if not db_model:
|
|
raise ValueError(f"AssistantModel with id {id} not found")
|
|
update_data = AssistantModel.model_validate(editor)
|
|
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, id: int = Path(description="ID")):
|
|
db_model = session.get(AssistantModel, id)
|
|
if not db_model:
|
|
raise ValueError(f"AssistantModel with id {id} not found")
|
|
session.delete(db_model)
|
|
session.commit()
|
|
dynamic_upgrade_cors(request=request, session=session)
|
|
|