feat: 支持 piico 设备国密加密

apps/common/utils/crypto.py

@@ -1,6 +1,7 @@
 import base64
 import logging
 import re
+
 from Cryptodome.Cipher import AES, PKCS1_v1_5
 from Cryptodome.Random import get_random_bytes
 from Cryptodome.PublicKey import RSA
@@ -11,6 +12,7 @@ from gmssl.sm4 import CryptSM4, SM4_ENCRYPT, SM4_DECRYPT
 from django.conf import settings
 from django.core.exceptions import ImproperlyConfigured
 
+from common.sdk.gm import piico
 
 secret_pattern = re.compile(r'password|secret|key|token', re.IGNORECASE)
 
@@ -64,6 +66,25 @@ class GMSM4EcbCrypto(BaseCrypto):
         return self.sm4_decryptor.crypt_ecb(data)
 
 
+class PiicoSM4EcbCrypto(BaseCrypto):
+
+    @staticmethod
+    def to_16(key):
+        while len(key) % 16 != 0:
+            key += b'\0'
+        return key  # 返回bytes
+
+    def __init__(self, key, device: piico.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:
+        return self.cipher.decrypt(data)
+
+
 class AESCrypto:
     """
     AES
@@ -164,6 +185,11 @@ def get_gm_sm4_ecb_crypto(key=None):
     return GMSM4EcbCrypto(key)
 
 
+def get_piico_gm_sm4_ecb_crypto(device, key=None):
+    key = key or settings.SECRET_KEY
+    return PiicoSM4EcbCrypto(key, device)
+
+
 aes_ecb_crypto = get_aes_crypto(mode='ECB')
 aes_crypto = get_aes_crypto(mode='GCM')
 gm_sm4_ecb_crypto = get_gm_sm4_ecb_crypto()
@@ -183,10 +209,16 @@ class Crypto:
         crypt_algo = settings.SECURITY_DATA_CRYPTO_ALGO
         if not crypt_algo:
             if settings.GMSSL_ENABLED:
-                crypt_algo = 'gm'
+                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)
+                    self.cryptor_map["piico_gm"] = get_piico_gm_sm4_ecb_crypto(device)
+                    crypt_algo = 'piico_gm'
+                else:
+                    crypt_algo = 'gm'
             else:
                 crypt_algo = 'aes'
-
         cryptor = self.cryptor_map.get(crypt_algo, None)
         if cryptor is None:
             raise ImproperlyConfigured(