from gmssl.sm4 import CryptSM4, SM4_ENCRYPT, SM4_DECRYPT
import binascii
from heapq import heappush, heappop
from collections import OrderedDict
def bytesToString(bs):
    return bytes.decode(bs,encoding='utf8')

class SM4:
    """
    国密sm4加解密
    """

    def __init__(self):
        self.crypt_sm4 = CryptSM4()

    def str_to_hexStr(self, hex_str):
        """
        字符串转hex
        :param hex_str: 字符串
        :return: hex
        """
        hex_data = hex_str.encode('utf-8')
        str_bin = binascii.unhexlify(hex_data)
        return str_bin.decode('utf-8')

    def encrypt(self, encrypt_key, value):
        """
        国密sm4加密
        :param encrypt_key: sm4加密key
        :param value: 待加密的字典
        :return: sm4加密后的hex值
        """
        value = self.toTreeMap(value)
        crypt_sm4 = self.crypt_sm4
        crypt_sm4.set_key(encrypt_key.encode(), SM4_ENCRYPT)
        encrypt_value = crypt_sm4.crypt_ecb(value.encode())  # bytes类型
        return encrypt_value.hex()

    def decrypt(self, decrypt_key, encrypt_value):
        """
        国密sm4解密
        :param decrypt_key:sm4加密key
        :param encrypt_value: 待解密的hex值
        :return: 原字符串
        """
        crypt_sm4 = self.crypt_sm4
        crypt_sm4.set_key(decrypt_key.encode(), SM4_DECRYPT)
        decrypt_value = crypt_sm4.crypt_ecb(bytes.fromhex(encrypt_value))  # bytes类型
        return self.str_to_hexStr(decrypt_value.hex())

    def toTreeMap(self, param_map):
        """
        将paramMap转换为java中的treeMap形式.将map的keys变为heapq.创建有序字典.
        :param param_map: 原字典
        :return: treeMap字典
        """
        param_map = dict(param_map)
        keys = param_map.keys()
        heap = []
        for item in keys:
            heappush(heap, item)
        sort = []
        while heap:
            sort.append(heappop(heap))
        resMap = OrderedDict()
        for key_ in sort:
            resMap[key_] = param_map.get(key_)
        key_list = [key_ for key_ in resMap.keys()]
        value_list = [value for value in resMap.values()]
        result = ""
        for i in range(len(key_list)):
            result = result + key_list[i] + value_list[i]

        return result


if __name__ == '__main__':
    str_data = {"code":"UIwoSt","timestamp":"20181024132419"}
    key = "d36fd6db915345548ebc62ef2f122403"
    SM4 = SM4()
    print("待加密内容：", str_data)
    encoding = SM4.encrypt(key, str_data)
    print("国密sm4加密后的结果：", encoding)
    print("国密sm4解密后的结果：", SM4.decrypt(key, encoding))