#!/usr/bin/env python3
"""
分析从狗上提取的Blowfish密钥dump
1. 在内存中搜索Blowfish P-array (key schedule后的运行时值)
2. 提取P-array和S-box,反推原始密钥
3. 用提取的P-array+S-box解密已知密文,验证正确性
"""
import os
import sys
import struct
import tarfile
import argparse
from pathlib import Path

# Blowfish标准初始常数 (从π推导)
P_INITIAL = [
    0x243f6a88, 0x85a308d3, 0x13198a2e, 0x03707344,
    0xa4093822, 0x299f31d0, 0x082efa98, 0xec4e6c89,
    0x452821e6, 0x38d01377, 0xbe5466cf, 0x34e90c6c,
    0xc0ac29b7, 0xc97c50dd, 0x3f84d5b5, 0xb5470917,
    0x9216d5d9, 0x8979fb1b
]

S0_INITIAL_FIRST = 0xd1310ba6
S1_INITIAL_FIRST = 0x4b7a70e9
S2_INITIAL_FIRST = 0xe93d5a68
S3_INITIAL_FIRST = 0x3a39ce37

# 已知明文-密文对 (用于验证)
KNOWN_PAIRS = [
    (bytes.fromhex('0000000000000000'), bytes.fromhex('12aa0354236b66e3')),
    (bytes.fromhex('feefff0004020305'), bytes.fromhex('5d9874ad7eee5851')),
    (bytes.fromhex('0802010000010900'), bytes.fromhex('c5ce0887ceb86f91')),
]


def find_pattern_locations(data, pattern_bytes):
    """在data中查找pattern的所有位置"""
    positions = []
    pos = 0
    while True:
        p = data.find(pattern_bytes, pos)
        if p < 0:
            break
        positions.append(p)
        pos = p + 1
    return positions


def looks_like_blowfish_buffer(data, offset):
    """
    检查offset处是否看起来像Blowfish state:
    布局假设: P[18] + S0[256] + S1[256] + S2[256] + S3[256]
    总大小: (18 + 4*256) * 4 = 4168 字节
    """
    BUF_SIZE = 4168
    if offset + BUF_SIZE > len(data):
        return False

    # 取出整个buffer
    buf = data[offset:offset+BUF_SIZE]

    # 解释为uint32列表
    words = struct.unpack(f'<{BUF_SIZE//4}I', buf)

    # 检查S-box是否包含至少部分初始值
    # 运行时,P和S-box已被key schedule修改,但所有1042个值的统计特性应该接近随机
    # 这里只能做粗略的"看起来像加密表"的检查

    # 唯一性检查: 1042个32位值应该几乎都不同
    unique = len(set(words))
    if unique < 1000:  # 太多重复说明不是表
        return False

    # 零值检查: 不应该有很多0
    zeros = sum(1 for w in words if w == 0)
    if zeros > 50:
        return False

    # 熵检查: 字节分布应该均匀
    byte_counts = [0] * 256
    for b in buf:
        byte_counts[b] += 1
    max_count = max(byte_counts)
    if max_count > len(buf) // 50:  # 任何字节出现频率超过2%可疑
        return False

    return True


def encrypt_with_state(plaintext_8, P, S0, S1, S2, S3, endian='<'):
    """
    用给定的P-array + S-box加密8字节明文
    endian: '<' (LE - Unitree style) 或 '>' (BE - standard)
    """
    pack_fmt = endian + 'II'
    L, R = struct.unpack(pack_fmt, plaintext_8)

    for i in range(16):
        L ^= P[i]
        # F(L)
        a = (L >> 24) & 0xff
        b = (L >> 16) & 0xff
        c = (L >> 8) & 0xff
        d = L & 0xff
        F = (((S0[a] + S1[b]) & 0xFFFFFFFF) ^ S2[c]) & 0xFFFFFFFF
        F = (F + S3[d]) & 0xFFFFFFFF
        R ^= F
        L, R = R, L

    L, R = R, L
    R ^= P[16]
    L ^= P[17]

    return struct.pack(pack_fmt, L, R)


def try_buffer_as_blowfish(buf_data, label="?"):
    """
    把一个4168字节的buffer解释为Blowfish state
    用已知明密文对验证
    """
    if len(buf_data) < 4168:
        return None

    # 假设布局1: P[0..17] + S0[0..255] + S1[0..255] + S2[0..255] + S3[0..255]
    words = struct.unpack('<1042I', buf_data[:4168])
    P = list(words[0:18])
    S0 = list(words[18:18+256])
    S1 = list(words[18+256:18+512])
    S2 = list(words[18+512:18+768])
    S3 = list(words[18+768:18+1024])

    # 测试已知对 - 同时尝试LE和BE两种字节序
    best_matches = 0
    best_endian = None
    for endian in ['<', '>']:
        matches = 0
        for pt, ct_expected in KNOWN_PAIRS:
            try:
                ct = encrypt_with_state(pt, P, S0, S1, S2, S3, endian=endian)
                if ct == ct_expected:
                    matches += 1
            except Exception:
                pass
        if matches > best_matches:
            best_matches = matches
            best_endian = endian

    return best_matches, best_endian


def search_blowfish_state(data, source_name, verbose=False, save_candidates=False):
    """在内存dump中搜索Blowfish状态"""
    print(f"\n=== 在 {source_name} 中搜索 Blowfish state ===")
    print(f"   大小: {len(data)} 字节")
    candidates_looking_like_bf = []

    # 策略1: 找S-box[0]首字 (运行时S-box会被修改,但偶尔可能未变)
    # 实际上key schedule后S-box几乎肯定变化,所以这条不会找到
    s0_pos = find_pattern_locations(data, struct.pack('<I', S0_INITIAL_FIRST))
    s1_pos = find_pattern_locations(data, struct.pack('<I', S1_INITIAL_FIRST))
    print(f"   S0初始首字位置数: {len(s0_pos)}")
    print(f"   S1初始首字位置数: {len(s1_pos)}")

    if s0_pos:
        # 这些是静态初始值,在二进制中也有 (说明是从binary mmap过来的)
        for p in s0_pos[:3]:
            print(f"     初始S0@0x{p:x}")

    # 策略2: 暴力搜索 - 假设 4168字节连续Blowfish state
    # 每4字节对齐位置都试
    print(f"   暴力搜索 Blowfish state (4字节对齐)...")
    found = []

    STEP = 4
    count = 0
    total = (len(data) - 4168) // STEP

    for offset in range(0, len(data) - 4168, STEP):
        count += 1
        if count % 50000 == 0:
            print(f"      进度: {count}/{total} ({100*count/total:.1f}%)")

        # 快速过滤: 必须looks_like
        if not looks_like_blowfish_buffer(data, offset):
            continue

        candidates_looking_like_bf.append(offset)
        if save_candidates and len(candidates_looking_like_bf) <= 20:
            cand_path = f'data/lowlevel_captures/candidate_{source_name}_{offset:08x}.bin'
            os.makedirs('data/lowlevel_captures', exist_ok=True)
            with open(cand_path, 'wb') as fp:
                fp.write(data[offset:offset+4168])
            print(f"   💾 保存候选 -> {cand_path}")

        result = try_buffer_as_blowfish(data[offset:offset+4168])
        if result is None:
            continue
        matches, endian = result
        if matches and matches > 0:
            endian_name = 'LE' if endian == '<' else 'BE'
            print(f"\n   ⭐ offset 0x{offset:x}: {matches}/3 已知对匹配! ({endian_name})")
            found.append((offset, matches, endian))
            if matches == 3:
                return found, data[offset:offset+4168]

    if candidates_looking_like_bf:
        print(f"   共找到 {len(candidates_looking_like_bf)} 个通过启发式过滤的候选位置")

    return found, None


def extract_key_from_state(state_buf):
    """从运行时P-array+S-box反推原始密钥 (这是Blowfish key schedule的逆过程)

    技术: Blowfish key schedule是
      1. P[i] = P_INITIAL[i] ^ key_word(i % keylen)
      2. 用此P和初始S-box加密 (0,0), 用结果替换P[0..1]
      3. 继续加密前一次结果,替换P[2..3], 直到P[16..17]
      4. 同样替换S-box

    要反推key: 由于步骤2-4都是确定性的,只要知道key就能复现
              而key只在步骤1出现, P[i] XOR P_INITIAL[i] = key_word(i % keylen)
    """
    words = struct.unpack('<1042I', state_buf[:4168])
    P_runtime = list(words[0:18])

    # 但是! P-array在步骤2-4也被修改了, 所以不能直接 P XOR P_INITIAL
    # 实际上, 步骤1后的P[i]立即在步骤2中被覆盖, 我们看到的是覆盖后的值
    # 这意味着无法从最终state直接反推key

    # 但我们已经有了完整的state, 不需要key就能加密/解密!
    print("\n=== 提取的密钥state ===")
    print(f"P-array (18 uint32):")
    for i in range(18):
        print(f"  P[{i:2d}] = 0x{P_runtime[i]:08x}")

    return P_runtime


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('input', help='blowfish_dump.tar.gz 或 包含.bin文件的目录')
    parser.add_argument('--verbose', '-v', action='store_true')
    parser.add_argument('--save-candidates', action='store_true',
                        help='保存所有"看起来像Blowfish state"的候选 (即使不能验证已知对)')
    args = parser.parse_args()

    input_path = Path(args.input)

    # 提取/找到所有bin文件
    bin_files = []
    if input_path.is_file() and input_path.suffix == '.gz':
        # tar.gz - 解压到临时目录
        import tempfile
        tmpdir = tempfile.mkdtemp(prefix='blowfish_analyze_')
        print(f"解压到 {tmpdir}...")
        with tarfile.open(input_path) as tf:
            tf.extractall(tmpdir)
        for f in Path(tmpdir).rglob('*.bin'):
            bin_files.append(f)
    elif input_path.is_dir():
        for f in input_path.rglob('*.bin'):
            bin_files.append(f)
    else:
        print(f"❌ 不支持的输入: {input_path}")
        return 1

    print(f"找到 {len(bin_files)} 个 .bin 文件")
    for f in bin_files:
        size = f.stat().st_size
        print(f"  {f.name} ({size} 字节)")

    # 按大小排序,先搜索小的 (更快)
    bin_files.sort(key=lambda f: f.stat().st_size)

    found_state = None
    for f in bin_files:
        with open(f, 'rb') as fp:
            data = fp.read()
        results, state = search_blowfish_state(data, f.name, verbose=args.verbose, save_candidates=args.save_candidates)
        if state is not None:
            found_state = state
            print(f"\n🎉 在 {f.name} 中找到完整匹配!")
            break
        elif results:
            print(f"\n⚠️  在 {f.name} 中找到 {len(results)} 个部分匹配位置")

    if found_state:
        P = extract_key_from_state(found_state)
        # 保存
        outpath = 'data/lowlevel_captures/blowfish_state.bin'
        os.makedirs('data/lowlevel_captures', exist_ok=True)
        with open(outpath, 'wb') as fp:
            fp.write(found_state)
        print(f"\n✅ Blowfish state已保存到 {outpath}")
        print(f"   可以用此state直接加密/解密通信数据")
        return 0
    else:
        print("\n❌ 未找到完整匹配的Blowfish state")
        print("   可能原因:")
        print("   1. 状态不在已dump的内存区域")
        print("   2. 状态被分散存储 (不连续)")
        print("   3. dump时机不对 (state还未初始化)")
        return 1


if __name__ == '__main__':
    sys.exit(main())