202 lines
6.2 KiB
Python
202 lines
6.2 KiB
Python
import json
|
|
import time
|
|
from typing import Optional, Dict, Any
|
|
from datetime import datetime, timedelta
|
|
from apitest.models.exceptions import AuthException
|
|
|
|
|
|
class AuthManager:
|
|
"""认证管理器"""
|
|
|
|
def __init__(self, base_url: str, credentials: Dict[str, str], logger):
|
|
"""初始化认证管理器
|
|
|
|
Args:
|
|
base_url: 基础URL
|
|
credentials: 认证凭据
|
|
logger: 日志记录器
|
|
"""
|
|
self.base_url = base_url.rstrip("/")
|
|
self.credentials = credentials
|
|
self.logger = logger
|
|
self._token: Optional[str] = None
|
|
self._token_expiry: Optional[datetime] = None
|
|
self._refresh_token: Optional[str] = None
|
|
self._login_endpoint: str = "/sys/auth/login"
|
|
|
|
def set_login_endpoint(self, endpoint: str):
|
|
"""设置登录端点
|
|
|
|
Args:
|
|
endpoint: 登录端点
|
|
"""
|
|
self._login_endpoint = endpoint
|
|
|
|
def login(self, login_endpoint: Optional[str] = None) -> Dict[str, Any]:
|
|
"""执行登录操作
|
|
|
|
Args:
|
|
login_endpoint: 登录端点，默认使用配置的端点
|
|
|
|
Returns:
|
|
登录响应数据
|
|
|
|
Raises:
|
|
AuthException: 登录失败时抛出
|
|
"""
|
|
endpoint = login_endpoint or self._login_endpoint
|
|
url = f"{self.base_url}{endpoint}"
|
|
|
|
self.logger.info(f"尝试登录: {url}")
|
|
|
|
try:
|
|
import requests
|
|
|
|
response = requests.post(
|
|
url,
|
|
json={
|
|
"username": self.credentials.get("username", ""),
|
|
"password": self.credentials.get("password", "")
|
|
},
|
|
timeout=10
|
|
)
|
|
|
|
if response.status_code == 200:
|
|
data = response.json()
|
|
|
|
if "data" in data and "token" in data["data"]:
|
|
self._token = data["data"]["token"]
|
|
self._refresh_token = data["data"].get("refreshToken")
|
|
|
|
expiry_seconds = data["data"].get("expiresIn", 3600)
|
|
self._token_expiry = datetime.now() + timedelta(seconds=expiry_seconds)
|
|
|
|
self.logger.info("登录成功")
|
|
return data
|
|
else:
|
|
raise AuthException("登录响应中未找到token")
|
|
else:
|
|
raise AuthException(f"登录失败: HTTP {response.status_code}")
|
|
|
|
except requests.RequestException as e:
|
|
raise AuthException(f"登录请求失败: {e}")
|
|
|
|
def get_token(self) -> Optional[str]:
|
|
"""获取当前token
|
|
|
|
Returns:
|
|
当前token，如果未登录则返回None
|
|
"""
|
|
return self._token
|
|
|
|
def set_token(self, token: str, expiry_seconds: int = 3600):
|
|
"""设置token
|
|
|
|
Args:
|
|
token: 认证令牌
|
|
expiry_seconds: 过期时间（秒），默认3600秒
|
|
"""
|
|
self._token = token
|
|
self._token_expiry = datetime.now() + timedelta(seconds=expiry_seconds)
|
|
self.logger.info("Token已设置")
|
|
|
|
def is_token_valid(self) -> bool:
|
|
"""检查token是否有效
|
|
|
|
Returns:
|
|
token是否有效
|
|
"""
|
|
if not self._token or not self._token_expiry:
|
|
return False
|
|
|
|
return datetime.now() < self._token_expiry
|
|
|
|
def refresh_token(self) -> bool:
|
|
"""刷新token
|
|
|
|
Returns:
|
|
刷新是否成功
|
|
"""
|
|
if not self._refresh_token:
|
|
self.logger.warning("没有可用的refresh token")
|
|
return False
|
|
|
|
try:
|
|
import requests
|
|
|
|
url = f"{self.base_url}/sys/auth/refresh"
|
|
response = requests.post(
|
|
url,
|
|
json={"refreshToken": self._refresh_token},
|
|
timeout=10
|
|
)
|
|
|
|
if response.status_code == 200:
|
|
data = response.json()
|
|
|
|
if "data" in data and "token" in data["data"]:
|
|
self._token = data["data"]["token"]
|
|
self._refresh_token = data["data"].get("refreshToken")
|
|
|
|
expiry_seconds = data["data"].get("expiresIn", 3600)
|
|
self._token_expiry = datetime.now() + timedelta(seconds=expiry_seconds)
|
|
|
|
self.logger.info("Token刷新成功")
|
|
return True
|
|
|
|
self.logger.warning("Token刷新失败，尝试重新登录")
|
|
return False
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Token刷新异常: {e}")
|
|
return False
|
|
|
|
def ensure_authenticated(self) -> str:
|
|
"""确保已认证，如果token过期则自动刷新或重新登录
|
|
|
|
Returns:
|
|
有效的token
|
|
|
|
Raises:
|
|
AuthException: 认证失败时抛出
|
|
"""
|
|
if not self._token:
|
|
self.login()
|
|
elif not self.is_token_valid():
|
|
if not self.refresh_token():
|
|
self.login()
|
|
|
|
if not self._token:
|
|
raise AuthException("无法获取有效的认证token")
|
|
|
|
return self._token
|
|
|
|
def logout(self):
|
|
"""登出"""
|
|
self._token = None
|
|
self._refresh_token = None
|
|
self._token_expiry = None
|
|
self.logger.info("已登出")
|
|
|
|
def get_auth_headers(self) -> Dict[str, str]:
|
|
"""获取认证请求头
|
|
|
|
Returns:
|
|
包含认证信息的请求头字典
|
|
"""
|
|
token = self.ensure_authenticated()
|
|
return {
|
|
"Authorization": f"Bearer {token}",
|
|
"Content-Type": "application/json"
|
|
}
|
|
|
|
def set_credentials(self, username: str, password: str):
|
|
"""设置认证凭据
|
|
|
|
Args:
|
|
username: 用户名
|
|
password: 密码
|
|
"""
|
|
self.credentials = {"username": username, "password": password}
|
|
self.logger.info("认证凭据已更新")
|