model-watchdog/notifier.py

"""告警出口：钉钉机器人 + 邮件群发，统一封装为"双重告警"。"""
import base64
import hashlib
import hmac
import logging
import time

import requests

from mailer import forward_notice, send_mail

log = logging.getLogger("notifier")

# 钉钉群机器人固定接口地址；access_token 作为 query 参数传入
DINGTALK_API = "https://oapi.dingtalk.com/robot/send"


def _dingtalk_sign(secret):
    """钉钉加签：毫秒时间戳 + "\\n" + secret，HMAC-SHA256 后 Base64。

    返回 (timestamp, sign)。sign 交给 requests 的 params 去做 URL 编码，
    不要自己 quote，否则容易出现二次编码导致验签失败。
    """
    ts = str(round(time.time() * 1000))
    digest = hmac.new(
        secret.encode("utf-8"),
        f"{ts}\n{secret}".encode("utf-8"),
        digestmod=hashlib.sha256,
    ).digest()
    return ts, base64.b64encode(digest).decode("utf-8")


def send_dingtalk(dt_cfg, title, text):
    """发送 markdown 消息到钉钉群机器人。未启用时直接跳过。

    凭证支持两种写法，二选一：
      access_token: "xxx"                      # 推荐，只填 token
      webhook: "https://...?access_token=xxx"  # 兼容整条 webhook
    配了 secret 就自动加签（对应机器人安全设置里的"加签"方式）。
    """
    if not dt_cfg.get("enabled"):
        log.info("[钉钉] 未启用，跳过")
        return

    url = (dt_cfg.get("webhook") or DINGTALK_API).strip()
    token = (dt_cfg.get("access_token") or "").strip()

    params = {}
    if "access_token=" not in url:
        if not token:
            log.warning("[钉钉] 已启用但未配置 access_token/webhook，跳过")
            return
        params["access_token"] = token

    secret = (dt_cfg.get("secret") or "").strip()
    if secret:
        params["timestamp"], params["sign"] = _dingtalk_sign(secret)

    at_mobiles = [str(m).strip() for m in (dt_cfg.get("at_mobiles") or []) if str(m).strip()]
    at_all = bool(dt_cfg.get("at_all"))
    # markdown 消息里必须出现 @手机号 字面量，@ 才会真的生效
    if at_mobiles:
        text = text + "\n\n" + " ".join(f"@{m}" for m in at_mobiles)

    payload = {"msgtype": "markdown", "markdown": {"title": title, "text": text}}
    if at_mobiles or at_all:
        payload["at"] = {"atMobiles": at_mobiles, "isAtAll": at_all}

    resp = requests.post(
        url,
        params=params,
        json=payload,
        headers={"Content-Type": "application/json"},
        timeout=dt_cfg.get("timeout_seconds", 15),
    )
    resp.raise_for_status()
    data = resp.json()
    if data.get("errcode") != 0:
        # errmsg 里通常直接写明原因（token 无效 / 验签失败 / 关键词不匹配 / 限流）
        raise RuntimeError(
            "钉钉返回错误 errcode=%s errmsg=%s" % (data.get("errcode"), data.get("errmsg"))
        )
    log.info("[钉钉] 告警已发送: %s", title)


def dual_forward(cfg, notice, dry_run=False):
    """下线通知双通道：邮件转发原文给通知列表 + 钉钉推送摘要。

    任一通道失败不影响另一个。返回 (转发是否成功, 钉钉是否成功)。
    """
    mid = notice["message_id"]
    if dry_run:
        log.info(
            "[DRY-RUN] 将转发下线通知\nMessage-ID: %s\n原发件人: %s\n原标题: %s",
            mid, notice["from"], notice["subject"],
        )
        return True, True

    fwd_ok = False
    try:
        forward_notice(cfg["mail"], notice)
        fwd_ok = True
    except Exception as e:  # noqa: BLE001
        log.error("[告警] 下线通知转发失败: %s", e)

    text = "\n".join([
        "检测到模型下线通知，已转发给通知列表。",
        "",
        f"原发件人: {notice['from']}",
        f"原发送时间: {notice['date']}",
        f"原邮件标题: {notice['subject']}",
        f"Message-ID: {mid}",
        "",
        "正文摘要:",
        (notice.get("body") or "")[:500],
    ])
    ding_ok = False
    try:
        send_dingtalk(cfg["dingtalk"], "【模型下线通知】" + notice["subject"], _to_markdown(
            "【模型下线通知】" + notice["subject"], text))
        ding_ok = True
    except Exception as e:  # noqa: BLE001
        log.error("[告警] 钉钉发送失败: %s", e)

    return fwd_ok, ding_ok


def dual_alert(cfg, subject, body, dry_run=False):
    """双重告警：邮件群发 + 钉钉推送。任一通道失败不影响另一个。

    返回 (邮件是否成功, 钉钉是否成功)。
    """
    if dry_run:
        log.info("[DRY-RUN] 将发送告警\n标题: %s\n正文:\n%s", subject, body)
        return True, True

    mail_ok = False
    try:
        send_mail(cfg["mail"], subject, body)
        mail_ok = True
    except Exception as e:  # noqa: BLE001
        log.error("[告警] 邮件发送失败: %s", e)

    ding_ok = False
    try:
        send_dingtalk(cfg["dingtalk"], subject, _to_markdown(subject, body))
        ding_ok = True
    except Exception as e:  # noqa: BLE001
        log.error("[告警] 钉钉发送失败: %s", e)

    return mail_ok, ding_ok


def _to_markdown(subject, body):
    """纯文本正文转钉钉 markdown：保留换行，避免被 md 合并成一行。"""
    lines = ["### " + subject, ""]
    for line in body.splitlines():
        lines.append(line.rstrip() + "  " if line.strip() else "")
    return "\n".join(lines)