# 伪代码示例
def initialize_state(key, nonce, counter):
    state = [0] * 16  # 16 个 32 位字
    
    # 位置 [0-3]: 固定常量
    state[0] = 0x61707865  # "expa"
    state[1] = 0x3320646e  # "nd 3"
    state[2] = 0x79622d32  # "2-by"
    state[3] = 0x6b206574  # "te k"
    
    # 位置 [4-11]: 256 位密钥（小端序）
    state[4:12] = bytes_to_words(key, 8)
    
    # 位置 [12]: 32 位计数器
    state[12] = counter
    
    # 位置 [13-15]: 96 位 Nonce
    state[13:16] = bytes_to_words(nonce, 3)
    
    return state

def quarter_round(state, a, b, c, d):
    """
    对状态矩阵的索引 a, b, c, d 位置执行 Quarter Round
    
    参数说明：
        state: 包含 16 个 32 位字的状态矩阵列表
        a, b, c, d: 状态矩阵的索引（0-15），指定要混合的 4 个字的位置
    
    示例：quarter_round(state, 0, 4, 8, 12)
        表示对 state[0], state[4], state[8], state[12] 这 4 个位置进行混合
        在 4×4 矩阵中，这对应第 1 列的所有元素
    
    所有运算都是模 2^32 进行
    """
    # 第一步：a += b; d ^= a; d <<<= 16
    state[a] = (state[a] + state[b]) & 0xFFFFFFFF
    state[d] ^= state[a]
    state[d] = ((state[d] << 16) | (state[d] >> 16)) & 0xFFFFFFFF
    
    # 第二步：c += d; b ^= c; b <<<= 12
    state[c] = (state[c] + state[d]) & 0xFFFFFFFF
    state[b] ^= state[c]
    state[b] = ((state[b] << 12) | (state[b] >> 20)) & 0xFFFFFFFF
    
    # 第三步：a += b; d ^= a; d <<<= 8
    state[a] = (state[a] + state[b]) & 0xFFFFFFFF
    state[d] ^= state[a]
    state[d] = ((state[d] << 8) | (state[d] >> 24)) & 0xFFFFFFFF
    
    # 第四步：c += d; b ^= c; b <<<= 7
    state[c] = (state[c] + state[d]) & 0xFFFFFFFF
    state[b] ^= state[c]
    state[b] = ((state[b] << 7) | (state[b] >> 25)) & 0xFFFFFFFF

# 示例：32 位数循环左移 16 位
原始值： 0b10101100_11110000_00001111_11001100
         ↓ 循环左移 16 位（<<<= 16）
结果值： 0b00001111_11001100_10101100_11110000
         └─────┬──────┘ └─────┬──────┘
              后16位      前16位（循环到后面）

# Python 实现
def rotate_left(value, shift, bits=32):
    mask = (1 << bits) - 1  # 0xFFFFFFFF for 32-bit
    return ((value << shift) | (value >> (bits - shift))) & mask

# 示例
x = 0xACF00FCC
y = rotate_left(x, 16)  # 循环左移 16 位
print(f"0x{y:08X}")  # 0x0FCCACF0

# 不好的选择示例
# 旋转 0 位：相当于没有旋转
# 旋转 32 位：回到原位，无意义
# 旋转 1/31 位：扩散太慢，需要更多轮次
# 旋转 4/8/16：2 的幂次，可能产生规律性模式

# ChaCha20 的选择经过验证
# 16: 快速跨半字
# 12: 不规则，增强混合
# 8:  跨字节
# 7:  最大化扩散

# 观察单个比特变化的扩散
def show_diffusion():
    state1 = [0] * 16
    state2 = [0] * 16
    state2[0] = 1  # 只改变 1 个比特
    
    # 执行 1 次 Quarter Round
    quarter_round(state1, 0, 4, 8, 12)
    quarter_round(state2, 0, 4, 8, 12)
    
    # 比较差异
    for i in [0, 4, 8, 12]:
        diff_bits = bin(state1[i] ^ state2[i]).count('1')
        print(f"位置 {i}: 变化了 {diff_bits} 个比特")
    
    # 输出示例：
    # 位置 0: 变化了 18 个比特  ← 从 1 位扩散到 18 位
    # 位置 4: 变化了 15 个比特
    # 位置 8: 变化了 16 个比特
    # 位置 12: 变化了 17 个比特

def chacha20_block(key, nonce, counter):
    # 初始化状态
    state = initialize_state(key, nonce, counter)
    working_state = state.copy()
    
    # 10 个双轮 = 20 轮
    for i in range(10):
        # 列轮（Column Round）
        quarter_round(working_state, 0, 4, 8, 12)   # 第 1 列
        quarter_round(working_state, 1, 5, 9, 13)   # 第 2 列
        quarter_round(working_state, 2, 6, 10, 14)  # 第 3 列
        quarter_round(working_state, 3, 7, 11, 15)  # 第 4 列
        
        # 对角轮（Diagonal Round）
        quarter_round(working_state, 0, 5, 10, 15)  # 对角线 1
        quarter_round(working_state, 1, 6, 11, 12)  # 对角线 2
        quarter_round(working_state, 2, 7, 8, 13)   # 对角线 3
        quarter_round(working_state, 3, 4, 9, 14)   # 对角线 4
    
    # 将工作状态与初始状态相加（防止逆向推导）
    for i in range(16):
        working_state[i] = (working_state[i] + state[i]) & 0xFFFFFFFF
    
    return working_state

# 最终输出
keystream_block = []
for i in range(16):
    output_word = (working_state[i] + initial_state[i]) & 0xFFFFFFFF
    keystream_block.append(output_word)

# 将 16 个 32 位字转换为 64 字节
keystream_bytes = words_to_bytes(keystream_block)

def chacha20_encrypt(key, nonce, plaintext, initial_counter=0):
    """
    ChaCha20 加密函数
    """
    ciphertext = bytearray()
    counter = initial_counter
    
    # 将明文分成 64 字节的块
    for i in range(0, len(plaintext), 64):
        # 生成当前块的密钥流
        keystream_block = chacha20_block(key, nonce, counter)
        keystream_bytes = words_to_bytes(keystream_block)
        
        # 获取当前明文块
        plaintext_block = plaintext[i:i+64]
        
        # XOR 操作：C = P ⊕ K
        for j in range(len(plaintext_block)):
            ciphertext.append(plaintext_block[j] ^ keystream_bytes[j])
        
        # 计数器递增
        counter += 1
    
    return bytes(ciphertext)

def chacha20_decrypt(key, nonce, ciphertext, initial_counter=0):
    """
    ChaCha20 解密函数（与加密完全相同）
    """
    # 流密码的对称性：解密 = 加密
    return chacha20_encrypt(key, nonce, ciphertext, initial_counter)

# 示例：明文长度 100 字节
# 第 1 块：使用密钥流的前 64 字节
# 第 2 块：只使用密钥流的前 36 字节（剩余 28 字节被丢弃）

plaintext_block = plaintext[64:100]  # 36 字节
keystream_block = chacha20_block(key, nonce, counter=1)  # 生成 64 字节
keystream_bytes = words_to_bytes(keystream_block)[:36]  # 只取前 36 字节

# XOR 操作
for j in range(36):
    ciphertext.append(plaintext_block[j] ^ keystream_bytes[j])

# 测试：改变 Nonce 的 1 比特
key = os.urandom(32)
nonce1 = b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00'
nonce2 = b'\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00'  # 第 8 字节改变 1 位

keystream1 = chacha20_block(key, nonce1, 0)
keystream2 = chacha20_block(key, nonce2, 0)

# 计算汉明距离
diff_bits = count_diff_bits(keystream1, keystream2)
print(f"改变比特数: {diff_bits} / 512")  # 通常约 256 位（50%）

# 攻击场景
C1 = P1 ⊕ Keystream(K, N, 0)
C2 = P2 ⊕ Keystream(K, N, 0)  # 重用了相同的 (K, N)

# 攻击者可以计算：
C1 ⊕ C2 = (P1 ⊕ Keystream) ⊕ (P2 ⊕ Keystream)
         = P1 ⊕ P2  # 密钥流被消除！

# 方法 1：递增计数器（推荐用于单一发送方）
class ChaCha20Encryptor:
    def __init__(self, key):
        self.key = key
        self.nonce_counter = 0
    
    def encrypt(self, plaintext):
        nonce = self.nonce_counter.to_bytes(12, 'little')
        self.nonce_counter += 1
        return chacha20_encrypt(self.key, nonce, plaintext)

# 方法 2：随机生成（96 位 Nonce 空间足够大）
import os
def encrypt_with_random_nonce(key, plaintext):
    nonce = os.urandom(12)  # 96 位随机 Nonce
    ciphertext = chacha20_encrypt(key, nonce, plaintext)
    return nonce + ciphertext  # Nonce 附加在密文前

# 攻击示例
C = P ⊕ Keystream
C' = C ⊕ mask  # 攻击者翻转某些位

# 解密后
P' = C' ⊕ Keystream = (C ⊕ mask) ⊕ Keystream = P ⊕ mask
# 攻击者成功修改了明文！

def chacha20_poly1305_encrypt(key, nonce, plaintext, associated_data=b''):
    """
    ChaCha20-Poly1305 认证加密
    """
    # 1. 使用 ChaCha20 加密
    ciphertext = chacha20_encrypt(key, nonce, plaintext, initial_counter=1)
    
    # 2. 生成 Poly1305 密钥（使用 counter=0）
    poly_key = chacha20_block(key, nonce, counter=0)[:32]
    
    # 3. 构造认证数据
    mac_data = associated_data + pad16(associated_data)
    mac_data += ciphertext + pad16(ciphertext)
    mac_data += len(associated_data).to_bytes(8, 'little')
    mac_data += len(ciphertext).to_bytes(8, 'little')
    
    # 4. 计算 MAC 标签
    tag = poly1305_mac(poly_key, mac_data)
    
    return ciphertext + tag  # 密文 + 16 字节 MAC 标签

# 伪代码：使用 SIMD 指令并行处理 4 个状态
def chacha20_simd_4x(keys, nonces, counters):
    # 同时处理 4 个独立的 ChaCha20 块
    states = [initialize_state(k, n, c) for k, n, c in zip(keys, nonces, counters)]
    
    for i in range(10):  # 10 个双轮
        # 使用 AVX2/NEON 指令并行执行 Quarter Round
        simd_quarter_round_4x(states, ...)
    
    return [words_to_bytes(s) for s in states]

# 安全的 ChaCha20 使用模板
from cryptography.hazmat.primitives.ciphers.aead import ChaCha20Poly1305
import os

# 1. 生成密钥（使用 CSPRNG）
key = ChaCha20Poly1305.generate_key()

# 2. 创建加密器
cipher = ChaCha20Poly1305(key)

# 3. 加密（自动生成 Nonce）
nonce = os.urandom(12)
associated_data = b"metadata"  # 可选的关联数据
ciphertext = cipher.encrypt(nonce, plaintext, associated_data)

# 4. 解密（验证 MAC）
try:
    plaintext = cipher.decrypt(nonce, ciphertext, associated_data)
except Exception:
    print("认证失败！密文被篡改或 Nonce/密钥错误")