287 lines
11 KiB
Python
287 lines
11 KiB
Python
from typing import Dict, Any, List, Optional, Tuple
|
||
import re
|
||
from loguru import logger
|
||
|
||
from base.plugin_common.message_plugin_interface import MessagePluginInterface
|
||
from base.plugin_common.plugin_interface import PluginStatus
|
||
from utils.decorator.plugin_decorators import plugin_stats_decorator
|
||
from utils.decorator.points_decorator import plugin_points_cost
|
||
from utils.robot_cmd.robot_command import Feature, PermissionStatus, GroupBotManager
|
||
from wechat_ipad import WechatAPIClient
|
||
|
||
|
||
class FanhaoSearchPlugin(MessagePluginInterface):
|
||
"""番号查询插件"""
|
||
|
||
FEATURE_KEY = "FANHAO"
|
||
FEATURE_DESCRIPTION = "🔎 番号查询功能 [番号]"
|
||
|
||
@property
|
||
def name(self) -> str:
|
||
return "番号查询"
|
||
|
||
@property
|
||
def version(self) -> str:
|
||
return "1.0.0"
|
||
|
||
@property
|
||
def description(self) -> str:
|
||
return "提供基于MongoDB的番号搜索功能，支持两个集合查询"
|
||
|
||
@property
|
||
def author(self) -> str:
|
||
return "liu.wei"
|
||
|
||
@property
|
||
def command_prefix(self) -> Optional[str]:
|
||
return ""
|
||
|
||
@property
|
||
def commands(self) -> List[str]:
|
||
return self._commands
|
||
|
||
@property
|
||
def feature_key(self) -> Optional[str]:
|
||
return self.FEATURE_KEY
|
||
|
||
@property
|
||
def feature_description(self) -> Optional[str]:
|
||
return self.FEATURE_DESCRIPTION
|
||
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.feature = self.register_feature()
|
||
self.mongo_client = None
|
||
self.mongo_db = None
|
||
|
||
def initialize(self, context: Dict[str, Any]) -> bool:
|
||
self.LOG = logger
|
||
self.LOG.info(f"正在初始化 {self.name} 插件...")
|
||
|
||
self.event_system = context.get("event_system")
|
||
|
||
cfg = self._config.get("FanhaoSearch", {})
|
||
self._commands = cfg.get("command", ["番号"]) # 例："番号 FNS-109"
|
||
self.command_format = cfg.get("command-format", "番号 番号编号 例如：番号 FNS-109")
|
||
self.enable = cfg.get("enable", True)
|
||
|
||
self.mongo_uri = cfg.get(
|
||
"mongo_uri",
|
||
"mongodb+srv://readonly:cS9NSuiJ1ebHnUL0@cluster0.8mosa.mongodb.net/sehuatang?retryWrites=true&w=majority",
|
||
)
|
||
self.mongo_db_name = cfg.get("db", "sehuatang")
|
||
self.collections = cfg.get(
|
||
"collections", ["hd_chinese_subtitles", "asia_codeless_originate", "asia_mosaic_originate", "4k_video"]
|
||
)
|
||
self.search_fields = cfg.get("search_fields", ["number"]) # 可能的字段名
|
||
|
||
# 延迟连接，在首次查询时连接，避免初始化阻塞
|
||
self.LOG.info(f"[{self.name}] 插件初始化完成，指令：{self._commands}")
|
||
return True
|
||
|
||
def start(self) -> bool:
|
||
self.LOG.info(f"[{self.name}] 插件已启动")
|
||
self.status = PluginStatus.RUNNING
|
||
return True
|
||
|
||
def stop(self) -> bool:
|
||
try:
|
||
if self.mongo_client:
|
||
self.mongo_client.close()
|
||
except Exception:
|
||
pass
|
||
self.LOG.info(f"[{self.name}] 插件已停止")
|
||
self.status = PluginStatus.STOPPED
|
||
return True
|
||
|
||
def can_process(self, message: Dict[str, Any]) -> bool:
|
||
if not self.enable:
|
||
return False
|
||
content = str(message.get("content", "")).strip()
|
||
command = content.split(" ")[0]
|
||
return command in self._commands
|
||
|
||
def _redact_mongo_uri(self, uri: str) -> str:
|
||
try:
|
||
# 隐藏用户名密码，仅保留协议和主机段
|
||
return re.sub(r"(mongodb\+srv://)(.*?@)", r"\\1***@", uri)
|
||
except Exception:
|
||
return "***"
|
||
|
||
def _ensure_mongo(self):
|
||
if self.mongo_client:
|
||
return
|
||
from pymongo import MongoClient
|
||
|
||
self.LOG.info(
|
||
f"[{self.name}] 准备连接MongoDB uri={self._redact_mongo_uri(self.mongo_uri)} db={self.mongo_db_name}"
|
||
)
|
||
try:
|
||
self.mongo_client = MongoClient(self.mongo_uri, serverSelectionTimeoutMS=5000)
|
||
# 探活
|
||
self.mongo_client.admin.command("ping")
|
||
self.mongo_db = self.mongo_client.get_database(self.mongo_db_name)
|
||
# 打印可见的数据库
|
||
try:
|
||
dbs = self.mongo_client.list_database_names()
|
||
self.LOG.info(f"[{self.name}] 可见数据库={dbs}")
|
||
except Exception as e:
|
||
self.LOG.warning(f"[{self.name}] 获取数据库列表失败: {e}")
|
||
try:
|
||
colls = self.mongo_db.list_collection_names()
|
||
except Exception as e:
|
||
colls = []
|
||
self.LOG.warning(f"[{self.name}] 获取集合列表失败: {e}")
|
||
self.LOG.info(f"[{self.name}] MongoDB连接成功，集合={colls}")
|
||
# 对配置集合进行计数探测
|
||
for cname in self.collections:
|
||
try:
|
||
c = self.mongo_db.get_collection(cname)
|
||
# 尝试快速计数（可能返回估算值，但足够判断可见性）
|
||
cnt = c.estimated_document_count()
|
||
self.LOG.info(f"[{self.name}] 集合探测 {self.mongo_db_name}.{cname} 文档数≈{cnt}")
|
||
except Exception as e:
|
||
self.LOG.warning(f"[{self.name}] 集合探测失败 {self.mongo_db_name}.{cname}: {e}")
|
||
except Exception as e:
|
||
self.LOG.error(f"[{self.name}] MongoDB连接失败: {e}")
|
||
raise
|
||
|
||
def _normalize_code(self, text: str) -> str:
|
||
# 1. 基础清理：判空、去首尾空格、转大写
|
||
text = (text or "").strip().upper()
|
||
# 用户输入 处理后 说明
|
||
# IPzz108 IPZZ-108 目标场景：自动补全了横杠
|
||
# ipzz108 IPZZ-108 全小写自动转大写并补全
|
||
# IPZZ-108 IPZZ-108 已经有横杠，正则不匹配（字母后是横杠不是数字），保持原样
|
||
# ipzz-108 IPZZ-108 小写带横杠，仅转大写，保持原样
|
||
# ipzz108 IPZZ-108 去除前后空格
|
||
# A1 A-1 极短代码也能兼容
|
||
|
||
# 2. 核心逻辑：使用正则查找“字母后面紧跟数字”的情况，并在中间插入横杠
|
||
# r'([A-Z])(\d)' 含义：捕获组1是任意大写字母，捕获组2是任意数字
|
||
# r'\1-\2' 含义：将匹配到的内容替换为“组1 + 横杠 + 组2”
|
||
return re.sub(r'([A-Z])(\d)', r'\1-\2', text)
|
||
|
||
def _build_queries(self, code_upper: str) -> List[Dict[str, Any]]:
|
||
# 精确匹配查询
|
||
or_exact = [{field: code_upper} for field in self.search_fields]
|
||
exact_query = {"$or": or_exact}
|
||
# 回退：大小写不敏感的等值匹配
|
||
or_regex = [
|
||
{field: {"$regex": f"^{re.escape(code_upper)}$", "$options": "i"}}
|
||
for field in self.search_fields
|
||
]
|
||
regex_query = {"$or": or_regex}
|
||
return [exact_query, regex_query]
|
||
|
||
def _query_collections(self, code_upper: str) -> Optional[Dict[str, Any]]:
|
||
self._ensure_mongo()
|
||
queries = self._build_queries(code_upper)
|
||
self.LOG.debug(f"[{self.name}] 标准化番号={code_upper}，查询字段={self.search_fields}")
|
||
for idx, query in enumerate(queries):
|
||
self.LOG.debug(f"[{self.name}] 执行查询({idx + 1}/{len(queries)}): {query}")
|
||
for coll_name in self.collections:
|
||
try:
|
||
coll = self.mongo_db.get_collection(coll_name)
|
||
self.LOG.debug(f"[{self.name}] 在集合 {coll_name} 查找…")
|
||
doc = coll.find_one(query)
|
||
if doc:
|
||
doc["_collection"] = coll_name
|
||
self.LOG.info(f"[{self.name}] 命中集合 {coll_name}")
|
||
return doc
|
||
else:
|
||
self.LOG.debug(f"[{self.name}] 集合 {coll_name} 未命中")
|
||
except Exception as e:
|
||
self.LOG.error(f"[{self.name}] 查询集合 {coll_name} 出错: {e}")
|
||
continue
|
||
return None
|
||
|
||
def _format_result(self, doc: Dict[str, Any]) -> str:
|
||
def pick(d: Dict[str, Any], keys: List[str]) -> str:
|
||
for k in keys:
|
||
v = d.get(k)
|
||
if v:
|
||
return str(v)
|
||
return ""
|
||
|
||
code = pick(doc, self.search_fields)
|
||
title = pick(doc, ["title", "name", "标题"]) or "未提供"
|
||
actress = pick(doc, ["actress", "actors", "performer", "女优", "演员"]) # 可为空
|
||
date_val = pick(doc, ["date", "publish_date", "发行日"]) # 例如：2025-09-10
|
||
post_time = pick(doc, ["post_time"]) # 例如：2025-09-10 10:42:04
|
||
|
||
magnet = pick(doc, ["magnet"]) # 磁力
|
||
magnet_115 = pick(doc, ["magnet_115"]) # 115专用磁力
|
||
|
||
lines = [
|
||
f"✅ 查询成功：{code}",
|
||
f"标题：{title}",
|
||
]
|
||
if actress:
|
||
lines.append(f"演员：{actress}")
|
||
|
||
if date_val and post_time:
|
||
lines.append(f"日期：{date_val}（发帖：{post_time}）")
|
||
elif date_val:
|
||
lines.append(f"日期：{date_val}")
|
||
elif post_time:
|
||
lines.append(f"发帖：{post_time}")
|
||
|
||
if magnet:
|
||
lines.append(f"磁力：{magnet}")
|
||
if magnet_115:
|
||
lines.append(f"115磁力：{magnet_115}")
|
||
|
||
return "\n".join(lines)
|
||
|
||
@plugin_stats_decorator(plugin_name="番号查询")
|
||
@plugin_points_cost(10, "番号查询消耗积分", FEATURE_KEY)
|
||
async def process_message(self, message: Dict[str, Any]) -> Tuple[bool, Optional[str]]:
|
||
content = str(message.get("content", "")).strip()
|
||
self.LOG.debug(f"插件执行： {self.name}：{content}")
|
||
|
||
command = content.split(" ")[0]
|
||
sender = message.get("sender")
|
||
roomid = message.get("roomid", "")
|
||
gbm: GroupBotManager = message.get("gbm")
|
||
bot: WechatAPIClient = message.get("bot")
|
||
|
||
# 参数检查
|
||
parts = content.split(" ")
|
||
if len(parts) < 2:
|
||
await bot.send_text_message((roomid if roomid else sender), f"❌命令格式错误！\n{self.command_format}",
|
||
sender)
|
||
return False, "命令格式错误"
|
||
|
||
if roomid and gbm.get_group_permission(roomid, self.feature) == PermissionStatus.DISABLED:
|
||
return False, "没有权限"
|
||
|
||
raw_code = content[len(command):].strip()
|
||
user_code = self._normalize_code(raw_code)
|
||
self.LOG.info(
|
||
f"[{self.name}] 收到查询 command={command} raw='{raw_code}' normalized='{user_code}' db={self.mongo_db_name} collections={self.collections}"
|
||
)
|
||
if not user_code:
|
||
await bot.send_text_message((roomid if roomid else sender), f"❌命令格式错误！\n{self.command_format}",
|
||
sender)
|
||
return False, "命令格式错误"
|
||
|
||
try:
|
||
doc = self._query_collections(user_code)
|
||
target = roomid if roomid else sender
|
||
if not doc:
|
||
self.LOG.warning(f"[{self.name}] 未找到番号：{user_code}")
|
||
await bot.send_text_message(target, f"未找到番号：{user_code}", sender)
|
||
return False, "未找到"
|
||
|
||
text = self._format_result(doc)
|
||
await bot.send_text_message(target, text, sender)
|
||
return True, "查询成功"
|
||
except Exception as e:
|
||
self.LOG.exception(f"处理番号查询出错: {e}")
|
||
return False, f"处理出错: {e}"
|
||
|
||
|
||
def get_plugin():
|
||
return FanhaoSearchPlugin()
|