抗量子算法，nist_pqc_apply，对str的操作，封装了对称加密，数字信封，签名，验签。

"""
抗量子算法，nist_pqc_apply，对str的操作，
封装了对称加密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 base64
import secrets

from quantcrypt import kdf

from utils.pqc.nist_pqc_alg import (gen_kyber_key_bytes, envelope_encrypt_bytes, envelope_decrypt_bytes,
                                    derive_key_bytes,
                                    gen_dss_sign_key_bytes, gen_sphincs_sign_key_bytes, sphincs_sign_bytes,
                                    dss_sign_bytes,
                                    verify_sphincs_sign_bytes, verify_dss_sign_bytes, krypton_encrypt_bytes,
                                    krypton_decrypt_bytes)


# 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()

code_tuple = ('Hex', 'base64')
mode_tuple = ('CBC', 'CTR')

def krypton_genkey():
    secret_key = secrets.token_bytes(64).hex().upper()
    return secret_key

def krypton_encrypt(secret_key,plain_text,code=code_tuple[0]):
    secret_key_bytes=bytes.fromhex(secret_key)
    plain_bytes=plain_text.encode()
    cipher_bytes,verify_dp_bytes=krypton_encrypt_bytes(secret_key_bytes,plain_bytes)

    if code.upper()==code_tuple[0].upper():
        cipher_text = cipher_bytes.hex().upper()
        verify_dp = verify_dp_bytes.hex().upper()
    else:
        cipher_text=base64.b64encode(cipher_bytes).decode()
        verify_dp=base64.b64encode(verify_dp_bytes).decode()

    return cipher_text,verify_dp

def krypton_decrypt(secret_key,cipher_text,verify_dp,code=code_tuple[0]):
    secret_key_bytes=bytes.fromhex(secret_key)
    if code.upper()==code_tuple[0].upper():
        cipher_bytes = bytes.fromhex(cipher_text)
        verify_dp_bytes = bytes.fromhex(verify_dp)
    else :
        cipher_bytes = base64.b64decode(cipher_text)
        verify_dp_bytes=base64.b64decode(verify_dp)

    decrypted_bytes=krypton_decrypt_bytes(secret_key_bytes,cipher_bytes,verify_dp_bytes)
    decrypted_text=decrypted_bytes.decode()
    return decrypted_text


def gen_kyber_key():
    public_key_bytes, secret_key_bytes = gen_kyber_key_bytes()
    public_key=public_key_bytes.hex().upper()
    secret_key=secret_key_bytes.hex().upper()
    return public_key, secret_key

def envelope_encrypt(public_key,plain_text):
    public_key_bytes=bytes.fromhex(public_key)
    plain_bytes=plain_text.encode()
    cipher_key_bytes,cipher_bytes=envelope_encrypt_bytes(public_key_bytes,plain_bytes)
    cipher_key=cipher_key_bytes.hex().upper()
    cipher_text=cipher_bytes.hex().upper()
    return cipher_key,cipher_text


def envelope_decrypt(secret_key,cipher_key,cipher_text):
    secret_key_bytes=bytes.fromhex(secret_key)
    cipher_key_bytes=bytes.fromhex(cipher_key)
    cipher_bytes=bytes.fromhex(cipher_text)
    shared_secret_bytes,dec_plain_bytes=envelope_decrypt_bytes(secret_key_bytes, cipher_key_bytes, cipher_bytes)
    dec_plain_text=dec_plain_bytes.decode()
    shared_secret_key=shared_secret_bytes.hex().upper()
    return shared_secret_key,dec_plain_text

def gen_dss_sign_key():
    public_key_bytes,secret_key_bytes = gen_dss_sign_key_bytes()
    public_key=public_key_bytes.hex().upper()
    secret_key=secret_key_bytes.hex().upper()
    return public_key,secret_key

def dss_sign(secret_key,message):
    secret_key_bytes=bytes.fromhex(secret_key)
    message_bytes=message.encode()
    signature_bytes = dss_sign_bytes(secret_key_bytes,message_bytes)
    signature=signature_bytes.hex().upper()
    return signature

def verify_dss_sign(public_key,message,signature):
    public_key_bytes=bytes.fromhex(public_key)
    message_bytes=message.encode()
    signature_bytes=bytes.fromhex(signature)
    verify_flag=verify_dss_sign_bytes(public_key_bytes,message_bytes,signature_bytes)
    return verify_flag

def gen_sphincs_sign_key():
    public_key_bytes,secret_key_bytes=gen_sphincs_sign_key_bytes()
    public_key=public_key_bytes.hex().upper()
    secret_key=secret_key_bytes.hex().upper()
    return public_key,secret_key

def sphincs_sign(secret_key,message):
    secret_key_bytes=bytes.fromhex(secret_key)
    message_bytes=message.encode()
    signature_bytes = sphincs_sign_bytes(secret_key_bytes,message_bytes)
    signature=signature_bytes.hex().upper()
    return signature

def verify_sphincs_sign(public_key,message,signature):
    public_key_bytes=bytes.fromhex(public_key)
    message_bytes=message.encode()
    signature_bytes=bytes.fromhex(signature)
    verify_flag=verify_sphincs_sign_bytes(public_key_bytes,message_bytes,signature_bytes)
    return verify_flag

def derive_key(master_key,key_len,num_keys=1):
    if num_keys> 65536:
        return b'error!num_keys<=65536'
    master_key_bytes=master_key.encode()
    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

def main():

    secret_key = krypton_genkey()
    plain_text = 'this is a test for Krypton'
    code='hex'

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

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

    decrypted_text=krypton_decrypt(secret_key,cipher_text,verify_dp,code)
    print(f'decrypted_text[{len(decrypted_text)}]={decrypted_text}')

    # print(f'Encryption')
    # public_key, secret_key =gen_kyber_key()
    # print(f'public_key[{len(public_key)}]={public_key}')
    # print('-' * 66)
    # print(f'secret_key[{len(secret_key)}]={secret_key}')
    # print('-'*66)
    # #plain_text='this a test for envelope_encrypt(AES[CTR])'
    # plain_text='1'
    # print(f'plain_text[{len(plain_text)}]={plain_text}')
    #
    # cipher_key,cipher_text=envelope_encrypt(public_key,plain_text)
    # print(f'cipher_key[{len(cipher_key)}]={cipher_key}')
    # print(f'cipher_text[{len(cipher_text)}]={cipher_text}')
    #
    # print('-' * 66)
    # shared_secret_key,dec_plain_text=envelope_decrypt(secret_key, cipher_key, cipher_text)
    # print(f'shared_secret_key[{len(shared_secret_key)}]={shared_secret_key}')
    # print(f'dec_plain_text[{len(dec_plain_text)}]={dec_plain_text}')
    # print('=' * 66)

    # print(f'DSS-Sign')
    # dss_public_key, dss_secret_key = gen_dss_sign_key()
    # print(f'dss_public_key[{len(dss_public_key)}]={dss_public_key}')
    # print('-' * 66)
    # print(f'dss_secret_key[{len(dss_secret_key)}]={dss_secret_key}')
    # print('-'*66)
    #
    # dss_message='1'
    # dss_signature=dss_sign(dss_secret_key,dss_message)
    # print(f'dss_signature[{len(dss_signature)}]={dss_signature}')
    # dss_verify_flag=verify_dss_sign(dss_public_key,dss_message,dss_signature)
    # print(f'dss_verify_flag={dss_verify_flag}')
    #
    #
    # print(f'FastSphincs-Sign')
    # sphincs_public_key, sphincs_secret_key = gen_sphincs_sign_key()
    # print(f'sphincs_public_key[{len(sphincs_public_key)}]={sphincs_public_key}')
    # print('-' * 66)
    # print(f'sphincs_secret_key[{len(sphincs_secret_key)}]={sphincs_secret_key}')
    # print('-'*66)
    # sphincs_message='1'
    # #sphincs_message=sha256(sphincs_message)
    # sphincs_signature=sphincs_sign(sphincs_secret_key, sphincs_message)
    # print(f'sphincs_signature[{len(sphincs_signature)}]={sphincs_signature}')
    # sphincs_verify_flag=verify_sphincs_sign(sphincs_public_key,sphincs_message,sphincs_signature)
    # print(f'sphincs_verify_flag={sphincs_verify_flag}')
    #
    # print('-' * 66)
    # print('derive_key')
    # master_key='12345678901234561234567890123456'
    # key_len=32
    # num_keys=3
    # tuple_three=derive_key(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()