mirror of
https://github.com/eigent-ai/eigent.git
synced 2026-04-28 11:40:25 +00:00
33829f2c1e
Co-authored-by: 4pmtong <web_chentong@163.com> Co-authored-by: Wendong-Fan <133094783+Wendong-Fan@users.noreply.github.com> Co-authored-by: Wendong-Fan <w3ndong.fan@gmail.com>
104 lines
3.6 KiB
Python
104 lines
3.6 KiB
Python
# ========= Copyright 2025-2026 @ Eigent.ai All Rights Reserved. =========
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# ========= Copyright 2025-2026 @ Eigent.ai All Rights Reserved. =========
|
|
|
|
from fastapi import Depends, Header
|
|
from fastapi_babel import _
|
|
from sqlmodel import Session, select
|
|
from app.component import code
|
|
from fastapi.security import OAuth2PasswordBearer
|
|
from app.component.database import session
|
|
from app.component.environment import env, env_not_empty
|
|
from datetime import timedelta, datetime
|
|
import jwt
|
|
from jwt.exceptions import InvalidTokenError
|
|
from app.model.mcp.proxy import ApiKey
|
|
from app.model.user.key import Key
|
|
from app.model.user.user import User
|
|
|
|
from app.exception.exception import (
|
|
NoPermissionException,
|
|
TokenException,
|
|
)
|
|
|
|
|
|
class Auth:
|
|
SECRET_KEY = env_not_empty("secret_key")
|
|
|
|
def __init__(self, id: int, expired_at: datetime):
|
|
self.id = id
|
|
self.expired_at = expired_at
|
|
self._user: User | None = None
|
|
|
|
@property
|
|
def user(self):
|
|
if self._user is None:
|
|
raise NoPermissionException("未查询到登录用户")
|
|
return self._user
|
|
|
|
@classmethod
|
|
def decode_token(cls, token: str):
|
|
try:
|
|
payload = jwt.decode(token, Auth.SECRET_KEY, algorithms=["HS256"])
|
|
id = payload["id"]
|
|
if payload["exp"] < int(datetime.now().timestamp()):
|
|
raise TokenException(code.token_expired, _("Validate credentials expired"))
|
|
except InvalidTokenError:
|
|
raise TokenException(code.token_invalid, _("Could not validate credentials"))
|
|
return Auth(id, payload["exp"])
|
|
|
|
@classmethod
|
|
def create_access_token(cls, user_id: int, expires_delta: timedelta | None = None):
|
|
to_encode: dict = {"id": user_id}
|
|
if expires_delta:
|
|
expire = datetime.now() + expires_delta
|
|
else:
|
|
expire = datetime.now() + timedelta(days=30)
|
|
to_encode.update({"exp": expire})
|
|
encoded_jwt = jwt.encode(to_encode, Auth.SECRET_KEY, algorithm="HS256")
|
|
return encoded_jwt
|
|
|
|
|
|
oauth2_scheme = OAuth2PasswordBearer(tokenUrl=f"{env('url_prefix', '')}/dev_login", auto_error=False)
|
|
|
|
|
|
async def auth(
|
|
token: str | None = Depends(oauth2_scheme),
|
|
session: Session = Depends(session),
|
|
) -> Auth | None:
|
|
if token is None:
|
|
return None
|
|
try:
|
|
model = Auth.decode_token(token)
|
|
user = session.get(User, model.id)
|
|
model._user = user
|
|
return model
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
async def auth_must(
|
|
token: str = Depends(oauth2_scheme),
|
|
session: Session = Depends(session),
|
|
) -> Auth:
|
|
model = Auth.decode_token(token)
|
|
user = session.get(User, model.id)
|
|
model._user = user
|
|
return model
|
|
|
|
|
|
async def key_must(headers: ApiKey = Header(), session: Session = Depends(session)):
|
|
model = session.exec(select(Key).where(Key.value == headers.api_key)).one_or_none()
|
|
if model is None:
|
|
raise TokenException(code.token_invalid, _(f"Could not validate key credentials: {headers.api_key}"))
|
|
return model
|