Merge branch 'refs/heads/main' into recyvan

# Conflicts:
#	recv/recv.py
#	sender/sender.py

.gitignore

@@ -1 +1,12 @@
 /.idea
+/tool/asymmetric/private.pem
+/tool/asymmetric/public.pem
+/tool/public.pem
+/tool/private.pem
+/private.pem
+/public.pem
+/sender/private.pem
+/recv/private.pem
+/recv/public.pem
+/sender/public.pem
+/recv/交给你了

config/config.py

@@ -1,2 +1,25 @@
+from enum import Enum, auto
 port = 8426
-priKeySavePath = "./priKey.txt"
+priKeySavePath = "./private.pem"
+pubKeySavePath = "./public.pem"
+
+class EncryptType(Enum):
+    AES_ECB = auto()
+    AES_CBC = auto()
+    AES_CFB = auto()
+    AES_OFB = auto()
+    SM4_ECB = auto()
+    SM4_CBC = auto()
+
+# 函数：通过文本获取枚举成员
+def getEncryptType(text):
+    try:
+        return EncryptType[text]
+    except KeyError:
+        raise ValueError(f"Invalid EncryptType name: {text}")
+
+# 函数：通过枚举成员获取文本
+def getEncryptTypeName(encryptType):
+    if not isinstance(encryptType, EncryptType):
+        raise TypeError("Expected a EncryptType enum member")
+    return encryptType.name

entity/Letter.py

@@ -1,10 +1,9 @@
 class Letter:
     sign = "计算得到"
-    encryptType = "SM4-GCM"
-    encryptKey = "计算获得"
+    encryptType = "SM4_GCM"
+    encryptKey = "计算获得" # recvPubKey 加密后的 对称加密秘钥 数据
     recvPubKey = ""
     senderPubKey = ""
 
     fileName = ""
-    fileSize = 0
     fileBase64 = ""

recv/test_recv.py

@@ -0,0 +1,13 @@
+from unittest.mock import patch
+
+import recv
+from sender import sender
+from tool import PriKeyHelper
+
+
+@patch('builtins.input', side_effect=['./public.pem', 'sm4', 'cbc'])
+def test_handleLetter(mock_input):
+    pki = PriKeyHelper.getUserKey()
+    letter = sender.main()
+    recv.handleLetter(letter)
+    pass

tool/PriKeyHelper.py

@@ -1,4 +1,24 @@
-def getUserKey() -> (str, str):
-    # 获取用户的公私钥对，若不存在，则生成之
-    # 返回 私钥，公钥 格式
-    pass
+import base64
+import os
+
+from config import config
+from tool.asymmetric import RSA
+
+
+def getUserKey() -> (str, str): # 返回base64编码
+    # 检查是否存在公私钥文件
+    if not os.path.exists(config.priKeySavePath) or not os.path.exists(config.pubKeySavePath):
+        # 生成新的密钥对
+        RSA.generate_keys(config.priKeySavePath,config.pubKeySavePath)
+
+    # 读取私钥
+    with open(config.priKeySavePath, "rb") as f:
+        data = f.read()
+        private_key = base64.b64encode(data).decode('utf-8')
+
+    # 读取公钥
+    with open(config.pubKeySavePath, "rb") as f:
+        data = f.read()
+        public_key = base64.b64encode(data).decode('utf-8')
+
+    return private_key, public_key

tool/asymmetric/RSA.py

@@ -0,0 +1,76 @@
+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
+
+
+def generate_keys(priKeySavePath, pubKeySavePath):
+    # 生成一个RSA密钥对象
+    key = RSA.generate(2048)
+
+    # 导出私钥
+    private_key = key.export_key()
+    with open(priKeySavePath, "wb") as f:
+        f.write(private_key)
+
+    # 导出公钥
+    public_key = key.publickey().export_key()
+    with open(pubKeySavePath, "wb") as f:
+        f.write(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)
+
+    return encrypted_message
+
+
+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(encrypted_message)
+
+    return decrypted_message
+
+
+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)
+    # 使用私钥对哈希值进行签名
+    signature = pkcs1_15.new(private_key).sign(hash_obj)
+
+    return binascii.hexlify(signature).decode('utf-8')
+
+
+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)
+    # 使用公钥验证签名
+    try:
+        pkcs1_15.new(public_key).verify(hash_obj, binascii.unhexlify(signature))
+        return True
+    except (ValueError, TypeError):
+        return False

tool/symmetric/AES.py

@@ -0,0 +1,50 @@
+# "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(32)
+
+    def encrypt(self, data: 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
+
+    def decrypt(self, encrypted_data: str, mode: str = 'ECB') -> str:
+        """解密数据"""
+        cipher = self._get_cipher(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"""
+        iv = None
+        if mode == 'CBC':
+            iv = get_random_bytes(AES.block_size)
+            return AES.new(self.key, AES.MODE_CBC, iv)
+        elif mode == 'CFB':
+            return AES.new(self.key, AES.MODE_CFB)
+        elif mode == 'OFB':
+            return AES.new(self.key, AES.MODE_OFB)
+        else:  # 默认是 ECB
+            return AES.new(self.key, AES.MODE_ECB)
+
+if __name__ == "__main__":
+    # 示例数据和密钥
+    data = "Hello, SM4!"
+    key = "1234567890abcdef"
+
+    enData , akey = AESUtils(key)
+
+    print(enData)
+

tool/symmetric/SM4.py

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

tool/test_PriKeyHelper.py

@@ -0,0 +1,39 @@
+import pytest
+
+from tool import PriKeyHelper
+from tool.asymmetric import RSA
+
+
+def test_encrypt():
+    key = PriKeyHelper.getUserKey()
+    message = b"Hello, this is a secret message."
+    # 打印keys
+    print(key[0])
+    print(key[1])
+
+    # 加密消息
+    encrypted = RSA.encrypt_message(message, key[1])
+    print(f"Encrypted: {encrypted}")
+
+    # 解密消息
+    decrypted = RSA.decrypt_message(encrypted, key[0])
+    print(f"Decrypted: {decrypted}")
+
+
+def test_sign():
+
+    key = PriKeyHelper.getUserKey()
+    message = b"This is a signed message."
+    # 打印keys
+    print(key[0])
+    print(key[1])
+
+
+    # 签名消息
+    signature = RSA.sign_message(message, key[0])
+    print(f"Signature: {signature}")
+
+    # 验证签名
+    is_valid = RSA.verify_signature(message, signature, key[1])
+    print(f"Is Signature Valid? {is_valid}")
+    assert is_valid == True