import sys
from PyQt5.QtWidgets import QMainWindow,QApplication,QFileDialog
import os
import socket
from Ui_untitled import Ui_MainWindow
from A_private_public import *
from digital_signature import *
from encrypt_message import *
from encrypt_decrypt_key import *
from B_private_public import *
from decrypt_ciphertext import *
from verify_digital_signature import *
from tcp_client import *
from tcp_server import *

class MyMainWindow(QMainWindow,Ui_MainWindow):
    def __init__(self,parent =None):
        super(MyMainWindow,self).__init__(parent)
        self.setupUi(self)

        self.send_file_path = None #初始化传输文件路径
        self.host = None#初始化ip和端口
        self.port = None
        self.host2 = None#初始化ip和端口
        self.port2 = None

        self.textEdit_2.setFontPointSize(20)
        self.textEdit_3.setFontPointSize(20)

        self.pushButton_9.clicked.connect(self.receive_BpubKey) #建立连接
        self.toolButton.clicked.connect(self.choose_file) #选择文件
        self.pushButton_3.clicked.connect(self.encrypt_message) #生成对称密钥并加密原文
        self.pushButton_8.clicked.connect(self.A_key) #生成A的密钥对
        self.pushButton_2.clicked.connect(self.digital_signature) # 执行数字签名
        self.pushButton.clicked.connect(self.encrypt_key) # 加密对称密钥
        self.pushButton_10.clicked.connect(self.update_zip) #发送加密包

        self.pushButton_12.clicked.connect(self.send_BpubKey) #等待连接
        self.pushButton_4.clicked.connect(self.B_key) # 生成B的密钥对
        self.pushButton_11.clicked.connect(self.download_zip) #下载加密包
        self.pushButton_5.clicked.connect(self.decrypt_key) # 解密对称密钥
        self.pushButton_6.clicked.connect(self.decrypt_ciphertext) # 解密密文
        self.pushButton_7.clicked.connect(self.verify_signature) # 验证数字签名

    def receive_BpubKey(self):
        try:
            s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            s.connect(('8.8.8.8', 80))
            ip = s.getsockname()[0]
        finally:
            s.close()
        try:
            self.text = self.plainTextEdit.toPlainText()
            self.host, self.port = self.text.split(':')
        except:
            # self.textEdit_2.setText('请输入正确的ip和端口')
            self.textEdit_2.setText("请输入正确的ip和端口")

        self.host = ip

        try:
            receive_Bpubkey(self.host,int(self.port))
            self.textEdit_2.setText('接收成功')
        except:
            self.textEdit_2.setText('接收失败')

    def update_zip(self):
        try:
            self.text = self.plainTextEdit.toPlainText()
            self.host, self.port = self.text.split(':')
        except:
            self.textEdit_2.setText('请输入正确的ip和端口')
        try:
            send_zip(self.host,int(self.port))
            self.textEdit_2.setText('发送成功')
        except:
            self.textEdit_2.setText('发送失败')


    def choose_file(self):
        file_dialog = QFileDialog()
        file_dialog.setFileMode(QFileDialog.AnyFile)
        file_dialog.setNameFilter("Text files (*.txt)")  # 过滤只显示文本文件
        file_dialog.setViewMode(QFileDialog.Detail)

        if file_dialog.exec_():
            file_paths = file_dialog.selectedFiles()
            self.textEdit_2.setText('\n'.join(file_paths))
            if file_paths:
                selected_file_path = file_paths[0]
                self.send_file_path = selected_file_path

    def encrypt_message(self):
        if self.send_file_path:
            current_directory = os.path.dirname(__file__)
            key_file_path = os.path.join(current_directory, 'symmetric_key.key')

            key = generate_and_save_key(key_file_path)
            encrypt_file(self.send_file_path, key)
            self.textEdit_2.setText('加密成功')
        else:
            self.textEdit_2.setText("请先选择文件进行加密")

    def A_key(self):
        A_key()
        self.textEdit_2.setText('成功生成发送方密钥对')

    def digital_signature(self):
        if self.send_file_path != None:  # 检查是否选择了文件
            hash_value = digital_signature(self.send_file_path)
            self.textEdit_2.setText(f"SHA-256摘要值: {hash_value}")
            self.textEdit_2.setText("数字签名已生成并保存！")
        else:
            self.textEdit_2.setText("请先选择文件进行数字签名")

    def encrypt_key(self):
        try:
            encrypt_key()
            self.textEdit_2.setText('对称密钥加密成功')
        except:
            self.textEdit_2.setText('加密失败，请检查公钥是否正确')


    def B_key(self):
        B_key()
        self.textEdit_3.setText('成功生成接收方密钥对')

    def decrypt_key(self):
        try:
            decrypt_key()
            self.textEdit_3.setText('对称密钥解密成功')
        except:
            self.textEdit_3.setText('解密失败，请检查私钥是否正确')

    def decrypt_ciphertext(self):
        try:
            current_directory = os.path.dirname(__file__)
            key_file_path = os.path.join(current_directory, 'decrypt_symmetric_key.key')
            encrypted_file_path = os.path.join(current_directory, 'file.txt.encrypted')

            key = load_key(key_file_path)
            decrypt_file(encrypted_file_path, key)
            self.textEdit_3.setText('密文解密成功')
        except:
            self.textEdit_3.setText('解密失败，请检查对称密钥是否正确')

    def verify_signature(self):
        try:
            verify_digital_signature()
            self.textEdit_3.setText('数字签名验证成功')
        except:
            self.textEdit_3.setText('请检查文件是否齐全且正确：发送方公钥、数字签名、对称密钥、密钥密文')

    def send_BpubKey(self):
        try:
            self.text = self.plainTextEdit_2.toPlainText()
            self.host2, self.port2 = self.text.split(':')
        except:
            self.textEdit_3.setText('请输入正确的ip和端口')
        try:
            send_Bpubkey(self.host2,int(self.port2))
            self.textEdit_3.setText('发送成功')
        except:
            self.textEdit_3.setText('发送失败')

    def download_zip(self):
        try:
            s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            s.connect(('8.8.8.8', 80))
            ip = s.getsockname()[0]
        finally:
            s.close()
        try:
            self.text = self.plainTextEdit_2.toPlainText()
            self.host2, self.port2 = self.text.split(':')
            self.host2 = ip
            # self.textEdit_3.setText(self.host2,self.port2)
        except:
            self.textEdit_3.setText('请输入正确的ip和端口')
        try:
            receive_zip(self.host2,int(self.port2))
            self.textEdit_3.setText('接收成功')
        except:
            self.textEdit_3.setText('接收失败')
 
if __name__ == "__main__":
    app = QApplication(sys.argv)
    myWin = MyMainWindow()
    myWin.show()
    sys.exit(app.exec_())