Files
KiloStar/kilostar/utils/access.py
T

131 lines
4.7 KiB
Python

# Copyright 2026 zhaoxi826
#
# 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.
from __future__ import annotations
import os
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING, Optional
import jwt
from fastapi import HTTPException, Request, status
from pydantic import BaseModel, ValidationError
from pwdlib import PasswordHash
if TYPE_CHECKING:
from kilostar.core.postgres_database.model.user import User
ALGORITHM = "HS256"
ACCESS_TOKEN_EXPIRE_MINUTES = 60 * 24
_INSECURE_SECRETS = {"secret", "114514", "changethiskey12345"}
class TokenData(BaseModel):
"""JWT 解码后的用户身份载荷。"""
user_id: str
username: Optional[str] = None
exp: Optional[int] = None
def _get_secret_key() -> str:
"""读取并校验 SECRET_KEY 环境变量。
校验在首次实际使用 JWT 时进行,避免在模块导入阶段抛错,
从而把"环境约束""模块加载"解耦。
"""
key = os.getenv("SECRET_KEY")
if not key or key in _INSECURE_SECRETS:
raise RuntimeError(
"未提供有效的 SECRET_KEY 或使用了不安全的默认值,请设置一个高熵的随机字符串"
)
return key
password_hasher = PasswordHash.recommended()
class Accessor:
"""封装认证与口令哈希相关的静态工具方法。"""
@staticmethod
def _decode_token(token: str) -> TokenData:
"""解码并校验 JWT,返回 TokenData;过期或无效时抛 401。"""
try:
payload = jwt.decode(token, _get_secret_key(), algorithms=[ALGORITHM])
return TokenData(**payload)
except jwt.ExpiredSignatureError:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="Token 已过期",
)
except (jwt.InvalidTokenError, ValidationError):
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="无效的认证凭证",
)
@staticmethod
def _create_access_token(data: dict) -> str:
"""根据 payload 生成带过期时间的 JWT 访问令牌。"""
to_encode = data.copy()
expire = datetime.now(timezone.utc) + timedelta(
minutes=ACCESS_TOKEN_EXPIRE_MINUTES
)
to_encode.update({"exp": int(expire.timestamp())})
return jwt.encode(to_encode, _get_secret_key(), algorithm=ALGORITHM)
@staticmethod
def verify_password(plain_password: str, hashed_password: str) -> bool:
"""校验明文口令是否匹配数据库中存储的哈希。"""
return password_hasher.verify(plain_password, hashed_password)
@staticmethod
def get_current_user(request: Request) -> TokenData:
"""从 Authorization Bearer 头解析当前请求的用户身份。"""
auth_header = request.headers.get("Authorization")
if not auth_header or not auth_header.startswith("Bearer "):
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="未提供认证头部",
)
token = auth_header.split(" ")[1]
return Accessor._decode_token(token)
@staticmethod
def login_hashed_password(user: "User", password: str) -> str:
"""完成登录核验:找不到用户或密码错误抛 401,否则签发新令牌。"""
if not user:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="用户不存在",
)
if not Accessor.verify_password(password, user.hashed_password):
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="用户名或密码错误",
)
token_payload = {"user_id": str(user.user_id), "username": user.user_name}
return Accessor._create_access_token(data=token_payload)
@staticmethod
def hash_password(password: str) -> str:
"""对明文口令做强哈希;空值或长度不足 6 位会抛 ValueError。"""
if not password:
raise ValueError("密码不能为空")
if len(password) < 6:
raise ValueError("密码长度不能小于 6 位")
return password_hasher.hash(password)