From 26890907007c737de5b7dda590ddcf60d551d629 Mon Sep 17 00:00:00 2001
From: Timmoc <timmoc@qq.com>
Date: Sun, 24 Nov 2024 20:42:16 +0800
Subject: [PATCH] =?UTF-8?q?=E6=9C=89=E5=B1=8E=E5=95=8A?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

修了一点okToCheck的代码，反正现在test_recv跑通了，剩下的都是小事了。

准备把这坨大的合并到main分支
---
 .gitignore                |  5 +++++
 recv/recv.py              | 14 +++++++-------
 recv/test_recv.py         | 11 ++++++++---
 sender/sender.py          | 17 ++++++++++-------
 tool/test_PriKeyHelper.py |  4 ++--
 5 files changed, 32 insertions(+), 19 deletions(-)

diff --git a/.gitignore b/.gitignore
index 56c7dcd..d23cf10 100644
--- a/.gitignore
+++ b/.gitignore
@@ -5,3 +5,8 @@
 /tool/private.pem
 /private.pem
 /public.pem
+/sender/private.pem
+/recv/private.pem
+/recv/public.pem
+/sender/public.pem
+/recv/交给你了
diff --git a/recv/recv.py b/recv/recv.py
index ac101ff..5bf5675 100644
--- a/recv/recv.py
+++ b/recv/recv.py
@@ -22,19 +22,19 @@ def handleLetter(letter:Letter):
     if pki[1] != letter.recvPubKey:
         raise Exception("信件不属于自己")
     # 用自己的私钥解密key 获得对称加密秘钥。
-    key = RSA.decrypt_message(letter.encryptKey,pki[0])
+    key = RSA.decrypt_message(base64.b64decode(letter.encryptKey),pki[0])
 
     # 根据不同的对称加密算法
     try:
         type = config.getEncryptType(letter.encryptType)
-    except KeyError:
-        raise KeyError("不支持的对称加密算法")
+    except ValueError:
+        raise ValueError("不支持的对称加密算法")
     # 进行解密fileBase64
-    data = None
+    data = b""
     if type == config.EncryptType.SM4_ECB:
-        SM4.decrypt_ecb(base64.b64decode(letter.fileBase64),key)
+        data = base64.b64decode(SM4.decrypt_ecb(letter.fileBase64,key))
     elif type == config.EncryptType.SM4_CBC:
-        SM4.decrypt_cbc_with_iv(base64.b64decode(letter.fileBase64),key)
+        data = base64.b64decode(SM4.decrypt_cbc_with_iv(letter.fileBase64,key))
     elif type == config.EncryptType.AES_ECB:
         raise NotImplementedError("未实现")
     elif type == config.EncryptType.AES_CBC:
@@ -50,5 +50,5 @@ def handleLetter(letter:Letter):
     # 保存文件
     with open(f"./{letter.fileName}","wb") as f:
         f.write(data)
-
+        print(f"签名验证有效，已将文件 {letter.fileName} 保存至当前目录下")
     return
\ No newline at end of file
diff --git a/recv/test_recv.py b/recv/test_recv.py
index 26601b1..ca0ad29 100644
--- a/recv/test_recv.py
+++ b/recv/test_recv.py
@@ -1,8 +1,13 @@
+from unittest.mock import patch
+
+import recv
 from sender import sender
 from tool import PriKeyHelper
 
 
-def test_handleLetter():
+@patch('builtins.input', side_effect=['./public.pem', 'sm4', 'cbc'])
+def test_handleLetter(mock_input):
     pki = PriKeyHelper.getUserKey()
-
-    pass
\ No newline at end of file
+    letter = sender.main()
+    recv.handleLetter(letter)
+    pass
diff --git a/sender/sender.py b/sender/sender.py
index d30a1b2..27127fc 100644
--- a/sender/sender.py
+++ b/sender/sender.py
@@ -1,4 +1,6 @@
 # 模式，文件，自己的公钥从哪里来，别人的公钥从哪里来
+import base64
+
 from entity.Letter import Letter
 from itsdangerous import base64_encode
 from tool import PriKeyHelper
@@ -22,14 +24,14 @@ def main():
     with open(path,"rb") as f:
         data = f.read()
     letter.fileName = "交给你了"
-    letter.encryptType = getEncryptType()
-    letter.encryptKey = getEncryptKey()
     letter.recvPubKey = getRecvPubKey()
     letter.senderPubKey = getSenderPubKey()
-    data, akey = SymEncryption(base64_encode(data).decode("utf-8"),letterSymKey)
+    letter.fileBase64, akey = SymEncryption(base64_encode(data).decode("utf-8"),letterSymKey)
     # data = "Hello, AES!"
+    letter.encryptKey = getEncryptKey()
+    letter.encryptType = getEncryptType()
 
-    letter.sign = getSign()
+    letter.sign = getSign(data)
 
 
 
@@ -43,6 +45,7 @@ def main():
     # 获取用户的公私钥对进行签名
     # 使用对方的公钥进行加密
     # 发送信件
+    return letter # 方便recv测试，以后可以删除。
     pass
 def selectFile() -> str:
     s = input("输入文件路径:")
@@ -83,6 +86,7 @@ def selectSymEncryptionChoice():
         elif encryWay == "sm4":
             encryMode = input("Choose the encryption mode (ecb/cbc): ").strip().lower()
             if encryMode in ["ecb", "cbc"]:
+                letterMode = encryMode
                 print(f"You have selected '{encryMode}' encryption mode.")
                 break  # 输入有效后退出循环
             else:
@@ -116,10 +120,9 @@ def SymEncryption(encryData, key: bytes = None):
         return encrypted_data, key
 
 # 获得签名
-def getSign():
+def getSign(document_bytes):
 
     #计算文件bytes
-    document_bytes = ""
     priKey, pubKey = PriKeyHelper.getUserKey()
     signDocuHash = RSA.sign_message(document_bytes, priKey)
 
@@ -134,7 +137,7 @@ def getEncryptType():
 def getEncryptKey():
     rsaEncrySymKey = RSA.encrypt_message(letterSymKey, getRecvPubKey())
 
-    return base64_encode(rsaEncrySymKey).decode("utf-8")
+    return base64.b64encode(rsaEncrySymKey).decode("utf-8")
 
 # 获得接收方的公钥
 def getRecvPubKey():
diff --git a/tool/test_PriKeyHelper.py b/tool/test_PriKeyHelper.py
index 8de5054..ed34aa4 100644
--- a/tool/test_PriKeyHelper.py
+++ b/tool/test_PriKeyHelper.py
@@ -6,7 +6,7 @@ from tool.asymmetric import RSA
 
 def test_encrypt():
     key = PriKeyHelper.getUserKey()
-    message = "Hello, this is a secret message."
+    message = b"Hello, this is a secret message."
     # 打印keys
     print(key[0])
     print(key[1])
@@ -23,7 +23,7 @@ def test_encrypt():
 def test_sign():
 
     key = PriKeyHelper.getUserKey()
-    message = "This is a signed message."
+    message = b"This is a signed message."
     # 打印keys
     print(key[0])
     print(key[1])