import binascii
import json
import os
import struct


class WechatAesKey:
    def __init__(self, pm, version, base):
        if not pm:
            raise ValueError("微信进程不存在，请确认微信是否已经启动！")
        self.versions = self.get_version_list()
        self.version = version
        self.base = base
        self.pm = pm

    def get_version_list(self):
        """微信获取版本号列表"""
        file_path = os.path.join(os.getcwd(), "version_list.json")

        # 检查文件是否存在
        if not os.path.exists(file_path):
            raise FileNotFoundError(
                f"微信版本号文件：“version_list.json”不存在，请到https://github.com/xaoyaoo/PyWxDump/blob/master/pywxdump/version_list.json下载最新文件，并放到工作目录"
            )

        # 读取 JSON 文件
        with open(file_path, "r") as file:
            try:
                data = json.load(file)
            except json.JSONDecodeError as e:
                raise ValueError(f"Error decoding JSON file {file_path}: {e}")

        # 检查文件是否为空
        if not data:
            raise ValueError(
                f"微信版本号文件：“version_list.json”为空，请到https://github.com/xaoyaoo/PyWxDump/blob/master/pywxdump/version_list.json下载最新文件，并放到工作目录"
            )

        return data

    def get_aes_key(self) -> str:
        """获取db密钥"""
        version = self.version
        base = self.base
        offset_array = self.versions.get(version, None)
        if not offset_array:
            raise ValueError(
                f"暂不支持版本 {version}，请到https://github.com/xaoyaoo/PyWxDump/blob/master/pywxdump/version_list.json下载最新文件，并放到工作目录"
            )
        for offset in offset_array:
            aesKey = self.getAesKey(self.pm, base, offset)
            if aesKey is not None:
                return aesKey

    def getAesKey(self, pm, base, offset):
        try:
            if self.is64Bit(pm):
                result = pm.read_bytes(base + offset, 8)  # 读取 AES Key 的地址
                addr = struct.unpack("<Q", result)[0]  # 地址为小端 8 字节整型
            else:
                result = pm.read_bytes(base + offset, 4)  # 读取 AES Key 的地址
                addr = struct.unpack("<I", result)[0]  # 地址为小端 4 字节整型

            aesKey = pm.read_bytes(addr, 0x20)  # 读取 AES Key
            result = binascii.b2a_hex(aesKey)  # 解码
            return result.decode()
        except Exception as e:
            return None

    def is64Bit(self, pm):
        exe_arch = self.getDllArch(list(pm.list_modules())[0].filename)
        if exe_arch == "x64":
            return True
        else:
            return False

    def getDllArch(self, dll_file):
        with open(dll_file, "rb") as f:
            doshdr = f.read(64)
            magic, padding, offset = struct.unpack("2s58si", doshdr)

            if magic != b"MZ":
                return None
            f.seek(offset, os.SEEK_SET)
            pehdr = f.read(6)

            # careful! H == unsigned short, x64 is negative with signed
            magic, padding, machine = struct.unpack("2s2sH", pehdr)

            if magic != b"PE":
                return None
            if machine == 0x014C:
                return "i386"
            if machine == 0x0200:
                return "IA64"
            if machine == 0x8664:
                return "x64"

            return "unknown"