mirror of
https://github.com/AstrBotDevs/AstrBot
synced 2026-07-17 09:50:31 +08:00
a9c16febf4
* test(skills): align sandbox cache tests with readonly behavior * ci(release): enforce core quality gate before publish * ci: enforce locked dependency installs in workflows * security: remove curl-pipe-shell installs * chore: align project python baseline to 3.12 * ci(dashboard): add explicit typecheck gate * chore(pre-commit): align ruff hook version with project * ci(codeql): add javascript-typescript analysis * chore(ruff): defer py312 migration lint rules * fix: resolve ruff violations without new ignores * fix: resolve ASYNC230 and ASYNC240 without ignores * fix(auth): replace utcnow with timezone-aware UTC now * fix: avoid blocking file read in file_to_base64
107 lines
3.5 KiB
Python
107 lines
3.5 KiB
Python
import asyncio
|
|
import os
|
|
import platform
|
|
import time
|
|
import uuid
|
|
from urllib.parse import unquote, urlparse
|
|
|
|
|
|
class FileTokenService:
|
|
"""维护一个简单的基于令牌的文件下载服务，支持超时和懒清除。"""
|
|
|
|
def __init__(self, default_timeout: float = 300) -> None:
|
|
self.lock = asyncio.Lock()
|
|
self.staged_files = {} # token: (file_path, expire_time)
|
|
self.default_timeout = default_timeout
|
|
|
|
async def _cleanup_expired_tokens(self) -> None:
|
|
"""清理过期的令牌"""
|
|
now = time.time()
|
|
expired_tokens = [
|
|
token for token, (_, expire) in self.staged_files.items() if expire < now
|
|
]
|
|
for token in expired_tokens:
|
|
self.staged_files.pop(token, None)
|
|
|
|
async def check_token_expired(self, file_token: str) -> bool:
|
|
async with self.lock:
|
|
await self._cleanup_expired_tokens()
|
|
return file_token not in self.staged_files
|
|
|
|
async def register_file(
|
|
self,
|
|
file_path: str,
|
|
timeout_seconds: float | None = None,
|
|
**kwargs,
|
|
) -> str:
|
|
"""向令牌服务注册一个文件。
|
|
|
|
Args:
|
|
file_path(str): 文件路径
|
|
timeout_seconds(float): 超时时间，单位秒（可选）
|
|
|
|
Returns:
|
|
str: 一个单次令牌
|
|
|
|
Raises:
|
|
FileNotFoundError: 当路径不存在时抛出
|
|
|
|
"""
|
|
# 处理 file:///
|
|
try:
|
|
parsed_uri = urlparse(file_path)
|
|
if parsed_uri.scheme == "file":
|
|
local_path = unquote(parsed_uri.path)
|
|
if platform.system() == "Windows" and local_path.startswith("/"):
|
|
local_path = local_path[1:]
|
|
else:
|
|
# 如果没有 file:/// 前缀，则认为是普通路径
|
|
local_path = file_path
|
|
except Exception:
|
|
# 解析失败时，按原路径处理
|
|
local_path = file_path
|
|
|
|
async with self.lock:
|
|
await self._cleanup_expired_tokens()
|
|
legacy_timeout = kwargs.pop("timeout", None)
|
|
if legacy_timeout is not None:
|
|
timeout_seconds = float(legacy_timeout)
|
|
|
|
if not await asyncio.to_thread(os.path.exists, local_path):
|
|
raise FileNotFoundError(
|
|
f"文件不存在: {local_path} (原始输入: {file_path})",
|
|
)
|
|
|
|
file_token = str(uuid.uuid4())
|
|
expire_time = time.time() + (
|
|
timeout_seconds if timeout_seconds is not None else self.default_timeout
|
|
)
|
|
# 存储转换后的真实路径
|
|
self.staged_files[file_token] = (local_path, expire_time)
|
|
return file_token
|
|
|
|
async def handle_file(self, file_token: str) -> str:
|
|
"""根据令牌获取文件路径，使用后令牌失效。
|
|
|
|
Args:
|
|
file_token(str): 注册时返回的令牌
|
|
|
|
Returns:
|
|
str: 文件路径
|
|
|
|
Raises:
|
|
KeyError: 当令牌不存在或已过期时抛出
|
|
FileNotFoundError: 当文件本身已被删除时抛出
|
|
|
|
"""
|
|
async with self.lock:
|
|
await self._cleanup_expired_tokens()
|
|
|
|
if file_token not in self.staged_files:
|
|
raise KeyError(f"无效或过期的文件 token: {file_token}")
|
|
|
|
file_path, _ = self.staged_files.pop(file_token)
|
|
if not await asyncio.to_thread(os.path.exists, file_path):
|
|
raise FileNotFoundError(f"文件不存在: {file_path}")
|
|
return file_path
|