from gmssl.sm4 import CryptSM4, SM4_ENCRYPT, SM4_DECRYPT
import binascii
from heapq import heappush, heappop
from collections import OrderedDict
class SM4():
    def __init__(self,key,mode):
        self.key=key.encode('utf-8')
        self.mode=mode
        self.iv = 'adsqqwertyuioplk'.encode('utf-8')
        self.crypt_sm4 = CryptSM4()
    def str_to_hexStr(self, hex_str):

        hex_data = hex_str.encode('utf-8')
        str_bin = binascii.unhexlify(hex_data)
        return str_bin.decode('utf-8')
    def encrypt(self,value):
        if self.mode=='ECB':
            crypt_sm4 = self.crypt_sm4
            crypt_sm4.set_key(self.key, SM4_ENCRYPT)
            encrypt_value = crypt_sm4.crypt_ecb(value.encode('utf-8')) #  bytes类型

        elif self.mode=='CBC':
            crypt_sm4 = self.crypt_sm4
            crypt_sm4.set_key(self.key, SM4_ENCRYPT)
            encrypt_value = crypt_sm4.crypt_cbc(self.iv , value.encode('utf-8')) #  bytes类型
        return encrypt_value.hex()

    def decrypt(self,encrypt_value):
        if self.mode=='ECB':
            crypt_sm4 = self.crypt_sm4
            crypt_sm4.set_key(self.key, SM4_DECRYPT)
            decrypt_value = crypt_sm4.crypt_ecb(bytes.fromhex(encrypt_value)) #  bytes类型

        elif self.mode=='CBC':
            crypt_sm4 = self.crypt_sm4
            crypt_sm4.set_key(self.key, SM4_DECRYPT)
            decrypt_value = crypt_sm4.crypt_cbc(self.iv , bytes.fromhex(encrypt_value)) #  bytes类型

        return self.str_to_hexStr(decrypt_value.hex())

str_data = "1234567890QWERTY"
key="3l5butlj26hvv313"
print("请输入加密模式")
mode=input()
SM4 = SM4(key,mode)
print("待加密内容：", str_data)
encoding = SM4.encrypt(str_data)
print("国密sm4加密后的结果：", encoding)
print("国密sm4解密后的结果：", SM4.decrypt(encoding))