python 的 sm4算法，sm4_alg.py ,bytes计算，代码过多，sm4_apply分开放

sm4_alg.py，注意最小集将部分代码注销，可以取消注销即可

"""
Author:tanglei
DateTime:2024-10-12
微信：ciss_cedar
欢迎一起学习
通过cryptography实现sm4，ecb，cbc，gcm的二次封装
并为应sm4_apply提供底层支持

"""
import os

from Cryptodome.Util.Padding import pad, unpad
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.ciphers import Cipher
from cryptography.hazmat.primitives.ciphers import algorithms, modes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC

block_size = 16
style = 'pkcs7'


def sm4_ecb_encrypt_bytes(key_bytes, plain_bytes):
    backend = default_backend()
    cipher = algorithms.SM4(key_bytes)
    mode = modes.ECB()  # 使用ECB模式
    padded_plaintext = pad(plain_bytes, block_size, style)
    encryptor = Cipher(cipher, mode, backend).encryptor()
    cipher_bytes = encryptor.update(padded_plaintext) + encryptor.finalize()
    return cipher_bytes


def sm4_ecb_decrypt_bytes(key_bytes, cipher_bytes):
    backend = default_backend()
    cipher = algorithms.SM4(key_bytes)
    mode = modes.ECB()  # 使用ECB模式
    decryptor = Cipher(cipher, mode, backend).decryptor()
    padded_plain_bytes = decryptor.update(cipher_bytes) + decryptor.finalize()
    plaintext = unpad(padded_plain_bytes, len(padded_plain_bytes), style)
    return plaintext


def sm4_cbc_encrypt_bytes(key_bytes, iv_bytes, plain_bytes):
    backend = default_backend()
    cipher = algorithms.SM4(key_bytes)
    mode = modes.CBC(iv_bytes)  # 使用ECB模式
    padded_plain_bytes = pad(plain_bytes, block_size, style)
    encryptor = Cipher(cipher, mode, backend).encryptor()
    cipher_bytes = encryptor.update(padded_plain_bytes) + encryptor.finalize()
    return cipher_bytes


def sm4_cbc_decrypt_bytes(key_bytes, iv_bytes, cipher_bytes):
    backend = default_backend()
    cipher = algorithms.SM4(key_bytes)
    mode = modes.CBC(iv_bytes)  # 使用CBC模式
    decryptor = Cipher(cipher, mode, backend).decryptor()
    padded_plain_bytes = decryptor.update(cipher_bytes) + decryptor.finalize()
    plain_bytes = unpad(padded_plain_bytes, len(padded_plain_bytes), style)
    return plain_bytes


def sm4_gcm_encrypt_bytes(key_bytes, plain_bytes, aad_bytes=None):
    backend = default_backend()
    cipher = algorithms.SM4(key_bytes)
    padded_aad = pad(aad_bytes, block_size, style)
    mode = modes.GCM(padded_aad or b'')
    encryptor = Cipher(cipher, mode, backend).encryptor()
    padded_plain_bytes = pad(plain_bytes, block_size, style)
    cipher_bytes = encryptor.update(padded_plain_bytes) + encryptor.finalize()
    tag_bytes = encryptor.tag
    return cipher_bytes, tag_bytes


def sm4_gcm_decrypt_bytes(key_bytes, cipher_bytes, tag_bytes, aad_bytes=None):
    backend = default_backend()
    cipher = algorithms.SM4(key_bytes)
    padded_aad_bytes = pad(aad_bytes, block_size, style)
    mode = modes.GCM(padded_aad_bytes or b'', tag_bytes)
    decryptor = Cipher(cipher, mode, backend).decryptor()
    padded_plain_bytes = decryptor.update(cipher_bytes) + decryptor.finalize()
    plain_bytes = unpad(padded_plain_bytes, len(padded_plain_bytes), style)
    return plain_bytes


def gen_random_bytes(length=16):
    # SM4 需要一个 128 位的密钥，即 16 字节
    # 使用软件算法，可以调用系统函数或者使用secrets生成随机数再进行派生
    # 当前使用32字符，16字节的随机数也是可以的。
    random_bytes = os.urandom(length)
    return random_bytes


def gen_bpe_key_bytes(password, salt, key_length=16):
    # key_length = 16
    # 使用一个密码（可以是任意字符串）和一个盐（随机生成）来生成密钥
    # password = b'20181016simkey00'  # 请使用一个更安全的密码
    # salt = os.urandom(16)  # 生成一个 16 字节的随机盐
    if key_length <= 32:

        kdf = PBKDF2HMAC(
            algorithm=hashes.SM3(),
            length=key_length,
            salt=salt,
            iterations=10000,  # 10000次循环
            backend=default_backend()
        )
        key = kdf.derive(password)

    elif key_length > 32 and key_length <= 64:

        kdf = PBKDF2HMAC(
            algorithm=hashes.SM3(),
            length=key_length,
            salt=salt,
            iterations=10000,  # 10000次循环
            backend=default_backend()
        )
        key = kdf.derive(password)
    else:
        return 'The key must be less than or equal to 64 bytes,please call gen_random_bytes(key_length=n)'

    return key


# 使用示例
if __name__ == '__main__':
    # length=16
    # str_random=gen_random_bytes(length)
    # print(f'str_random{[length]}={str_random.hex().upper()}')
    # password='simkey.1'
    # password=password.encode()
    # salt='01020304050607080102030405060708'
    # salt = salt.encode()
    # length=33
    # key=gen_bpe_key_bytes(password,salt,length)
    # print(f'key{[length]}={key.hex().upper()}')
    # key = b'secret_sm4_key00'
    # plaintext = b'this is a secret message11111111111111111111111111111111111'
    # aad=b'secret_sm4_key00'
    # ciphertext, tag = sm4_gcm_encrypt_bytes(key,plaintext,aad)
    # print(ciphertext.hex().upper(),tag.hex().upper())
    # recovered_plaintext = sm4_gcm_decrypt_bytes(key,ciphertext,tag ,aad)
    # print(recovered_plaintext)
    # assert plaintext == recovered_plaintext, "Decryption failed"

    key = b'secret_sm4_key00'
    plaintext = b'this is a secret message11111111111111111111111111111111111'
    aad = b'secret_sm4_key00'
    ciphertext, tag = sm4_gcm_encrypt_bytes(key, plaintext, aad)
    print(f'key={key}')
    print(f'plaintext={plaintext}')
    print(f'ciphertext={ciphertext.hex().upper()}')
    print(f'tag={tag.hex().upper()}')

    recovered_plaintext = sm4_gcm_decrypt_bytes(key, ciphertext, tag, aad)
    print(f'recovered_plaintext={recovered_plaintext}')
    assert plaintext == recovered_plaintext, "Decryption failed"