first
This commit is contained in:
319
bot_notifier.py
Normal file
319
bot_notifier.py
Normal file
@@ -0,0 +1,319 @@
|
||||
# ==================== Telegram 通知模块 ====================
|
||||
# 用于将日志和状态变更推送到 Telegram
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from typing import Callable, List, Optional, Dict
|
||||
from telegram import Bot, Message
|
||||
from telegram.error import TelegramError
|
||||
|
||||
from config import (
|
||||
TELEGRAM_ADMIN_CHAT_IDS,
|
||||
TELEGRAM_NOTIFY_ON_COMPLETE,
|
||||
TELEGRAM_NOTIFY_ON_ERROR,
|
||||
)
|
||||
|
||||
|
||||
def make_progress_bar(current: int, total: int, width: int = 10) -> str:
|
||||
"""生成文本进度条
|
||||
|
||||
Args:
|
||||
current: 当前进度
|
||||
total: 总数
|
||||
width: 进度条宽度 (字符数)
|
||||
|
||||
Returns:
|
||||
进度条字符串，如 "████████░░ 80%"
|
||||
"""
|
||||
if total <= 0:
|
||||
return "░" * width + " 0%"
|
||||
|
||||
percent = min(current / total, 1.0)
|
||||
filled = int(width * percent)
|
||||
empty = width - filled
|
||||
|
||||
bar = "█" * filled + "░" * empty
|
||||
percent_text = f"{int(percent * 100)}%"
|
||||
|
||||
return f"{bar} {percent_text}"
|
||||
|
||||
|
||||
class ProgressTracker:
|
||||
"""进度跟踪器 - 用于实时更新 Telegram 消息显示进度"""
|
||||
|
||||
def __init__(self, bot: Bot, chat_ids: List[int], team_name: str, total: int):
|
||||
self.bot = bot
|
||||
self.chat_ids = chat_ids
|
||||
self.team_name = team_name
|
||||
self.total = total
|
||||
self.current = 0
|
||||
self.success = 0
|
||||
self.failed = 0
|
||||
self.current_account = ""
|
||||
self.current_step = ""
|
||||
self.messages: Dict[int, Message] = {} # chat_id -> Message
|
||||
self._last_update = 0
|
||||
self._update_interval = 2 # 最小更新间隔 (秒)
|
||||
self._loop: asyncio.AbstractEventLoop = None
|
||||
|
||||
def _get_progress_text(self) -> str:
|
||||
"""生成进度消息文本"""
|
||||
bar = make_progress_bar(self.current, self.total, 12)
|
||||
|
||||
lines = [
|
||||
f"<b>Processing: {self.team_name}</b>",
|
||||
"",
|
||||
f"Progress: {bar}",
|
||||
f"Accounts: {self.current}/{self.total}",
|
||||
f"Success: {self.success} | Failed: {self.failed}",
|
||||
]
|
||||
|
||||
if self.current_account:
|
||||
lines.append("")
|
||||
lines.append(f"Current: <code>{self.current_account}</code>")
|
||||
|
||||
if self.current_step:
|
||||
lines.append(f"Step: {self.current_step}")
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
async def _send_initial_message(self):
|
||||
"""发送初始进度消息"""
|
||||
text = self._get_progress_text()
|
||||
for chat_id in self.chat_ids:
|
||||
try:
|
||||
msg = await self.bot.send_message(
|
||||
chat_id=chat_id,
|
||||
text=text,
|
||||
parse_mode="HTML"
|
||||
)
|
||||
self.messages[chat_id] = msg
|
||||
except TelegramError:
|
||||
pass
|
||||
|
||||
async def _update_messages(self):
|
||||
"""更新所有进度消息"""
|
||||
text = self._get_progress_text()
|
||||
for chat_id, msg in self.messages.items():
|
||||
try:
|
||||
await self.bot.edit_message_text(
|
||||
chat_id=chat_id,
|
||||
message_id=msg.message_id,
|
||||
text=text,
|
||||
parse_mode="HTML"
|
||||
)
|
||||
except TelegramError:
|
||||
pass
|
||||
|
||||
def _schedule_update(self):
|
||||
"""调度消息更新 (限流)"""
|
||||
now = time.time()
|
||||
if now - self._last_update < self._update_interval:
|
||||
return
|
||||
|
||||
self._last_update = now
|
||||
|
||||
if self._loop:
|
||||
asyncio.run_coroutine_threadsafe(self._update_messages(), self._loop)
|
||||
|
||||
def start(self, loop: asyncio.AbstractEventLoop):
|
||||
"""启动进度跟踪"""
|
||||
self._loop = loop
|
||||
asyncio.run_coroutine_threadsafe(self._send_initial_message(), loop)
|
||||
|
||||
def update(self, current: int = None, account: str = None, step: str = None):
|
||||
"""更新进度 (供同步代码调用)"""
|
||||
if current is not None:
|
||||
self.current = current
|
||||
if account is not None:
|
||||
self.current_account = account
|
||||
if step is not None:
|
||||
self.current_step = step
|
||||
|
||||
self._schedule_update()
|
||||
|
||||
def account_done(self, email: str, success: bool):
|
||||
"""标记账号处理完成"""
|
||||
self.current += 1
|
||||
if success:
|
||||
self.success += 1
|
||||
else:
|
||||
self.failed += 1
|
||||
self.current_account = ""
|
||||
self.current_step = ""
|
||||
self._schedule_update()
|
||||
|
||||
def finish(self):
|
||||
"""完成进度跟踪，发送最终状态"""
|
||||
self.current_step = "Completed!"
|
||||
if self._loop:
|
||||
asyncio.run_coroutine_threadsafe(self._update_messages(), self._loop)
|
||||
|
||||
|
||||
class BotNotifier:
|
||||
"""Telegram 通知推送器"""
|
||||
|
||||
def __init__(self, bot: Bot, chat_ids: List[int] = None):
|
||||
self.bot = bot
|
||||
self.chat_ids = chat_ids or TELEGRAM_ADMIN_CHAT_IDS
|
||||
self._message_queue: asyncio.Queue = None
|
||||
self._worker_task: asyncio.Task = None
|
||||
self._current_progress: Optional[ProgressTracker] = None
|
||||
self._loop: asyncio.AbstractEventLoop = None
|
||||
|
||||
async def start(self):
|
||||
"""启动消息发送队列"""
|
||||
self._message_queue = asyncio.Queue()
|
||||
self._worker_task = asyncio.create_task(self._message_worker())
|
||||
self._loop = asyncio.get_event_loop()
|
||||
|
||||
async def stop(self):
|
||||
"""停止消息发送队列"""
|
||||
if self._worker_task:
|
||||
self._worker_task.cancel()
|
||||
try:
|
||||
await self._worker_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
async def _message_worker(self):
|
||||
"""后台消息发送工作线程"""
|
||||
while True:
|
||||
try:
|
||||
message, level = await self._message_queue.get()
|
||||
await self._send_to_all(message)
|
||||
self._message_queue.task_done()
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
async def _send_to_all(self, message: str, parse_mode: str = "HTML"):
|
||||
"""发送消息到所有管理员"""
|
||||
for chat_id in self.chat_ids:
|
||||
try:
|
||||
await self.bot.send_message(
|
||||
chat_id=chat_id,
|
||||
text=message,
|
||||
parse_mode=parse_mode
|
||||
)
|
||||
except TelegramError:
|
||||
pass
|
||||
|
||||
def queue_message(self, message: str, level: str = "info"):
|
||||
"""将消息加入发送队列 (非阻塞)"""
|
||||
if self._message_queue:
|
||||
try:
|
||||
self._message_queue.put_nowait((message, level))
|
||||
except asyncio.QueueFull:
|
||||
pass
|
||||
|
||||
async def notify(self, message: str, level: str = "info"):
|
||||
"""直接发送通知 (阻塞)"""
|
||||
await self._send_to_all(message)
|
||||
|
||||
def create_progress(self, team_name: str, total: int) -> ProgressTracker:
|
||||
"""创建进度跟踪器"""
|
||||
self._current_progress = ProgressTracker(
|
||||
self.bot, self.chat_ids, team_name, total
|
||||
)
|
||||
if self._loop:
|
||||
self._current_progress.start(self._loop)
|
||||
return self._current_progress
|
||||
|
||||
def get_progress(self) -> Optional[ProgressTracker]:
|
||||
"""获取当前进度跟踪器"""
|
||||
return self._current_progress
|
||||
|
||||
def clear_progress(self):
|
||||
"""清除进度跟踪器"""
|
||||
if self._current_progress:
|
||||
self._current_progress.finish()
|
||||
self._current_progress = None
|
||||
|
||||
async def notify_task_started(self, team_name: str):
|
||||
"""通知任务开始"""
|
||||
await self.notify(f"<b>Task Started</b>\nTeam: {team_name}")
|
||||
|
||||
async def notify_task_completed(self, team_name: str, success: int, failed: int):
|
||||
"""通知任务完成"""
|
||||
if not TELEGRAM_NOTIFY_ON_COMPLETE:
|
||||
return
|
||||
status = "All Success" if failed == 0 else f"{failed} Failed"
|
||||
await self.notify(
|
||||
f"<b>Task Completed</b>\n"
|
||||
f"Team: {team_name}\n"
|
||||
f"Success: {success}\n"
|
||||
f"Status: {status}"
|
||||
)
|
||||
|
||||
async def notify_error(self, message: str, details: str = ""):
|
||||
"""通知错误"""
|
||||
if not TELEGRAM_NOTIFY_ON_ERROR:
|
||||
return
|
||||
text = f"<b>Error</b>\n{message}"
|
||||
if details:
|
||||
text += f"\n<code>{details[:500]}</code>"
|
||||
await self.notify(text)
|
||||
|
||||
async def notify_account_status(self, email: str, status: str, team_name: str = ""):
|
||||
"""通知账号状态变更"""
|
||||
icon = {
|
||||
"completed": "OK",
|
||||
"authorized": "AUTH",
|
||||
"registered": "REG",
|
||||
"failed": "FAIL",
|
||||
}.get(status, status.upper())
|
||||
|
||||
text = f"[{icon}] {email}"
|
||||
if team_name:
|
||||
text = f"<b>{team_name}</b>\n{text}"
|
||||
self.queue_message(text, "info")
|
||||
|
||||
|
||||
# 全局通知器实例 (在 telegram_bot.py 中初始化)
|
||||
_notifier: Optional[BotNotifier] = None
|
||||
|
||||
|
||||
def set_notifier(notifier: BotNotifier):
|
||||
"""设置全局通知器"""
|
||||
global _notifier
|
||||
_notifier = notifier
|
||||
|
||||
|
||||
def get_notifier() -> Optional[BotNotifier]:
|
||||
"""获取全局通知器"""
|
||||
return _notifier
|
||||
|
||||
|
||||
def notify_sync(message: str, level: str = "info"):
|
||||
"""同步方式发送通知 (供非异步代码使用)"""
|
||||
if _notifier:
|
||||
_notifier.queue_message(message, level)
|
||||
|
||||
|
||||
# ==================== 进度更新接口 (供 run.py 使用) ====================
|
||||
|
||||
def progress_start(team_name: str, total: int) -> Optional[ProgressTracker]:
|
||||
"""开始进度跟踪"""
|
||||
if _notifier:
|
||||
return _notifier.create_progress(team_name, total)
|
||||
return None
|
||||
|
||||
|
||||
def progress_update(account: str = None, step: str = None):
|
||||
"""更新当前进度"""
|
||||
if _notifier and _notifier.get_progress():
|
||||
_notifier.get_progress().update(account=account, step=step)
|
||||
|
||||
|
||||
def progress_account_done(email: str, success: bool):
|
||||
"""标记账号完成"""
|
||||
if _notifier and _notifier.get_progress():
|
||||
_notifier.get_progress().account_done(email, success)
|
||||
|
||||
|
||||
def progress_finish():
|
||||
"""完成进度跟踪"""
|
||||
if _notifier:
|
||||
_notifier.clear_progress()
|
||||
Reference in New Issue
Block a user