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

import base64
import os


from entity.Letter import Letter
from tool import PriKeyHelper
from tool.PriKeyHelper import getUserKey
from tool.asymmetric import RSA
from tool.symmetric.AES import AESUtils
from tool.symmetric.SM4 import encrypt_ecb, encrypt_cbc_with_iv
from Crypto.Random import get_random_bytes


# 三个全局变量，记录信封
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.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)


    # 获取用户的公私钥对进行签名
    # 使用对方的公钥进行加密
    # 发送信件
    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():
    rsaEncrySymKey = RSA.encrypt_message(letterSymKey, getRecvPubKey())

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

# 获得接收方的公钥
def getRecvPubKey():
    # recPubKey = input("plz input Receiver's Public Key: ")
    # 在某某地方获得对方的公钥，然后保存到某个地方，输入路径uoqu
    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()