返回 douyin-downloader
notifier.py
根目录 / utils / notifier.py
1 """多渠道下载完成通知。
2
3 支持:
4 - Bark(iOS 推送)
5 - Telegram Bot
6 - 通用 Webhook(POST JSON)
7 - (企业微信可通过 Webhook 模式接入)
8
9 用法:
10 notifier = build_notifier(config)
11 await notifier.send(
12 title="下载完成",
13 body="成功 12 / 失败 0 / 跳过 3",
14 level="success",
15 )
16
17 配置示例(default_config.py):
18 "notifications": {
19 "enabled": False,
20 "on_success": True,
21 "on_failure": True,
22 "providers": [
23 {"type": "bark", "url": "https://api.day.app/<device_key>"},
24 {"type": "telegram", "bot_token": "...", "chat_id": "..."},
25 {"type": "webhook", "url": "https://example.com/hook",
26 "headers": {"Authorization": "Bearer ..."}},
27 ],
28 }
29 """
30
31 from __future__ import annotations
32
33 import asyncio
34 import copy
35 from typing import Any, Dict, List
36 from urllib.parse import parse_qsl, quote, urlencode, urlsplit, urlunsplit
37
38 import aiohttp
39
40 from utils.logger import setup_logger
41
42 logger = setup_logger("Notifier")
43
44
45 def _mask_credential(value: str) -> str:
46 """Mask a credential: keep first 4 + last 4 chars, replace middle with ``***``.
47
48 If ``len(value) < 8`` (too short to usefully partial-mask), return ``***``.
49 The input is coerced to ``str`` to tolerate non-string inputs coming from
50 user config.
51 """
52 if value is None:
53 return "***"
54 text = value if isinstance(value, str) else str(value)
55 if len(text) >= 8:
56 return text[:4] + "***" + text[-4:]
57 return "***"
58
59
60 def _mask_url_query(url: str) -> str:
61 """Return ``url`` with every query-string value masked via ``_mask_credential``.
62
63 Scheme, host, path, and fragment are preserved. Parameter names are kept as-is;
64 only the values are masked. An empty / non-string / unparseable URL is returned
65 unchanged.
66 """
67 if not isinstance(url, str) or not url:
68 return url
69 try:
70 parts = urlsplit(url)
71 except ValueError:
72 return url
73 if not parts.query:
74 return url
75 masked_pairs = [
76 (key, _mask_credential(val)) for key, val in parse_qsl(parts.query, keep_blank_values=True)
77 ]
78 new_query = urlencode(masked_pairs)
79 return urlunsplit((parts.scheme, parts.netloc, parts.path, new_query, parts.fragment))
80
81
82 def _masked_config_for_log(provider_type: str, config: Dict[str, Any]) -> Dict[str, Any]:
83 """Return a deep-copied ``config`` with sensitive fields masked per provider type.
84
85 - ``bark``: mask ``device_key``
86 - ``telegram``: mask ``bot_token``
87 - ``webhook``: mask every query-string value in ``url`` (host/path preserved)
88
89 Unknown types and non-dict configs are returned as a deep copy without
90 modification so callers can always log the result safely.
91 """
92 if not isinstance(config, dict):
93 # Return a deep copy of whatever was passed so callers never mutate the
94 # original object via the returned reference.
95 return copy.deepcopy(config)
96
97 masked = copy.deepcopy(config)
98 ptype = (provider_type or "").strip().lower()
99 if ptype == "bark":
100 if "device_key" in masked:
101 masked["device_key"] = _mask_credential(masked.get("device_key") or "")
102 elif ptype == "telegram":
103 if "bot_token" in masked:
104 masked["bot_token"] = _mask_credential(masked.get("bot_token") or "")
105 elif ptype == "webhook":
106 if "url" in masked and isinstance(masked.get("url"), str):
107 masked["url"] = _mask_url_query(masked["url"])
108 return masked
109
110
111 class _BaseProvider:
112 def __init__(self, settings: Dict[str, Any]):
113 self.settings = settings or {}
114
115 async def send(self, session: aiohttp.ClientSession, title: str, body: str, level: str) -> bool:
116 raise NotImplementedError
117
118
119 class BarkProvider(_BaseProvider):
120 """Bark 推送,URL 形如 https://api.day.app/<device_key>。
121
122 参考:https://bark.day.app/
123 """
124
125 async def send(self, session: aiohttp.ClientSession, title: str, body: str, level: str) -> bool:
126 base_url = str(self.settings.get("url") or "").rstrip("/")
127 if not base_url:
128 logger.warning("Bark notification skipped: missing url")
129 return False
130 sound = str(self.settings.get("sound") or "")
131 # Bark 以 URL path 传参:/{device_key}/{title}/{body}
132 url = f"{base_url}/{quote(title, safe='')}/{quote(body, safe='')}"
133 params: Dict[str, str] = {}
134 if sound:
135 params["sound"] = sound
136 try:
137 async with session.get(url, params=params) as resp:
138 ok = resp.status == 200
139 if not ok:
140 logger.warning("Bark notification HTTP %s", resp.status)
141 return ok
142 except Exception as exc:
143 logger.warning("Bark notification failed: %s", exc)
144 return False
145
146
147 class TelegramProvider(_BaseProvider):
148 """Telegram Bot 推送。需要配置 bot_token 与 chat_id。"""
149
150 async def send(self, session: aiohttp.ClientSession, title: str, body: str, level: str) -> bool:
151 bot_token = str(self.settings.get("bot_token") or "")
152 chat_id = str(self.settings.get("chat_id") or "")
153 if not bot_token or not chat_id:
154 logger.warning("Telegram notification skipped: missing bot_token/chat_id")
155 return False
156 url = f"https://api.telegram.org/bot{bot_token}/sendMessage"
157 text = f"*{title}*\n{body}" if title else body
158 payload = {
159 "chat_id": chat_id,
160 "text": text,
161 "parse_mode": "Markdown",
162 }
163 try:
164 async with session.post(url, json=payload) as resp:
165 ok = resp.status == 200
166 if not ok:
167 logger.warning("Telegram notification HTTP %s", resp.status)
168 return ok
169 except Exception as exc:
170 logger.warning("Telegram notification failed: %s", exc)
171 return False
172
173
174 class WebhookProvider(_BaseProvider):
175 """通用 Webhook:POST JSON {title, body, level}。可用于接企业微信/飞书/钉钉 bot。"""
176
177 async def send(self, session: aiohttp.ClientSession, title: str, body: str, level: str) -> bool:
178 url = str(self.settings.get("url") or "")
179 if not url:
180 logger.warning("Webhook notification skipped: missing url")
181 return False
182 extra_headers = {str(k): str(v) for k, v in (self.settings.get("headers") or {}).items()}
183 payload: Dict[str, Any] = {"title": title, "body": body, "level": level}
184 # 允许通过 extra_body 合并到 payload,便于适配某些平台(企业微信 msgtype 等)。
185 extra_body = self.settings.get("extra_body")
186 if isinstance(extra_body, dict):
187 payload.update(extra_body)
188 try:
189 async with session.post(url, json=payload, headers=extra_headers) as resp:
190 ok = resp.status < 400
191 if not ok:
192 logger.warning("Webhook notification HTTP %s", resp.status)
193 return ok
194 except Exception as exc:
195 logger.warning("Webhook notification failed: %s", exc)
196 return False
197
198
199 _PROVIDER_REGISTRY = {
200 "bark": BarkProvider,
201 "telegram": TelegramProvider,
202 "webhook": WebhookProvider,
203 }
204
205
206 class Notifier:
207 """聚合通知器,并发分发至所有启用的 provider。"""
208
209 def __init__(
210 self,
211 providers: List[_BaseProvider],
212 *,
213 on_success: bool = True,
214 on_failure: bool = True,
215 ):
216 self.providers = providers
217 self.on_success = on_success
218 self.on_failure = on_failure
219
220 @property
221 def enabled(self) -> bool:
222 return bool(self.providers)
223
224 async def send(
225 self,
226 title: str,
227 body: str,
228 *,
229 level: str = "info",
230 ) -> Dict[str, bool]:
231 """发送通知,返回 {provider_name: ok} 映射。"""
232 if not self.providers:
233 return {}
234
235 is_failure = level in {"failure", "error"}
236 is_success = level in {"success", "info"}
237 if is_failure and not self.on_failure:
238 return {}
239 if is_success and not self.on_success:
240 return {}
241
242 async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=10)) as session:
243 tasks = [p.send(session, title=title, body=body, level=level) for p in self.providers]
244 results = await asyncio.gather(*tasks, return_exceptions=True)
245
246 summary: Dict[str, bool] = {}
247 for provider, result in zip(self.providers, results):
248 name = type(provider).__name__
249 if isinstance(result, Exception):
250 logger.warning("Provider %s crashed: %s", name, result)
251 summary[name] = False
252 else:
253 summary[name] = bool(result)
254 return summary
255
256
257 def build_notifier(config_source: Any) -> Notifier:
258 """从 ConfigLoader 或 dict 构造 Notifier。"""
259 cfg: Any
260 if hasattr(config_source, "get"):
261 cfg = config_source.get("notifications", {}) or {}
262 elif isinstance(config_source, dict):
263 cfg = config_source.get("notifications", {}) or {}
264 else:
265 cfg = {}
266
267 # 用户可能误写成 `notifications: on` 等 scalar:防御性降级为 disabled。
268 if not isinstance(cfg, dict):
269 logger.warning(
270 "notifications config must be a dict, got %s; treating as disabled.",
271 type(cfg).__name__,
272 )
273 return Notifier(providers=[])
274
275 if not cfg.get("enabled", False):
276 return Notifier(providers=[])
277
278 providers: List[_BaseProvider] = []
279 for entry in cfg.get("providers") or []:
280 if not isinstance(entry, dict):
281 continue
282 ptype = str(entry.get("type") or "").strip().lower()
283 cls = _PROVIDER_REGISTRY.get(ptype)
284 if cls is None:
285 logger.warning("Unknown notification provider type: %s", ptype)
286 continue
287 providers.append(cls(entry))
288
289 return Notifier(
290 providers=providers,
291 on_success=bool(cfg.get("on_success", True)),
292 on_failure=bool(cfg.get("on_failure", True)),
293 )
294
294 lines PYTHON