FileSecureTransfer/senders/sender.py

# 模式，文件，自己的公钥从哪里来，别人的公钥从哪里来
import base64
import os

import pyfiglet
from Crypto.Random import get_random_bytes

from entity.Letter import Letter
from tool import PriKeyHelper
from tool.PriKeyHelper import getUserKey
from tool.asymmetric import RSA
from tool.hash import Segwit
from tool.symmetric.AES import AESUtils
from tool.symmetric.SM4 import encrypt_ecb, encrypt_cbc_with_iv

# 三个全局变量，记录信封
letterWay = ""
letterMode = ""
letterSymKey = b""


def main():
    # greet
    print("")
    print(">>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>")
    greet = pyfiglet.figlet_format("File Secure Transfer", font="slant", width=250)
    print(greet)
    print("                                                             <For Secure And Fast File Transfer>")
    author = "                                                       <-Made By Li-Nuo-Cheng Tan-Jun-Wen Ren-Qing-Yu->"
    print(author)
    print("<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<")

    letter = Letter()
    # 用户输入各种数据填充letter字段
    path = selectFile()

    with open(path, "rb") as f:
        data = f.read()

    letter.fileName = getFileName(path)
    letter.recvPubKey = getRecvPubKey()
    letter.senderPubKey = getSenderPubKey()
    letter.fileBase64, akey = SymEncryption(base64.b64encode(data).decode("utf-8"), letterSymKey)
    letter.encryptKey = getEncryptKey(letter.recvPubKey)
    letter.encryptType = getEncryptType()

    letter.sign = getSign(data)

    print(letter.fileName)
    print(letter.fileBase64)
    print(letter.sign)
    print(letter.encryptType)
    print(letter.encryptKey)
    print(letter.recvPubKey)
    print(letter.senderPubKey)

    # by Timmoc
    print(
        f"已以 {Segwit.encodeSegwit(letter.senderPubKey.encode('utf-8'))} 的身份向"
        f" {Segwit.encodeSegwit(letter.recvPubKey.encode('utf-8'))} 生成加密信件")
    # by Timmoc

    # 获取用户的公私钥对进行签名
    # 使用对方的公钥进行加密
    # 发送信件
    return letter  # 方便recv测试，以后可以删除。
    pass


def selectFile() -> str:
    # s = "public.pem"
    while True:
        try:
            s = input("输入文件路径:")
            with open(s, "rb") as _:
                pass
            return s
        except FileNotFoundError:
            print("错误: 文件未找到，请输入正确的文件路径。")
        except PermissionError:
            print("错误: 无法访问该文件，请检查权限。")
        except Exception as e:
            print(f"发生未知错误: {e}")


# 获得文件名
def getFileName(fName: str) -> str:
    filePath = os.path.split(fName)
    return filePath[-1]


def sendLetter(letter: Letter, target="192.168.195.162:8426"):
    # 向目标ip和端口发送指定的信件
    pass


# 选择对称加密的方法和模式(aes/sm4)
def selectSymEncryptionChoice():
    global letterWay, letterMode
    encryWay = ""
    encryMode = ""

    # 选择加密算法
    while True:
        encryWay = input("选择加密算法 (aes/sm4): ").strip().lower()  # 统一转成小写
        if encryWay in ["aes", "sm4"]:
            letterWay = encryWay
            print(f"已经选择 '{encryWay}'.")
            break  # 输入有效后退出循环
        else:
            print("非法选择，请输入 'aes' 或 'sm4'.")

    # 选择加密算法的模式
    while True:
        if encryWay == "aes":
            encryMode = input("选择加密算法模式 (ecb/cbc): ").strip().lower()
            if encryMode in ["ecb", "cbc"]:
                letterMode = encryMode
                print(f"已选择 '{encryMode}' 加密模式.")
                break  # 输入有效后退出循环
            else:
                print("非法输入，请输入 ecb/cbc")

        elif encryWay == "sm4":
            encryMode = input("选择加密模式 (ecb/cbc): ").strip().lower()
            if encryMode in ["ecb", "cbc"]:
                letterMode = encryMode
                print(f"已选择 '{encryMode}' 加密模式.")
                break  # 输入有效后退出循环
            else:
                print("非法输入，请输入 enter ecb/cbc")

    # 返回加密方法和加密方法的模式
    return encryWay, encryMode


# 使用对称加密，返回加密后的数据和随机生成的密钥
def SymEncryption(encryData, key: bytes = None):
    global letterSymKey
    # 获得加密的方法和加密的模式
    way, mode = selectSymEncryptionChoice()

    if way == "aes":
        aesUtils = AESUtils(key)
        encryptedData, trueKey = aesUtils.encrypt(encryData, mode=mode)  # 这里encryData要改为文件内容
        letterSymKey = trueKey

        return encryptedData, trueKey

    if way == "sm4":
        key = get_random_bytes(16)
        if mode == "ecb":
            encrypted_data = encrypt_ecb(encryData, key)
        # mode 为 cbc
        else:
            encrypted_data = encrypt_cbc_with_iv(encryData, key)
        letterSymKey = key
        return encrypted_data, key


# 获得签名
def getSign(document_bytes):
    # 计算文件bytes
    priKey, pubKey = PriKeyHelper.getUserKey()
    signDocuHash = RSA.sign_message(document_bytes, priKey)

    return signDocuHash


# 获得加密的方法和模式，封装信封
def getEncryptType():
    encryType = f"{letterWay}_{letterMode}".upper()

    return encryType


# 对称密钥，返回的是使用接收方公钥加密后的对称密钥
def getEncryptKey(getRecvPubKey):
    rsaEncrySymKey = RSA.encrypt_message(letterSymKey, getRecvPubKey)

    return base64.b64encode(rsaEncrySymKey).decode("utf-8")


# 获得接收方的公钥
def getRecvPubKey():
    # recPubKey = input(" plz input Receiver's Public Key: ")
    # 在某某地方获得对方的公钥，然后保存到某个地方，输入路径
    recPubKeyPath = input("请输入接受方的公钥文件路径：")
    with open(recPubKeyPath, "rb") as f:
        data = f.read()
        recPubKey = base64.b64encode(data).decode('utf-8')
    return recPubKey


# 获得发送方的公钥
def getSenderPubKey():
    privateKey, publicKey = getUserKey()
    return publicKey


if __name__ == "__main__":
    # 示例数据和密钥
    # data = "Hello, AES!"
    #
    # data, akey = SymEncryption(data)

    # print(data,akey)
    #
    # encryptType = f"{letterWay}_{letterMode}".upper()
    # print(encryptType)
    main()