python 的sm2的封装，包括数字信封等

"""
Author: tanglei
DateTime:2024-11-18 完成
微信：ciss_cedar
欢迎一起学习
"""

from gmssl import sm2, func
from sm2_genkey import SM2_Key
from sm3_apply import sm3_hash
from sm4_apply import sm4_ecb_encrypt,sm4_ecb_decrypt


def sm2_digital_envelope(public_key,sm4_key,source):
    """
    数字信封
    :param public_key:str 公钥
    :param sm4_key:str sm4 ecb 模式 对称密钥
    :param source:str 数据源
    :return:enc_sm4_key 加密后的sm4密钥,enc_source sm4 ecb加密后的数据
    """
    enc_sm4_key=sm2_encrypt(public_key,sm4_key)
    enc_source=sm4_ecb_encrypt(sm4_key,source)
    return enc_sm4_key,enc_source

def sm2_digital_envelope_decrypt(private_key,enc_sm4_key,enc_source):
    """
    sm2数字信封解密
    :param private_key: str 私钥
    :param enc_sm4_key: str 加密后的sm4 ecb密钥
    :param enc_source: sm4 ecb 模式加密后的数据
    :return: str key,source 数据源
    """
    sm4_key=sm2_decrypt(private_key,enc_sm4_key)
    source=sm4_ecb_decrypt(sm4_key,enc_source)
    return sm4_key,source

# # 生成SM2公私钥对
def sm2_keypair_generate():
    """
    :return: public_key, private_key 公钥，私钥
    """
    private_key = ''
    public_key = ''
    ecc_table =SM2_Key.default_ecc_table
    my_sm2 = SM2_Key(private_key=private_key, public_key=public_key, ecc_table=ecc_table, mode=1)
    public_key, private_key = my_sm2.generate_keypair()
    return public_key, private_key

def sm2_encrypt(public_key,source):
    """
    SM2 加密 mode=1 默认c1c3c2
    :param public_key: str 公钥
    :param private_key: str 私钥
    :param source: str 数据源
    :return: 加密结果 16进制字符串
    """
    private_key=''
    sm2_crypt = sm2.CryptSM2(
        public_key=public_key, private_key=private_key, mode=1)
    source = source.encode()
    enc_source = sm2_crypt.encrypt(source)
    return enc_source.hex().upper()

def sm2_decrypt(private_key,source):
    """
    sm2 解密
    :param private_key: str 私钥
    :param source: sm2解密结果
    :return: 解密结果
    """
    public_key=''
    sm2_crypt = sm2.CryptSM2(
        public_key=public_key, private_key=private_key, mode=1)

    source = bytes.fromhex(source)
    dec_source =sm2_crypt.decrypt(source)
    return dec_source.decode()


def sm2_sign(public_key,private_key,source):
    """非标准签名,对原文签名没有进行SM3，不建议使用
    :param public_key: str 公钥
    :param private_key: str 私钥
    :param source: 需要签名的原始数据
    :return: 签名结果 str 16进制
    """
    data=source.encode()
    # 对接java 时验签失败可以使用
    # sm2_crypt = sm2.CryptSM2(
    #     public_key=public_key, private_key=private_key, asn1=True)

    sm2_crypt = sm2.CryptSM2(
        public_key=public_key, private_key=private_key)
    random_hex_str = func.random_hex(sm2_crypt.para_len)
    sign = sm2_crypt.sign(data, random_hex_str)  # 16进制
    return sign.upper()

def sm2_sign_verify(public_key,sign,source):
    """
     非标准签名的验签，不建议使用
    :param public_key: 公钥 str
    :param private_key: 私钥 str
    :param sign: 签名值 str 16进制
    :param source: 签名原始数据 str
    :return: bool
    """

    data=source.encode()
    #sign=bytes.fromhex(sign)
    # 对接java 时验签失败可以使用
    # sm2_crypt = sm2.CryptSM2(
    #     public_key=public_key, private_key=private_key, asn1=True)
    # SM2签名算法的结果通常包含两个大整数R和S，这两个整数是无符号的。但在DER编码过程中，如果R或S的最高位是1（即表示负数，因为在计算机中整数通常采用补码表示），则需要在这些整数的最前面填充一个00字节，以使其成为一个正整数进行表示。这个填充过程会增加签名结果的长度。
    # 当R和S的最高位都是0时，不需要填充，签名结果会包含6个字节的标志字节（用于表示整数类型和长度），以及R和S的值，累计70字节。
    # 当R或S中只有一个的最高位是1时，需要填充一个00字节，签名结果会包含7个字节的标志字节（其中一个用于填充），以及R和S的值，累计71字节。
    # 当R和S的最高位都是1时，都需要填充00字节，签名结果会包含8个字节的标志字节（其中两个用于填充），以及R和S的值，累计72字节。
    #
    private_key=''
    sm2_crypt = sm2.CryptSM2(
        public_key=public_key, private_key=private_key)
    source_verify_flag=sm2_crypt.verify(sign, data)
    return source_verify_flag


def sm2_sign_with_sm3(public_key,private_key,source):
    """
    标准签名,签名会用到公钥信息
    :param public_key: 公钥 str
    :param private_key: 私钥 str
    :param source: 需要签名的数据源 str
    :return: 签名结果 str 16进制
    """
    data = source.encode()  # bytes类型
    # 对接java 时验签失败可以使用
    # sm2_crypt = sm2.CryptSM2(
    #     public_key=public_key, private_key=private_key, asn1=True)

    sm2_crypt = sm2.CryptSM2(
        public_key=public_key, private_key=private_key)
    sign = sm2_crypt.sign_with_sm3(data)  # 16进制
    return sign.upper()

def sm2_sign_with_sm3_verify(public_key,sign,source):
    """
    标准验签
    :param public_key:公钥 str
    :param private_key: 私钥 str
    :param sign: 签名值
    :param source: 用于签名的原始数据
    :return: 验签结果 bool
    """
    private_key=''
    data = source.encode()  # bytes类型
    # 对接java 时验签失败可以使用
    # sm2_crypt = sm2.CryptSM2(
    #     public_key=public_key, private_key=private_key, asn1=True)

    sm2_crypt = sm2.CryptSM2(
        public_key=public_key, private_key=private_key)
    std_verify_flag=sm2_crypt.verify_with_sm3(sign, data)  # 16进制
    return std_verify_flag


if __name__ == '__main__':

    public_key = 'B9C9A6E04E9C91F7BA880429273747D7EF5DDEB0BB2FF6317EB00BEF331A83081A6994B8993F3F5D6EADDDB81872266C87C018FB4162F5AF347B483E24620207'
    private_key = 'B9AB0B828FF68872F21A837FC303668428DEA11DCD1B24429D0C99E24EED83D5'
    source = 'a'
    sign=sm2_sign(public_key,private_key,source)
    print(f'source={source}')
    print(f'sign={sign}')
    verify_flag=sm2_sign_verify(public_key,sign,source)
    print(f'sign verify={verify_flag}')
    print('=' * 66)
    sm3_value=sm3_hash(source)
    sign_sm3_value = sm2_sign(public_key, private_key, sm3_value)
    print(f'sm3_value={sm3_value}')
    print(f'sign_sm3_value={sign_sm3_value}')
    verify_flag = sm2_sign_verify(public_key, sign_sm3_value, sm3_value)
    print(f'sign verify={verify_flag}')
    print('-' * 66)
    print(f'source={source}')
    sign_sm3 = sm2_sign_with_sm3(public_key, private_key, source)
    print(f'sm3_value={sm3_value}')
    print(f'sign_sm3={sign_sm3}')
    verify_flag = sm2_sign_with_sm3_verify(public_key, sign_sm3, source)
    print(f'sign_sm3 verify={verify_flag}')
    print('=' * 66)
    python_ciphertext = sm2_encrypt(public_key,source)
    csharp_ciphertext='6909BDB065901478065130E53A72EA9832CE543B5CFE26E27DF6DCEAF924B9CB9981A1124DBA972B17CF481C0C604230B604CE3E9DD24E0188285EC728F07DE798E60DD8765767CD9C519A509E1639A36EF7FF632F7BB3CAC686C4B04BF9C29847'
    python_plaintext = sm2_decrypt(private_key,python_ciphertext)
    csharp_plaintext = sm2_decrypt(private_key,csharp_ciphertext)
    print(f'source={source}')
    print(f'public_key={public_key}')
    print(f'private_key={private_key}')
    print(f'python_ciphertext={python_ciphertext}')
    print(f'csharp_ciphertext={csharp_ciphertext}')
    print(f'python_plaintext={python_plaintext}')
    print(f'csharp_plaintext={csharp_plaintext}')
    print(f'source==python_plaintext={source==python_plaintext}')
    print(f'source==csharp_plaintext={source == csharp_plaintext}')
    print('-' * 66)

    sm4_key='2934412A66B7A186DC35DC40E926F9EE'
    enc_sm4_key,enc_source=sm2_digital_envelope(public_key,sm4_key,source)
    print(f'enc_sm4_key={enc_sm4_key}')
    print(f'enc_source={enc_source}')
    print(f'sm4_key={sm4_key}')
    dec_sm4_key=sm2_decrypt(private_key,enc_sm4_key)
    print(f'dec_sm4_key={dec_sm4_key}')
    print(f'sm4_key==dec_sm4_key={sm4_key == dec_sm4_key}')
    dec_source=sm2_digital_envelope_decrypt(private_key,enc_sm4_key,enc_source)
    print(f'dec_source={dec_source}')
    print(f'source==dec_source={source==dec_source}')
    print('=' * 66)