FileSecureTransfer/recvs/recv.py

import base64
import logging
import os.path

from config import config
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, AES


def getLetter():
    # 阻塞自身 从指定端口获取信件
    #
    # 解析信件
    letter = None
    handleLetter(letter)
    pass


def handleLetter(letter: Letter):
    # 解析信件 确认收信人
    # 获取自身key
    pki = PriKeyHelper.getUserKey()
    print("pki is: ",pki[1])
    print("letter.recvPubKey is: ",letter.recvPubKey)
    if pki[1] != letter.recvPubKey:
        raise Exception("信件不属于自己")
    # 用自己的私钥解密key 获得对称加密秘钥。
    key = RSA.decrypt_message(base64.b64decode(letter.encryptKey), pki[0])

    # 根据不同的对称加密算法
    try:
        type = config.getEncryptType(letter.encryptType)
    except ValueError:
        raise ValueError("不支持的对称加密算法")
    # 进行解密fileBase64
    data = b""
    if type == config.EncryptType.SM4_ECB:
        data = base64.b64decode(SM4.decrypt_ecb(letter.fileBase64, key))
    elif type == config.EncryptType.SM4_CBC:
        data = base64.b64decode(SM4.decrypt_cbc_with_iv(letter.fileBase64, key))
    elif type == config.EncryptType.AES_ECB:
        data = base64.b64decode(AES.AESUtils().decrypt(key, letter.fileBase64, "ecb"))
    elif type == config.EncryptType.AES_CBC:
        data = base64.b64decode(AES.AESUtils().decrypt(key, letter.fileBase64,"cbc"))
    else:
        raise KeyError("不支持的对称加密算法")

    # 用发信人的公钥验签摘要
    result = RSA.verify_signature(data, letter.sign, letter.senderPubKey)
    if not result:
        raise Exception("签名验证失败，文件不可信")

    # 保存文件

    # 默认下载目录
    download_dir = DownloadPathTool.get_download_directory()
    filename = letter.fileName
    base_name, ext = os.path.splitext(filename)  # 分离文件名和扩展名
    newName = None
    count = 0
    while True:
        if count:
            newName = f"{base_name}({count}){ext}"
        else:
            newName = filename
        path = f"{download_dir}/{newName}"
        if os.path.exists(path):
            logging.debug("文件已存在，自动避免覆盖写入")
            count += 1
        else:
            break
    with open(path, "wb") as f:
        f.write(data)
        print(f"确认收到来自 {Segwit.encodeSegwit(letter.senderPubKey.encode('utf-8'))} 的文件")
        print(f"签名验证有效，已将文件 {newName} 保存至 {download_dir} 下")
    return