From a7e3cb67e88d045defdac8bf5fee2dac79d29088 Mon Sep 17 00:00:00 2001
From: UniDarkstars <1830018670@qq.com>
Date: Tue, 26 Nov 2024 14:14:48 +0800
Subject: [PATCH] =?UTF-8?q?=E6=9B=B4=E6=94=B9=E4=BA=86recv=E5=92=8CAES?=
 =?UTF-8?q?=EF=BC=8CAES=E4=BD=BF=E7=94=A8=E4=BF=A1=E5=B0=81=E5=86=85?=
 =?UTF-8?q?=E7=9A=84=E5=AF=86=E9=92=A5=E8=BF=9B=E8=A1=8C=E8=A7=A3=E5=AF=86?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 recv/recv.py          |  4 ++--
 tool/symmetric/AES.py | 55 ++++++++++++++++++++++++++-----------------
 2 files changed, 36 insertions(+), 23 deletions(-)

diff --git a/recv/recv.py b/recv/recv.py
index ac9d829..82802a8 100644
--- a/recv/recv.py
+++ b/recv/recv.py
@@ -40,9 +40,9 @@ def handleLetter(letter: Letter):
     elif type == config.EncryptType.SM4_CBC:
         data = base64.b64decode(SM4.decrypt_cbc_with_iv(letter.fileBase64, key))
     elif type == config.EncryptType.AES_ECB:
-        data = AES.AESUtils(key).decrypt(letter.fileBase64)
+        data = AES.AESUtils(key).decrypt(key, letter.fileBase64, "ecb")
     elif type == config.EncryptType.AES_CBC:
-        data = AES.AESUtils(key).decrypt(letter.fileBase64,"cbc")
+        data = AES.AESUtils(key).decrypt(key, letter.fileBase64,"cbc")
     else:
         raise KeyError("不支持的对称加密算法")
 
diff --git a/tool/symmetric/AES.py b/tool/symmetric/AES.py
index 87e9589..7197e2e 100644
--- a/tool/symmetric/AES.py
+++ b/tool/symmetric/AES.py
@@ -1,33 +1,29 @@
-# "ecb", "cbc", "cfb", "ofb"
-
 from Crypto.Cipher import AES
 from Crypto.Util.Padding import pad, unpad
 from Crypto.Random import get_random_bytes
 import base64
-from fontTools.misc.eexec import encrypt
-
 
 class AESUtils:
     def __init__(self, key: bytes = None):
         """生成一个随机密钥"""
-        self.key = get_random_bytes(16)
+        self.key = key if key else get_random_bytes(16)
 
-    def encrypt(self, data: str, mode: str = 'ECB') -> (str,str):
+    def encrypt(self, data1: str, mode: str = 'ECB') -> (str, str):
         """加密数据"""
-        cipher = self._get_cipher(mode)
-        data = pad(data.encode(), AES.block_size)  # 填充数据
-        ciphertext = cipher.encrypt(data)
-        return base64.b64encode(ciphertext).decode(), self.key
+        cipher = self._get_encipher(mode)
+        data1 = pad(data1.encode('utf-8'), AES.block_size)  # 填充数据
+        ciphertext = cipher.encrypt(data1)
+        return base64.b64encode(ciphertext).decode('utf-8'), self.key
 
-    def decrypt(self, encrypted_data: str, mode: str = 'ECB') -> str:
+    def decrypt(self, dekey, encrypted_data: str, mode: str = 'ECB') -> str:
         """解密数据"""
-        cipher = self._get_cipher(mode)
-        encrypted_data = base64.b64decode(encrypted_data)
+        cipher = self._get_decipher(dekey, mode)
+        encrypted_data = base64.b64decode(encrypted_data)  # 解码密文
         plaintext = unpad(cipher.decrypt(encrypted_data), AES.block_size)
         return plaintext.decode()
 
-    def _get_cipher(self, mode: str) -> AES:
-        """根据模式返回相应的 AES cipher"""
+    def _get_encipher(self, mode: str) -> AES:
+        """根据模式返回相应的加密 AES cipher"""
         iv = None
         if mode == 'CBC':
             iv = get_random_bytes(AES.block_size)
@@ -39,12 +35,29 @@ class AESUtils:
         else:  # 默认是 ECB
             return AES.new(self.key, AES.MODE_ECB)
 
-if __name__ == "__main__":
-    # 示例数据和密钥
-    data = "Hello, SM4!"
-    key = "1234567890abcdef"
+    def _get_decipher(self, dekey, mode: str) -> AES:
+        """根据模式返回相应的解密 AES cipher"""
+        iv = None
+        if mode == 'CBC':
+            iv = get_random_bytes(AES.block_size)
+            return AES.new(dekey, AES.MODE_CBC, iv)
+        elif mode == 'CFB':
+            return AES.new(dekey, AES.MODE_CFB)
+        elif mode == 'OFB':
+            return AES.new(dekey, AES.MODE_OFB)
+        else:  # 默认是 ECB
+            return AES.new(dekey, AES.MODE_ECB)
 
-    enData , akey = AESUtils(key)
+if __name__ == '__main__':
+    key = b'1234567890abcdef'  # AES 密钥必须是 16 字节
+    aes = AESUtils(key)
 
-    print(enData)
+    # 加密
+    plaintext = "Hello, AES!"
+    enData, akey = aes.encrypt(plaintext, mode='ECB')
+    print("Encrypted Data:", enData)
+    print("Key:", akey)
 
+    # 解密
+    decrypted = aes.decrypt(aes.key, enData, mode='ECB')
+    print("Decrypted:", decrypted)
\ No newline at end of file