SM4的message使用str，key使用bytes类型

tool/symmetric/SM4.py

@@ -1,23 +1,30 @@
+import base64
 import os
 
 from gmssl import sm4
-
-
-def encrypt_ecb(data, key):
+"""
+输入：
+消息或者加密后消息：字符串类型
+key：字节类型
+输出：
+消息或者加密后消息：base64的字符串类型
+"""
+
+def encrypt_ecb(data: str, key: bytes) -> str:
     cipher = sm4.CryptSM4()
     cipher.set_key(key, sm4.SM4_ENCRYPT)
-    encrypted_data = cipher.crypt_ecb(data)
-    return encrypted_data
+    encrypted_data = cipher.crypt_ecb(data.encode('utf-8'))
+    return base64.b64encode(encrypted_data).decode('utf-8')
 
 
-def decrypt_ecb(encrypted_data, key):
+def decrypt_ecb(encrypted_hex: str, key: bytes) -> str:
     cipher = sm4.CryptSM4()
     cipher.set_key(key, sm4.SM4_DECRYPT)
-    decrypted_data = cipher.crypt_ecb(encrypted_data)
-    return decrypted_data
+    decrypted_data = cipher.crypt_ecb(bytes.fromhex(encrypted_hex))
+    return decrypted_data.decode('utf-8')
 
 
-def encrypt_cbc_with_iv(data, key):
+def encrypt_cbc_with_iv(data: str, key: bytes) -> str:
     cipher = sm4.CryptSM4()
     cipher.set_key(key, sm4.SM4_ENCRYPT)
 
@@ -25,30 +32,31 @@ def encrypt_cbc_with_iv(data, key):
     iv = os.urandom(16)
 
     # 加密数据
-    encrypted_data = cipher.crypt_cbc(iv, data)
+    encrypted_data = cipher.crypt_cbc(iv, data.encode('utf-8'))
 
     # 将IV和加密后的数据拼接在一起
-    return iv + encrypted_data
+    return base64.b64encode(iv + encrypted_data).decode('utf-8')
 
 
-def decrypt_cbc_with_iv(encrypted_data, key):
+def decrypt_cbc_with_iv(encrypted_bytes_base64: str, key) -> str:
+    encrypted_bytes = base64.b64decode(encrypted_bytes_base64)
     cipher = sm4.CryptSM4()
     cipher.set_key(key, sm4.SM4_DECRYPT)
 
     # 提取IV
-    iv = encrypted_data[:16]
+    iv = encrypted_bytes[:16]
 
     # 提取加密后的数据
-    encrypted_data = encrypted_data[16:]
+    encrypted_data = encrypted_bytes[16:]
 
     # 解密数据
     decrypted_data = cipher.crypt_cbc(iv, encrypted_data)
-    return decrypted_data
+    return base64.b64decode(decrypted_data).decode('utf-8')
 
 
 if __name__ == "__main__":
     # 示例数据和密钥
-    data = b"Hello, SM4!"
+    data = "Hello, SM4!"
     key = b"1234567890abcdef"
 
     # 加密
@@ -60,7 +68,7 @@ if __name__ == "__main__":
     print(f"Decrypted: {decrypted_data}")
 
     # 示例数据和密钥
-    data = b"Hello, SM4 CBC with random IV!"
+    data = "Hello, SM4 CBC with random IV!"
     key = b"1234567890abcdef"
 
     # 加密