Codex
144 lines
5.2 KiB
Python
144 lines
5.2 KiB
Python
#!/usr/bin/env python
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
企业微信回调加解密工具
|
||
|
||
实现企业微信事件订阅所需的签名验证和消息加解密。
|
||
基于企业微信官方文档：https://developer.work.weixin.qq.com/document/path/90968
|
||
|
||
加密方案：AES-CBC，PKCS#7 padding，key 由 EncodingAESKey + "=" base64 解码得到
|
||
签名方案：SHA1(sort(token, timestamp, nonce, encrypt))
|
||
"""
|
||
import base64
|
||
import hashlib
|
||
import secrets
|
||
import struct
|
||
import time
|
||
from typing import Dict
|
||
from xml.etree import ElementTree
|
||
|
||
from cryptography.hazmat.backends import default_backend
|
||
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
|
||
|
||
|
||
class WecomCrypto:
|
||
"""企业微信回调消息加解密"""
|
||
|
||
def __init__(self, token: str, aes_key: str, corp_id: str):
|
||
self.token = token
|
||
self.corp_id = corp_id
|
||
self.aes_key_bytes = base64.b64decode(aes_key + "=")
|
||
|
||
def _generate_nonce(self, length: int = 16) -> str:
|
||
return secrets.token_hex(length // 2)
|
||
|
||
def _pkcs7_pad(self, data: bytes, block_size: int = 32) -> bytes:
|
||
padding_len = block_size - (len(data) % block_size)
|
||
return data + bytes([padding_len] * padding_len)
|
||
|
||
def _pkcs7_unpad(self, data: bytes) -> bytes:
|
||
padding_len = data[-1]
|
||
return data[:-padding_len]
|
||
|
||
def _sign(self, token: str, timestamp: str, nonce: str, encrypt: str) -> str:
|
||
"""计算签名 SHA1(sort(token, timestamp, nonce, encrypt))"""
|
||
parts = sorted([token, timestamp, nonce, encrypt])
|
||
raw = "".join(parts).encode("utf-8")
|
||
return hashlib.sha1(raw).hexdigest()
|
||
|
||
def encrypt(self, plaintext: str) -> Dict[str, str]:
|
||
"""
|
||
加密明文消息，返回回调响应体
|
||
|
||
返回: {"msg_signature": ..., "timeStamp": ..., "nonce": ..., "encrypt": ...}
|
||
"""
|
||
nonce = self._generate_nonce()
|
||
timestamp = str(int(time.time()))
|
||
|
||
random_bytes = secrets.token_bytes(16)
|
||
text_bytes = plaintext.encode("utf-8")
|
||
corp_id_bytes = self.corp_id.encode("utf-8")
|
||
content = random_bytes + struct.pack("!I", len(text_bytes)) + text_bytes + corp_id_bytes
|
||
|
||
padded = self._pkcs7_pad(content)
|
||
iv = self.aes_key_bytes[:16]
|
||
cipher = Cipher(algorithms.AES(self.aes_key_bytes), modes.CBC(iv), backend=default_backend())
|
||
encryptor = cipher.encryptor()
|
||
encrypted = encryptor.update(padded) + encryptor.finalize()
|
||
|
||
encrypt_str = base64.b64encode(encrypted).decode("utf-8")
|
||
signature = self._sign(self.token, timestamp, nonce, encrypt_str)
|
||
|
||
return {
|
||
"msg_signature": signature,
|
||
"timeStamp": timestamp,
|
||
"nonce": nonce,
|
||
"encrypt": encrypt_str,
|
||
}
|
||
|
||
def decrypt(self, encrypt_str: str) -> str:
|
||
"""解密密文，返回明文 XML 字符串，并校验 receiveid"""
|
||
encrypted = base64.b64decode(encrypt_str)
|
||
|
||
iv = self.aes_key_bytes[:16]
|
||
cipher = Cipher(algorithms.AES(self.aes_key_bytes), modes.CBC(iv), backend=default_backend())
|
||
decryptor = cipher.decryptor()
|
||
decrypted = decryptor.update(encrypted) + decryptor.finalize()
|
||
|
||
unpadded = self._pkcs7_unpad(decrypted)
|
||
|
||
msg_len = struct.unpack("!I", unpadded[16:20])[0]
|
||
plaintext = unpadded[20:20 + msg_len].decode("utf-8")
|
||
|
||
receive_id = unpadded[20 + msg_len:].decode("utf-8")
|
||
if receive_id != self.corp_id:
|
||
raise ValueError(f"receiveid 校验失败: 期望 {self.corp_id}，实际 {receive_id}")
|
||
|
||
return plaintext
|
||
|
||
def verify_signature(self, msg_signature: str, timestamp: str, nonce: str, encrypt: str) -> bool:
|
||
"""验证回调签名"""
|
||
calculated = self._sign(self.token, timestamp, nonce, encrypt)
|
||
return calculated == msg_signature
|
||
|
||
def handle_callback(
|
||
self,
|
||
msg_signature: str,
|
||
timestamp: str,
|
||
nonce: str,
|
||
post_data: str,
|
||
) -> str:
|
||
"""
|
||
处理 POST 回调：从 XML 中提取 Encrypt，验签 + 解密
|
||
|
||
返回解密后的明文 XML 字符串
|
||
"""
|
||
root = ElementTree.fromstring(post_data)
|
||
encrypt = root.find("Encrypt").text
|
||
|
||
if not self.verify_signature(msg_signature, timestamp, nonce, encrypt):
|
||
raise ValueError("签名验证失败")
|
||
return self.decrypt(encrypt)
|
||
|
||
def handle_verify(self, msg_signature: str, timestamp: str, nonce: str, echostr: str) -> str:
|
||
"""
|
||
处理 GET 验证请求（URL 验证）
|
||
|
||
验签后解密 echostr，返回明文用于响应
|
||
"""
|
||
if not self.verify_signature(msg_signature, timestamp, nonce, echostr):
|
||
raise ValueError("签名验证失败")
|
||
return self.decrypt(echostr)
|
||
|
||
def generate_success_response(self) -> str:
|
||
"""生成回调成功响应（加密 "success"，返回 XML）"""
|
||
result = self.encrypt("success")
|
||
return (
|
||
f'<xml>'
|
||
f'<Encrypt><![CDATA[{result["encrypt"]}]]></Encrypt>'
|
||
f'<MsgSignature><![CDATA[{result["msg_signature"]}]]></MsgSignature>'
|
||
f'<TimeStamp>{result["timeStamp"]}</TimeStamp>'
|
||
f'<Nonce><![CDATA[{result["nonce"]}]]></Nonce>'
|
||
f'</xml>'
|
||
)
|