from base64 import b64decode
from Crypto.Cipher import AES
import json

class WXBizDataCrypt:
    def __init__(self, appid, session_key):
        self.appid = appid
        self.session_key = b64decode(session_key)

    def decrypt(self, encrypted_data, iv):
        # AES解密
        try:
            cipher = AES.new(self.session_key, AES.MODE_CBC, b64decode(iv))
            decrypted = json.loads(self._unpad(cipher.decrypt(b64decode(encrypted_data))))
            if decrypted['watermark']['appid'] != self.appid:
                raise Exception('Invalid Buffer')
            return decrypted
        except Exception as e:
            raise Exception('Failed to decrypt data')

    def _unpad(self, s):
        try:
            return s[:-ord(s[len(s) - 1:])]
        except Exception as e:
            print("Error during unpadding:", e)
            return s  # 返回原始数据以便进一步分析错误