修复SM4和RSA的bug

tool/asymmetric/RSA.py

@@ -1,9 +1,10 @@
 import base64
 import binascii
+
 from Crypto.Cipher import PKCS1_OAEP
+from Crypto.Hash import SHA256
 from Crypto.PublicKey import RSA
 from Crypto.Signature import pkcs1_15
-from Crypto.Hash import SHA256
 
 
 def generate_keys(priKeySavePath, pubKeySavePath):
@@ -21,45 +22,46 @@ def generate_keys(priKeySavePath, pubKeySavePath):
         f.write(public_key)
 
 
-def encrypt_message(message, public_key_base64):
+def encrypt_message(message: bytes, public_key_base64):
     # 加载公钥
     public_key_bin = base64.b64decode(public_key_base64)
     public_key = RSA.import_key(public_key_bin)
     # 使用公钥加密消息
     cipher_rsa = PKCS1_OAEP.new(public_key)
-    encrypted_message = cipher_rsa.encrypt(message.encode('utf-8'))
+    encrypted_message = cipher_rsa.encrypt(message)
 
-    return binascii.hexlify(encrypted_message).decode('utf-8')
+    return encrypted_message
 
 
-def decrypt_message(encrypted_message, private_key_base64):
+def decrypt_message(encrypted_message: bytes, private_key_base64):
     # 加载私钥
     private_key_bin = base64.b64decode(private_key_base64)
     private_key = RSA.import_key(private_key_bin)
     # 使用私钥解密消息
     cipher_rsa = PKCS1_OAEP.new(private_key)
-    decrypted_message = cipher_rsa.decrypt(binascii.unhexlify(encrypted_message))
+    decrypted_message = cipher_rsa.decrypt(encrypted_message)
+
+    return decrypted_message
 
-    return decrypted_message.decode('utf-8')
 
-def sign_message(message, private_key_base64):
+def sign_message(message: bytes, private_key_base64):
     # 加载私钥
     private_key_bin = base64.b64decode(private_key_base64)
     private_key = RSA.import_key(private_key_bin)
     # 计算消息的哈希值
-    hash_obj = SHA256.new(message.encode('utf-8'))
+    hash_obj = SHA256.new(message)
     # 使用私钥对哈希值进行签名
     signature = pkcs1_15.new(private_key).sign(hash_obj)
 
     return binascii.hexlify(signature).decode('utf-8')
 
 
-def verify_signature(message, signature, public_key_base64):
+def verify_signature(message: bytes, signature, public_key_base64):
     # 加载公钥
     public_key_bin = base64.b64decode(public_key_base64)
     public_key = RSA.import_key(public_key_bin)
     # 计算消息的哈希值
-    hash_obj = SHA256.new(message.encode('utf-8'))
+    hash_obj = SHA256.new(message)
     # 使用公钥验证签名
     try:
         pkcs1_15.new(public_key).verify(hash_obj, binascii.unhexlify(signature))

tool/symmetric/SM4.py

@@ -2,49 +2,54 @@ import os
 
 from gmssl import sm4
 
+
 def encrypt_ecb(data, key):
     cipher = sm4.CryptSM4()
-    cipher.set_key(key.encode('utf-8'), sm4.SM4_ENCRYPT)
+    cipher.set_key(key, sm4.SM4_ENCRYPT)
-    encrypted_data = cipher.crypt_ecb(data.encode('utf-8'))
+    encrypted_data = cipher.crypt_ecb(data)
-    return encrypted_data.hex()
+    return encrypted_data
+
 
-def decrypt_ecb(encrypted_hex, key):
+def decrypt_ecb(encrypted_data, key):
     cipher = sm4.CryptSM4()
-    cipher.set_key(key.encode('utf-8'), sm4.SM4_DECRYPT)
+    cipher.set_key(key, sm4.SM4_DECRYPT)
-    decrypted_data = cipher.crypt_ecb(bytes.fromhex(encrypted_hex))
+    decrypted_data = cipher.crypt_ecb(encrypted_data)
-    return decrypted_data.decode('utf-8')
+    return decrypted_data
+
 
 def encrypt_cbc_with_iv(data, key):
     cipher = sm4.CryptSM4()
-    cipher.set_key(key.encode('utf-8'), sm4.SM4_ENCRYPT)
+    cipher.set_key(key, sm4.SM4_ENCRYPT)
 
     # 生成随机的16字节IV
     iv = os.urandom(16)
 
     # 加密数据
-    encrypted_data = cipher.crypt_cbc(iv, data.encode('utf-8'))
+    encrypted_data = cipher.crypt_cbc(iv, data)
 
     # 将IV和加密后的数据拼接在一起
     return iv + encrypted_data
 
-def decrypt_cbc_with_iv(encrypted_bytes, key):
+
+def decrypt_cbc_with_iv(encrypted_data, key):
     cipher = sm4.CryptSM4()
-    cipher.set_key(key.encode('utf-8'), sm4.SM4_DECRYPT)
+    cipher.set_key(key, sm4.SM4_DECRYPT)
 
     # 提取IV
-    iv = encrypted_bytes[:16]
+    iv = encrypted_data[:16]
 
     # 提取加密后的数据
-    encrypted_data = encrypted_bytes[16:]
+    encrypted_data = encrypted_data[16:]
 
     # 解密数据
     decrypted_data = cipher.crypt_cbc(iv, encrypted_data)
-    return decrypted_data.decode('utf-8')
+    return decrypted_data
+
 
 if __name__ == "__main__":
     # 示例数据和密钥
-    data = "Hello, SM4!"
+    data = b"Hello, SM4!"
-    key = "1234567890abcdef"
+    key = b"1234567890abcdef"
 
     # 加密
     encrypted_data = encrypt_ecb(data, key)
@@ -55,8 +60,8 @@ if __name__ == "__main__":
     print(f"Decrypted: {decrypted_data}")
 
     # 示例数据和密钥
-    data = "Hello, SM4 CBC with random IV!"
+    data = b"Hello, SM4 CBC with random IV!"
-    key = "1234567890abcdef"
+    key = b"1234567890abcdef"
 
     # 加密
     encrypted_data = encrypt_cbc_with_iv(data, key)