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()