import os import base64 from django.conf import settings from gmssl.sm4 import CryptSM4, SM4_ENCRYPT, SM4_DECRYPT from gmssl import sm2, sm3, func import logging from common.sdk.gm import Device from .base import BaseCrypto, BaseCryptoSuite from .common import padding_key class GMSM4EcbCrypto(BaseCrypto): def __init__(self, key): self.key = padding_key(key, 16) self.sm4_encryptor = CryptSM4() self.sm4_encryptor.set_key(self.key, SM4_ENCRYPT) self.sm4_decryptor = CryptSM4() self.sm4_decryptor.set_key(self.key, SM4_DECRYPT) def _encrypt(self, data: bytes) -> bytes: return self.sm4_encryptor.crypt_ecb(data) def _decrypt(self, data: bytes) -> bytes: return self.sm4_decryptor.crypt_ecb(data) class GMDeviceSM4EcbCrypto(BaseCrypto): @staticmethod def to_16(key): while len(key) % 16 != 0: key += b'\0' return key # 返回bytes def __init__(self, key, device: Device): key = padding_key(key, 16) self.cipher = device.new_sm4_ebc_cipher(key) def _encrypt(self, data: bytes) -> bytes: return self.cipher.encrypt(self.to_16(data)) def _decrypt(self, data: bytes) -> bytes: bs = self.cipher.decrypt(data) return bs.rstrip(b'\0') class Sm3Hasher: name = 'sm3' block_size = 64 digest_size = 32 def __init__(self, data): self.__data = data def hexdigest(self): return sm3.sm3_hash(func.bytes_to_list(self.__data)) def digest(self): return bytes.fromhex(self.hexdigest()) @staticmethod def hash(cls, msg=b''): return cls(msg) def update(self, data): self.__data += data @classmethod def copy(cls, data): return cls(data) class GmCryptoSuite(BaseCryptoSuite): crypto = None _key = None def __init__(self, key, mode: str = 'ecb'): self.mode = mode if key: self.key = key @property def key(self): return self._key @key.setter def key(self, value): self._key = value self.crypto = GMSM4EcbCrypto(value) def encrypt(self, text): return self.crypto.encrypt(text) def decrypt(self, cipher_text): return self.crypto.decrypt(cipher_text) def gen_key_pair(self): private_key = os.urandom(32).hex() sm2_crypt = sm2.CryptSM2(public_key='', private_key=private_key) public_key = sm2_crypt._kg(int(private_key, 16), sm2.default_ecc_table['g']) return private_key, public_key def encrypt_with_key_pair(self, plain_text, public_key): sm2_crypt = sm2.CryptSM2(public_key=public_key, private_key='') message = sm2_crypt.encrypt(plain_text.encode()) cipher_text = base64.b64encode(message).decode() return cipher_text def decrypt_with_key_pair(self, cipher_text, private_key): try: message = base64.b64decode(cipher_text.encode()) sm2_crypt = sm2.CryptSM2(public_key='', private_key=private_key) return sm2_crypt.decrypt(message).decode() except Exception as e: logging.error(f"Decrypt with key pair error: {e}") return cipher_text def hash(self, msg): return Sm3Hasher.hash(msg) def get_gm_sm4_ecb_crypto(key=None): key = key or settings.SECRET_KEY return GMSM4EcbCrypto(key) def get_gm_device_sm4_ecb_crypto(device, key=None): key = key or settings.SECRET_KEY return GMDeviceSM4EcbCrypto(key, device) def gen_gm_key_pair(): private_key = os.urandom(32).hex() sm2_crypt = sm2.CryptSM2(public_key='', private_key=private_key) public_key = sm2_crypt._kg(int(private_key, 16), sm2.default_ecc_table['g']) return private_key, public_key gm_sm4_ecb_crypto = get_gm_sm4_ecb_crypto()