#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2020/1/3 2:02
# @Author  : leon yan
# @Site    : 
# @File    : Encryption.py
# @Software: PyCharm

from Cryptodome import Random
from Cryptodome.Cipher import AES
from PyQt5.Qt import QWidget, QDateTime, QDate, QMessageBox, QFileDialog, QApplication
from PyQt5.QtCore import QDataStream, QFile, QDir, QIODevice, pyqtSlot

from resource.Encryption_ui import Ui_Form


class EncryptionForm(QWidget, Ui_Form):
    def __init__(self, parent=None, *args, **kwargs):
        super().__init__(parent, *args, **kwargs)
        self.setupUi(self)
        self.fileLab.setText("需打开或新建一个加密文件")
        curDateTime = QDateTime.currentDateTime()
        self.dateEdit.setDate(curDateTime.date())
        self.key = b'aasdw&65JH*&^JHL'
        self.__testFileName = ""
        self.ciphertext = ""

    def __del__(self):
        print("Encryption is deleted")

    def __iniWrite(self):  ##开始写文件操作
        self.fileDevice = QFile(self.__testFileName)  # 创建文件对象
        if not self.fileDevice.open(QIODevice.WriteOnly):
            del self.fileDevice  # 删除对象
            return False
        self.fileStream = QDataStream(self.fileDevice)
        self.fileStream.setVersion(QDataStream.Qt_5_12)  # 设置流版本号，写入和读取的版本号要兼容
        return True

    def __iniRead(self):  ##开始读文件操作
        if not QFile.exists(self.__testFileName):
            QMessageBox.critical(self, "错误", "文件不存在")
            return False
        self.fileDevice = QFile(self.__testFileName)  # 创建文件对象
        if not self.fileDevice.open(QIODevice.ReadOnly):
            del self.fileDevice  # 删除对象
            return False

        self.fileStream = QDataStream(self.fileDevice)
        self.fileStream.setVersion(QDataStream.Qt_5_12)  # 设置版本号，写入和读取的版本号要兼容
        return True

    def __delFileStream(self):  ##结束写文件操作
        self.fileDevice.close()
        del self.fileStream
        del self.fileDevice

    @pyqtSlot()
    def on_btnWrite_clicked(self):
        curPath = QDir.currentPath()  # 获取系统当前目录
        title = "选择文件"  # 对话框标题
        filt = "License文件(*.lic)"  # 文件过滤器
        fileName, flt = QFileDialog.getSaveFileName(self, title, curPath, filt)
        if (fileName == ""):
            return
        self.__testFileName = fileName  # lICENSE用文件
        self.fileLab.setText(fileName)

        '''加密'''
        self.do_encryption()

    @pyqtSlot()
    def on_btnRead_clicked(self):
        curPath = QDir.currentPath()  # 获取系统当前目录
        title = "选择文件"  # 对话框标题
        filt = "License文件(*.lic)"  # 文件过滤器
        fileName, flt = QFileDialog.getOpenFileName(self, title, curPath, filt)
        # fileName, flt = QFileDialog.getSaveFileName(self, title, curPath, filt)
        if (fileName == ""):
            return
        self.__testFileName = fileName  # lICENSE用文件
        self.fileLab.setText(fileName)

        if self.__iniRead():
            try:
                fsize = self.fileDevice.size()
                self.ciphertext = self.fileStream.readRawData(fsize)
                self.do_decrypt()
            except Exception as e:
                QMessageBox.critical(self, "读String过程出现错误", str(e))
            finally:
                self.__delFileStream()

    '''私钥解密'''

    def do_decrypt(self):
        # 解密的话要用key和iv生成新的AES对象
        mydecrypt = AES.new(self.key, AES.MODE_CFB, self.ciphertext[:16])
        # 使用新生成的AES对象，将加密的密文解密
        decrypttext = mydecrypt.decrypt(self.ciphertext[16:])
        self.read2panel(decrypttext)

    def read2panel(self, steam):
        aa = steam.decode()
        aa2 = aa.split('$$$')
        self.company_le.setText(aa2[0])
        self.machine_le.setText(aa2[1])
        self.account_le.setText(aa2[2])
        self.psw_le.setText(aa2[3])
        aa3 = QDate.fromString(aa2[4], "yyyy年M月d日")
        self.dateEdit.setDate(aa3)
        if aa2[5] == "True":
            self.checkBox.setChecked(True)

    '''公钥加密'''

    def do_encryption(self):
        company = str(self.company_le.text())
        machine = str(self.machine_le.text())
        account = str(self.account_le.text())
        psw = str(self.psw_le.text())
        dataEdit = str(self.dateEdit.text())
        if self.checkBox.isChecked():
            superuser = True
        else:
            superuser = None
        theStr = company + '$$$' + machine + '$$$' + account + '$$$' + psw + '$$$' + dataEdit + '$$$' + str(superuser)

        # 密钥key 长度必须为16（AES-128）、24（AES-192）、或32（AES-256）Bytes 长度.
        # 目前AES-128足够用
        # 生成长度等于AES块大小的不可重复的密钥向量
        iv = Random.new().read(AES.block_size)

        # 使用key和iv初始化AES对象, 使用MODE_CFB模式
        mycipher = AES.new(self.key, AES.MODE_CFB, iv)
        # 加密的明文长度必须为16的倍数，如果长度不为16的倍数，则需要补足为16的倍数
        # 将iv（密钥向量）加到加密的密文开头，一起传输
        self.ciphertext = iv + mycipher.encrypt(theStr.encode())
        '''保存至文件'''
        if self.__iniWrite():
            self.fileStream.writeRawData(self.ciphertext)
            self.__delFileStream()


if __name__ == '__main__':
    # pass
    import sys

    app = QApplication(sys.argv)
    Encryption_Form = EncryptionForm()

    Encryption_Form.show()
    sys.exit(app.exec_())