from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad
import binascii

class PDCPEncryption:
    def __init__(self, key, count, bearer, direction):
        """
        初始化PDCP加密上下文
        :param key: 256-bit加密密钥 (32字节)
        :param count: PDCP COUNT值 (32位)
        :param bearer: 承载ID (5位)
        :param direction: 方向 (0=上行, 1=下行)
        """
        self.key = key
        self.count = count
        self.bearer = bearer
        self.direction = direction
        
        # 初始化加密算法 (5G使用AES-CTR)
        self.cipher = None
        self._update_cipher()
    
    def _update_cipher(self):
        """更新加密算法的初始向量(IV)"""
        # 5G PDCP IV构造 (3GPP TS 33.501)
        iv = bytearray(16)
        
        # IV[0..4] = COUNT (32位) + BEARER (5位) + DIRECTION (1位)
        iv[0] = (self.count >> 24) & 0xFF
        iv[1] = (self.count >> 16) & 0xFF
        iv[2] = (self.count >> 8) & 0xFF
        iv[3] = self.count & 0xFF
        iv[4] = ((self.bearer & 0x1F) << 3) | ((self.direction & 0x01) << 2)
        
        # IV[5..15] = 0 (保留位)
        for i in range(5, 16):
            iv[i] = 0x00
            
        # 创建AES-CTR加密器
        self.cipher = AES.new(self.key, AES.MODE_CTR, nonce=b'', initial_value=bytes(iv))
    
    def encrypt(self, plaintext):
        """
        PDCP数据加密
        :param plaintext: 明文数据
        :return: 加密后的数据
        """
        # 加密数据
        ciphertext = self.cipher.encrypt(plaintext)
        
        # 更新COUNT值 (在实际系统中需要处理COUNT回绕)
        self.count += 1
        self._update_cipher()
        
        return ciphertext
    
    def decrypt(self, ciphertext):
        """
        PDCP数据解密
        :param ciphertext: 密文数据
        :return: 解密后的数据
        """
        # 解密数据
        plaintext = self.cipher.decrypt(ciphertext)
        
        # 更新COUNT值
        self.count += 1
        self._update_cipher()
        
        return plaintext

# 示例用法
if __name__ == "__main__":
    # 示例参数
    key = b'\x00\x01\x02\x03\x04\x05\x06\x07\x08\x09\x0A\x0B\x0C\x0D\x0E\x0F' \
          b'\x10\x11\x12\x13\x14\x15\x16\x17\x18\x19\x1A\x1B\x1C\x1D\x1E\x1F'  # 256-bit密钥
    count = 0x12345678  # PDCP COUNT
    bearer = 0x1A       # 承载ID (5位)
    direction = 0       # 上行
    
    # 创建PDCP加密上下文
    pdcp = PDCPEncryption(key, count, bearer, direction)
    
    # 要加密的数据
    plaintext = b'This is a test message for 5G PDCP encryption!'
    print(f"原始明文: {plaintext}")
    print(f"十六进制: {binascii.hexlify(plaintext)}")
    
    # 加密
    ciphertext = pdcp.encrypt(plaintext)
    print(f"\n加密后的密文: {binascii.hexlify(ciphertext)}")
    
    # 重置COUNT以解密
    pdcp.count = count
    pdcp._update_cipher()
    
    # 解密
    decrypted = pdcp.decrypt(ciphertext)
    print(f"\n解密后的明文: {decrypted}")
    print(f"十六进制: {binascii.hexlify(decrypted)}")