notify.py

#!/usr/bin/env python3
"""
轻量邮件通知套件。供定时任务 / watcher / 一次性脚本复用。

配置: 同目录 notify_config.json（集中存放 SMTP 账号/授权码/默认收件人）。
缺失时回退到环境变量 NOTIFY_SMTP_USER / NOTIFY_SMTP_PASS / NOTIFY_TO。

用法 A — Python import:
    from notify import send
    send("训练完成", "PSNR 34.1 dB")
    send("带附件", body, attachments=["/path/fig.png"])

用法 B — 命令行 / 管道（cron、shell 直接用，无需写 Python）:
    python3 notify.py "主题" "正文"
    echo "正文" | python3 notify.py "主题"
    tail -20 train.log | python3 notify.py "训练进度"

退出码: 0 成功 / 1 失败（已重试）。失败信息打到 stderr，不抛栈。
"""
import sys, os, json, time, ssl, smtplib
from email.mime.text import MIMEText
from email.mime.multipart import MIMEMultipart
from email.mime.application import MIMEApplication

_CFG_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "notify_config.json")


def _load_cfg():
    cfg = {
        "smtp_host": "smtp.qq.com",
        "smtp_port": 465,
        "smtp_user": os.environ.get("NOTIFY_SMTP_USER", ""),
        "smtp_pass": os.environ.get("NOTIFY_SMTP_PASS", ""),
        "to": os.environ.get("NOTIFY_TO", ""),
    }
    if os.path.exists(_CFG_PATH):
        try:
            with open(_CFG_PATH, "r", encoding="utf-8") as f:
                cfg.update({k: v for k, v in json.load(f).items() if v != ""})
        except Exception as e:
            print(f"[notify] 配置读取失败({e})，回退环境变量", file=sys.stderr)
    return cfg


def send(subject, body, to=None, attachments=None, retries=3, retry_wait=10):
    """发送邮件。成功返回 True，失败（重试耗尽）返回 False。"""
    cfg = _load_cfg()
    to_addr = to or cfg["to"]
    if not (cfg["smtp_user"] and cfg["smtp_pass"] and to_addr):
        print("[notify] 缺少 smtp_user/smtp_pass/to，无法发送", file=sys.stderr)
        return False

    recipients = [a.strip() for a in (to_addr if isinstance(to_addr, str) else ",".join(to_addr)).split(",") if a.strip()]

    msg = MIMEMultipart()
    msg["From"] = cfg["smtp_user"]
    msg["To"] = ", ".join(recipients)
    msg["Subject"] = subject
    msg.attach(MIMEText(body, "plain", "utf-8"))

    for path in (attachments or []):
        try:
            with open(path, "rb") as f:
                part = MIMEApplication(f.read(), Name=os.path.basename(path))
            part["Content-Disposition"] = f'attachment; filename="{os.path.basename(path)}"'
            msg.attach(part)
        except Exception as e:
            print(f"[notify] 附件跳过 {path}: {e}", file=sys.stderr)

    last_err = None
    for attempt in range(1, retries + 1):
        try:
            ctx = ssl.create_default_context()
            with smtplib.SMTP_SSL(cfg["smtp_host"], int(cfg["smtp_port"]), context=ctx, timeout=30) as s:
                s.login(cfg["smtp_user"], cfg["smtp_pass"])
                s.sendmail(cfg["smtp_user"], recipients, msg.as_string())
            return True
        except Exception as e:
            last_err = e
            print(f"[notify] 第 {attempt}/{retries} 次发送失败: {e}", file=sys.stderr)
            if attempt < retries:
                time.sleep(retry_wait)
    print(f"[notify] 放弃，最后错误: {last_err}", file=sys.stderr)
    return False


def _main(argv):
    if len(argv) < 2:
        print('用法: notify.py "主题" ["正文"]   (正文缺省时从 stdin 读)', file=sys.stderr)
        return 2
    subject = argv[1]
    if len(argv) >= 3:
        body = argv[2]
    else:
        body = sys.stdin.read() if not sys.stdin.isatty() else ""
    return 0 if send(subject, body) else 1


if __name__ == "__main__":
    sys.exit(_main(sys.argv))