diff --git a/backend/apps/system/api/assistant.py b/backend/apps/system/api/assistant.py index ab8f5419..b9657e21 100644 --- a/backend/apps/system/api/assistant.py +++ b/backend/apps/system/api/assistant.py @@ -1,11 +1,13 @@ from datetime import timedelta from fastapi import APIRouter, FastAPI, Request from sqlmodel import Session, select -from apps.system.crud.user import get_user_by_account +from apps.system.crud.assistant import get_assistant_info from apps.system.models.system_model import AssistantModel +from apps.system.schemas.auth import CacheName, CacheNamespace from apps.system.schemas.system_schema import AssistantBase, AssistantDTO, AssistantValidator from common.core.deps import SessionDep from common.core.security import create_access_token +from common.core.sqlbot_cache import clear_cache from common.utils.time import get_timestamp from starlette.middleware.cors import CORSMiddleware from common.core.config import settings @@ -13,13 +15,15 @@ router = APIRouter(tags=["system/assistant"], prefix="/system/assistant") @router.get("/validator/{id}", response_model=AssistantValidator) async def info(session: SessionDep, id: int): - db_model = session.get(AssistantModel, id) + db_model = get_assistant_info(session, id) if not db_model: return AssistantValidator() access_token_expires = timedelta(minutes=settings.ACCESS_TOKEN_EXPIRE_MINUTES) - user = get_user_by_account(session=session, account='admin') + assistantDict = { + "id": 1, "account": 'admin', "oid": 1, "assistant_id": id + } access_token = create_access_token( - user.to_dict(), expires_delta=access_token_expires + assistantDict, expires_delta=access_token_expires ) return AssistantValidator(True, True, True, access_token) @@ -39,6 +43,7 @@ async def add(request: Request, session: SessionDep, creator: AssistantBase): @router.put("") +@clear_cache(namespace=CacheNamespace.EMBEDDED_INFO, cacheName=CacheName.ASSISTANT_INFO, keyExpression="editor.id") async def update(request: Request, session: SessionDep, editor: AssistantDTO): id = editor.id db_model = session.get(AssistantModel, id) @@ -52,12 +57,13 @@ async def update(request: Request, session: SessionDep, editor: AssistantDTO): @router.get("/{id}", response_model=AssistantModel) async def get_one(session: SessionDep, id: int): - db_model = session.get(AssistantModel, id) + db_model = get_assistant_info(session, id) if not db_model: raise ValueError(f"AssistantModel with id {id} not found") return db_model -@router.delete("/{id}") +@router.delete("/{id}") +@clear_cache(namespace=CacheNamespace.EMBEDDED_INFO, cacheName=CacheName.ASSISTANT_INFO, keyExpression="id") async def delete(request: Request, session: SessionDep, id: int): db_model = session.get(AssistantModel, id) if not db_model: diff --git a/backend/apps/system/crud/assistant.py b/backend/apps/system/crud/assistant.py new file mode 100644 index 00000000..0eeeba8c --- /dev/null +++ b/backend/apps/system/crud/assistant.py @@ -0,0 +1,12 @@ + + +from sqlmodel import Session +from apps.system.models.system_model import AssistantModel +from apps.system.schemas.auth import CacheName, CacheNamespace +from common.core.sqlbot_cache import cache + + +@cache(namespace=CacheNamespace.EMBEDDED_INFO, cacheName=CacheName.ASSISTANT_INFO, keyExpression="assistant_id") +async def get_assistant_info(*, session: Session, assistant_id: int) -> AssistantModel | None: + db_model = session.get(AssistantModel, assistant_id) + return db_model \ No newline at end of file diff --git a/backend/apps/system/middleware/auth.py b/backend/apps/system/middleware/auth.py index eff32618..e821e9c7 100644 --- a/backend/apps/system/middleware/auth.py +++ b/backend/apps/system/middleware/auth.py @@ -1,28 +1,88 @@ -from fastapi import Depends +from typing import Optional +from fastapi import Request from fastapi.responses import JSONResponse +import jwt +from sqlmodel import Session from starlette.middleware.base import BaseHTTPMiddleware +from common.core.db import engine +from apps.system.crud.assistant import get_assistant_info +from apps.system.crud.user import get_user_info +from apps.system.schemas.system_schema import UserInfoDTO +from common.core import security from common.core.config import settings -# from common.core.deps import get_current_user +from common.core.schemas import TokenPayload from common.utils.whitelist import whiteUtils - +from fastapi.security.utils import get_authorization_scheme_param class TokenMiddleware(BaseHTTPMiddleware): + + def __init__(self, app): super().__init__(app) async def dispatch(self, request, call_next): - tokenkey = settings.TOKEN_KEY + if self.is_options(request) or whiteUtils.is_whitelisted(request.url.path): return await call_next(request) + assistantTokenKey = settings.ASSISTANT_TOKEN_KEY + assistantToken = request.headers.get(assistantTokenKey) + #if assistantToken and assistantToken.lower().startswith("assistant "): + if assistantToken: + validate_pass, data, assistant = await self.validateAssistant(assistantToken) + if validate_pass: + request.state.current_user = data + request.state.assistant = assistant + return await call_next(request) + return JSONResponse({"error": f"Unauthorized:[{data}]"}, status_code=401) + #validate pass + tokenkey = settings.TOKEN_KEY token = request.headers.get(tokenkey) - if not token or not token.startswith("Bearer "): - return JSONResponse({"error": "Unauthorized"}, status_code=401) - """ user = await get_current_user() - request.state.user = user """ - return await call_next(request) + validate_pass, data = await self.validateToken(token) + if validate_pass: + request.state.current_user = data + return await call_next(request) + return JSONResponse({"error": f"Unauthorized:[{data}]"}, status_code=401) - def is_options(self, request): + def is_options(self, request: Request): return request.method == "OPTIONS" - \ No newline at end of file + async def validateToken(self, token: Optional[str]): + 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!" + payload = jwt.decode( + param, settings.SECRET_KEY, algorithms=[security.ALGORITHM] + ) + token_data = TokenPayload(**payload) + try: + 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) + return True, session_user + except Exception as e: + return False, e + + + async def validateAssistant(self, assistantToken: Optional[str]): + if not assistantToken: + return False, f"Miss Token[{settings.TOKEN_KEY}]!" + schema, param = get_authorization_scheme_param(assistantToken) + 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!" + try: + 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) + assistant_info = get_assistant_info(session, payload['assistant_id']) + return True, session_user, assistant_info + except Exception as e: + return False, e \ No newline at end of file diff --git a/backend/apps/system/schemas/auth.py b/backend/apps/system/schemas/auth.py index fd50737d..70b87d74 100644 --- a/backend/apps/system/schemas/auth.py +++ b/backend/apps/system/schemas/auth.py @@ -8,10 +8,11 @@ class LocalLoginSchema(BaseModel): class CacheNamespace(Enum): AUTH_INFO = "sqlbot:auth" + EMBEDDED_INFO = "sqlbot:embedded" def __str__(self): return self.value class CacheName(Enum): USER_INFO = "user:info" - + ASSISTANT_INFO = "assistant:info" def __str__(self): return self.value \ No newline at end of file diff --git a/backend/common/core/config.py b/backend/common/core/config.py index db0496ce..99c7fbc7 100644 --- a/backend/common/core/config.py +++ b/backend/common/core/config.py @@ -57,8 +57,9 @@ class Settings(BaseSettings): POSTGRES_PASSWORD: str = "" POSTGRES_DB: str = "" - TOKEN_KEY: str - DEFAULT_PWD: str + TOKEN_KEY: str = "X-SQLBOT-TOKEN" + DEFAULT_PWD: str = "SQLBot@123456" + ASSISTANT_TOKEN_KEY: str = "X-SQLBOT-ASSISTANT-TOKEN" @computed_field # type: ignore[prop-decorator] @property diff --git a/backend/common/core/deps.py b/backend/common/core/deps.py index 3480886b..85942b0f 100644 --- a/backend/common/core/deps.py +++ b/backend/common/core/deps.py @@ -1,51 +1,21 @@ from typing import Annotated -import jwt -from fastapi import Depends, HTTPException, Request, status -# from fastapi.security import OAuth2PasswordBearer -from jwt.exceptions import InvalidTokenError -from pydantic import ValidationError +from fastapi import Depends, Request from sqlmodel import Session -from apps.system.crud.user import get_user_info from apps.system.schemas.system_schema import UserInfoDTO -from common.core.schemas import TokenPayload, XOAuth2PasswordBearer -from common.core import security -from common.core.config import settings from common.core.db import get_session from common.utils.locale import I18n -reusable_oauth2 = XOAuth2PasswordBearer( - tokenUrl=f"{settings.API_V1_STR}/login/access-token" -) - - SessionDep = Annotated[Session, Depends(get_session)] -TokenDep = Annotated[str, Depends(reusable_oauth2)] i18n = I18n() async def get_i18n(request: Request): return i18n(request) Trans = Annotated[I18n, Depends(get_i18n)] -async def get_current_user(session: SessionDep, token: TokenDep) -> UserInfoDTO: - try: - payload = jwt.decode( - token, settings.SECRET_KEY, algorithms=[security.ALGORITHM] - ) - token_data = TokenPayload(**payload) - except (InvalidTokenError, ValidationError): - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail="Could not validate credentials", - ) - session_user = await get_user_info(session = session, user_id = token_data.id) - session_user = UserInfoDTO.model_validate(session_user) - if not session_user: - raise HTTPException(status_code=404, detail="User not found") - - if session_user.status != 1: - raise HTTPException(status_code=400, detail="Inactive user") - return session_user +async def get_current_user(request: Request) -> UserInfoDTO: + return request.state.current_user + CurrentUser = Annotated[UserInfoDTO, Depends(get_current_user)] diff --git a/backend/common/core/schemas.py b/backend/common/core/schemas.py index 30746692..6d9d75fa 100644 --- a/backend/common/core/schemas.py +++ b/backend/common/core/schemas.py @@ -18,8 +18,11 @@ class Token(SQLModel): class XOAuth2PasswordBearer(OAuth2PasswordBearer): async def __call__(self, request: Request) -> Optional[str]: authorization = request.headers.get(settings.TOKEN_KEY) + if request.headers.get(settings.ASSISTANT_TOKEN_KEY): + authorization = request.headers.get(settings.ASSISTANT_TOKEN_KEY) scheme, param = get_authorization_scheme_param(authorization) - if not authorization or scheme.lower() != "bearer": + + if not authorization or scheme.lower() not in ["bearer", "assistant"]: if self.auto_error: raise HTTPException( status_code=HTTP_401_UNAUTHORIZED, diff --git a/frontend/src/stores/assistant.ts b/frontend/src/stores/assistant.ts new file mode 100644 index 00000000..99bc8eb2 --- /dev/null +++ b/frontend/src/stores/assistant.ts @@ -0,0 +1,31 @@ +import { defineStore } from 'pinia' +import { store } from './index' + +interface AssistantState { + token: string +} + +export const AssistantStore = defineStore('assistant', { + state: (): AssistantState => { + return { + token: '', + } + }, + getters: { + getToken(): string { + return this.token + }, + }, + actions: { + setToken(token: string) { + this.token = token + }, + clear() { + this.$reset() + }, + }, +}) + +export const useAssistantStore = () => { + return AssistantStore(store) +} diff --git a/frontend/src/utils/request.ts b/frontend/src/utils/request.ts index b3d4f213..5665172c 100644 --- a/frontend/src/utils/request.ts +++ b/frontend/src/utils/request.ts @@ -10,6 +10,8 @@ import axios, { import { useCache } from '@/utils/useCache' import { getLocale } from './utils' +import { useAssistantStore } from '@/stores/assistant' +const assistantStore = useAssistantStore() const { wsCache } = useCache() // Response data structure export interface ApiResponse { @@ -70,6 +72,10 @@ class HttpService { if (token && config.headers) { config.headers['X-SQLBOT-TOKEN'] = `Bearer ${token}` } + if (assistantStore.getToken) { + config.headers['X-SQLBOT-ASSISTANT-TOKEN'] = `Assistant ${assistantStore.getToken}` + if (config.headers['X-SQLBOT-TOKEN']) config.headers.delete('X-SQLBOT-TOKEN') + } const locale = getLocale() if (locale) { /* const mapping = { diff --git a/frontend/src/views/embedded/assistant.vue b/frontend/src/views/embedded/assistant.vue index fb9f59e6..c66a92d1 100644 --- a/frontend/src/views/embedded/assistant.vue +++ b/frontend/src/views/embedded/assistant.vue @@ -62,8 +62,9 @@ import IconOpeEdit from '@/assets/svg/operate/ope-edit.svg' import IconOpeDelete from '@/assets/svg/operate/ope-delete.svg' import { useRoute } from 'vue-router' import { assistantApi } from '@/api/assistant' -import { useUserStore } from '@/stores/user' -const userStore = useUserStore() +import { useAssistantStore } from '@/stores/assistant' + +const assistantStore = useAssistantStore() const route = useRoute() const chatRef = ref() @@ -96,7 +97,7 @@ const loading = ref(true) onBeforeMount(async () => { const assistantId = route.params.id validator.value = await assistantApi.validate(assistantId) - userStore.setToken(validator.value.token) + assistantStore.setToken(validator.value.token) loading.value = false })