yidaima_tools/wechat_publisher.py

from __future__ import annotations

import json
import time
from pathlib import Path
from typing import Optional

import requests


class WeChatPublisher:
    """微信公众号草稿箱发布器"""

    def __init__(self, appid: str, appsecret: str):
        self.appid = appid
        self.appsecret = appsecret
        self.access_token: Optional[str] = None
        self.token_expires_at: float = 0

    def _get_access_token(self) -> str:
        """获取或刷新 access_token"""
        if self.access_token and time.time() < self.token_expires_at - 300:
            return self.access_token

        url = "https://api.weixin.qq.com/cgi-bin/token"
        params = {
            "grant_type": "client_credential",
            "appid": self.appid,
            "secret": self.appsecret
        }

        response = requests.get(url, params=params, timeout=30)
        data = response.json()

        if "access_token" not in data:
            errcode = data.get("errcode")
            errmsg = data.get("errmsg", "")

            if errcode == 40164 or "not in whitelist" in errmsg.lower() or "ip" in errmsg.lower():
                import re
                ip_match = re.search(r'(\d+\.\d+\.\d+\.\d+)', errmsg)
                current_ip = ip_match.group(1) if ip_match else "当前IP"

                raise RuntimeError(
                    f"获取 access_token 失败: IP 不在白名单中\n"
                    f"当前 IP: {current_ip}\n"
                    f"解决方法: 登录微信公众平台 → 开发 → 基本配置 → IP白名单 → 添加以上IP地址"
                )
            elif errcode == 40013:
                raise RuntimeError(f"获取 access_token 失败: AppID 无效")
            elif errcode == 40125:
                raise RuntimeError(f"获取 access_token 失败: AppSecret 无效")
            else:
                raise RuntimeError(f"获取 access_token 失败: {data}")

        self.access_token = data["access_token"]
        self.token_expires_at = time.time() + data.get("expires_in", 7200)

        return self.access_token

    def upload_thumb(self, image_path: str) -> str:
        """
        上传封面图片获取 thumb_media_id

        Args:
            image_path: 图片文件路径

        Returns:
            thumb_media_id
        """
        token = self._get_access_token()
        url = f"https://api.weixin.qq.com/cgi-bin/material/add_material?access_token={token}&type=image"

        with open(image_path, 'rb') as f:
            files = {'media': f}
            response = requests.post(url, files=files, timeout=30)

        result = response.json()

        if result.get("errcode"):
            raise RuntimeError(f"上传封面图片失败: {result}")

        media_id = result.get("media_id")
        if not media_id:
            raise RuntimeError(f"上传封面图片失败，未返回 media_id: {result}")

        return media_id

    def upload_thumb_from_url(self, image_url: str) -> str:
        """
        从URL下载图片并上传获取 thumb_media_id

        Args:
            image_url: 图片URL

        Returns:
            thumb_media_id
        """
        import tempfile

        # 下载图片
        response = requests.get(image_url, timeout=30)
        response.raise_for_status()

        # 保存到临时文件
        suffix = Path(image_url).suffix or '.jpg'
        with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as f:
            f.write(response.content)
            temp_path = f.name

        try:
            return self.upload_thumb(temp_path)
        finally:
            # 清理临时文件
            Path(temp_path).unlink(missing_ok=True)

    def add_draft(
        self,
        title: str,
        content: str,
        thumb_media_id: str,
        author: str = "",
        digest: str = "",
        content_source_url: str = ""
    ) -> dict:
        """
        添加草稿到公众号草稿箱

        Args:
            title: 文章标题
            content: 文章内容（HTML格式）
            thumb_media_id: 封面图片的 media_id（必需）
            author: 作者
            digest: 摘要
            content_source_url: 阅读原文链接

        Returns:
            API 返回的原始数据
        """
        token = self._get_access_token()
        url = f"https://api.weixin.qq.com/cgi-bin/draft/add?access_token={token}"

        # 构建图文消息
        article = {
            "title": title,
            "content": content,
            "thumb_media_id": thumb_media_id,
            "show_cover_pic": 1,
            "need_open_comment": 0,
            "only_fans_can_comment": 0
        }

        if author:
            article["author"] = author
        if digest:
            article["digest"] = digest
        if content_source_url:
            article["content_source_url"] = content_source_url

        payload = {"articles": [article]}

        # 调试信息
        payload_str = json.dumps(payload, ensure_ascii=False)
        print(f"[WeChat] 发送内容长度: {len(payload_str)} 字符")

        response = requests.post(
            url,
            data=payload_str.encode('utf-8'),
            headers={"Content-Type": "application/json; charset=utf-8"},
            timeout=30
        )

        result = response.json()

        if result.get("errcode", 0) != 0:
            raise RuntimeError(f"添加草稿失败: {result}")

        return result


def publish_to_wechat(
    title: str,
    content: str,
    appid: str,
    appsecret: str,
    thumb_image_path: Optional[str] = None,
    thumb_image_url: Optional[str] = None,
    author: str = ""
) -> bool:
    """
    便捷函数：发布文章到微信公众号草稿箱

    Args:
        title: 文章标题
        content: 文章内容（HTML格式）
        appid: 微信公众号 AppID
        appsecret: 微信公众号 AppSecret
        thumb_image_path: 封面图片本地路径（可选，与 thumb_image_url 二选一）
        thumb_image_url: 封面图片URL（可选，与 thumb_image_path 二选一）
        author: 作者

    Returns:
        是否成功
    """
    try:
        publisher = WeChatPublisher(appid, appsecret)

        # 获取封面图片的 media_id
        if thumb_image_path:
            print(f"[WeChat] 正在上传封面图片: {thumb_image_path}")
            thumb_media_id = publisher.upload_thumb(thumb_image_path)
        elif thumb_image_url:
            print(f"[WeChat] 正在下载并上传封面图片: {thumb_image_url}")
            thumb_media_id = publisher.upload_thumb_from_url(thumb_image_url)
        else:
            raise RuntimeError("必须提供封面图片路径或URL")

        print(f"[WeChat] 封面图片上传成功，media_id: {thumb_media_id}")

        # 添加草稿
        result = publisher.add_draft(
            title=title,
            content=content,
            thumb_media_id=thumb_media_id,
            author=author
        )

        print(f"[WeChat] 草稿添加成功，media_id: {result.get('media_id')}")
        return True

    except Exception as e:
        print(f"[WeChat] 发布失败: {e}")
        return False


def main():
    """测试发布功能"""
    from config_loader import get_config

    config = get_config()
    appid = config.wechat_appid
    appsecret = config.wechat_appsecret

    if not appid or not appsecret:
        print("[ERROR] 未配置微信公众号，请编辑 config.yaml 文件:")
        print("  wechat:")
        print("    appid: \"your_appid\"")
        print("    appsecret: \"your_appsecret\"")
        return

    # 测试内容
    title = "测试文章"
    content = "<p>这是一篇测试文章</p>"

    # 使用默认封面图片
    default_thumb_path = r"c:\Users\南音\Desktop\yidaima\assets\img\bg.jpg"

    success = publish_to_wechat(
        title=title,
        content=content,
        appid=appid,
        appsecret=appsecret,
        thumb_image_path=default_thumb_path
    )

    if success:
        print("[INFO] 发布成功！")
    else:
        print("[ERROR] 发布失败！")


if __name__ == "__main__":
    main()