424 lines
17 KiB
Python
424 lines
17 KiB
Python
# 群清单管理
|
||
# 群功能管理
|
||
# 0.加入或者关闭群机器人 #启用群机器人 #关闭群机器人
|
||
# 1.每日新闻自动播报 #启用每日新闻播报 #关闭每日新闻播报
|
||
# 2.每日群发言总结 #启用群发言 #关闭群发言
|
||
# 3.群AI能力 #启用群AI #关闭群AI
|
||
# 4.群总结能力 #启用群总结 #关闭群总结
|
||
# 5.sehuatang PDF能力 #启用pdf #关闭pdf
|
||
import os
|
||
from enum import Enum
|
||
from typing import List
|
||
|
||
import yaml
|
||
from loguru import logger
|
||
|
||
from db.connection import DBConnectionManager
|
||
|
||
|
||
class PermissionStatus(Enum):
|
||
"""权限状态枚举"""
|
||
ENABLED = "enabled"
|
||
DISABLED = "disabled"
|
||
|
||
|
||
# 在Feature枚举类中添加新的功能权限
|
||
|
||
class Feature(Enum):
|
||
"""功能权限枚举，带序号"""
|
||
ROBOT = 1, "🔧 群机器人 [总开关]"
|
||
DAILY_NEWS = 2, "📰 每日新闻自动播报 [每日8:30定时发送]"
|
||
PDF_CAPABILITY = 3, "📄 sehuatang PDF能力 [无]"
|
||
EPIC = 4, "📊 EPIC自动播报 [每周五自动发送]" # 新增的功能
|
||
DAILY_SUMMARY = 5, "🕤 每日群发言总结 [每日9:30定时发送]"
|
||
|
||
# DAILY_SUMMARY = 3, "🕤 每日群发言总结 [每日9:30定时发送]"
|
||
# AI_CAPABILITY = 4, "🤖 AI对话 [ai, 聊天, AI] 用法：ai 如何写一个机器人？"
|
||
# SUMMARY_CAPABILITY = 5, "📝 群总结能力 [#总结]"
|
||
# PIC = 8, "🖼️ 图来能力 [图来, 秀人]"
|
||
# TASK_GAME = 9, "📚 百科答题 [/t, /s, /a 任务ID 答案]"
|
||
# MUSIC = 10, "🎵 点歌功能 [点歌, 音乐, 音乐点播, 点播音乐, 音乐点歌]"
|
||
# SIGNIN = 11, "✅ 签到功能 [签到, 每日签到, qd, Qd, QD, 上班, 牛马]"
|
||
# POINT_TRADE = 12, "🎁 积分赠送 [积分赠送 1 @XX, 积分排行, 打劫 @XX, 保释 @XX]"
|
||
# BEAUTY_LEG = 13, "🦵 腿来能力 [美腿, 腿来]"
|
||
# VIDEO = 14, "🎥 黑丝视频 [黑丝视频, 黑丝, 来个黑丝, 搞个黑丝]"
|
||
# VIDEO_MAN = 15, "💪 肌肉视频 [猛男, 肌肉, 帅哥]"
|
||
# # GROUP_ADD = 16, "加群提醒"
|
||
# DOUYIN_PARSER = 17, "🎥 抖音链接转视频"
|
||
# GROUP_MEMBER_CHANGE = 18, "👥 群成员变更提醒 [自动触发]"
|
||
# # KID_PHOTO_EXTRACT = 19, "儿童照片提取转发功能" # 小朋友照片提取功能
|
||
# NEWS = 20, "🌍 全球政治经济新闻"
|
||
# WEATHER = 21, "🌤️ 天气查询 [上海天气, 天气上海]"
|
||
# JD_TOKEN = 22, "🔑 JD_京豆token设置 [设置京东 pt_key=xxx;pt_pin=xxx; 备注名称]"
|
||
# AI_AUTO = 23, "💬 仿真对话"
|
||
# GUESS_MUSIC = 24, "🎤 猜歌名游戏 [猜歌名 - 开始 | 猜歌名 歌手名 - 指定歌手 | 猜歌名 歌名 - 提交答案]"
|
||
|
||
def __new__(cls, value, description):
|
||
obj = object.__new__(cls)
|
||
obj._value_ = value
|
||
obj.description = description
|
||
return obj
|
||
|
||
def __str__(self):
|
||
return self.description
|
||
|
||
@classmethod
|
||
def get_max_value(cls) -> int:
|
||
"""获取当前最大的枚举值"""
|
||
return max([member.value for member in cls])
|
||
|
||
@classmethod
|
||
def register_feature(cls, key: str, description: str) -> 'Feature':
|
||
"""注册新的功能权限
|
||
|
||
Args:
|
||
key: 功能键名
|
||
value: 权限值
|
||
description: 权限描述
|
||
|
||
Returns:
|
||
Feature: 新注册的功能权限
|
||
"""
|
||
# 检查value是否已存在
|
||
for feature in cls:
|
||
if feature.name == key:
|
||
logger.warning(f"功能权限值 {key} 已存在，将被覆盖")
|
||
feature.description = description
|
||
return feature
|
||
|
||
# 创建新的功能权限
|
||
feature = cls._member_map_.get(key)
|
||
if feature is None:
|
||
feature = object.__new__(cls)
|
||
feature._value_ = cls.get_max_value() + 1
|
||
feature._name_ = key # 设置枚举成员名称
|
||
feature.description = description
|
||
cls._member_map_[key] = feature
|
||
cls._member_names_.append(key)
|
||
|
||
return feature
|
||
|
||
@classmethod
|
||
def get_feature(cls, key: str) -> 'Feature':
|
||
"""获取已注册的功能
|
||
|
||
Args:
|
||
key: 功能键名
|
||
|
||
Returns:
|
||
Feature: 功能枚举实例
|
||
"""
|
||
return cls._member_map_.get(key)
|
||
|
||
@classmethod
|
||
def get_all_features(cls) -> List['Feature']:
|
||
"""获取所有功能
|
||
|
||
Returns:
|
||
List[Feature]: 所有功能列表
|
||
"""
|
||
return list(cls)
|
||
|
||
@classmethod
|
||
def _missing_(cls, value):
|
||
"""处理未找到的枚举值"""
|
||
if isinstance(value, int):
|
||
for member in cls:
|
||
if member.value == value:
|
||
return member
|
||
return None
|
||
|
||
|
||
def get_redis_connection():
|
||
# 初始化时加载本地缓存
|
||
db_manager: DBConnectionManager = DBConnectionManager.get_instance()
|
||
return db_manager.get_redis_connection()
|
||
|
||
|
||
class GroupBotManager:
|
||
"""群机器人管理，支持本地缓存"""
|
||
|
||
# 本地缓存作为类级别静态属性
|
||
local_cache = {
|
||
"group_permissions": {}, # 用于缓存群组功能权限
|
||
"group_list": set() # 用于缓存 group:list
|
||
}
|
||
|
||
# 管理员列表缓存
|
||
admin_list = [] # 用于缓存管理员列表
|
||
|
||
@staticmethod
|
||
def display_menu_status(group_id):
|
||
"""显示所有功能列表及其在指定群组中的当前状态，带emoji"""
|
||
menu = []
|
||
for feature in Feature:
|
||
status = GroupBotManager.get_group_permission(group_id, feature)
|
||
status_emoji = "✅" if status == PermissionStatus.ENABLED else "❌"
|
||
status_str = "启用" if status == PermissionStatus.ENABLED else "关闭"
|
||
menu.append(f"{status_emoji} {status_str}-{feature.value}-{feature.description}")
|
||
return "\n".join(menu)
|
||
|
||
@staticmethod
|
||
def load_local_cache():
|
||
r = get_redis_connection()
|
||
"""从 Redis 加载数据到本地缓存"""
|
||
# 加载管理员列表
|
||
group_list = r.smembers("group:list")
|
||
# 输出group_list 列表
|
||
logger.debug(f"group_list: {group_list}")
|
||
GroupBotManager.local_cache["group_list"] = set(group_list)
|
||
|
||
# 加载群组权限
|
||
for group_id in GroupBotManager.local_cache["group_list"]:
|
||
key = f'group:{group_id}:permissions'
|
||
GroupBotManager.local_cache["group_permissions"][group_id] = {}
|
||
for feature in Feature:
|
||
status_value = r.hget(key, feature.name)
|
||
if status_value:
|
||
GroupBotManager.local_cache["group_permissions"][group_id][feature] = PermissionStatus(status_value)
|
||
else:
|
||
GroupBotManager.local_cache["group_permissions"][group_id][feature] = PermissionStatus.DISABLED
|
||
|
||
@staticmethod
|
||
def save_to_redis():
|
||
r = get_redis_connection()
|
||
"""将本地缓存保存回 Redis"""
|
||
# 保存 group:list 到 Redis
|
||
r.sadd("group:list", *GroupBotManager.local_cache["group_list"])
|
||
|
||
# 保存每个群组的权限到 Redis
|
||
for group_id, permissions in GroupBotManager.local_cache["group_permissions"].items():
|
||
key = f'group:{group_id}:permissions'
|
||
for feature, status in permissions.items():
|
||
r.hset(key, feature.name, status.value)
|
||
|
||
@staticmethod
|
||
def set_group_permission(group_id, feature: Feature, status: PermissionStatus):
|
||
r = get_redis_connection()
|
||
"""设置群组功能权限并更新本地缓存"""
|
||
# 更新本地缓存
|
||
if group_id not in GroupBotManager.local_cache["group_permissions"]:
|
||
GroupBotManager.local_cache["group_permissions"][group_id] = {}
|
||
|
||
GroupBotManager.local_cache["group_permissions"][group_id][feature] = status
|
||
|
||
# 同步到 Redis
|
||
key = f'group:{group_id}:permissions'
|
||
r.hset(key, feature.name, status.value)
|
||
# 输出保存到redis成功
|
||
logger.info(f"set_group_permission({group_id}, {feature.name}, {status.value})")
|
||
|
||
@staticmethod
|
||
def get_group_permission(group_id, feature: Feature):
|
||
"""获取群组某个功能的权限状态"""
|
||
# 先从本地缓存获取
|
||
if group_id in GroupBotManager.local_cache["group_permissions"]:
|
||
return GroupBotManager.local_cache["group_permissions"][group_id].get(feature, PermissionStatus.DISABLED)
|
||
else:
|
||
return PermissionStatus.DISABLED
|
||
|
||
@staticmethod
|
||
def check_permission(group_id, feature: Feature):
|
||
"""检查某个功能是否启用，若未启用则返回提示信息"""
|
||
status = GroupBotManager.get_group_permission(group_id, feature)
|
||
if status == PermissionStatus.DISABLED:
|
||
return f"该功能未启用，请开启 {feature.description}。"
|
||
return None # 如果已启用，则返回 None，不做处理
|
||
|
||
@staticmethod
|
||
def handle_command(group_id, command_str):
|
||
"""统一处理群功能指令"""
|
||
# print(f"PermissionStatus robot_menu command_str: {command_str}")
|
||
# 命令解析
|
||
command_parts = command_str.strip().split("-")
|
||
|
||
# 如果是GROUP_LIST指令，返回 group:list 清单
|
||
if command_str.strip().upper() == "群列表":
|
||
return GroupBotManager.get_group_list()
|
||
|
||
# 如果是清除群指令
|
||
if command_str.strip().startswith("清除群-"):
|
||
target_group_id = command_str.strip().split("-")[1]
|
||
return GroupBotManager.remove_group(target_group_id)
|
||
|
||
if len(command_parts) < 2:
|
||
return None
|
||
|
||
feature_str = command_parts[0]
|
||
action = command_parts[1]
|
||
|
||
# 如果第一个参数是序号，则转化为对应的功能
|
||
if feature_str.isdigit():
|
||
feature_num = int(feature_str)
|
||
# 通过遍历查找匹配的枚举值
|
||
feature = None
|
||
for f in Feature:
|
||
if f.value == feature_num:
|
||
feature = f
|
||
break
|
||
if not feature:
|
||
return f"无效的功能序号: {feature_num}"
|
||
else:
|
||
try:
|
||
feature = Feature[feature_str] # 通过枚举名称获取功能枚举
|
||
except KeyError:
|
||
return f"无效功能名称: {feature_str}"
|
||
|
||
r = get_redis_connection()
|
||
# 处理群机器人的启用和关闭（特别操作：更新 group:list）
|
||
if feature == Feature.ROBOT:
|
||
if action == "启用":
|
||
GroupBotManager.set_group_permission(group_id, feature, PermissionStatus.ENABLED)
|
||
# 启用群机器人时，将 group_id 加入 group:list
|
||
GroupBotManager.local_cache["group_list"].add(group_id)
|
||
# 同步到 Redis
|
||
r.sadd("group:list", group_id)
|
||
return f"BOT已启用，群组 {group_id} 已加入"
|
||
elif action == "关闭":
|
||
GroupBotManager.set_group_permission(group_id, feature, PermissionStatus.DISABLED)
|
||
# 关闭群机器人时，从 group:list 中删除 group_id
|
||
GroupBotManager.local_cache["group_list"].remove(group_id)
|
||
# 同步到 Redis
|
||
r.srem("group:list", group_id)
|
||
return f"BOT已关闭，群组 {group_id} 已移除"
|
||
else:
|
||
return "无效操作，仅支持启用或关闭"
|
||
|
||
# 先检查群机器人权限
|
||
robot_status = GroupBotManager.get_group_permission(group_id, Feature.ROBOT)
|
||
if robot_status != PermissionStatus.ENABLED and feature != Feature.ROBOT:
|
||
return "群机器人未启用，无法执行其他功能操作"
|
||
|
||
# 根据不同的操作启用或禁用功能
|
||
if action == "启用":
|
||
GroupBotManager.set_group_permission(group_id, feature, PermissionStatus.ENABLED)
|
||
return f"功能启用：\n {feature.description} 已启用"
|
||
elif action == "关闭":
|
||
GroupBotManager.set_group_permission(group_id, feature, PermissionStatus.DISABLED)
|
||
return f"功能关闭： \n {feature.description} 已关闭"
|
||
else:
|
||
return "无效操作，仅支持启用或关闭"
|
||
|
||
@staticmethod
|
||
def list_group_permissions(group_id):
|
||
"""列出群组所有功能及其状态"""
|
||
permissions = {}
|
||
for feature in Feature:
|
||
status = GroupBotManager.get_group_permission(group_id, feature)
|
||
permissions[feature] = status if status else PermissionStatus.DISABLED # 默认为禁用状态
|
||
return permissions
|
||
|
||
@staticmethod
|
||
def get_group_list():
|
||
"""返回所有启用了群机器人的群组清单，格式为集合"""
|
||
return list(GroupBotManager.local_cache["group_list"])
|
||
|
||
@staticmethod
|
||
def remove_group(group_id):
|
||
"""一键清除某个群的所有设置，用于退群或关闭群时处理
|
||
|
||
Args:
|
||
group_id: 群ID
|
||
|
||
Returns:
|
||
str: 操作结果信息
|
||
"""
|
||
# 检查群是否在列表中
|
||
if group_id not in GroupBotManager.local_cache["group_list"]:
|
||
return f"群 {group_id} 不在机器人管理列表中"
|
||
|
||
r = get_redis_connection()
|
||
# 从本地缓存中移除群组
|
||
GroupBotManager.local_cache["group_list"].remove(group_id)
|
||
if group_id in GroupBotManager.local_cache["group_permissions"]:
|
||
del GroupBotManager.local_cache["group_permissions"][group_id]
|
||
|
||
# 从Redis中移除群组
|
||
r.srem("group:list", group_id)
|
||
r.delete(f'group:{group_id}:permissions')
|
||
|
||
return f"已成功清除群 {group_id} 的所有设置"
|
||
|
||
@staticmethod
|
||
def get_admin_list() -> List[str]:
|
||
"""获取管理员列表
|
||
返回系统管理员的微信ID列表
|
||
直接返回缓存的 admin_list，性能更好
|
||
"""
|
||
return GroupBotManager.admin_list
|
||
|
||
@staticmethod
|
||
def is_admin(user_id: str) -> bool:
|
||
"""检查用户是否为管理员"""
|
||
# admin_list 现在是 GroupBotManager 的类属性，可以直接访问
|
||
admin_list = GroupBotManager.get_admin_list()
|
||
return user_id in admin_list
|
||
|
||
|
||
# 示例命令
|
||
def simulate_commands():
|
||
# 加载本地缓存
|
||
GroupBotManager.load_local_cache()
|
||
|
||
group_id = "49571962306@chatroom"
|
||
|
||
print(GroupBotManager.get_group_permission(group_id, Feature.AI_CAPABILITY) == PermissionStatus.DISABLED)
|
||
|
||
print(GroupBotManager.get_group_permission(group_id, Feature.SUMMARY_CAPABILITY) == PermissionStatus.ENABLED)
|
||
|
||
|
||
if __name__ == '__main__':
|
||
# 测试新增功能
|
||
print("当前功能列表：")
|
||
for feature in Feature.get_all_features():
|
||
print(f"{feature.name}: {feature.value} - {feature.description}")
|
||
print(f"最大VALUE：{Feature.get_max_value()}")
|
||
# 测试注册新功能
|
||
print("\n注册新功能...")
|
||
new_feature = Feature.register_feature("TEST_FEATURE", "🧪 测试功能 [测试]")
|
||
print(f"新功能: {new_feature.name} = {new_feature.value} - {new_feature.description}")
|
||
|
||
# 验证新功能是否添加成功
|
||
print("\n更新后的功能列表：")
|
||
for feature in Feature.get_all_features():
|
||
print(f"{feature.name}: {feature.value} - {feature.description}")
|
||
|
||
# 测试重复注册
|
||
print("\n测试重复注册...")
|
||
same_feature = Feature.register_feature("TEST_FEATURE", "🧪 测试功能 [测试]")
|
||
print(f"重复注册结果: {same_feature.name} = {same_feature.value} - {same_feature.description}")
|
||
|
||
# 测试获取功能
|
||
print("\n测试获取功能...")
|
||
retrieved_feature = Feature.get_feature("TEST_FEATURE")
|
||
print(f"获取到的功能: {retrieved_feature.name} = {retrieved_feature.value} - {retrieved_feature.description}")
|
||
|
||
# 测试通过值查找功能
|
||
# print("\n测试通过值查找功能...")
|
||
# feature_by_value = Feature(1, "") # 查找 ROBOT
|
||
# print(f"通过值1查找: {feature_by_value.name} = {feature_by_value.value} - {feature_by_value.description}")
|
||
# feature_by_value = Feature(2, "") # 查找 TEST_FEATURE
|
||
# print(f"通过值2查找: {feature_by_value.name} = {feature_by_value.value} - {feature_by_value.description}")
|
||
|
||
# 测试枚举属性访问
|
||
print("\n测试枚举属性访问...")
|
||
print(f"ROBOT.name: {Feature.ROBOT.name}")
|
||
print(f"ROBOT.value: {Feature.ROBOT.value}")
|
||
print(f"ROBOT.description: {Feature.ROBOT.description}")
|
||
print(f"TEST_FEATURE.name: {Feature.TEST_FEATURE.name}")
|
||
print(f"TEST_FEATURE.value: {Feature.TEST_FEATURE.value}")
|
||
print(f"TEST_FEATURE.description: {Feature.TEST_FEATURE.description}")
|
||
|
||
# 测试枚举比较
|
||
print("\n测试枚举比较...")
|
||
print(f"ROBOT == TEST_FEATURE: {Feature.ROBOT == Feature.TEST_FEATURE}")
|
||
print(f"ROBOT != TEST_FEATURE: {Feature.ROBOT != Feature.TEST_FEATURE}")
|
||
print(f"ROBOT == ROBOT: {Feature.ROBOT == Feature.ROBOT}")
|
||
print(f"TEST_FEATURE == TEST_FEATURE: {Feature.TEST_FEATURE == Feature.TEST_FEATURE}")
|
||
|
||
# 测试枚举迭代
|
||
print("\n测试枚举迭代...")
|
||
print("所有功能:")
|
||
for feature in Feature:
|
||
print(f" {feature.name}: {feature.value} - {feature.description}")
|