H3ConuMS-v2/backend/app/services/wechat_service.py

"""企业微信 (WeChat Work) 应用消息服务"""
import time
import hashlib
import base64
import socket
import struct
import urllib.parse
import xml.etree.ElementTree as ET
import logging
import requests
from Crypto.Cipher import AES
from app.core.config import settings

logger = logging.getLogger(__name__)


def _get_db_setting(key: str, default: str = "") -> str:
    """从数据库读取 SystemSetting，失败时返回 default"""
    try:
        from app.core.database import SessionLocal
        from app.models.setting import SystemSetting
        db = SessionLocal()
        try:
            row = db.query(SystemSetting).filter_by(key=key).first()
            return row.value if row else default
        finally:
            db.close()
    except Exception:
        return default


class WeChatService:
    """企业微信应用消息服务"""

    def __init__(self):
        # 优先读环境变量，fallback 到数据库配置
        self.corpid = settings.WECHAT_CORPID or _get_db_setting("wechat_corpid")
        self.corpsecret = settings.WECHAT_CORPSECRET or _get_db_setting("wechat_corpsecret")
        self.agentid = settings.WECHAT_AGENTID or _get_db_setting("wechat_agentid")
        self.token = settings.WECHAT_TOKEN
        self.encoding_aes_key = settings.WECHAT_ENCODING_AES_KEY
        self.use_proxy = settings.WECHAT_USE_PROXY
        self.proxy_api_url = settings.WECHAT_PROXY_API_URL
        self._access_token = None
        self._token_expires_at = 0

    # ── access_token 管理 ───────────────────────────────────────────────────

    def get_access_token(self) -> str | None:
        """获取企业微信 access_token，自动缓存和续期"""
        now = time.time()
        if self._access_token and now < self._token_expires_at:
            return self._access_token

        if not self.corpid or not self.corpsecret:
            logger.warning("WECHAT_CORPID 或 WECHAT_CORPSECRET 未配置")
            return None

        try:
            if self.use_proxy:
                url = f"{self.proxy_api_url}/cgi-bin/gettoken?corpid={self.corpid}&corpsecret={self.corpsecret}"
            else:
                url = f"https://qyapi.weixin.qq.com/cgi-bin/gettoken?corpid={self.corpid}&corpsecret={self.corpsecret}"

            resp = requests.get(url, timeout=10)
            result = resp.json()

            if result.get("errcode") == 0:
                self._access_token = result.get("access_token")
                expires_in = result.get("expires_in", 7200) - 300  # 提前5分钟过期
                self._token_expires_at = now + expires_in
                logger.info(f"企业微信 access_token 获取成功，过期时间: {time.strftime('%H:%M:%S', time.localtime(self._token_expires_at))}")
                return self._access_token
            elif result.get("errcode") == 60020 and not self.use_proxy:
                logger.info("IP受限，尝试使用代理获取 access_token")
                self.use_proxy = True
                return self.get_access_token()
            else:
                logger.error(f"获取 access_token 失败: {result}")
                return None
        except Exception as e:
            logger.error(f"获取 access_token 异常: {e}")
            return None

    # ── 消息分条 ────────────────────────────────────────────────────────────

    @staticmethod
    def _split_long_message(content: str, max_chars: int = 1800) -> list[str]:
        """将长消息按换行边界拆分为多条，避免企微截断"""
        if len(content) <= max_chars:
            return [content]
        chunks = []
        lines = content.split('\n')
        current = ''
        for line in lines:
            if len(current) + len(line) + 1 > max_chars and current:
                chunks.append(current.strip())
                current = line
            else:
                current += ('\n' + line) if current else line
        if current.strip():
            chunks.append(current.strip())
        return chunks

    # ── 发送消息 ────────────────────────────────────────────────────────────

    def _send_markdown_single(self, content: str, to_user: str) -> bool:
        """发送单条 Markdown 消息（内部方法）"""
        access_token = self.get_access_token()
        if not access_token:
            return False

        data = {
            "touser": to_user, "toparty": "", "totag": "",
            "msgtype": "markdown",
            "agentid": int(self.agentid),
            "markdown": {"content": content}
        }

        if self.use_proxy:
            url = f"{self.proxy_api_url}/cgi-bin/message/send?access_token={access_token}"
        else:
            url = f"https://qyapi.weixin.qq.com/cgi-bin/message/send?access_token={access_token}"

        resp = requests.post(url, json=data, timeout=15)
        result = resp.json()
        errcode = result.get("errcode", -1)
        if errcode == 0:
            return True
        elif errcode == 40014:
            self._access_token = None
            self._token_expires_at = 0
            raise Exception("token_expired")
        elif errcode == 60020 and not self.use_proxy:
            self.use_proxy = True
            raise Exception("ip_restricted")
        else:
            logger.warning(f"企业微信消息发送失败: {result.get('errmsg')} (errcode={errcode})")
            return False

    def send_markdown(self, content: str, to_user: str = "@all") -> bool:
        """发送 Markdown 消息，自动分条"""
        if not self.corpid or not self.corpsecret or not self.agentid:
            logger.warning("企业微信未配置，跳过发送")
            return False

        chunks = self._split_long_message(content)
        success = True
        for i, chunk in enumerate(chunks):
            prefix = f"({i+1}/{len(chunks)})\n" if len(chunks) > 1 else ""
            for attempt in range(3):
                try:
                    ok = self._send_markdown_single(prefix + chunk, to_user)
                    if ok:
                        break
                    if attempt < 2:
                        time.sleep(1)
                except Exception as e:
                    logger.warning(f"发送分片 {i+1}/{len(chunks)} 异常: {e}")
                    if attempt < 2:
                        time.sleep(1)
                        continue
                    success = False
            if i < len(chunks) - 1:
                time.sleep(0.5)  # 避免频率限制
        return success


    # ── 发送文本消息 ────────────────────────────────────────────────────────

    def send_text_message(self, content: str, to_user: str = "@all") -> bool:
        """通过应用消息 API 发送文本消息，自动分条"""
        if not self.corpid or not self.corpsecret or not self.agentid:
            return False

        chunks = self._split_long_message(content, max_chars=1800)
        success = True
        for i, chunk in enumerate(chunks):
            prefix = f"({i+1}/{len(chunks)})\n" if len(chunks) > 1 else ""
            try:
                access_token = self.get_access_token()
                if not access_token:
                    return False
                data = {
                    "touser": to_user, "toparty": "", "totag": "",
                    "msgtype": "text",
                    "agentid": int(self.agentid),
                    "text": {"content": prefix + chunk}
                }
                url = f"{self.proxy_api_url}/cgi-bin/message/send?access_token={access_token}" if self.use_proxy else f"https://qyapi.weixin.qq.com/cgi-bin/message/send?access_token={access_token}"
                resp = requests.post(url, json=data, timeout=15)
                if resp.json().get("errcode") != 0:
                    success = False
            except Exception as e:
                logger.error(f"发送文本消息分片 {i+1}/{len(chunks)} 失败: {e}")
                success = False
            if i < len(chunks) - 1:
                time.sleep(0.5)
        return success

    # ── 菜单管理 ────────────────────────────────────────────────────────────

    def create_menu(self, menu_data: dict) -> bool:
        """创建/更新企业微信应用菜单"""
        try:
            access_token = self.get_access_token()
            if not access_token:
                return False
            if self.use_proxy:
                url = f"{self.proxy_api_url}/cgi-bin/menu/create?access_token={access_token}&agentid={self.agentid}"
            else:
                url = f"https://qyapi.weixin.qq.com/cgi-bin/menu/create?access_token={access_token}&agentid={self.agentid}"
            resp = requests.post(url, json=menu_data, timeout=15)
            result = resp.json()
            if result.get("errcode") == 0:
                logger.info("企业微信菜单创建成功")
                return True
            logger.warning(f"创建菜单失败: {result}")
            return False
        except Exception as e:
            logger.error(f"创建菜单异常: {e}")
            return False

    # ── 获取用户信息 ─────────────────────────────────────────────────────────

    def get_user_info(self, userid: str) -> dict:
        """获取企业微信用户信息"""
        try:
            access_token = self.get_access_token()
            if not access_token:
                return {"errcode": -1, "errmsg": "无 access_token"}
            if self.use_proxy:
                url = f"{self.proxy_api_url}/cgi-bin/user/get?access_token={access_token}&userid={userid}"
            else:
                url = f"https://qyapi.weixin.qq.com/cgi-bin/user/get?access_token={access_token}&userid={userid}"
            resp = requests.get(url, timeout=10)
            return resp.json()
        except Exception as e:
            logger.error(f"获取用户信息失败: {e}")
            return {"errcode": -1, "errmsg": str(e)}

    # ── URL 验证（GET 回调）──────────────────────────────────────────────────

    def verify_url(self, msg_signature: str, timestamp: str, nonce: str, echostr: str) -> str | None:
        """验证企业微信回调 URL"""
        try:
            echostr = urllib.parse.unquote(echostr)
            temp_list = [self.token, timestamp, nonce, echostr]
            temp_list.sort()
            temp_str = ''.join(temp_list)
            hash_str = hashlib.sha1(temp_str.encode('utf-8')).hexdigest()

            if hash_str != msg_signature:
                logger.error(f"URL验证签名不匹配: expected={msg_signature}, got={hash_str}")
                return None

            if self.encoding_aes_key:
                return self._decrypt_echostr(echostr)
            return echostr
        except Exception as e:
            logger.error(f"URL验证异常: {e}")
            return None

    def _decrypt_echostr(self, echostr: str) -> str | None:
        """解密 echostr"""
        try:
            aes_key = base64.b64decode(self.encoding_aes_key + '=')
            encrypted = base64.b64decode(echostr)
            cipher = AES.new(aes_key, AES.MODE_CBC, aes_key[:16])
            decrypted = cipher.decrypt(encrypted)
            decrypted = decrypted[:-decrypted[-1]]  # PKCS7 unpad
            content = decrypted[16:]
            xml_len = socket.ntohl(struct.unpack("I", content[:4])[0])
            xml_content = content[4:xml_len + 4]
            received_id = content[xml_len + 4:].decode('utf-8')
            if received_id != self.corpid:
                logger.error(f"企业ID验证失败: {received_id} != {self.corpid}")
                return None
            return xml_content.decode('utf-8')
        except Exception as e:
            logger.error(f"解密echostr失败: {e}")
            return None

    # ── 消息解密（POST 回调）─────────────────────────────────────────────────

    def parse_message(self, xml_data: bytes) -> dict | None:
        """解析企业微信回调的加密 XML 消息"""
        try:
            root = ET.fromstring(xml_data)
            msg = {child.tag: child.text for child in root}

            if 'Encrypt' in msg:
                decrypted = self._decrypt_message(msg['Encrypt'])
                decrypted_root = ET.fromstring(decrypted)
                msg = {child.tag: child.text for child in decrypted_root}

            return msg
        except Exception as e:
            logger.error(f"解析消息失败: {e}")
            return None

    def _decrypt_message(self, encrypted_msg: str) -> str:
        """解密企业微信推送的消息"""
        aes_key = base64.b64decode(self.encoding_aes_key + '=')
        encrypted = base64.b64decode(encrypted_msg)
        cipher = AES.new(aes_key, AES.MODE_CBC, aes_key[:16])
        decrypted = cipher.decrypt(encrypted)
        decrypted = decrypted[:-decrypted[-1]]  # PKCS7 unpad
        content = decrypted[16:]
        xml_len = socket.ntohl(struct.unpack("I", content[:4])[0])
        xml_content = content[4:xml_len + 4]
        received_id = content[xml_len + 4:].decode('utf-8')
        if received_id != self.corpid:
            raise Exception(f"企业ID验证失败: {received_id} != {self.corpid}")
        return xml_content.decode('utf-8')


# 模块级单例和便捷函数
_service: WeChatService | None = None


def get_wechat_service() -> WeChatService:
    global _service
    if _service is None:
        _service = WeChatService()
    return _service


def send_wechat_markdown(content: str) -> bool:
    """便捷函数：发送企业微信 Markdown 消息"""
    return get_wechat_service().send_markdown(content)