diff --git a/recv/recv.py b/recv/recv.py index 8e085d6..98679c3 100644 --- a/recv/recv.py +++ b/recv/recv.py @@ -7,7 +7,7 @@ from entity.Letter import Letter from tool import PriKeyHelper, DownloadPathTool from tool.asymmetric import RSA from tool.hash import Segwit -from tool.symmetric import SM4 +from tool.symmetric import SM4, AES def getLetter(): @@ -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: - raise NotImplementedError("未实现") + data = AES.AESUtils(key).decrypt(letter.fileBase64) elif type == config.EncryptType.AES_CBC: - raise NotImplementedError("未实现") + data = AES.AESUtils(key).decrypt(letter.fileBase64,"cbc") else: raise KeyError("不支持的对称加密算法") diff --git a/recv/test_recv.py b/recv/test_recv.py index ca0ad29..1106e10 100644 --- a/recv/test_recv.py +++ b/recv/test_recv.py @@ -2,12 +2,31 @@ from unittest.mock import patch import recv from sender import sender -from tool import PriKeyHelper + + +@patch('builtins.input', side_effect=['./public.pem', 'aes', 'ecb']) +def test_handleLetter_aes_ecb(mock_input): + letter = sender.main() + recv.handleLetter(letter) + pass + + +@patch('builtins.input', side_effect=['./public.pem', 'aes', 'cbc']) +def test_handleLetter_aes_cbc(mock_input): + letter = sender.main() + recv.handleLetter(letter) + pass + + +@patch('builtins.input', side_effect=['./public.pem', 'sm4', 'ecb']) +def test_handleLetter_sm4_ecb(mock_input): + letter = sender.main() + recv.handleLetter(letter) + pass @patch('builtins.input', side_effect=['./public.pem', 'sm4', 'cbc']) -def test_handleLetter(mock_input): - pki = PriKeyHelper.getUserKey() +def test_handleLetter_sm4_cbc(mock_input): letter = sender.main() recv.handleLetter(letter) pass