graylog-ops/graylog_ops/notify.py

"""Webhook notifications for newly convicted malicious IPs.

Supported payload flavours: WeCom (企业微信), DingTalk, Feishu, Slack and a
generic JSON body. Delivery failures are reported but never break detection.
"""
import base64
import hashlib
import re
import hmac
import json
import time
import urllib.error
import urllib.parse
import urllib.request
from datetime import datetime

from .geo import flag

DEFAULT_TITLE = "新增恶意 IP"


class Notifier:
    def __init__(self, config):
        self.kind = (config.get("type") or "generic").lower()
        self.url = config.get("url") or ""
        if not self.url:
            raise ValueError("notify config is missing 'url'")
        self.secret = config.get("secret") or ""
        self.link_prefix = config.get("link_prefix") or ""
        self.timeout = int(config.get("timeout") or 10)
        self.title = config.get("title") or DEFAULT_TITLE

    def link_for(self, ip):
        if not self.link_prefix:
            return ""
        return self.link_prefix.rstrip("/") + "/#/ip/" + ip

    @staticmethod
    def local_time(value):
        try:
            text = str(value).replace("Z", "+00:00")
            return datetime.fromisoformat(text).astimezone().strftime("%m-%d %H:%M")
        except (ValueError, TypeError):
            return str(value)

    @staticmethod
    def risk_label(score):
        if score >= 90:
            return "🔴 高危"
        if score >= 70:
            return "🟠 中危"
        return "🟡 关注"

    def render(self, entries):
        count = len(entries)
        heading = self.title if count == 1 else "新增 %d 个恶意 IP" % count
        lines = ["### 🚨 %s" % heading, ""]
        for index, entry in enumerate(entries, 1):
            ip = entry.get("ip", "?")
            geo = " · ".join(part for part in
                             (entry.get("country_name"), entry.get("region")) if part)
            flag_text = flag(entry.get("country_code"))
            place = ("%s %s" % (flag_text, geo)).strip() if geo else ""
            score = float(entry.get("score") or 0)
            lead = ("**%d. %s**" % (index, ip)) if count > 1 else ("**%s**" % ip)
            lines.append("%s%s" % (lead, ("  " + place) if place else ""))
            lines.append("评分 **%.1f**（%s）｜请求 **%s** 次"
                         % (score, self.risk_label(score), entry.get("total") or 0))
            if entry.get("first_seen"):
                lines.append("首次攻击 **%s**" % self.local_time(entry["first_seen"]))
            rules = entry.get("rules_text")
            if rules:
                lines.append("")
                lines.append("**命中类型**")
                for part in rules.split("；"):
                    part = part.strip()
                    if not part:
                        continue
                    part = re.sub(r"\s+(\d+)\s*次$", r" × \1", part)
                    lines.append("- " + part)
            link = self.link_for(ip)
            if link:
                lines.append("")
                lines.append("[🔗 查看详情与证据](%s)" % link)
            if index != count:
                lines.append("")
                lines.append("---")
                lines.append("")
        return "\n".join(lines)

    def signed_url(self, timestamp=None):
        """DingTalk custom robots with 加签 need timestamp + HMAC-SHA256."""
        if self.kind != "dingtalk" or not self.secret:
            return self.url
        stamp = str(timestamp if timestamp is not None
                    else int(time.time() * 1000))
        string_to_sign = stamp + "\n" + self.secret
        digest = hmac.new(self.secret.encode("utf-8"),
                          string_to_sign.encode("utf-8"),
                          hashlib.sha256).digest()
        sign = urllib.parse.quote_plus(base64.b64encode(digest).decode("ascii"))
        separator = "&" if "?" in self.url else "?"
        return "%s%stimestamp=%s&sign=%s" % (self.url, separator, stamp, sign)

    def payload(self, text, title=None):
        title = title or self.title
        if self.kind == "wecom":
            return {"msgtype": "markdown", "markdown": {"content": text}}
        if self.kind == "dingtalk":
            return {"msgtype": "markdown", "markdown": {"title": title, "text": text}}
        if self.kind == "feishu":
            return {"msg_type": "text", "content": {"text": text}}
        if self.kind == "slack":
            return {"text": text}
        return {"title": title, "text": text}

    def post(self, text, title=None):
        body = json.dumps(self.payload(text, title), ensure_ascii=False).encode("utf-8")
        request = urllib.request.Request(
            self.signed_url(), data=body, method="POST",
            headers={"Content-Type": "application/json; charset=utf-8"})
        try:
            with urllib.request.urlopen(request, timeout=self.timeout) as response:
                reply = response.read().decode("utf-8", "replace")
        except urllib.error.HTTPError as exc:
            raise RuntimeError("webhook HTTP %s: %s"
                               % (exc.code, exc.read()[:200].decode("utf-8", "replace")))
        except urllib.error.URLError as exc:
            raise RuntimeError("webhook unreachable: %s" % exc.reason)
        if reply and self.kind in ("wecom", "dingtalk", "feishu"):
            try:
                parsed = json.loads(reply)
            except ValueError:
                return reply
            code = parsed.get("errcode", parsed.get("code", 0))
            if code:
                raise RuntimeError("webhook rejected: %s" % reply[:200])
        return reply

    def send(self, entries):
        if not entries:
            return ""
        return self.post(self.render(entries))

    def send_test(self):
        return self.post("**%s**\n\n这是一条测试推送，说明通知通道工作正常。"
                         % self.title, title=self.title + " · 测试")


def load_notifier(path):
    with open(path, "r", encoding="utf-8") as handle:
        return Notifier(json.load(handle))


def notifier_from_env():
    """Build a notifier from NOTIFY_* environment variables (container mode)."""
    import os
    url = (os.environ.get("NOTIFY_URL") or "").strip()
    if not url:
        return None
    return Notifier({
        "type": os.environ.get("NOTIFY_TYPE", "generic"),
        "url": url,
        "secret": os.environ.get("NOTIFY_SECRET", ""),
        "link_prefix": os.environ.get("NOTIFY_LINK_PREFIX", ""),
        "title": os.environ.get("NOTIFY_TITLE", "新增恶意 IP"),
    })