Files
abot/message_sign/main.py
2025-03-19 09:55:08 +08:00

206 lines
8.0 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
from datetime import datetime, timedelta
import logging
import mysql.connector.pooling
import tomllib
import pytz
import redis
from typing import Optional
from wcferry import Wcf, WxMsg
from message_util import MessageUtil
from robot_cmd.robot_command import GroupBotManager, Feature, PermissionStatus
from db.connection import DBConnectionManager
from db.sign_in import SignInDB
from db.sign_in_redis import SignInRedisDB
# 创建表的SQL语句保留在这里用于初始化表
CREATE_TABLE_SQL = """
CREATE TABLE IF NOT EXISTS t_sign_record (
id BIGINT AUTO_INCREMENT PRIMARY KEY,
wx_id VARCHAR(100) NOT NULL,
group_id VARCHAR(100) NOT NULL,
wx_nick_name VARCHAR(100) NOT NULL,
points INT DEFAULT 0,
sign_stat DATETIME,
signin_streak INT DEFAULT 0,
create_time DATETIME DEFAULT CURRENT_TIMESTAMP,
update_time DATETIME DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
UNIQUE KEY unique_sign (wx_id, group_id)
)
"""
class SignInSystem:
def __init__(self, wcf: Wcf, gbm: GroupBotManager, all_contacts: dict,
db_pool: mysql.connector.pooling.MySQLConnectionPool, redis_pool: redis.ConnectionPool,
message_util: MessageUtil):
# 读取配置文件
with open('message_sign/config.toml', 'rb') as f:
self.config = tomllib.load(f)['SignIn']
self.LOG = logging.getLogger(__name__)
if not self.config['enable']:
raise Exception("签到功能未启用")
self.wcf = wcf
self.gbm = gbm
self.message_util = message_util
self.all_contacts = all_contacts
# 初始化数据库连接管理器
self.db_manager = DBConnectionManager.get_instance()
# 初始化数据库操作类
self.sign_in_db = SignInDB(self.db_manager)
self.sign_in_redis = SignInRedisDB(self.db_manager)
self.command = self.config['command']
self.min_point = self.config['min-point']
self.max_point = self.config['max-point']
self.streak_cycle = self.config['streak-cycle']
self.max_streak_point = self.config['max-streak-point']
# 时区设置
self.timezone = 'Asia/Shanghai'
# 从 Redis 初始化签到数据
self.today_signin_count = self.sign_in_redis.load_signin_count()
last_reset_date = self.sign_in_redis.get_last_reset_date()
if last_reset_date:
self.last_reset_date = last_reset_date
else:
self.last_reset_date = datetime.now(tz=pytz.timezone(self.timezone)).date()
self.sign_in_redis.save_last_reset_date(self.last_reset_date)
self.LOG.info(f"[签到] 组件初始化完成 {self.command_format}")
@property
def command_format(self):
return ','.join(self.command)
@property
def enable(self):
return self.config['enable']
def initialize_table(self):
"""初始化数据库表"""
self.sign_in_db.execute_update(CREATE_TABLE_SQL)
def reset_today_count_if_needed(self):
"""检查并重置每日签到计数"""
current_date = datetime.now(tz=pytz.timezone(self.timezone)).date()
if current_date != self.last_reset_date:
self.today_signin_count.clear()
self.sign_in_redis.reset_daily_counts()
self.last_reset_date = current_date
self.sign_in_redis.save_last_reset_date(self.last_reset_date)
self.LOG.info(f"[签到] 已重置每日签到计数,日期更新为 {current_date}")
def get_today_signin_count(self, group_id: str) -> int:
"""获取群内今日签到人数(使用缓存)"""
self.reset_today_count_if_needed()
return self.today_signin_count.get(group_id, 0)
def get_user_record(self, wx_id: str, group_id: str) -> Optional[dict]:
"""获取用户签到记录"""
return self.sign_in_db.get_user_record(wx_id, group_id)
def calculate_points(self, streak: int) -> int:
"""根据连续签到天数计算积分"""
base_points = self.min_point
extra_points = min(streak // self.streak_cycle, self.max_streak_point)
total_points = base_points + extra_points
return min(total_points, self.max_point)
def member_sign_in(self, message: WxMsg):
"""会员签到功能"""
if not self.enable:
return
content = str(message.content).strip()
command = content.split(" ")
if not len(command) or command[0] not in self.command:
return
if self.gbm.get_group_permission(message.roomid, Feature.SIGNIN) == PermissionStatus.DISABLED:
return
# 获取当前时间,带有时区信息
current_time = datetime.now(tz=pytz.timezone(self.timezone))
# 获取当天零点的时间
today_start = current_time.replace(hour=0, minute=0, second=0, microsecond=0)
# 获取昨天的时间
yesterday = today_start - timedelta(days=1)
# 获取用户的签到记录
user_record = self.get_user_record(message.sender, message.roomid)
wx_nick_name = self.all_contacts.get(message.sender, message.sender)
# 判断用户是否已经签到过
if user_record and user_record.get('sign_stat'):
sign_stat = user_record['sign_stat']
# 确保 sign_stat 和 today_start 是同一时区对象
if isinstance(sign_stat, datetime) and sign_stat.tzinfo is None:
sign_stat = pytz.timezone(self.timezone).localize(sign_stat)
# 如果 sign_stat 已经大于或等于今天的零点,则认为用户已经签到过了
if sign_stat >= today_start:
self.message_util.send_text_msg(f"您今天已经签到过了!当前积分:{user_record['points']}", message.roomid,
message.sender)
return
streak = 0
streak_broken = False
old_streak = 1
if user_record and user_record['sign_stat']:
last_sign_date = user_record['sign_stat'].replace(hour=0, minute=0, second=0, microsecond=0)
# 确保 sign_stat 和 today_start 是同一时区对象
if isinstance(last_sign_date, datetime) and last_sign_date.tzinfo is None:
last_sign_date = pytz.timezone(self.timezone).localize(last_sign_date)
self.LOG.info(
f"last_sign_date: {last_sign_date}, yesterday: {yesterday}, user_streak: {user_record['signin_streak']}")
if last_sign_date == yesterday:
streak = user_record['signin_streak'] + 1
old_streak = streak
streak_broken = False
else:
streak = 1
streak_broken = True
else:
streak = 1
today_signin_rank = self.get_today_signin_count(message.roomid) + 1
self.today_signin_count[message.roomid] = today_signin_rank
self.sign_in_redis.save_signin_count(message.roomid, today_signin_rank)
points_to_add = self.calculate_points(streak)
# 使用数据库操作类更新或创建签到记录
if user_record:
self.sign_in_db.update_sign_record(
message.sender, message.roomid, wx_nick_name,
points_to_add, current_time, streak
)
else:
self.sign_in_db.create_sign_record(
message.sender, message.roomid, wx_nick_name,
points_to_add, current_time, streak
)
output = f"签到成功,加[{points_to_add}]分,第[{today_signin_rank}]个!"
if streak_broken and old_streak > 0: # 只有在真的断签且之前有签到记录时才显示
output += f"断开了 {old_streak} 天连签!"
elif streak > 1:
output += f"连签 {streak} 天!"
self.message_util.send_text_msg(output, message.roomid,
message.sender)
def __del__(self):
"""连接池由外部管理,不需要手动关闭"""
pass