import os
import binascii
from Crypto.Cipher import AES
from Crypto.Util import Counter

class PDCPNEA2Cipher:
    def __init__(self, key, count, bearer, direction):
        """
        初始化PDCP NEA2加密器
        
        :param key: 16字节加密密钥（十六进制字符串或字节）
        :param count: 32位PDCP COUNT值
        :param bearer: 5位承载标识
        :param direction: 1位方向标识(0:上行, 1:下行)
        """
        # 处理密钥输入
        if isinstance(key, str):
            # 如果是十六进制字符串，转换为字节
            self.key = binascii.unhexlify(key)
        else:
            self.key = key
            
        if len(self.key) != 16:
            raise ValueError("密钥必须是16字节长度")
        
        self.count = count
        self.bearer = bearer
        self.direction = direction
        
    def _generate_iv(self):
        """生成初始化向量(IV)"""
        # 根据3GPP TS 33.401规范构建IV
        # IV = COUNT[32] | BEARER[5] | DIRECTION[1] | 0[26]
        iv = bytearray(16)
        
        # 设置COUNT (4字节，大端序)
        iv[0] = (self.count >> 24) & 0xFF
        iv[1] = (self.count >> 16) & 0xFF
        iv[2] = (self.count >> 8) & 0xFF
        iv[3] = self.count & 0xFF
        
        # 设置BEARER和DIRECTION
        # 第4字节: BEARER[5] + DIRECTION[1] + 0[2]
        iv[4] = ((self.bearer & 0x1F) << 3) | ((self.direction & 0x01) << 2)
        
        # 剩余字节保持为0
        return bytes(iv)
    
    def encrypt(self, plaintext, hex_input=False, hex_output=False):
        """
        加密数据
        
        :param plaintext: 明文数据（字节或十六进制字符串）
        :param hex_input: 输入是否为十六进制字符串
        :param hex_output: 是否输出十六进制字符串
        :return: 加密后的数据
        """
        # 处理输入数据
        if hex_input:
            # 如果是十六进制字符串，转换为字节
            data = binascii.unhexlify(plaintext)
        else:
            data = plaintext
        
        # 生成IV
        iv = self._generate_iv()
        
        # 创建CTR模式的计数器
        # NEA2使用完整的16字节IV作为初始计数器值
        ctr = Counter.new(128, initial_value=int.from_bytes(iv, byteorder='big'))
        
        # 创建AES-CTR密码器
        cipher = AES.new(self.key, AES.MODE_CTR, counter=ctr)
        
        # 加密数据
        ciphertext = cipher.encrypt(data)
        
        # 处理输出格式
        if hex_output:
            return binascii.hexlify(ciphertext).decode('utf-8')
        else:
            return ciphertext
    
    def decrypt(self, ciphertext, hex_input=False, hex_output=False):
        """
        解密数据
        
        :param ciphertext: 密文数据（字节或十六进制字符串）
        :param hex_input: 输入是否为十六进制字符串
        :param hex_output: 是否输出十六进制字符串
        :return: 解密后的数据
        """
        # 处理输入数据
        if hex_input:
            # 如果是十六进制字符串，转换为字节
            data = binascii.unhexlify(ciphertext)
        else:
            data = ciphertext
        
        # 解密过程与加密相同（CTR模式对称）
        # 生成IV
        iv = self._generate_iv()
        
        # 创建CTR模式的计数器
        ctr = Counter.new(128, initial_value=int.from_bytes(iv, byteorder='big'))
        
        # 创建AES-CTR密码器
        cipher = AES.new(self.key, AES.MODE_CTR, counter=ctr)
        
        # 解密数据
        plaintext = cipher.decrypt(data)
        
        # 处理输出格式
        if hex_output:
            return binascii.hexlify(plaintext).decode('utf-8')
        else:
            return plaintext

# 示例使用
if __name__ == "__main__":
    # 测试参数 - 使用十六进制格式
    key_hex = "00112233445566778899aabbccddeeff"  # 128位密钥（十六进制字符串）
    count = 0x12345678    # 32位COUNT值
    bearer = 1         # 承载标识 (最大31)，BR_ID - 1（DRB2,bearer则为1）
    direction = 0         # 数据方向，0为上行，1为下行
    
    # 创建加密器实例
    cipher = PDCPNEA2Cipher(key_hex, count, bearer, direction)
    
    # 示例明文数据（十六进制字符串）
    plaintext_hex = "48656c6c6f2050444350204e45413220656e6372797074696f6e21"
    
    # 加密（输入和输出都使用十六进制）
    ciphertext_hex = cipher.encrypt(plaintext_hex, hex_input=True, hex_output=True)
    print(f"明文(hex): {plaintext_hex}")
    print(f"密文(hex): {ciphertext_hex}")
    
    # 解密（输入和输出都使用十六进制）
    decrypted_hex = cipher.decrypt(ciphertext_hex, hex_input=True, hex_output=True)
    print(f"解密后(hex): {decrypted_hex}")
    
    # 验证加解密是否正确
    assert plaintext_hex == decrypted_hex, "加解密测试失败!"
    print("加解密测试成功!")
    
    # 也可以将十六进制结果转换回可读文本
    decrypted_text = binascii.unhexlify(decrypted_hex).decode('utf-8')
    print(f"解密文本: {decrypted_text}")