outlook.py

#!/usr/bin/env python3
"""Outlook 整理工具 — Microsoft Graph API

用法:
  python3 outlook.py /me/messages?$top=5                # 直接调 Graph 任意路径(兼容旧用法)
  python3 outlook.py whoami                             # 账户信息
  python3 outlook.py folders                            # 列出所有邮件文件夹
  python3 outlook.py stats [--folder 收件箱] [--limit 5000] [--top 30]   # 按发件域名统计
  python3 outlook.py organize --dry-run                 # 预览整理方案(不移动)
  python3 outlook.py organize [--confirm]               # 执行整理(按规则移动到文件夹)
  python3 outlook.py move <message_id> <文件夹名>        # 移动单封邮件
  python3 outlook.py rules                              # 显示当前整理规则

权限: Mail.Read / Mail.ReadWrite / User.Read / offline_access
Token 自动刷新, 存储于脚本同目录 .outlook-token.json (已被 .gitignore 忽略)
整理规则从同目录 rules.json 加载(未提供时用内置通用示例)
"""
import json, os, sys, time, argparse, urllib.request, urllib.error, urllib.parse, threading
from collections import Counter
from concurrent.futures import ThreadPoolExecutor, as_completed

HOME = os.path.dirname(os.path.abspath(__file__))
TOKEN_PATH = os.path.join(HOME, ".outlook-token.json")
CLIENT_ID = "14d82eec-204b-4c2f-b7e8-296a70dab67e"  # Microsoft Graph PowerShell
TOKEN_URL = "https://login.microsoftonline.com/consumers/oauth2/v2.0/token"
GRAPH = "https://graph.microsoft.com/v1.0"
SCOPES = "https://graph.microsoft.com/Mail.Read https://graph.microsoft.com/Mail.ReadWrite https://graph.microsoft.com/Mail.Send https://graph.microsoft.com/User.Read offline_access"

# ── 整理规则 ────────────────────────────────────────────────────────────
# 优先从同目录 rules.json 加载: [{"keyword": "github.com", "folder": "GitHub"}, ...]
# 域名匹配为子串匹配(如 "hsbc" 可匹配 notification.hsbc.com.hk)
# 未提供 rules.json 时使用下列通用示例
DEFAULT_RULES = [
    {"keyword": "github.com",       "folder": "GitHub"},
    {"keyword": "gitlab.com",      "folder": "GitHub"},
    {"keyword": "jetbrains.com",   "folder": "JetBrains"},
    {"keyword": "stackoverflow",   "folder": "GitHub"},
    {"keyword": "linkedin.com",    "folder": "社交通知"},
    {"keyword": "facebook",        "folder": "社交通知"},
    {"keyword": "twitter.com",     "folder": "社交通知"},
    {"keyword": "instagram.com",   "folder": "社交通知"},
    {"keyword": "newsletter",      "folder": "订阅与优惠"},
    {"keyword": "marketing",       "folder": "订阅与优惠"},
    {"keyword": "promo",           "folder": "订阅与优惠"},
]


def load_rules():
    """加载整理规则: rules.json(同目录) > 内置示例。返回 [(keyword, folder)]。"""
    rules_path = os.path.join(HOME, "rules.json")
    if os.path.exists(rules_path):
        with open(rules_path) as f:
            data = json.load(f)
        return [(r["keyword"], r["folder"]) for r in data]
    return [(r["keyword"], r["folder"]) for r in DEFAULT_RULES]


RULES = load_rules()


class TokenError(Exception):
    pass


def load_token():
    with open(TOKEN_PATH) as f:
        return json.load(f)


PENDING_PATH = os.path.join(HOME, ".auth-pending.json")


def _device_code_request():
    """发起 device code 请求, 返回响应 dict。"""
    import urllib.parse
    body = urllib.parse.urlencode({
        "client_id": CLIENT_ID,
        "scope": SCOPES,
    }).encode()
    req = urllib.request.Request("https://login.microsoftonline.com/consumers/oauth2/v2.0/devicecode", data=body)
    with urllib.request.urlopen(req, timeout=30) as r:
        return json.load(r)


def _poll_once(device_code):
    """单次轮询 token 端点。返回 (status, detail)：
    status: success(带 token) | pending | declined | expired | error"""
    import urllib.parse
    body = urllib.parse.urlencode({
        "grant_type": "urn:ietf:params:oauth:grant-type:device_code",
        "client_id": CLIENT_ID,
        "device_code": device_code,
    }).encode()
    try:
        req = urllib.request.Request("https://login.microsoftonline.com/consumers/oauth2/v2.0/token", data=body)
        with urllib.request.urlopen(req, timeout=30) as r:
            tok = json.load(r)
            tok["acquired_at"] = time.time()
            save_token(tok)
            return "success", tok
    except urllib.error.HTTPError as e:
        err = json.load(e)
        code = err.get("error")
        if code in ("authorization_pending",):
            return "pending", None
        if code == "slow_down":
            return "slow_down", None
        if code == "authorization_declined":
            return "declined", None
        if code == "expired_token":
            return "expired", None
        return "error", err.get("error_description", code)
    except Exception as e:
        return "error", f"{type(e).__name__}: {str(e)[:100]}"


def cmd_authorize(args=None):
    """OAuth 2.0 device code 授权。

    AI 友好两阶段用法:
      1. python3 outlook.py authorize --json    # 发起, 立即返回 user_code/URL, 不阻塞
      2. python3 outlook.py authorize --poll    # 轮询一次, 输出 JSON 状态
         (AI 可循环调用 --poll, 直到 status=success)
    人用交互模式: python3 outlook.py authorize  # 发起并阻塞轮询直到完成
    """
    if getattr(args, "poll", False):
        # ── 阶段 2: 轮询 ──
        try:
            with open(PENDING_PATH) as f:
                pending = json.load(f)
        except (FileNotFoundError, json.JSONDecodeError):
            print(json.dumps({"status": "error", "message": "没有待处理的授权，请先运行 authorize --json"},
                             ensure_ascii=False))
            return 1
        status, detail = _poll_once(pending["device_code"])
        if status == "success":
            try:
                os.remove(PENDING_PATH)
            except OSError:
                pass
            print(json.dumps({"status": "success",
                              "message": "授权成功，token 已保存",
                              "expires_in": detail.get("expires_in"),
                              "scope": detail.get("scope")}, ensure_ascii=False))
            return 0
        if status == "pending":
            print(json.dumps({"status": "pending",
                              "message": "用户尚未完成授权，请稍后重试 --poll"}, ensure_ascii=False))
            return 2
        if status == "slow_down":
            print(json.dumps({"status": "pending", "message": "轮询过频，请稍候再试"}, ensure_ascii=False))
            return 2
        if status == "declined":
            os.remove(PENDING_PATH)
            print(json.dumps({"status": "declined", "message": "用户拒绝了授权"}, ensure_ascii=False))
            return 1
        if status == "expired":
            os.remove(PENDING_PATH)
            print(json.dumps({"status": "expired", "message": "授权码已过期，请重新 authorize --json"}, ensure_ascii=False))
            return 1
        print(json.dumps({"status": "error", "message": detail}, ensure_ascii=False))
        return 1

    # ── 发起阶段 ──
    d = _device_code_request()
    if getattr(args, "json", False):
        # AI 模式: 只发起, 立即返回结构化结果
        with open(PENDING_PATH, "w") as f:
            json.dump({"device_code": d["device_code"],
                       "expires_at": time.time() + d.get("expires_in", 900)}, f)
        print(json.dumps({
            "status": "pending_user_action",
            "verification_uri": d["verification_uri"],
            "user_code": d["user_code"],
            "expires_in": d.get("expires_in"),
            "poll_hint": "python3 outlook.py authorize --poll",
        }, ensure_ascii=False))
        return 0

    # 人用交互模式: 发起 + 阻塞轮询
    print(f"\n请打开浏览器访问: {d['verification_uri']}")
    print(f"输入代码: {d['user_code']}")
    print("登录并确认授权，本程序会自动继续...\n")
    deadline = time.time() + d.get("expires_in", 900) - 30
    while time.time() < deadline:
        status, detail = _poll_once(d["device_code"])
        if status == "success":
            print("✓ 授权成功，token 已保存到", TOKEN_PATH)
            return 0
        if status == "pending" or status == "slow_down":
            time.sleep(5)
            continue
        if status == "declined":
            print("✗ 授权被拒绝")
            return 1
        if status == "expired":
            print("✗ 授权码已过期，请重试")
            return 1
        print("✗ 授权错误:", detail)
        return 1
    print("✗ 授权超时，请重新运行")
    return 1


def save_token(tok):
    tmp = TOKEN_PATH + ".tmp"
    with open(tmp, "w") as f:
        f.write(json.dumps(tok))
    os.replace(tmp, TOKEN_PATH)
    try:
        os.chmod(TOKEN_PATH, 0o600)
    except OSError:
        pass


def refresh_token(tok):
    body = urllib.parse.urlencode({
        "grant_type": "refresh_token",
        "client_id": CLIENT_ID,
        "refresh_token": tok["refresh_token"],
        "scope": SCOPES,
    }).encode()
    req = urllib.request.Request(TOKEN_URL, data=body)
    with urllib.request.urlopen(req, timeout=30) as r:
        new = json.load(r)
    new["acquired_at"] = time.time()
    if "refresh_token" not in new:
        new["refresh_token"] = tok["refresh_token"]
    save_token(new)
    return new


_token_lock = threading.Lock()


def get_token():
    with _token_lock:
        tok = load_token()
        if time.time() - tok.get("acquired_at", 0) > tok.get("expires_in", 3600) - 300:
            return refresh_token(tok)
        return tok


def graph(path, method="GET", body=None, retries=3):
    """调用 Graph v1.0。path 可以是 /me/... 或完整 URL(nextLink)。"""
    for attempt in range(retries):
        tok = get_token()
        url = path if path.startswith("http") else GRAPH + path
        req = urllib.request.Request(url, method=method)
        req.add_header("Authorization", "Bearer " + tok["access_token"])
        if body is not None:
            req.add_header("Content-Type", "application/json")
            req.data = json.dumps(body).encode()
        try:
            with urllib.request.urlopen(req, timeout=30) as r:
                raw = r.read().decode()
                return r.status, (json.loads(raw) if raw else {})
        except urllib.error.HTTPError as e:
            if e.code == 429:  # 限流, 退避重试
                time.sleep(2 ** attempt * 2)
                continue
            try:
                return e.code, json.loads(e.read().decode())
            except Exception:
                return e.code, {"error": str(e)}
        except Exception as e:
            if attempt < retries - 1:
                time.sleep(2)
                continue
            raise
    raise RuntimeError("Graph 调用重试失败")


def paginate(path, select=None, limit=5000):
    """分页拉取, yield 每条记录; select 为逗号分隔字段名。"""
    q = "?" + urllib.parse.urlencode({"$top": 50, "$select": select}) if select else "?$top=50"
    url = path + q
    count = 0
    while url and count < limit:
        status, data = graph(url)
        if status >= 400:
            raise RuntimeError(f"Graph 错误 {status}: {data}")
        for item in data.get("value", []):
            yield item
            count += 1
            if count >= limit:
                return
        url = data.get("@odata.nextLink")
        time.sleep(0.1)


# ── 子命令实现 ─────────────────────────────────────────────────────────────

def cmd_whoami(args=None):
    status, me = graph("/me")
    if status >= 400:
        print("错误:", me)
        return 1
    print(f"账户: {me.get('userPrincipalName')}  ({me.get('displayName')})")
    return 0


def cmd_folders(args=None):
    status, data = graph("/me/mailFolders?$top=200&$select=displayName,totalItemCount,unreadItemCount")
    if status >= 400:
        print("错误:", data)
        return 1
    for f in sorted(data.get("value", []), key=lambda x: -x.get("totalItemCount", 0)):
        print(f"  {f['displayName']:<20} 总数={f.get('totalItemCount', 0):<6} 未读={f.get('unreadItemCount', 0)}")
    return 0


def cmd_stats(args):
    folders = get_folders()
    folder = next((f for f in folders if f["displayName"] == args.folder), None)
    if not folder:
        print(f"未找到文件夹: {args.folder}")
        return 1
    domain_count = Counter()
    total = 0
    for m in paginate(f"/me/mailFolders/{folder['id']}/messages", select="from", limit=args.limit):
        addr = (m.get("from") or {}).get("emailAddress", {}).get("address", "?")
        domain = addr.split("@")[-1] if "@" in addr else "?"
        domain_count[domain] += 1
        total += 1
    print(f"共统计 [{args.folder}] {total} 封\n")
    for d, c in domain_count.most_common(args.top):
        print(f"  {d:<45} {c}")
    return 0


def get_folders():
    status, data = graph("/me/mailFolders?$top=200&$select=displayName,id,totalItemCount")
    if status >= 400:
        raise RuntimeError(f"获取文件夹失败: {data}")
    return data.get("value", [])


def match_rule(domain):
    for keyword, folder in RULES:
        if keyword in domain:
            return folder
    return None


def cmd_rules(args=None):
    print("当前整理规则 (域名关键字 -> 文件夹):")
    for keyword, folder in RULES:
        print(f"  {keyword:<25} -> {folder}")
    return 0


def cmd_organize(args):
    """按规则把邮件移动到目标文件夹。--dry-run 只预览。"""
    folders = get_folders()
    by_name = {f["displayName"]: f for f in folders}
    inbox = by_name.get("收件箱") or by_name.get("Inbox")
    if not inbox:
        print("未找到收件箱")
        return 1

    # 目标文件夹存在性检查
    missing = set(f for _, f in RULES) - set(by_name)
    if missing:
        print(f"⚠ 以下目标文件夹不存在, 相关规则将跳过: {missing}")

    plan = Counter()       # folder -> 数量
    skipped = Counter()    # 跳过原因
    moves = []             # (message_id, subject, 源域, 目标文件夹)
    for m in paginate(f"/me/mailFolders/{inbox['id']}/messages", select="from,subject,id", limit=args.limit):
        addr = (m.get("from") or {}).get("emailAddress", {}).get("address", "?")
        domain = addr.split("@")[-1] if "@" in addr else "?"
        target = match_rule(domain)
        if not target:
            skipped["无匹配规则"] += 1
            continue
        if target not in by_name:
            skipped[f"文件夹不存在:{target}"] += 1
            continue
        plan[target] += 1
        moves.append((m["id"], m.get("subject", "")[:50], domain, target))

    print(f"\n整理预览 ({'DRY-RUN, 未执行任何移动' if args.dry_run else '实际执行'}):")
    print(f"收件箱扫描: 匹配 {len(moves)} 封, 未匹配 {skipped.get('无匹配规则', 0)} 封")
    for folder, n in plan.most_common():
        print(f"  -> {folder:<12} {n} 封")
    if skipped:
        print("跳过:", dict(skipped))

    if args.dry_run:
        print("\n示例(每目标前 5 封):")
        shown = Counter()
        for mid, subj, domain, target in moves:
            if shown[target] >= 5:
                continue
            shown[target] += 1
            print(f"  [{target}] {subj}  ({domain})")
        print(f"\n共 {len(moves)} 封待移动。执行请加 --confirm")
        return 0

    if not args.confirm:
        print("\n这是实际执行! 确认无误请加 --confirm 重新运行")
        return 0

    # 实际执行 (多线程并发移动)
    ok = fail = 0
    log_lock = threading.Lock()
    workers = args.workers

    def do_move(item):
        mid, subj, domain, target = item
        status, res = graph(f"/me/messages/{mid}/move", method="POST",
                            body={"destinationId": by_name[target]["id"]})
        return (status < 300, subj, target, res)

    t0 = time.time()
    with ThreadPoolExecutor(max_workers=workers) as ex:
        futs = [ex.submit(do_move, it) for it in moves]
        for fut in as_completed(futs):
            ok_flag, subj, target, res = fut.result()
            with log_lock:
                if ok_flag:
                    ok += 1
                    print(f"  ✓ -> {target}: {subj}")
                else:
                    fail += 1
                    print(f"  ✗ -> {target}: {subj} ({res.get('error', {}).get('message', res)})")
    print(f"\n完成: 成功 {ok}, 失败 {fail}, 用时 {int(time.time()-t0)} 秒")
    return 0


def cmd_move(args):
    folders = get_folders()
    target = next((f for f in folders if f["displayName"] == args.folder), None)
    if not target:
        print(f"未找到文件夹: {args.folder}")
        return 1
    status, res = graph(f"/me/messages/{args.message_id}/move", method="POST",
                        body={"destinationId": target["id"]})
    if status < 300:
        print(f"✓ 已移动到 {args.folder}")
        return 0
    print("失败:", res)
    return 1


def main():
    args = sys.argv[1:]
    if not args:
        print(__doc__)
        return 0
    if args[0].startswith("/") or args[0].startswith("http"):
        status, data = graph(args[0])
        print("HTTP", status)
        print(json.dumps(data, indent=1, ensure_ascii=False)[:8000])
        return 0

    p = argparse.ArgumentParser(prog="outlook.py")
    sub = p.add_subparsers(dest="cmd")
    sub.add_parser("whoami")
    a = sub.add_parser("authorize")
    a.add_argument("--json", action="store_true", help="AI 模式: 发起授权, 输出 JSON, 不阻塞")
    a.add_argument("--poll", action="store_true", help="AI 模式: 轮询一次授权状态, 输出 JSON")
    sub.add_parser("folders")
    sub.add_parser("rules")
    s = sub.add_parser("stats")
    s.add_argument("--folder", default="收件箱")
    s.add_argument("--limit", type=int, default=5000)
    s.add_argument("--top", type=int, default=30)
    o = sub.add_parser("organize")
    o.add_argument("--dry-run", action="store_true")
    o.add_argument("--confirm", action="store_true")
    o.add_argument("--limit", type=int, default=5000)
    o.add_argument("--workers", type=int, default=8, help="并发移动线程数(默认 8)")
    m = sub.add_parser("move")
    m.add_argument("message_id")
    m.add_argument("folder")
    opts = p.parse_args(args)

    handlers = {
        "whoami": cmd_whoami, "authorize": cmd_authorize, "folders": cmd_folders,
        "rules": cmd_rules,
        "stats": cmd_stats, "organize": cmd_organize, "move": cmd_move,
    }
    if opts.cmd in handlers:
        try:
            return handlers[opts.cmd](opts)
        except FileNotFoundError:
            print("未找到授权 token (.outlook-token.json)。")
            print("请先运行: python3 outlook.py authorize  完成 OAuth 授权。")
            return 1
        except TokenError as e:
            print("token 错误:", e)
            return 1
    print(__doc__)
    return 0


if __name__ == "__main__":
    sys.exit(main())