from __future__ import annotations import base64 import hashlib import hmac import json import logging import os import threading import time from collections import defaultdict from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field from io import BytesIO from typing import Any from urllib.error import HTTPError, URLError from urllib.request import Request, urlopen from PIL import Image, ImageDraw, ImageFont from .storage import DeliveryStatsStore, utc_now_iso LOGGER = logging.getLogger(__name__) ALERT_CLASSES = ("fire", "smoke") CHANNEL_NAMES = {"wechat": "企业微信", "feishu": "飞书"} @dataclass(slots=True) class SessionState: consecutive: dict[str, int] = field( default_factory=lambda: defaultdict(int) ) last_alert_at: dict[str, float] = field( default_factory=lambda: defaultdict(float) ) last_seen_at: float = field(default_factory=time.monotonic) class AlertManager: def __init__( self, webhook_url: str | None = None, wechat_webhook_url: str | None = None, feishu_webhook_url: str | None = None, feishu_secret: str | None = None, confirm_frames: int = 3, cooldown_seconds: float = 60.0, session_ttl_seconds: float = 3600.0, stats_store: DeliveryStatsStore | None = None, ) -> None: self.wechat_webhook_url = ( wechat_webhook_url or webhook_url or os.getenv("WECHAT_WEBHOOK_URL") ) self.feishu_webhook_url = ( feishu_webhook_url or os.getenv("FEISHU_WEBHOOK_URL") ) self.feishu_secret = feishu_secret or os.getenv("FEISHU_SECRET") self.confirm_frames = max(1, confirm_frames) self.cooldown_seconds = max(0.0, cooldown_seconds) self.session_ttl_seconds = max(60.0, session_ttl_seconds) self.stats_store = stats_store self._states: dict[str, SessionState] = {} self._lock = threading.Lock() self._delivery_stats = { channel: { "delivered": 0, "failed": 0, "tests": 0, "last_status": "idle", "last_delivery_at": None, "last_test_at": None, "last_error": None, } for channel in CHANNEL_NAMES } self._executor = ThreadPoolExecutor( max_workers=4, thread_name_prefix="alert-dispatch", ) @property def channels(self) -> dict[str, bool]: return { "wechat": bool(self.wechat_webhook_url), "feishu": bool(self.feishu_webhook_url), } @property def enabled(self) -> bool: return any(self.channels.values()) @property def channel_status(self) -> list[dict[str, Any]]: configured = self.channels if self.stats_store is not None: saved = self.stats_store.get_all() stats = { channel: { "delivered": 0, "failed": 0, "tests": 0, "last_status": "idle", "last_delivery_at": None, "last_test_at": None, "last_error": None, **saved.get(channel, {}), } for channel in CHANNEL_NAMES } else: with self._lock: stats = { channel: dict(values) for channel, values in self._delivery_stats.items() } return [ { "id": channel, "name": CHANNEL_NAMES[channel], "configured": configured[channel], "signed": channel == "feishu" and bool(self.feishu_secret), "capabilities": ( ["Markdown 告警", "检测截图"] if channel == "wechat" else ["交互卡片", "签名校验"] ), **stats[channel], } for channel in CHANNEL_NAMES ] def test_channel(self, channel: str) -> dict[str, Any]: if channel not in CHANNEL_NAMES: raise ValueError(f"Unsupported robot channel: {channel}") if not self.channels[channel]: raise ValueError(f"{CHANNEL_NAMES[channel]}机器人尚未配置") try: if channel == "wechat": self._send_wechat_test() else: self._send_feishu_test() except (HTTPError, URLError, TimeoutError, ValueError) as error: self._record_delivery( channel, success=False, error=str(error), test=True, ) raise RuntimeError(str(error)) from error self._record_delivery(channel, success=True, test=True) return next( status for status in self.channel_status if status["id"] == channel ) def evaluate( self, session_id: str, detections: list[dict[str, Any]], image: Image.Image, ) -> dict[str, Any]: now = time.monotonic() present = { detection["class"] for detection in detections if detection.get("class") in ALERT_CLASSES } triggered: list[str] = [] with self._lock: self._prune_sessions(now) state = self._states.setdefault(session_id, SessionState()) state.last_seen_at = now for class_name in ALERT_CLASSES: state.consecutive[class_name] = ( state.consecutive[class_name] + 1 if class_name in present else 0 ) ready = state.consecutive[class_name] >= self.confirm_frames cooldown_elapsed = ( now - state.last_alert_at[class_name] >= self.cooldown_seconds ) if ready and cooldown_elapsed: state.last_alert_at[class_name] = now triggered.append(class_name) consecutive = dict(state.consecutive) if triggered and self.enabled: self._executor.submit( self._send_alerts, annotate_image(image, detections), triggered, detections, ) return { "triggered": bool(triggered), "classes": triggered, "confirmed_frames": self.confirm_frames, "consecutive": consecutive, "cooldown_seconds": self.cooldown_seconds, "notification_enabled": self.enabled, "notification_channels": self.channels, } def reset(self, session_id: str) -> None: with self._lock: self._states.pop(session_id, None) def close(self) -> None: """Shut down the alert dispatch executor (call on service shutdown).""" self._executor.shutdown(wait=True, cancel_futures=True) def _prune_sessions(self, now: float) -> None: expired = [ session_id for session_id, state in self._states.items() if now - state.last_seen_at > self.session_ttl_seconds ] for session_id in expired: del self._states[session_id] def _send_alerts( self, image: Image.Image, triggered: list[str], detections: list[dict[str, Any]], ) -> None: if self.wechat_webhook_url: try: self._send_wechat_alert(image, triggered, detections) self._record_delivery("wechat", success=True) except (HTTPError, URLError, TimeoutError, ValueError) as error: self._record_delivery("wechat", success=False, error=str(error)) LOGGER.error("WeChat alert failed: %s", error) if self.feishu_webhook_url: try: self._send_feishu_alert(triggered, detections) self._record_delivery("feishu", success=True) except (HTTPError, URLError, TimeoutError, ValueError) as error: self._record_delivery("feishu", success=False, error=str(error)) LOGGER.error("Feishu alert failed: %s", error) def _record_delivery( self, channel: str, *, success: bool, error: str | None = None, test: bool = False, ) -> None: if self.stats_store is not None: self.stats_store.record( channel, success=success, error=error, test=test, ) return now = utc_now_iso() with self._lock: stats = self._delivery_stats[channel] stats["last_status"] = "success" if success else "failed" stats["last_error"] = None if success else error stats["delivered" if success else "failed"] += 1 if test: stats["tests"] += 1 stats["last_test_at"] = now else: stats["last_delivery_at"] = now def _send_wechat_test(self) -> None: if not self.wechat_webhook_url: return post_json( self.wechat_webhook_url, { "msgtype": "markdown", "markdown": { "content": ( "燧安机器人连接测试\n" "> 消息通道工作正常\n" f"> 测试时间:{time.strftime('%Y-%m-%d %H:%M:%S')}" ) }, }, ) def _send_feishu_test(self) -> None: if not self.feishu_webhook_url: return payload: dict[str, Any] = { "msg_type": "interactive", "card": { "header": { "template": "green", "title": { "tag": "plain_text", "content": "燧安机器人连接测试", }, }, "elements": [ { "tag": "markdown", "content": ( "**消息通道工作正常**\n" f"测试时间:{time.strftime('%Y-%m-%d %H:%M:%S')}" ), } ], }, } if self.feishu_secret: timestamp = str(int(time.time())) payload["timestamp"] = timestamp payload["sign"] = feishu_signature( timestamp, self.feishu_secret, ) post_json(self.feishu_webhook_url, payload) def _send_wechat_alert( self, image: Image.Image, triggered: list[str], detections: list[dict[str, Any]], ) -> None: if not self.wechat_webhook_url: return target_text, max_confidence = alert_summary(triggered, detections) message = ( f"🔥 烟火检测告警\n" f"> 检测目标:{target_text}\n" f"> 最高置信度:{max_confidence:.1%}\n" f"> 请及时查看现场视频。" ) post_json( self.wechat_webhook_url, {"msgtype": "markdown", "markdown": {"content": message}}, ) image_bytes = encode_jpeg(image) post_json( self.wechat_webhook_url, { "msgtype": "image", "image": { "base64": base64.b64encode(image_bytes).decode("ascii"), "md5": hashlib.md5( image_bytes, usedforsecurity=False, ).hexdigest(), }, }, ) def _send_feishu_alert( self, triggered: list[str], detections: list[dict[str, Any]], ) -> None: if not self.feishu_webhook_url: return target_text, max_confidence = alert_summary(triggered, detections) payload: dict[str, Any] = { "msg_type": "interactive", "card": { "config": {"wide_screen_mode": True}, "header": { "template": "red", "title": { "tag": "plain_text", "content": "烟火检测告警", }, }, "elements": [ { "tag": "markdown", "content": ( f"**检测目标:** {target_text}\n" f"**最高置信度:** {max_confidence:.1%}\n" f"**告警时间:** " f"{time.strftime('%Y-%m-%d %H:%M:%S')}\n" f"请及时查看现场视频。" ), } ], }, } if self.feishu_secret: timestamp = str(int(time.time())) payload["timestamp"] = timestamp payload["sign"] = feishu_signature( timestamp, self.feishu_secret, ) post_json(self.feishu_webhook_url, payload) def alert_summary( triggered: list[str], detections: list[dict[str, Any]], ) -> tuple[str, float]: labels = {"fire": "火焰", "smoke": "烟雾"} target_text = "、".join(labels[name] for name in triggered) max_confidence = max( ( float(detection.get("confidence", 0.0)) for detection in detections if detection.get("class") in triggered ), default=0.0, ) return target_text, max_confidence def feishu_signature(timestamp: str, secret: str) -> str: string_to_sign = f"{timestamp}\n{secret}".encode() digest = hmac.new(string_to_sign, digestmod=hashlib.sha256).digest() return base64.b64encode(digest).decode("ascii") def post_json(url: str, payload: dict[str, Any]) -> dict[str, Any]: request = Request( url, data=json.dumps(payload, ensure_ascii=False).encode("utf-8"), headers={"Content-Type": "application/json"}, method="POST", ) with urlopen(request, timeout=10) as response: result = json.loads(response.read().decode("utf-8")) error_code = result.get( "errcode", result.get("code", result.get("StatusCode", 0)), ) if error_code != 0: raise ValueError( result.get("errmsg") or result.get("msg") or result.get("StatusMessage") or "Unknown robot webhook error" ) return result def annotate_image( image: Image.Image, detections: list[dict[str, Any]], ) -> Image.Image: annotated = image.copy() draw = ImageDraw.Draw(annotated) font = ImageFont.load_default() colors = {"fire": "#ff3b30", "smoke": "#00a89b"} for detection in detections: box = detection.get("box") if not box or len(box) != 4: continue class_name = str(detection.get("class", "target")) confidence = float(detection.get("confidence", 0.0)) color = colors.get(class_name, "#ffd166") coordinates = tuple(round(value) for value in box) draw.rectangle(coordinates, outline=color, width=4) draw.text( (coordinates[0] + 4, max(0, coordinates[1] - 16)), f"{class_name} {confidence:.0%}", fill=color, font=font, ) return annotated def encode_jpeg(image: Image.Image, max_bytes: int = 1_900_000) -> bytes: working = image.convert("RGB") quality = 88 while True: output = BytesIO() working.save(output, format="JPEG", quality=quality, optimize=True) payload = output.getvalue() if len(payload) <= max_bytes: return payload if quality > 55: quality -= 10 continue width, height = working.size working = working.resize( (max(1, int(width * 0.8)), max(1, int(height * 0.8))) ) quality = 75