Simon
377 lines
14 KiB
Python
377 lines
14 KiB
Python
# =============================================================================
|
||
# 企微IT智能服务台 — 企微消息 AES 加解密工具
|
||
# =============================================================================
|
||
# 说明：实现企微回调消息的加解密和签名验证
|
||
# 参考企微官方加解密库逻辑，使用 Python cryptography 库重写
|
||
# 加密模式：AES-CBC-256，PKCS7 填充
|
||
# 签名算法：SHA1(sort(token, timestamp, nonce, encrypt))
|
||
# =============================================================================
|
||
|
||
import base64
|
||
import hashlib
|
||
import json
|
||
import logging
|
||
import secrets
|
||
import string
|
||
import struct
|
||
import time
|
||
import xml.etree.ElementTree as ET
|
||
from typing import Dict, Optional, Tuple
|
||
|
||
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
|
||
from cryptography.hazmat.backends import default_backend
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class WecomCrypto:
|
||
"""企微消息加解密工具类。
|
||
|
||
实现企微回调消息的完整加解密流程：
|
||
1. 签名验证：SHA1(sort(token, timestamp, nonce, encrypt))
|
||
2. AES 解密：CBC 模式，PKCS7 去填充
|
||
3. AES 加密：CBC 模式，PKCS7 填充
|
||
|
||
Attributes:
|
||
token: 企微回调配置的 Token
|
||
aes_key: 解码后的 AES 密钥（32 字节）
|
||
corp_id: 企业ID
|
||
"""
|
||
|
||
def __init__(self, token: str, encoding_aes_key: str, corp_id: str):
|
||
"""初始化加解密工具。
|
||
|
||
Args:
|
||
token: 企微回调配置的 Token（用于签名验证）
|
||
encoding_aes_key: 企微回调配置的 EncodingAESKey（43 位 Base64 字符串）
|
||
corp_id: 企业ID（用于消息体中的校验）
|
||
"""
|
||
self.token = token
|
||
# EncodingAESKey 是 43 位 Base64 编码字符串
|
||
# 解码时需要先补上 1 个 "=" 使其成为合法的 Base64 字符串
|
||
# 解码后得到 32 字节的 AES 密钥（AES-256）
|
||
self.aes_key = base64.b64decode(encoding_aes_key + "=")
|
||
# AES CBC 模式的 IV（初始化向量）取密钥的前 16 字节
|
||
self.iv = self.aes_key[:16]
|
||
self.corp_id = corp_id
|
||
|
||
# --------------------------------------------------------------------------
|
||
# 签名验证
|
||
# --------------------------------------------------------------------------
|
||
def verify_signature(
|
||
self, signature: str, timestamp: str, nonce: str, encrypt: str
|
||
) -> bool:
|
||
"""验证企微回调签名。
|
||
|
||
企微的签名算法：SHA1(sort([token, timestamp, nonce, encrypt]))
|
||
将 token、timestamp、nonce、encrypt 四个字符串字典序排列后拼接，计算 SHA1。
|
||
|
||
Args:
|
||
signature: 企微传来的签名（msg_signature 参数）
|
||
timestamp: 时间戳
|
||
nonce: 随机数
|
||
encrypt: 加密的消息内容
|
||
|
||
Returns:
|
||
bool: 签名是否验证通过
|
||
"""
|
||
# 将四个参数按字典序排列
|
||
sort_list = sorted([self.token, timestamp, nonce, encrypt])
|
||
# 拼接为一个字符串
|
||
concat_str = "".join(sort_list)
|
||
# 计算 SHA1 哈希值
|
||
computed = hashlib.sha1(concat_str.encode("utf-8")).hexdigest()
|
||
# 与传入的签名比较
|
||
is_valid = computed == signature
|
||
if not is_valid:
|
||
logger.warning(
|
||
f"签名验证失败: computed={computed}, received={signature}"
|
||
)
|
||
return is_valid
|
||
|
||
def generate_signature(
|
||
self, timestamp: str, nonce: str, encrypt: str
|
||
) -> str:
|
||
"""生成签名（用于加密响应消息时）。
|
||
|
||
Args:
|
||
timestamp: 时间戳
|
||
nonce: 随机数
|
||
encrypt: 加密后的消息内容
|
||
|
||
Returns:
|
||
str: SHA1 签名字符串
|
||
"""
|
||
sort_list = sorted([self.token, timestamp, nonce, encrypt])
|
||
concat_str = "".join(sort_list)
|
||
return hashlib.sha1(concat_str.encode("utf-8")).hexdigest()
|
||
|
||
# --------------------------------------------------------------------------
|
||
# AES 解密
|
||
# --------------------------------------------------------------------------
|
||
def decrypt(self, encrypted_text: str) -> str:
|
||
"""AES-CBC 解密企微加密消息。
|
||
|
||
解密流程：
|
||
1. Base64 解码得到密文
|
||
2. AES-CBC 解密
|
||
3. 去除 PKCS7 填充
|
||
4. 提取明文内容（格式：random(16) + msg_len(4) + msg + corp_id）
|
||
|
||
Args:
|
||
encrypted_text: Base64 编码的加密消息
|
||
|
||
Returns:
|
||
str: 解密后的明文 XML 字符串
|
||
|
||
Raises:
|
||
ValueError: 解密失败时抛出
|
||
"""
|
||
try:
|
||
# 1. Base64 解码得到密文字节
|
||
encrypted_data = base64.b64decode(encrypted_text)
|
||
|
||
# 2. AES-CBC 解密
|
||
cipher = Cipher(
|
||
algorithms.AES(self.aes_key),
|
||
modes.CBC(self.iv),
|
||
backend=default_backend(),
|
||
)
|
||
decryptor = cipher.decryptor()
|
||
decrypted_data = decryptor.update(encrypted_data) + decryptor.finalize()
|
||
|
||
# 3. 去除 PKCS7 填充
|
||
# PKCS7: 填充字节的值等于填充的字节数
|
||
# 例如填充了 5 个字节，则每个字节的值都是 0x05
|
||
pad_len = decrypted_data[-1]
|
||
# 校验填充是否合法（填充值应等于填充长度）
|
||
if pad_len < 1 or pad_len > 32:
|
||
raise ValueError(f"无效的 PKCS7 填充值: {pad_len}")
|
||
# 检查所有填充字节是否一致
|
||
for i in range(pad_len):
|
||
if decrypted_data[-(i + 1)] != pad_len:
|
||
raise ValueError("PKCS7 填充校验失败")
|
||
# 去除填充部分
|
||
plaintext_data = decrypted_data[:-pad_len]
|
||
|
||
# 4. 提取明文内容
|
||
# 企微加密格式：random(16字节) + msg_len(4字节网络序) + msg + corp_id
|
||
# 跳过前 16 字节随机串
|
||
msg_len = struct.unpack("!I", plaintext_data[16:20])[0]
|
||
msg_content = plaintext_data[20 : 20 + msg_len].decode("utf-8")
|
||
from_corp_id = plaintext_data[20 + msg_len :].decode("utf-8")
|
||
|
||
# 5. 校验 corp_id 是否匹配
|
||
if from_corp_id != self.corp_id:
|
||
raise ValueError(
|
||
f"corp_id 不匹配: expected={self.corp_id}, got={from_corp_id}"
|
||
)
|
||
|
||
logger.debug(f"AES 解密成功, 明文长度: {len(msg_content)}")
|
||
return msg_content
|
||
|
||
except Exception as e:
|
||
logger.error(f"AES 解密失败: {e}")
|
||
raise ValueError(f"消息解密失败: {e}") from e
|
||
|
||
# --------------------------------------------------------------------------
|
||
# AES 加密
|
||
# --------------------------------------------------------------------------
|
||
def encrypt(self, plaintext: str) -> str:
|
||
"""AES-CBC 加密消息（用于被动回复）。
|
||
|
||
加密流程：
|
||
1. 构造明文：random(16) + msg_len(4) + msg + corp_id
|
||
2. PKCS7 填充
|
||
3. AES-CBC 加密
|
||
4. Base64 编码
|
||
|
||
Args:
|
||
plaintext: 要加密的明文字符串
|
||
|
||
Returns:
|
||
str: Base64 编码的加密结果
|
||
"""
|
||
try:
|
||
# 1. 构造明文字节
|
||
# 16 字节随机串（增加密文随机性，防止相同明文产生相同密文）
|
||
random_str = secrets.token_bytes(16)
|
||
# 消息内容字节
|
||
msg_bytes = plaintext.encode("utf-8")
|
||
# 消息长度（4 字节网络序，大端）
|
||
msg_len = struct.pack("!I", len(msg_bytes))
|
||
# corp_id 字节
|
||
corp_id_bytes = self.corp_id.encode("utf-8")
|
||
# 拼接
|
||
plaintext_data = random_str + msg_len + msg_bytes + corp_id_bytes
|
||
|
||
# 2. PKCS7 填充
|
||
# 块大小为 32 字节（AES-256 的块大小实际是 16 字节，
|
||
# 但企微官方实现使用 32 字节块做 PKCS7 填充）
|
||
block_size = 32
|
||
pad_len = block_size - (len(plaintext_data) % block_size)
|
||
# 填充字节的值等于填充的字节数
|
||
plaintext_data += bytes([pad_len] * pad_len)
|
||
|
||
# 3. AES-CBC 加密
|
||
cipher = Cipher(
|
||
algorithms.AES(self.aes_key),
|
||
modes.CBC(self.iv),
|
||
backend=default_backend(),
|
||
)
|
||
encryptor = cipher.encryptor()
|
||
encrypted_data = encryptor.update(plaintext_data) + encryptor.finalize()
|
||
|
||
# 4. Base64 编码
|
||
result = base64.b64encode(encrypted_data).decode("utf-8")
|
||
logger.debug(f"AES 加密成功, 密文长度: {len(result)}")
|
||
return result
|
||
|
||
except Exception as e:
|
||
logger.error(f"AES 加密失败: {e}")
|
||
raise ValueError(f"消息加密失败: {e}") from e
|
||
|
||
# --------------------------------------------------------------------------
|
||
# 解密回调消息（完整流程）
|
||
# --------------------------------------------------------------------------
|
||
def decrypt_message(
|
||
self,
|
||
xml_body: str,
|
||
msg_signature: str,
|
||
timestamp: str,
|
||
nonce: str,
|
||
) -> Dict[str, str]:
|
||
"""解密企微回调消息的完整流程。
|
||
|
||
流程：
|
||
1. 从 XML 中提取 Encrypt 字段
|
||
2. 验证签名
|
||
3. AES 解密
|
||
4. 解析明文 XML，返回消息内容字典
|
||
|
||
Args:
|
||
xml_body: 企微 POST 的 XML 请求体
|
||
msg_signature: 企微传来的签名
|
||
timestamp: 时间戳
|
||
nonce: 随机数
|
||
|
||
Returns:
|
||
Dict[str, str]: 解密后的消息内容，包含 from_user_name, content, msg_type 等
|
||
|
||
Raises:
|
||
ValueError: 签名验证失败或解密失败
|
||
"""
|
||
# 1. 解析 XML，提取 Encrypt 字段
|
||
try:
|
||
root = ET.fromstring(xml_body)
|
||
encrypt_node = root.find("Encrypt")
|
||
if encrypt_node is None:
|
||
raise ValueError("XML 中未找到 Encrypt 字段")
|
||
encrypt = encrypt_node.text or ""
|
||
except ET.ParseError as e:
|
||
raise ValueError(f"XML 解析失败: {e}") from e
|
||
|
||
# 2. 验证签名
|
||
if not self.verify_signature(msg_signature, timestamp, nonce, encrypt):
|
||
raise ValueError("签名验证失败")
|
||
|
||
# 3. AES 解密
|
||
plaintext_xml = self.decrypt(encrypt)
|
||
|
||
# 4. 解析明文 XML
|
||
try:
|
||
plain_root = ET.fromstring(plaintext_xml)
|
||
result = {}
|
||
# 提取常见字段
|
||
for child in plain_root:
|
||
result[child.tag] = child.text or ""
|
||
logger.info(f"消息解密成功: from={result.get('FromUserName', 'unknown')}")
|
||
return result
|
||
except ET.ParseError as e:
|
||
raise ValueError(f"明文 XML 解析失败: {e}") from e
|
||
|
||
# --------------------------------------------------------------------------
|
||
# 加密响应消息（完整流程）
|
||
# --------------------------------------------------------------------------
|
||
def encrypt_message(
|
||
self, reply_content: str, nonce: Optional[str] = None
|
||
) -> str:
|
||
"""加密响应消息的完整流程（用于被动回复）。
|
||
|
||
流程：
|
||
1. AES 加密消息内容
|
||
2. 生成签名
|
||
3. 构造响应 XML
|
||
|
||
Args:
|
||
reply_content: 要回复的消息内容
|
||
nonce: 随机数（可选，默认自动生成）
|
||
|
||
Returns:
|
||
str: 加密后的 XML 响应字符串
|
||
"""
|
||
# 生成时间戳和随机数
|
||
timestamp = str(int(time.time()))
|
||
if nonce is None:
|
||
nonce = "".join(
|
||
secrets.choice(string.ascii_letters + string.digits)
|
||
for _ in range(10)
|
||
)
|
||
|
||
# AES 加密
|
||
encrypt = self.encrypt(reply_content)
|
||
|
||
# 生成签名
|
||
signature = self.generate_signature(timestamp, nonce, encrypt)
|
||
|
||
# 构造响应 XML
|
||
response_xml = (
|
||
f"<xml>"
|
||
f"<Encrypt><![CDATA[{encrypt}]]></Encrypt>"
|
||
f"<MsgSignature><![CDATA[{signature}]]></MsgSignature>"
|
||
f"<TimeStamp>{timestamp}</TimeStamp>"
|
||
f"<Nonce><![CDATA[{nonce}]]></Nonce>"
|
||
f"</xml>"
|
||
)
|
||
|
||
logger.debug("响应消息加密完成")
|
||
return response_xml
|
||
|
||
# --------------------------------------------------------------------------
|
||
# 验证 URL 有效性（GET 请求解密 echostr）
|
||
# --------------------------------------------------------------------------
|
||
def decrypt_echostr(
|
||
self,
|
||
msg_signature: str,
|
||
timestamp: str,
|
||
nonce: str,
|
||
echostr: str,
|
||
) -> str:
|
||
"""解密企微回调验证的 echostr。
|
||
|
||
企微配置回调 URL 时会发送 GET 请求，需要：
|
||
1. 验证签名
|
||
2. 解密 echostr
|
||
3. 返回明文 echostr
|
||
|
||
Args:
|
||
msg_signature: 企微传来的签名
|
||
timestamp: 时间戳
|
||
nonce: 随机数
|
||
echostr: 加密的验证字符串
|
||
|
||
Returns:
|
||
str: 解密后的 echostr 明文
|
||
|
||
Raises:
|
||
ValueError: 签名验证失败或解密失败
|
||
"""
|
||
# 验证签名
|
||
if not self.verify_signature(msg_signature, timestamp, nonce, echostr):
|
||
raise ValueError("回调URL验证签名失败")
|
||
|
||
# 解密 echostr
|
||
plaintext = self.decrypt(echostr)
|
||
logger.info("回调URL验证成功，echostr 解密完成")
|
||
return plaintext
|