forked from mttu-developers/konabot
150 lines
5.7 KiB
Python
150 lines
5.7 KiB
Python
import asyncio
|
|
from dataclasses import dataclass
|
|
from hashlib import md5
|
|
import time
|
|
|
|
from nonebot import logger
|
|
from nonebot_plugin_apscheduler import driver
|
|
from konabot.common.path import DATA_PATH
|
|
import os
|
|
from PIL import Image
|
|
|
|
IMAGE_PATH = DATA_PATH / "temp" / "images"
|
|
|
|
@dataclass
|
|
class ImageResource:
|
|
name: str
|
|
expire: int
|
|
|
|
@dataclass
|
|
class StorageImage:
|
|
name: str
|
|
resources: dict[str,
|
|
dict[str,ImageResource]] # {群号: {QQ号: ImageResource}}
|
|
|
|
class ImageStorager:
|
|
images_pool: dict[str,StorageImage] = {}
|
|
|
|
max_storage: int = 10 * 1024 * 1024 # 最大存储10MB
|
|
max_image_count: int = 200 # 最大存储图片数量
|
|
|
|
@staticmethod
|
|
def init():
|
|
if not IMAGE_PATH.exists():
|
|
IMAGE_PATH.mkdir(parents=True, exist_ok=True)
|
|
|
|
@staticmethod
|
|
def delete_path_image(name: str):
|
|
resource_path = IMAGE_PATH / name
|
|
if resource_path.exists():
|
|
os.remove(resource_path)
|
|
|
|
@classmethod
|
|
async def clear_expire_image(cls):
|
|
# 清理过期的图片资源，将未被删除的放入列表中，如果超过最大数量则删除最早过期的
|
|
remaining_images = []
|
|
current_time = time.time()
|
|
for name, storage_image in list(ImageStorager.images_pool.items()):
|
|
for group_id, resources in list(storage_image.resources.items()):
|
|
for qq_id, resource in list(resources.items()):
|
|
if resource.expire < current_time:
|
|
del storage_image.resources[group_id][qq_id]
|
|
cls.delete_path_image(name)
|
|
else:
|
|
remaining_images.append((name, group_id, qq_id, resource.expire))
|
|
if not storage_image.resources:
|
|
del ImageStorager.images_pool[name]
|
|
# 如果剩余图片超过最大数量，按过期时间排序并删除最早过期的
|
|
if len(remaining_images) > ImageStorager.max_image_count:
|
|
remaining_images.sort(key=lambda x: x[3]) # 按过期时间排序
|
|
to_delete = len(remaining_images) - ImageStorager.max_image_count
|
|
for i in range(to_delete):
|
|
name, group_id, qq_id, _ = remaining_images[i]
|
|
resource = ImageStorager.images_pool[name].resources[group_id][qq_id]
|
|
del ImageStorager.images_pool[name].resources[group_id][qq_id]
|
|
cls.delete_path_image(name)
|
|
|
|
@classmethod
|
|
def _add_to_pool(cls, image: bytes, name: str, group_id: str, qq_id: str, expire: int = 36000):
|
|
expire_time = time.time() + expire
|
|
if name not in cls.images_pool:
|
|
cls.images_pool[name] = StorageImage(name=name,resources={})
|
|
if group_id not in cls.images_pool[name].resources:
|
|
cls.images_pool[name].resources[group_id] = {}
|
|
cls.images_pool[name].resources[group_id][qq_id] = ImageResource(name=name, expire=expire_time)
|
|
|
|
@classmethod
|
|
def save_image(cls, image: bytes, name: str, group_id: str, qq_id: str) -> None:
|
|
"""
|
|
以哈希值命名保存图片，并返回图片资源信息
|
|
"""
|
|
# 检测图像大小，不得超过 10 MB
|
|
if len(image) > cls.max_storage:
|
|
raise ValueError("图片大小超过 10 MB 限制")
|
|
hash_name = md5(image).hexdigest()
|
|
ext = os.path.splitext(name)[1]
|
|
file_name = f"{hash_name}{ext}"
|
|
full_path = IMAGE_PATH / file_name
|
|
with open(full_path, "wb") as f:
|
|
f.write(image)
|
|
# 将文件写入 images_pool
|
|
cls._add_to_pool(image, file_name, group_id, qq_id)
|
|
|
|
@classmethod
|
|
def load_image(cls, name: str, group_id: str, qq_id: str) -> Image:
|
|
if name not in cls.images_pool:
|
|
return None
|
|
if group_id not in cls.images_pool[name].resources:
|
|
return None
|
|
# 寻找对应 QQ 号 的资源，如果没有就返回相同群下的第一个资源
|
|
if qq_id not in cls.images_pool[name].resources[group_id]:
|
|
first_qq_id = next(iter(cls.images_pool[name].resources[group_id]))
|
|
qq_id = first_qq_id
|
|
resource = cls.images_pool[name].resources[group_id][qq_id]
|
|
resource_path = IMAGE_PATH / resource.name
|
|
return Image.open(resource_path)
|
|
|
|
class ImageStoragerManager:
|
|
def __init__(self, interval: int = 300): # 默认 5 分钟执行一次
|
|
self.interval = interval
|
|
self._clear_task = None
|
|
self._running = False
|
|
|
|
async def start_auto_clear(self):
|
|
"""启动自动任务"""
|
|
self._running = True
|
|
self._clear_task = asyncio.create_task(self._auto_clear_loop())
|
|
|
|
logger.info(f"自动清理任务已启动，间隔: {self.interval}秒")
|
|
|
|
async def stop_auto_clear(self):
|
|
"""停止自动清理任务"""
|
|
if self._clear_task:
|
|
self._running = False
|
|
self._clear_task.cancel()
|
|
try:
|
|
await self._clear_task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
logger.info("自动清理任务已停止")
|
|
else:
|
|
logger.warning("没有正在运行的自动清理任务")
|
|
|
|
async def _auto_clear_loop(self):
|
|
"""自动清理循环"""
|
|
while self._running:
|
|
try:
|
|
await asyncio.sleep(self.interval)
|
|
await ImageStorager.clear_expire_image()
|
|
except asyncio.CancelledError:
|
|
break
|
|
except Exception as e:
|
|
logger.error(f"定时清理失败: {e}")
|
|
|
|
image_manager = ImageStoragerManager(interval=300) # 每5分钟清理一次
|
|
|
|
@driver.on_startup
|
|
async def init_image_storage():
|
|
ImageStorager.init()
|
|
# 启用定时任务清理过期图片
|
|
await image_manager.start_auto_clear() |