130 lines
4.3 KiB
Python
130 lines
4.3 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import os
|
|
import re
|
|
import time
|
|
import urllib.parse
|
|
import urllib.request
|
|
from dataclasses import dataclass
|
|
|
|
|
|
@dataclass
|
|
class NotifierConfig:
|
|
bot_token: str
|
|
chat_id: str
|
|
min_level: str = "INFO"
|
|
rate_limit_per_sec: float = 5.0
|
|
|
|
|
|
class Notifier:
|
|
LEVELS = {"INFO": 10, "WARNING": 20, "CRITICAL": 30}
|
|
|
|
def __init__(self, config: NotifierConfig | None = None) -> None:
|
|
self.config = config
|
|
self.enabled = config is not None and bool(config.bot_token) and bool(config.chat_id)
|
|
self._queue: asyncio.PriorityQueue[tuple[int, float, str, str]] = asyncio.PriorityQueue(maxsize=1000)
|
|
self._worker: asyncio.Task | None = None
|
|
self._stopping = False
|
|
self._last_send_ts = 0.0
|
|
|
|
@classmethod
|
|
def from_env(cls) -> "Notifier":
|
|
token = (
|
|
os.getenv("GAMMA_TELEGRAM_TOKEN", "").strip()
|
|
or os.getenv("TELEGRAM_BOT_TOKEN", "").strip()
|
|
)
|
|
chat_id = (
|
|
os.getenv("GAMMA_TELEGRAM_CHAT_ID", "").strip()
|
|
or os.getenv("TELEGRAM_CHAT_ID", "").strip()
|
|
)
|
|
min_level = (
|
|
os.getenv("GAMMA_TELEGRAM_MIN_LEVEL", "").strip().upper()
|
|
or os.getenv("TELEGRAM_MIN_LEVEL", "INFO").strip().upper()
|
|
or "INFO"
|
|
)
|
|
if not token or not chat_id:
|
|
return cls(config=None)
|
|
return cls(config=NotifierConfig(bot_token=token, chat_id=chat_id, min_level=min_level))
|
|
|
|
async def start(self) -> None:
|
|
if not self.enabled or self._worker is not None:
|
|
return
|
|
self._stopping = False
|
|
self._worker = asyncio.create_task(self._worker_loop(), name="strategy32-telegram-notifier")
|
|
|
|
async def stop(self) -> None:
|
|
self._stopping = True
|
|
if self._worker is None:
|
|
return
|
|
self._worker.cancel()
|
|
try:
|
|
await self._worker
|
|
except asyncio.CancelledError:
|
|
pass
|
|
except Exception:
|
|
pass
|
|
self._worker = None
|
|
|
|
async def send(self, level: str, message: str) -> bool:
|
|
level = level.upper().strip()
|
|
if not self.enabled:
|
|
return False
|
|
if self.LEVELS.get(level, 0) < self.LEVELS.get(self.config.min_level.upper(), 10):
|
|
return False
|
|
try:
|
|
priority = 0 if level == "CRITICAL" else 1
|
|
self._queue.put_nowait((priority, time.time(), level, _mask_sensitive(message)))
|
|
return True
|
|
except asyncio.QueueFull:
|
|
return False
|
|
|
|
async def _worker_loop(self) -> None:
|
|
while not self._stopping:
|
|
_priority, _ts, level, message = await self._queue.get()
|
|
text = f"[{level}] {message}"
|
|
try:
|
|
await self._send_telegram(text)
|
|
except Exception:
|
|
continue
|
|
|
|
async def _send_telegram(self, text: str) -> None:
|
|
min_interval = 1.0 / max(1e-6, float(self.config.rate_limit_per_sec))
|
|
elapsed = time.time() - self._last_send_ts
|
|
if elapsed < min_interval:
|
|
await asyncio.sleep(min_interval - elapsed)
|
|
|
|
encoded = urllib.parse.urlencode(
|
|
{
|
|
"chat_id": self.config.chat_id,
|
|
"text": text,
|
|
"disable_web_page_preview": "true",
|
|
}
|
|
).encode("utf-8")
|
|
url = f"https://api.telegram.org/bot{self.config.bot_token}/sendMessage"
|
|
|
|
def _call() -> None:
|
|
req = urllib.request.Request(
|
|
url,
|
|
data=encoded,
|
|
method="POST",
|
|
headers={"Content-Type": "application/x-www-form-urlencoded"},
|
|
)
|
|
with urllib.request.urlopen(req, timeout=10) as resp:
|
|
raw = resp.read().decode("utf-8")
|
|
data = json.loads(raw)
|
|
if not isinstance(data, dict) or not bool(data.get("ok")):
|
|
raise RuntimeError("telegram_send_failed")
|
|
|
|
await asyncio.to_thread(_call)
|
|
self._last_send_ts = time.time()
|
|
|
|
|
|
def _mask_sensitive(text: str) -> str:
|
|
out = str(text)
|
|
out = re.sub(r"0x[a-fA-F0-9]{64}", "0x***MASKED_PRIVATE_KEY***", out)
|
|
out = re.sub(r"0x[a-fA-F0-9]{40}", "0x***MASKED_ADDRESS***", out)
|
|
out = re.sub(r"\b\d{7,}:[A-Za-z0-9_-]{20,}\b", "***MASKED_TOKEN***", out)
|
|
return out
|