抗量子算法，nist_pqc_alg，对bytes的操作，封装了对称加密，数字信封，签名和验签

"""
抗量子算法，nist_pqc_alg，对bytes的操作，
封装了对称加密krypton，数字信封envelope_，签名dss_sign，sphincs_sign，
验签verify_dss_sign，verify_sphincs_sign。
加密和签名的密钥需要分别生成。ml-kem-1024 封装密钥长度:   1568 字节，
解封装密钥长度: 3168 字节;共享密钥长度为 32 字节。
krypton 是二次封装了库中提供对称加密，其验证数据长度为160字节【16进制表示为320字符】

Author: tanglei
DateTime:2025-01-05 完成
微信：ciss_cedar
欢迎一起学习

"""
import secrets

from quantcrypt import kem,dss,cipher,kdf

from hashlib import sha256,sha3_256,sha3_384,sha3_512
from utils.algorithm.symmetric.symmetric_alg import MySymCipher,AlgName,CipherMode

# kem,     # Key Encapsulation Mechanism algorithms - public-key cryptography
# dss,     # Digital Signature Scheme algorithms - secret-key signatures
# cipher,  # The Krypton Cipher - symmetric cipher based on AES-256
# kdf,     # Argon2 helpers + KMAC-KDF - key derivation functions
# errors,  # All errors QuantCrypt may raise - also available from other modules
# utils    # Helper utilities from all modules - gathered into one module
# 加密的密钥只能用于加密，签名的密钥只能用于签名 和RSA，SM2不一样
#my_kdf=kdf.KKDF()

def krypton_genkey_bytes():
     key_bytes=secrets.token_bytes(64)
     return key_bytes

def krypton_encrypt_bytes(secret_key_bytes,plain_bytes):
    encryption = cipher.Krypton(secret_key_bytes)
    encryption.begin_encryption()
    cipher_bytes = encryption.encrypt(plain_bytes)
    verify_dp_bytes = encryption.finish_encryption()
    return cipher_bytes,verify_dp_bytes

def krypton_decrypt_bytes(secret_key_bytes,cipher_bytes,verify_dp_bytes):
    encryption = cipher.Krypton(secret_key_bytes)
    encryption.begin_decryption(verify_dp_bytes)
    decrypted_bytes = encryption.decrypt(cipher_bytes)
    encryption.finish_decryption()
    return decrypted_bytes

def gen_kyber_key_bytes():
    my_Kyber = kem.Kyber()
    public_key_bytes, secret_key_bytes = my_Kyber.keygen()
    return public_key_bytes, secret_key_bytes

def envelope_encrypt_bytes(public_key_bytes,plain_bytes):
    # 加密操作
    my_Kyber = kem.Kyber()
    cipher_key_bytes,shared_secret_bytes = my_Kyber.encaps(public_key_bytes)
    master_key = shared_secret_bytes
    key_len = 32
    key, iv = kdf.KKDF(master_key, key_len, 2)
    iv = iv[0:16]
    #shared_secret_bytes
    #print(f'key={key},iv={iv}')
    #key=shared_secret_bytes
    #iv = b'0102030405060708'
    alg_name = AlgName.AES.value
    mode = CipherMode.CTR.value
    my_alg = MySymCipher(key, iv, mode, alg_name)
    cipher_bytes = my_alg.encrypt_bytes(plain_bytes)
    return cipher_key_bytes,cipher_bytes


def envelope_decrypt_bytes(secret_key_bytes,cipher_key_bytes,cipher_bytes):
    my_Kyber = kem.Kyber()
    shared_secret_bytes = my_Kyber.decaps(secret_key_bytes,cipher_key_bytes)
    master_key = shared_secret_bytes
    key_len = 32
    key, iv = kdf.KKDF(master_key, key_len, 2)
    iv = iv[0:16]
    #iv=b'0102030405060708'
    #
    # print(f'key={key},iv={iv}')
    # shared_secret_bytes
    alg_name = AlgName.AES.value
    mode = CipherMode.CTR.value
    my_alg = MySymCipher(key, iv, mode, alg_name)
    dec_plain_bytes = my_alg.decrypt_bytes(cipher_bytes)
    return shared_secret_bytes,dec_plain_bytes

def gen_dss_sign_key_bytes():
    my_dss=dss.Dilithium()
    public_key_bytes, secret_key_bytes = my_dss.keygen()
    return public_key_bytes,secret_key_bytes

def dss_sign_bytes(secret_key_bytes,message_bytes):
    my_dss = dss.Dilithium()
    signature_bytes = my_dss.sign(secret_key_bytes, message_bytes)
    return signature_bytes

def verify_dss_sign_bytes(public_key_bytes,message_bytes,signature_bytes):
    my_dss = dss.Dilithium()
    verify_flag=my_dss.verify(public_key_bytes,message_bytes,signature_bytes)
    return verify_flag

def gen_sphincs_sign_key_bytes():
    my_FastSphincs = dss.FastSphincs()
    public_key_bytes, secret_key_bytes = my_FastSphincs.keygen()
    return public_key_bytes,secret_key_bytes

def sphincs_sign_bytes(secret_key_bytes,message_bytes):
    my_FastSphincs=dss.FastSphincs()
    signature_bytes = my_FastSphincs.sign(secret_key_bytes, message_bytes)
    return signature_bytes

def verify_sphincs_sign_bytes(public_key_bytes,message_bytes,signature_bytes):
    my_FastSphincs=dss.FastSphincs()
    verify_flag=my_FastSphincs.verify(public_key_bytes,message_bytes,signature_bytes)
    return verify_flag

def derive_key_bytes(master_key_bytes,key_len,num_keys=1):
    if num_keys> 65536:
        return b'error!num_keys<=65536'
    key_tuple_bytes = kdf.KKDF(master_key_bytes, key_len, num_keys=num_keys)

    #key_tuple = tuple(key_bytes.hex().upper() for key_bytes in key_tuple_bytes)

    return key_tuple_bytes

def main():

    secret_key_bytes = krypton_genkey_bytes()
    plain_bytes = b'this is a test for Krypton'

    # 创建加密实例并开始加密 AES(EAX)
    # AES（Advanced Encryption Standard）的 EAX 模式（Encrypt - then - Authenticate - then - Translate）
    # 是一种同时提供保密性和完整性认证的加密模式。它结合了加密和认证的功能，主要用于保护数据在传输或存储过程中的安全性
    # 和完整性
    print(f'secret_key_bytes[{len(secret_key_bytes)}]={secret_key_bytes}')
    print(f'plain_bytes[{len(plain_bytes)}]={plain_bytes}')
    cipher_bytes,verify_dp=krypton_encrypt_bytes(secret_key_bytes,plain_bytes)

    print(f'cipher_bytes[{len(cipher_bytes)}]={cipher_bytes}')
    print(f'verify_dp[{len(verify_dp)}]={verify_dp}')

    decrypted_bytes=krypton_decrypt_bytes(secret_key_bytes,cipher_bytes,verify_dp)
    print(f'decrypted_bytes[{len(decrypted_bytes)}]={decrypted_bytes}')


    # print(f'Encryption')
    # public_key_bytes, secret_key_bytes =gen_kyber_key_bytes()
    # print(f'public_key_bytes[{len(public_key_bytes)}]={public_key_bytes}')
    # print('-' * 66)
    # print(f'secret_key_bytes[{len(secret_key_bytes)}]={secret_key_bytes}')
    # print('-'*66)
    # #plain_bytes=b'this a test for envelope_encrypt(AES[CTR])'
    # plain_bytes=b'1'
    # print(f'plain_bytes[{len(plain_bytes)}]={plain_bytes}')
    # cipher_key_bytes,cipher_bytes=envelope_encrypt_bytes(public_key_bytes,plain_bytes)
    # print(f'cipher_key_bytes[{len(cipher_key_bytes)}]={cipher_key_bytes}')
    # print(f'cipher_bytes[{len(cipher_bytes)}]={cipher_bytes}')
    # print('-' * 66)
    # shared_secret_bytes,dec_plain_bytes=envelope_decrypt_bytes(secret_key_bytes, cipher_key_bytes, cipher_bytes)
    # print(f'shared_secret_bytes[{len(shared_secret_bytes)}]={shared_secret_bytes}')
    # print(f'dec_plain_bytes[{len(dec_plain_bytes)}]={dec_plain_bytes}')
    # print('=' * 66)

    # print(f'DSS-Sign')
    # dss_public_key_bytes, dss_secret_key_bytes = gen_dss_sign_key_bytes()
    # print(f'dss_public_key_bytes[{len(dss_public_key_bytes)}]={dss_public_key_bytes}')
    # print('-' * 66)
    # print(f'dss_secret_key_bytes[{len(dss_secret_key_bytes)}]={dss_secret_key_bytes}')
    # print('-'*66)
    # dss_message=b'1'
    # dss_signature_bytes=dss_sign_bytes(dss_secret_key_bytes, dss_message)
    # print(f'dss_signature_bytes[{len(dss_signature_bytes)}]={dss_signature_bytes}')
    # dss_verify_flag=verify_dss_sign_bytes(dss_public_key_bytes,dss_message,dss_signature_bytes)
    # print(f'dss_verify_flag={dss_verify_flag}')
    #
    #
    # print(f'FastSphincs-Sign')
    # sphincs_public_key_bytes, sphincs_secret_key_bytes = gen_sphincs_sign_key_bytes()
    # print(f'sphincs_public_key_bytes[{len(sphincs_public_key_bytes)}]={sphincs_public_key_bytes}')
    # print('-' * 66)
    # print(f'sphincs_secret_key_bytes[{len(sphincs_secret_key_bytes)}]={sphincs_secret_key_bytes}')
    # print('-'*66)
    # sphincs_message=b'1'
    # #sphincs_message=sha256(sphincs_message)
    # sphincs_signature_bytes=sphincs_sign_bytes(sphincs_secret_key_bytes, sphincs_message)
    # print(f'sphincs_signature_bytes[{len(sphincs_signature_bytes)}]={sphincs_signature_bytes}')
    # sphincs_verify_flag=verify_sphincs_sign_bytes(sphincs_public_key_bytes,sphincs_message,sphincs_signature_bytes)
    # print(f'sphincs_verify_flag={sphincs_verify_flag}')
    # print('-' * 66)
    # print('derive_key')
    # master_key=b'12345678901234561234567890123456'
    # key_len=32
    # num_keys=3
    # tuple_three=derive_key_bytes(master_key,key_len,num_keys=num_keys)
    # print(f'master_key={master_key}')
    # print(f'tuple_three[0]={tuple_three[0]}')
    # print(f'tuple_three[1]={tuple_three[1]}')
    # print(f'tuple_three[2]={tuple_three[2]}')


if __name__ == '__main__':

    main()