206 lines
8.0 KiB
Python
206 lines
8.0 KiB
Python
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
|