# -*- coding: utf-8 -*-
"""
微信桌面版Cookie自动捕获工具
监控微信Cookie数据库，检测到停车缴费域名Cookie后自动推送到系统API

使用方法：
  python wechat_cookie_capture.py

在微信桌面版中打开停车缴费二维码后，脚本会自动检测并推送Cookie。
按 Ctrl+C 停止。
"""

import sys, os, json, time, hashlib, sqlite3, shutil, tempfile
sys.stdout.reconfigure(encoding='utf-8')

import ctypes, ctypes.wintypes, base64
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
import urllib.request

# ===== 配置 =====
PARKING_API = 'https://parking-auto.pages.dev/api/push-cookie'
PARKING_PASSWORD = '012345jaqct'
TARGET_DOMAINS = ['zkzbgz.com']  # 只采集支付域名，QR入口域名(zkzbkj.com)不需要Cookie
CHECK_INTERVAL = 5  # 秒

# ===== 微信路径 =====
WECHAT_BASE = os.path.join(os.environ['APPDATA'], 'Tencent', 'xwechat', 'radium', 'web')
LOCAL_STATE_PATH = os.path.join(WECHAT_BASE, 'Local State')
PROFILES_DIR = os.path.join(WECHAT_BASE, 'profiles')


def get_decryption_key():
    """从微信Local State读取加密密钥并用DPAPI解密"""
    with open(LOCAL_STATE_PATH, 'r') as f:
        ls = json.load(f)
    enc_key = base64.b64decode(ls['os_crypt']['encrypted_key'])[5:]  # 跳过DPAPI前缀

    class DATA_BLOB(ctypes.Structure):
        _fields_ = [('cbData', ctypes.wintypes.DWORD),
                     ('pbData', ctypes.POINTER(ctypes.c_char))]

    blob_in = DATA_BLOB(len(enc_key), ctypes.create_string_buffer(enc_key, len(enc_key)))
    blob_out = DATA_BLOB()
    if not ctypes.windll.crypt32.CryptUnprotectData(
            ctypes.byref(blob_in), None, None, None, None, 0, ctypes.byref(blob_out)):
        raise RuntimeError('DPAPI解密失败')
    return ctypes.string_at(blob_out.pbData, blob_out.cbData)


def find_cookie_dbs():
    """查找所有微信Cookie数据库"""
    dbs = []
    for profile in os.listdir(PROFILES_DIR):
        db_path = os.path.join(PROFILES_DIR, profile, 'Network', 'Cookies')
        if os.path.exists(db_path):
            dbs.append(db_path)
    return dbs


def read_parking_cookies(db_path, aes_key):
    """从指定数据库读取停车域名Cookie"""
    tmp = os.path.join(tempfile.gettempdir(), f'wc_cap_{os.getpid()}.db')
    try:
        shutil.copy2(db_path, tmp)
        db = sqlite3.connect(tmp)
        where = ' OR '.join(f"host_key LIKE '%{d}%'" for d in TARGET_DOMAINS)
        rows = db.execute(
            f"SELECT host_key, name, encrypted_value FROM cookies WHERE {where}"
        ).fetchall()
        db.close()
    finally:
        try: os.unlink(tmp)
        except: pass

    aesgcm = AESGCM(aes_key)
    cookies = []
    for host, name, enc_val in rows:
        if enc_val[:3] == b'v10':
            try:
                dec = aesgcm.decrypt(enc_val[3:15], enc_val[15:], None).decode('utf-8', errors='replace')
                # Strip decryption padding: find actual cookie content
                # ASP.NET Core session cookies start with 'CfDJ8'
                cf_idx = dec.find('CfDJ8')
                if cf_idx >= 0:
                    dec = dec[cf_idx:]
                else:
                    # Fallback: strip all leading non-printable chars
                    dec = dec.lstrip(''.join(chr(c) for c in range(0, 0x20)) + chr(0x7f))
                cookies.append({'domain': host, 'name': name, 'value': dec})
            except Exception:
                pass
    return cookies


def dedup_cookies(cookies):
    """去重：同一个cookie name只保留最后一个值（多profile场景）"""
    seen = {}
    for c in cookies:
        seen[c['name']] = c
    return list(seen.values())


def push_cookies(cookie_str):
    """推送Cookie到停车系统API"""
    data = json.dumps({'cookie': cookie_str}).encode('utf-8')
    req = urllib.request.Request(PARKING_API, data=data, method='POST', headers={
        'Content-Type': 'application/json',
        'Authorization': f'Bearer {PARKING_PASSWORD}',
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
    })
    try:
        with urllib.request.urlopen(req, timeout=10) as resp:
            result = json.loads(resp.read())
            return result
    except Exception as e:
        return {'success': False, 'error': str(e)}


def main():
    print('=' * 44)
    print(' 微信Cookie自动捕获工具')
    print('=' * 44)
    print()

    if not os.path.exists(LOCAL_STATE_PATH):
        print('[错误] 未找到微信Local State:', LOCAL_STATE_PATH)
        print('  请确认微信桌面版已安装并打开过')
        return

    dbs = find_cookie_dbs()
    if not dbs:
        print('[错误] 未找到Cookie数据库')
        return

    print(f'[信息] 找到 {len(dbs)} 个Cookie数据库')

    print('[信息] 正在获取解密密钥...')
    try:
        aes_key = get_decryption_key()
        print(f'[成功] 密钥获取成功（{len(aes_key)}字节AES-256）')
    except Exception as e:
        print(f'[错误] 密钥获取失败: {e}')
        return

    print()
    print(f'[监听] 每{CHECK_INTERVAL}秒检查一次')
    print(f'[监听] 目标域名: {", ".join(TARGET_DOMAINS)}')
    print('[监听] 请在微信中打开停车缴费二维码...')
    print('[监听] 按 Ctrl+C 停止')
    print('-' * 44)

    last_hash = ''

    # 初始检查
    all_cookies = []
    for db in dbs:
        all_cookies.extend(read_parking_cookies(db, aes_key))
    all_cookies = dedup_cookies(all_cookies)

    if all_cookies:
        cookie_str = '; '.join(f"{c['name']}={c['value']}" for c in all_cookies)
        last_hash = hashlib.md5(cookie_str.encode()).hexdigest()
        print(f'[信息] 已有 {len(all_cookies)} 个停车Cookie，自动推送')
        result = push_cookies(cookie_str)
        print(f'[{"成功" if result.get("success") else "失败"}] {result.get("message", result.get("error", ""))}')

    while True:
        try:
            time.sleep(CHECK_INTERVAL)
            all_cookies = []
            for db in dbs:
                all_cookies.extend(read_parking_cookies(db, aes_key))
            all_cookies = dedup_cookies(all_cookies)

            if all_cookies:
                cookie_str = '; '.join(f"{c['name']}={c['value']}" for c in all_cookies)
                h = hashlib.md5(cookie_str.encode()).hexdigest()
                if h != last_hash:
                    last_hash = h
                    ts = time.strftime('%H:%M:%S')
                    print(f'\n[检测到] Cookie已更新! {ts}')
                    for c in all_cookies:
                        v = c['value'][:50] + '...' if len(c['value']) > 50 else c['value']
                        print(f'  {c["domain"]} | {c["name"]} = {v}')
                    result = push_cookies(cookie_str)
                    print(f'[{"成功" if result.get("success") else "失败"}] {result.get("message", result.get("error", ""))}')
        except KeyboardInterrupt:
            print('\n\n[停止] 已退出')
            break
        except Exception as e:
            print(f'[错误] {e}')


if __name__ == '__main__':
    main()