from ctypes.wintypes import MAX_PATH
import os
import winreg
import hmac
import ctypes
import hashlib
from Crypto.Cipher import AES


class WechatDB:
    def __init__(self, aes_key):
        self.aes_key = bytes.fromhex(aes_key)
        self.save_path_directory = os.path.join(os.getcwd(), "wechat_db")
        if not os.path.exists(self.save_path_directory):
            os.makedirs(self.save_path_directory)

    def get_wechat_user_directory(self) -> str:
        """
        获取微信用户数据库保存文件目录
        """
        try:
            path = r"Software\Tencent\WeChat"
            with winreg.OpenKey(winreg.HKEY_CURRENT_USER, path) as registry_key:
                # 读取注册表项中的值
                value, _ = winreg.QueryValueEx(registry_key, "FileSavePath")
                if value == "MyDocument:":
                    dll = ctypes.windll.shell32
                    buf = ctypes.create_unicode_buffer(MAX_PATH + 1)
                    if dll.SHGetSpecialFolderPathW(None, buf, 0x0005, False):
                        return os.path.join(buf.value, "WeChat Files")
                    else:
                        raise Exception("未找到微信用户数据路径,请确保电脑上安装了微信")
                return value
        except FileNotFoundError:
            raise Exception("未找到微信用户数据路径,请确保电脑上安装了微信")

    def decrypt_db(self, path):
        KEY_SIZE = 32
        DEFAULT_ITER = 64000
        DEFAULT_PAGESIZE = 4096  # 4048数据 + 16IV + 20 HMAC + 12
        SQLITE_FILE_HEADER = bytes("SQLite format 3", encoding="ASCII") + bytes(
            1
        )  # SQLite 文件头

        with open(path, "rb") as f:
            # TODO: 优化，考虑超大文件
            blist = f.read()

        salt = blist[:16]  # 前16字节为盐
        key = hashlib.pbkdf2_hmac(
            "sha1", self.aes_key, salt, DEFAULT_ITER, KEY_SIZE
        )  # 获得Key

        page1 = blist[16:DEFAULT_PAGESIZE]  # 丢掉salt

        mac_salt = bytes([x ^ 0x3A for x in salt])
        mac_key = hashlib.pbkdf2_hmac("sha1", key, mac_salt, 2, KEY_SIZE)

        hash_mac = hmac.new(mac_key, digestmod="sha1")
        hash_mac.update(page1[:-32])
        hash_mac.update(bytes(ctypes.c_int(1)))

        if hash_mac.digest() != page1[-32:-12]:
            raise RuntimeError("密码错误！")

        pages = [
            blist[i : i + DEFAULT_PAGESIZE]
            for i in range(DEFAULT_PAGESIZE, len(blist), DEFAULT_PAGESIZE)
        ]
        pages.insert(0, page1)  # 把第一页补上

        return SQLITE_FILE_HEADER, key, pages

    def save_db(self, header, key, pages, path):
        with open(path, "wb") as f:
            f.write(header)  # 写入文件头

            for i in pages:
                t = AES.new(key, AES.MODE_CBC, i[-48:-32])
                f.write(t.decrypt(i[:-48]))
                f.write(i[-48:])


if __name__ == "__main__":
    wxdb = WechatDB("")
    print(wxdb.get_wechat_user_directory())