269 lines
9.3 KiB
Python
269 lines
9.3 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""Telegram Bot 通知与回调。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import html
|
||
import json
|
||
import urllib.error
|
||
import urllib.parse
|
||
import urllib.request
|
||
from typing import Any, Mapping, Protocol, Sequence
|
||
|
||
import httpx
|
||
|
||
MAX_CANDIDATES_PER_MESSAGE = 3
|
||
|
||
|
||
class TelegramError(RuntimeError):
|
||
pass
|
||
|
||
|
||
class TelegramTransport(Protocol):
|
||
def post(self, url: str, data: Mapping[str, Any], timeout: float = 30) -> Any: ...
|
||
|
||
|
||
class UrllibTelegramTransport:
|
||
def post(self, url: str, data: Mapping[str, Any], timeout: float = 30) -> Any:
|
||
request = urllib.request.Request(
|
||
url,
|
||
data=urllib.parse.urlencode(data).encode("utf-8"),
|
||
method="POST",
|
||
)
|
||
try:
|
||
with urllib.request.urlopen(request, timeout=timeout) as response:
|
||
return json.loads(response.read().decode("utf-8", "replace"))
|
||
except urllib.error.HTTPError as exc:
|
||
# Telegram 对业务错误也返回 4xx + JSON 描述；保留描述便于定位。
|
||
body = exc.read().decode("utf-8", "replace")
|
||
try:
|
||
return json.loads(body)
|
||
except json.JSONDecodeError as parse_error:
|
||
raise TelegramError(
|
||
f"Telegram HTTP {exc.code}: {body[:200]}"
|
||
) from parse_error
|
||
|
||
|
||
def format_duration(seconds: int) -> str:
|
||
seconds = max(0, int(seconds or 0))
|
||
hours, remainder = divmod(seconds, 3600)
|
||
minutes, secs = divmod(remainder, 60)
|
||
if hours:
|
||
return f"{hours}:{minutes:02d}:{secs:02d}"
|
||
return f"{minutes}:{secs:02d}"
|
||
|
||
|
||
def _result_line(result: Any) -> str:
|
||
state = getattr(result, "state", "")
|
||
title = html.escape(str(getattr(result, "title", "") or ""))
|
||
message = str(getattr(result, "message", "") or "")
|
||
aid = getattr(result, "aid", None)
|
||
suffix = f"《{title}》" if title else ""
|
||
ref = f"（aid {aid}）" if aid else ""
|
||
icon = {
|
||
"added": "✅", "recovered": "♻️", "exists": "ℹ️",
|
||
"error": "⚠️", "missing": "❔",
|
||
}.get(state, "•")
|
||
body = message or "完成"
|
||
return f"{icon} {body}{suffix}{ref}".strip()
|
||
|
||
|
||
def format_ingest_result(result: Any) -> str:
|
||
"""候选确认后的回执（候选确认 / 手动入库单条）。"""
|
||
return _result_line(result)
|
||
|
||
|
||
def format_manual_reply(manual_result: Any) -> str:
|
||
"""私聊文本手动入库的汇总回执。"""
|
||
refs = int(getattr(manual_result, "refs", 0) or 0)
|
||
if refs <= 0:
|
||
return "未识别到 B 站视频引用，请发送视频链接或 BV 号，例如 BV1x8x166EzY。"
|
||
results = list(getattr(manual_result, "results", []) or [])
|
||
lines = [_result_line(item) for item in results]
|
||
header = f"共识别 {refs} 个引用，入库 {int(getattr(manual_result, 'added', 0) or 0)} 条："
|
||
return "\n".join([header, *lines])
|
||
|
||
|
||
HELP_TEXT = (
|
||
"发送 B 站视频链接、BV 号或 av 号即可手动入库并开始追踪。\n"
|
||
"例如：https://www.bilibili.com/video/BV1x8x166EzY 或 BV1x8x166EzY"
|
||
)
|
||
|
||
|
||
def format_candidate(candidate: Mapping[str, Any]) -> str:
|
||
title = html.escape(str(candidate.get("title") or "(无标题)"))
|
||
up_name = html.escape(str(candidate.get("up_name") or "未知UP"))
|
||
duration = format_duration(int(candidate.get("duration") or 0))
|
||
play = int(candidate.get("play") or 0)
|
||
like = int(candidate.get("like") or 0)
|
||
return (
|
||
f"🎬 <b>{title}</b>\n"
|
||
f"UP: {up_name}\n"
|
||
f"时长 {duration} · 播放 {play} · 点赞 {like}（索引值）"
|
||
)
|
||
|
||
|
||
class HttpxTelegramTransport:
|
||
"""基于 httpx 连接池的 Telegram 传输。"""
|
||
|
||
def __init__(
|
||
self,
|
||
*,
|
||
timeout: float = 30,
|
||
client: httpx.Client | None = None,
|
||
) -> None:
|
||
self.timeout = timeout
|
||
self._client = client or httpx.Client(
|
||
timeout=httpx.Timeout(timeout, connect=10),
|
||
limits=httpx.Limits(
|
||
max_keepalive_connections=5,
|
||
keepalive_expiry=30,
|
||
),
|
||
follow_redirects=False,
|
||
trust_env=True,
|
||
)
|
||
|
||
def post(self, url: str, data: Mapping[str, Any], timeout: float = 30) -> Any:
|
||
try:
|
||
response = self._client.post(url, data=dict(data), timeout=timeout)
|
||
except httpx.TransportError as exc:
|
||
raise TelegramError(str(exc)) from exc
|
||
try:
|
||
payload = response.json()
|
||
except ValueError as exc:
|
||
raise TelegramError(
|
||
f"Telegram HTTP {response.status_code}: {response.text[:200]}"
|
||
) from exc
|
||
if response.status_code >= 400 and not (
|
||
isinstance(payload, Mapping) and "ok" in payload
|
||
):
|
||
raise TelegramError(
|
||
f"Telegram HTTP {response.status_code}: {str(payload)[:200]}"
|
||
)
|
||
return payload
|
||
|
||
def close(self) -> None:
|
||
self._client.close()
|
||
|
||
|
||
class TelegramClient:
|
||
def __init__(
|
||
self,
|
||
token: str,
|
||
chat_id: int,
|
||
*,
|
||
transport: TelegramTransport | None = None,
|
||
timeout: float = 30,
|
||
poll_timeout: int = 25,
|
||
logger=None,
|
||
) -> None:
|
||
if not token:
|
||
raise ValueError("telegram.bot_token 不能为空")
|
||
if not chat_id:
|
||
raise ValueError("telegram.chat_id 不能为空")
|
||
self.token = token
|
||
self.chat_id = int(chat_id)
|
||
self.transport = transport or HttpxTelegramTransport()
|
||
self.timeout = timeout
|
||
self.poll_timeout = poll_timeout
|
||
self.logger = logger
|
||
|
||
def close(self) -> None:
|
||
close = getattr(self.transport, "close", None)
|
||
if callable(close):
|
||
close()
|
||
|
||
def _call(
|
||
self,
|
||
method: str,
|
||
params: Mapping[str, Any] | None = None,
|
||
*,
|
||
timeout: float | None = None,
|
||
) -> Any:
|
||
url = f"https://api.telegram.org/bot{self.token}/{method}"
|
||
payload = self.transport.post(
|
||
url, params or {}, timeout=self.timeout if timeout is None else timeout
|
||
)
|
||
if not isinstance(payload, Mapping):
|
||
raise TelegramError("Telegram 响应不是对象")
|
||
if not payload.get("ok"):
|
||
raise TelegramError(str(payload.get("description") or "Telegram API 错误"))
|
||
return payload.get("result")
|
||
|
||
def send_text(
|
||
self, text: str, *, reply_markup: Mapping[str, Any] | None = None
|
||
) -> Any:
|
||
params: dict[str, Any] = {
|
||
"chat_id": self.chat_id,
|
||
"text": text,
|
||
"parse_mode": "HTML",
|
||
"disable_web_page_preview": True,
|
||
}
|
||
if reply_markup is not None:
|
||
params["reply_markup"] = json.dumps(reply_markup, ensure_ascii=False)
|
||
return self._call("sendMessage", params)
|
||
|
||
def send_candidates(self, candidates: Sequence[Mapping[str, Any]]) -> bool:
|
||
# 通知器边界再做一次 aid 去重，避免调用方重复入队导致消息主体和按钮重复。
|
||
items: list[Mapping[str, Any]] = []
|
||
seen: set[int] = set()
|
||
for candidate in candidates:
|
||
aid = candidate.get("aid")
|
||
if aid in seen:
|
||
continue
|
||
seen.add(aid)
|
||
items.append(candidate)
|
||
if not items:
|
||
return True
|
||
for index in range(0, len(items), MAX_CANDIDATES_PER_MESSAGE):
|
||
chunk = items[index:index + MAX_CANDIDATES_PER_MESSAGE]
|
||
text = "\n\n".join(format_candidate(candidate) for candidate in chunk)
|
||
keyboard = {
|
||
"inline_keyboard": [
|
||
[
|
||
{"text": "📥 入库", "callback_data": f"add:{candidate['aid']}"},
|
||
{"text": "🙈 忽略", "callback_data": f"ign:{candidate['aid']}"},
|
||
{"text": "🚫 拉黑", "callback_data": f"blk:{candidate['aid']}"},
|
||
]
|
||
for candidate in chunk
|
||
]
|
||
}
|
||
self.send_text(text, reply_markup=keyboard)
|
||
return True
|
||
|
||
def get_updates(
|
||
self, offset: int | None = None, *, timeout: int | None = None
|
||
) -> list[dict[str, Any]]:
|
||
params: dict[str, Any] = {
|
||
"timeout": self.poll_timeout if timeout is None else int(timeout),
|
||
"allowed_updates": json.dumps(["callback_query", "message"]),
|
||
}
|
||
if offset is not None:
|
||
params["offset"] = int(offset)
|
||
# 长轮询本身会挂起 poll_timeout 秒，HTTP 读超时必须留出更大余量。
|
||
result = self._call(
|
||
"getUpdates", params, timeout=self.poll_timeout + 15
|
||
)
|
||
return result if isinstance(result, list) else []
|
||
|
||
def answer_callback(self, callback_id: str, text: str = "") -> Any:
|
||
return self._call(
|
||
"answerCallbackQuery",
|
||
{"callback_query_id": callback_id, "text": text[:200]},
|
||
)
|
||
|
||
def edit_reply_markup(
|
||
self, message_id: int, *, reply_markup: Mapping[str, Any] | None = None
|
||
) -> Any:
|
||
"""移除或替换原消息的按钮；空 keyboard 表示移除。"""
|
||
return self._call(
|
||
"editMessageReplyMarkup",
|
||
{
|
||
"chat_id": self.chat_id,
|
||
"message_id": int(message_id),
|
||
"reply_markup": json.dumps(
|
||
reply_markup or {"inline_keyboard": []}, ensure_ascii=False
|
||
),
|
||
},
|
||
)
|