import collections

# 定义S盒，将4位输入映射到4位输出
s_box_def = {
    0x0: 0x6, 0x1: 0x4, 0x2: 0xc, 0x3: 0x5,
    0x4: 0x0, 0x5: 0x7, 0x6: 0x2, 0x7: 0xe,
    0x8: 0x1, 0x9: 0xf, 0xA: 0x3, 0xB: 0xd,
    0xC: 0x8, 0xD: 0xa, 0xE: 0x9, 0xF: 0xb,
}

# 计算逆S盒，通过将原S盒的键值对反转得到
inv_s_box_def = {v: k for k, v in s_box_def.items()}

def split_to_nibbles(word16):
    """将16位字分割成4个4位的nibble（从最高位开始）"""
    return [(word16 >> 12) & 0xF, (word16 >> 8) & 0xF, (word16 >> 4) & 0xF, word16 & 0xF]

def join_nibbles(nibbles_array):
    """将4个4位nibble合并为一个16位字"""
    if len(nibbles_array) != 4:
        raise ValueError("Nibble数组必须包含4个nibble")
    return (nibbles_array[0] << 12) | (nibbles_array[1] << 8) | \
           (nibbles_array[2] << 4) | nibbles_array[3]

def apply_inv_s_box_layer_to_nibbles(state_nibbles):
    """对每个nibble应用逆S盒"""
    return [inv_s_box_def[n] for n in state_nibbles]

def load_plain_cipher_pairs(filename):
    """从文件加载明文-密文对"""
    pairs = {}
    try:
        with open(filename, 'r') as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                try:
                    # 解析每行的明文和密文（十六进制格式）
                    pt_hex, ct_hex = line.split(',')
                    pt = int(pt_hex, 16)
                    ct = int(ct_hex, 16)
                    pairs[pt] = ct
                except ValueError:
                    print(f"警告：跳过格式错误行: {line}")
    except FileNotFoundError:
        print(f"错误：文件未找到 {filename}")
        return None
    return pairs

def recover_last_round_key_k4(pt_ct_map, cipher_input_diff, 
                              expected_s_layer3_output_diffs, assumed_s_layer3_input_diff):
    """
    使用差分密码分析恢复最后一轮密钥k4
    pt_ct_map: 明文到密文的映射字典
    cipher_input_diff: 明文的差分（例如0x0020）
    expected_s_layer3_output_diffs: 预期的第三轮S层输出差分集合
    assumed_s_layer3_input_diff: 假设的第三轮S层输入差分（例如0x0020）
    """
    k4_candidate_counts = collections.defaultdict(int)  # 存储k4候选及其出现次数
    
    num_pairs_processed = 0  # 处理的明文对总数
    num_good_delta_c_pairs = 0  # 符合预期的密文差分对数

    # 遍历所有明文，生成具有指定差分的明文对
    plaintexts_in_map = sorted(list(pt_ct_map.keys()))  # 排序以确保确定性

    for m1 in plaintexts_in_map:
        c1 = pt_ct_map[m1]
        m2 = m1 ^ cipher_input_diff  # 计算m2 = m1 XOR 差分
        
        if m2 in pt_ct_map:
            c2 = pt_ct_map[m2]
            num_pairs_processed += 1
            
            delta_c = c1 ^ c2  # 密文差分
            
            
            # 仅处理符合预期密文差分的对
            if delta_c in expected_s_layer3_output_diffs:
                print(delta_c)
                num_good_delta_c_pairs += 1
                c1_nibbles = split_to_nibbles(c1)
                c2_nibbles = split_to_nibbles(c2)

                # 遍历所有可能的k4候选（16位）
                for k4_guess_val in range(1 << 16):  # 0x0000到0xFFFF
                    k4_guess_nibbles = split_to_nibbles(k4_guess_val)
                    
                    # 计算第三轮S层的输出V3 = C XOR k4_guess
                    v3_c1_nibbles = [cn ^ k4n for cn, k4n in zip(c1_nibbles, k4_guess_nibbles)]
                    v3_c2_nibbles = [cn ^ k4n for cn, k4n in zip(c2_nibbles, k4_guess_nibbles)]
                    
                    # 计算第三轮S层的输入U3 = S^{-1}(V3)
                    try:
                        u3_c1_nibbles = apply_inv_s_box_layer_to_nibbles(v3_c1_nibbles)
                        u3_c2_nibbles = apply_inv_s_box_layer_to_nibbles(v3_c2_nibbles)
                    except KeyError: 
                        # 处理无效的S盒输入（理论上不会发生，因为S盒是双射）
                        continue 
                        
                    u3_c1_val = join_nibbles(u3_c1_nibbles)
                    u3_c2_val = join_nibbles(u3_c2_nibbles)
                    
                    # 检查第三轮S层输入差分是否符合假设
                    if (u3_c1_val ^ u3_c2_val) == assumed_s_layer3_input_diff:
                        k4_candidate_counts[k4_guess_val] += 1  # 候选计数增加
    
    # 输出统计信息
    print(f"明文映射中的唯一明文数量: {len(plaintexts_in_map)}")
    print(f"处理的总明文对数: {num_pairs_processed}")
    print(f"符合预期delta_c的对数: {num_good_delta_c_pairs}")

    if not k4_candidate_counts:
        print("未找到k4候选。可能原因：数据不足、假设错误或差分特征不符合。")
        return None

    # 按出现次数排序候选（降序），次数相同则按k4值升序
    sorted_candidates = sorted(k4_candidate_counts.items(), key=lambda item: (-item[1], item[0]))
    
    best_k4 = sorted_candidates[0][0]  # 最可能的k4
    
    print("\n前5个k4候选（k4值: 出现次数）:")
    for i in range(min(5, len(sorted_candidates))):
        print(f"  0x{sorted_candidates[i][0]:04x}: {sorted_candidates[i][1]}")
        
    return best_k4

# 主程序
if __name__ == "__main__":
    # 注意：替换为实际文件路径
    ciphertext_file_path = "C:\\Users\\zywoo\\Desktop\\ciphertextB_13.txt" 
    
    pt_ct_map = load_plain_cipher_pairs(ciphertext_file_path)

    if pt_ct_map is None or not pt_ct_map:
        print(f"无法从文件加载数据: {ciphertext_file_path}。程序退出。")
    else:
        print(f"成功加载 {len(pt_ct_map)} 对明文-密文。")

        # 根据题目设定参数
        cipher_input_difference = 0x0020  # 明文差分
        assumed_s_layer3_input_difference = 0x0020  # 假设第三轮S层输入差分
        expected_s_layer3_output_differences = { 0x0020}  # 预期密文差分

        print(f"\n开始恢复k4，参数如下：")
        print(f"  明文输入差分 (delta_m): 0x{cipher_input_difference:04x}")
        print(f"  假设的第三轮S层输入差分 (delta_U3): 0x{assumed_s_layer3_input_difference:04x}")
        print(f"  预期的第三轮S层输出差分 (delta_V3 = delta_c): {[f'0x{val:04x}' for val in sorted(list(expected_s_layer3_output_differences))]}")
        
        # 调用函数恢复k4
        recovered_k4 = recover_last_round_key_k4(
            pt_ct_map,
            cipher_input_difference,
            expected_s_layer3_output_differences,
            assumed_s_layer3_input_difference
        )

        if recovered_k4 is not None:
            print(f"\n最可能的k4: 0x{recovered_k4:04x}")
        else:
            print("\n恢复k4失败，无有效候选。")    高概率差分0020到0020，为甚么密文匹配的才64个，这不对吧