ForcePilot/backend/package/yuxi/channel/extensions/wechat_miniprogram/crypto.py

55 lines
1.8 KiB
Python
Raw Normal View History

import base64
import hashlib
import os
import socket
import struct
from Crypto.Cipher import AES
class WeChatCrypto:
BLOCK_SIZE = 32
def __init__(self, token: str, encoding_aes_key: str, app_id: str):
self.token = token
self.app_id = app_id
self.aes_key = base64.b64decode(encoding_aes_key + "=")
def verify_signature(
self,
signature: str,
timestamp: str,
nonce: str,
msg_encrypt: str = "",
) -> bool:
computed = self.compute_signature(timestamp, nonce, msg_encrypt)
return computed == signature
def compute_signature(self, timestamp: str, nonce: str, msg_encrypt: str = "") -> str:
params = sorted([self.token, timestamp, nonce, msg_encrypt])
return hashlib.sha1("".join(params).encode()).hexdigest()
def decrypt(self, msg_encrypt: str) -> tuple[str, str]:
cipher = AES.new(self.aes_key, AES.MODE_CBC, iv=self.aes_key[:16])
raw = base64.b64decode(msg_encrypt)
decrypted = cipher.decrypt(raw)
pad = decrypted[-1]
content = decrypted[:-pad]
msg_len = socket.ntohl(struct.unpack("I", content[16:20])[0])
xml = content[20:20 + msg_len].decode("utf-8")
receiver_app_id = content[20 + msg_len:].decode("utf-8")
return xml, receiver_app_id
def encrypt(self, msg: str) -> str:
random_bytes = os.urandom(16)
msg_bytes = msg.encode("utf-8")
msg_len = socket.htonl(len(msg_bytes))
raw = random_bytes + struct.pack("I", msg_len) + msg_bytes + self.app_id.encode("utf-8")
pad_len = self.BLOCK_SIZE - (len(raw) % self.BLOCK_SIZE)
raw = raw + bytes([pad_len]) * pad_len
cipher = AES.new(self.aes_key, AES.MODE_CBC, iv=self.aes_key[:16])
return base64.b64encode(cipher.encrypt(raw)).decode()