From ede7e688e0f6f9147409cb28792cdfa8d38cce93 Mon Sep 17 00:00:00 2001 From: halo Date: Fri, 6 Mar 2026 17:16:42 +0800 Subject: [PATCH] feat: enhance key handling in cipher classes --- apps/common/sdk/gm/piico/__init__.py | 8 +++- apps/common/sdk/gm/piico/cipher.py | 19 ++++------ apps/common/sdk/gm/piico/session_mixin.py | 46 +++++++++++------------ apps/common/utils/crypto.py | 4 +- 4 files changed, 37 insertions(+), 40 deletions(-) diff --git a/apps/common/sdk/gm/piico/__init__.py b/apps/common/sdk/gm/piico/__init__.py index 936a89ca6..e21516a7e 100644 --- a/apps/common/sdk/gm/piico/__init__.py +++ b/apps/common/sdk/gm/piico/__init__.py @@ -1,7 +1,13 @@ +from django.conf import settings + from .device import Device +DEFAULT_DRIVER_PATH = "./lib/libpiico_ccmu.so" -def open_piico_device(driver_path) -> Device: + +def open_piico_device(driver_path=None) -> Device: + if driver_path is None: + driver_path = settings.PIICO_DRIVER_PATH or DEFAULT_DRIVER_PATH d = Device() d.open(driver_path) return d diff --git a/apps/common/sdk/gm/piico/cipher.py b/apps/common/sdk/gm/piico/cipher.py index 13333900d..e2602ce6f 100644 --- a/apps/common/sdk/gm/piico/cipher.py +++ b/apps/common/sdk/gm/piico/cipher.py @@ -22,21 +22,17 @@ class EBCCipher: def __init__(self, session, key_val): self._session = session - self._key = self.__get_key(key_val) + self._key = key_val if isinstance(key_val, bytes) else bytes(key_val, encoding='utf-8') self._alg = "sm4_ebc" self._iv = None - def __get_key(self, key_val): - key_val = self.__padding(key_val) - return self._session.import_key(key_val) - @staticmethod - def __padding(val): - # padding - val = bytes(val) - while len(val) == 0 or len(val) % 16 != 0: - val += b'\0' - return val + def __padding(data): + if not isinstance(data, bytes): + data = bytes(data, encoding='utf-8') + while len(data) == 0 or len(data) % 16 != 0: + data += b'\0' + return data def encrypt(self, plain_text): plain_text = self.__padding(plain_text) @@ -48,7 +44,6 @@ class EBCCipher: return bytes(plain_text) def destroy(self): - self._session.destroy_cipher_key(self._key) self._session.close() diff --git a/apps/common/sdk/gm/piico/session_mixin.py b/apps/common/sdk/gm/piico/session_mixin.py index d625f627f..fc645ad0b 100644 --- a/apps/common/sdk/gm/piico/session_mixin.py +++ b/apps/common/sdk/gm/piico/session_mixin.py @@ -1,4 +1,3 @@ - from .ecc import * from .exception import PiicoError @@ -13,9 +12,8 @@ class BaseMixin: class SM2Mixin(BaseMixin): def ecc_encrypt(self, public_key, plain_text, alg_id): - pos = 1 + pos = 0 k1 = bytes([0] * 32) + bytes(public_key[pos:pos + 32]) - k1 = (c_ubyte * len(k1))(*k1) pos += 32 k2 = bytes([0] * 32) + bytes(public_key[pos:pos + 32]) @@ -88,21 +86,6 @@ class SM3Mixin(BaseMixin): class SM4Mixin(BaseMixin): - def import_key(self, key_val): - # to c lang - key_val = (c_ubyte * len(key_val))(*key_val) - - key = c_void_p() - ret = self._driver.SDF_ImportKey(self._session, key_val, c_int(len(key_val)), pointer(key)) - if ret != 0: - raise PiicoError("import key failed", ret) - return key - - def destroy_cipher_key(self, key): - ret = self._driver.SDF_DestroyKey(self._session, key) - if ret != 0: - raise Exception("destroy key failed") - def encrypt(self, plain_text, key, alg, iv=None): return self.__do_cipher_action(plain_text, key, alg, iv, True) @@ -110,20 +93,35 @@ class SM4Mixin(BaseMixin): return self.__do_cipher_action(cipher_text, key, alg, iv, False) def __do_cipher_action(self, text, key, alg, iv=None, encrypt=True): - text = (c_ubyte * len(text))(*text) + text_buf = (c_ubyte * len(text))(*text) + key_buf = (c_ubyte * len(key))(*key) + + iv_buf = None + iv_len = 0 if iv is not None: - iv = (c_ubyte * len(iv))(*iv) + iv_buf = (c_ubyte * len(iv))(*iv) + iv_len = len(iv) temp_data = (c_ubyte * len(text))() temp_data_length = c_int() if encrypt: - ret = self._driver.SDF_Encrypt(self._session, key, c_int(alg), iv, text, c_int(len(text)), temp_data, - pointer(temp_data_length)) + ret = self._driver.SPII_EncryptEx( + self._session, text_buf, c_int(len(text)), + key_buf, c_int(len(key)), + iv_buf, c_int(iv_len), + c_int(alg), + temp_data, pointer(temp_data_length), + ) if ret != 0: raise PiicoError("encrypt failed", ret) else: - ret = self._driver.SDF_Decrypt(self._session, key, c_int(alg), iv, text, c_int(len(text)), temp_data, - pointer(temp_data_length)) + ret = self._driver.SPII_DecryptEx( + self._session, text_buf, c_int(len(text)), + key_buf, c_int(len(key)), + iv_buf, c_int(iv_len), + c_int(alg), + temp_data, pointer(temp_data_length), + ) if ret != 0: raise PiicoError("decrypt failed", ret) return temp_data[:temp_data_length.value] diff --git a/apps/common/utils/crypto.py b/apps/common/utils/crypto.py index 0a0721e50..4d0db7cbf 100644 --- a/apps/common/utils/crypto.py +++ b/apps/common/utils/crypto.py @@ -210,9 +210,7 @@ class Crypto: if not crypt_algo: if settings.GMSSL_ENABLED: if settings.PIICO_DEVICE_ENABLE: - piico_driver_path = settings.PIICO_DRIVER_PATH if settings.PIICO_DRIVER_PATH \ - else "./lib/libpiico_ccmu.so" - device = piico.open_piico_device(piico_driver_path) + device = piico.open_piico_device() self.cryptor_map["piico_gm"] = get_piico_gm_sm4_ecb_crypto(device) crypt_algo = 'piico_gm' else: