pysafe_gui.py

# -*- coding: utf-8 -*-

# Form implementation generated from reading ui file 'ui/pysafe.ui'
#
# Created by: PyQt5 UI code generator 5.15.6
#
# WARNING: Any manual changes made to this file will be lost when pyuic5 is
# run again.  Do not edit this file unless you know what you are doing.


from PyQt5 import QtCore, QtGui, QtWidgets
from PyQt5.QtWidgets import QApplication, QInputDialog, QLineEdit, QFileDialog
from PyQt5.QtGui import QIcon
import base64
import getpass
import os
import pathlib
import secrets
import sys
import cryptography
from cryptography.fernet import Fernet
from cryptography.hazmat.primitives.kdf.scrypt import Scrypt
import qdarkstyle
import glob

class Ui_MainWindow(object):
    def setupUi(self, MainWindow):
        MainWindow.setObjectName("MainWindow")
        MainWindow.resize(700, 520)
        MainWindow.setMinimumSize(QtCore.QSize(700, 520))
        MainWindow.setMaximumSize(QtCore.QSize(700, 520))
        icon = QtGui.QIcon()
        icon.addPixmap(QtGui.QPixmap("./img/pysafe.ico"), QtGui.QIcon.Normal, QtGui.QIcon.Off)
        MainWindow.setWindowIcon(icon)
        self.centralwidget = QtWidgets.QWidget(MainWindow)
        self.centralwidget.setObjectName("centralwidget")
        self.pysafe_label = QtWidgets.QLabel(self.centralwidget)
        self.pysafe_label.setGeometry(QtCore.QRect(270, 10, 141, 41))
        self.pysafe_label.setStyleSheet("font: bold 25pt \"Arial\";")
        self.pysafe_label.setObjectName("pysafe_label")
        self.progressBar = QtWidgets.QProgressBar(self.centralwidget)
        self.progressBar.setGeometry(QtCore.QRect(10, 460, 311, 23))
        self.progressBar.setProperty("value", 0)
        self.progressBar.setObjectName("progressBar")
        self.files_text = QtWidgets.QTextBrowser(self.centralwidget)
        self.files_text.setGeometry(QtCore.QRect(10, 90, 311, 351))
        self.files_text.setObjectName("files_text")
        self.start_btn = QtWidgets.QPushButton(self.centralwidget)
        self.start_btn.setGeometry(QtCore.QRect(610, 460, 75, 23))
        self.start_btn.setObjectName("start_btn")
        self.start_btn.setDisabled(True)
        self.files_label = QtWidgets.QLabel(self.centralwidget)
        self.files_label.setGeometry(QtCore.QRect(120, 60, 53, 21))
        self.files_label.setStyleSheet("font: 75  14pt \"Arial\";")
        self.files_label.setObjectName("files_label")
        self.line = QtWidgets.QFrame(self.centralwidget)
        self.line.setGeometry(QtCore.QRect(340, 90, 3, 350))
        self.line.setFrameShape(QtWidgets.QFrame.VLine)
        self.line.setFrameShadow(QtWidgets.QFrame.Sunken)
        self.line.setObjectName("line")
        self.status_text = QtWidgets.QTextBrowser(self.centralwidget)
        self.status_text.setGeometry(QtCore.QRect(360, 90, 331, 351))
        self.status_text.setObjectName("status_text")
        self.status_lebel = QtWidgets.QLabel(self.centralwidget)
        self.status_lebel.setGeometry(QtCore.QRect(500, 55, 67, 31))
        self.status_lebel.setStyleSheet("font: 75  14pt \"Arial\";")
        self.status_lebel.setObjectName("status_lebel")
        MainWindow.setCentralWidget(self.centralwidget)
        self.menubar = QtWidgets.QMenuBar(MainWindow)
        self.menubar.setGeometry(QtCore.QRect(0, 0, 700, 22))
        self.menubar.setObjectName("menubar")
        self.menuFile = QtWidgets.QMenu(self.menubar)
        self.menuFile.setObjectName("menuFile")
        self.menuAbout = QtWidgets.QMenu(self.menubar)
        self.menuAbout.setObjectName("menuAbout")
        MainWindow.setMenuBar(self.menubar)
        self.actionEncrypt = QtWidgets.QAction(MainWindow)
        self.actionEncrypt.setObjectName("actionEncrypt")
        self.actionDecrypt = QtWidgets.QAction(MainWindow)
        self.actionDecrypt.setObjectName("actionDecrypt")
        self.actionAbout = QtWidgets.QAction(MainWindow)
        self.actionAbout.setObjectName("actionAbout")
        self.menuFile.addAction(self.actionEncrypt)
        self.menuFile.addAction(self.actionDecrypt)
        self.menuAbout.addAction(self.actionAbout)
        self.menubar.addAction(self.menuFile.menuAction())
        self.menubar.addAction(self.menuAbout.menuAction())

        self.retranslateUi(MainWindow)
        QtCore.QMetaObject.connectSlotsByName(MainWindow)

    def retranslateUi(self, MainWindow):
        _translate = QtCore.QCoreApplication.translate
        MainWindow.setWindowTitle(_translate("MainWindow", "PySafe"))
        self.pysafe_label.setText(_translate("MainWindow", "PYSAFE"))
        self.start_btn.setText(_translate("MainWindow", "START"))
        self.files_label.setText(_translate("MainWindow", "Files"))
        self.status_lebel.setText(_translate("MainWindow", "Status"))
        self.menuFile.setTitle(_translate("MainWindow", "File"))
        self.menuAbout.setTitle(_translate("MainWindow", "Help"))
        self.actionEncrypt.setText(_translate("MainWindow", "Encrypt"))
        self.actionDecrypt.setText(_translate("MainWindow", "Decrypt"))
        self.actionAbout.setText(_translate("MainWindow", "Github : SAJAD-net/Pysafe"))
        
        self.actionEncrypt.triggered.connect(lambda : self.select_path(action="encrypt"))
        self.actionDecrypt.triggered.connect(lambda : self.select_path(action="decrypt"))
        self.start_btn.clicked.connect(self.start)

    
    def start(self):
        password = self.passwordDialog()
        if password:
            self.key = self.generate_key(password)
        else:
            return

        self.status_text.append(f"Total files: {self.count}")
        self.progressBar.setValue(0)
        self.num = 1
        
        if self.action == "encrypt":
            self.encrypt(self.files)
            self.progressBar.setValue(100)
        elif self.action == "decrypt":
            self.decrypt(self.files)
            self.progressBar.setValue(100)
            

    def passwordDialog(self):
        """Small input dialog for password"""
        password, ok = QInputDialog.getText(None , 'Input Dialog', 'Password:',echo=QLineEdit.Password)
        if ok:
            return password
        else:
            return False
            
    def select_path(self, action):
        """Select a directory to encrypt/decrypt all the containing files recursively"""
        self.path = QFileDialog.getExistingDirectory(None, 'Select Folder')
        fold_list = glob.glob(f"{self.path}/**/*", recursive=True)
        self.action = action
        self.count = 0
        self.files = []
        
        for file in fold_list:
            if os.path.isfile(file):
                if "pysalt.salt" in file:
                    continue
                fname = file.split("/")[-1]
                self.files_text.append(f"[+] {fname}")
                self.files.append(file)
                self.count += 1
        self.files_text.append("-"*44)
        self.start_btn.setEnabled(True)

    
    def generate_salt(self, size=16):
        """Generate the salt used for key derivation"""
        return secrets.token_bytes(size)


    def derive_key(self, salt, password):
        """Derive the key from the `password` using the passed `salt`"""
        kdf = Scrypt(salt=salt, length=32, n=2**14, r=8, p=1)
        return kdf.derive(password.encode())


    def load_salt(self):
        # load salt from pysalt.salt file
        if os.path.exists(f"{self.path}/conf/pysalt.salt"):    
            self.salt = open(f"{self.path}/conf/pysalt.salt", "rb").read()
        else:
            self.salt = self.generate_salt()
            if not os.path.exists(f"{self.path}/conf"):
                os.mkdir(f"{self.path}/conf")
            with open(f"{self.path}/conf/pysalt.salt", "wb") as salt_file:
                salt_file.write(self.salt)

    
    def generate_key(self, password, salt_size=16):
        """
        Generates a key from a `password` and the salt.
        """    
        
        # load existing salt
        self.load_salt()
        
        # generate the key from the salt and the password
        derived_key = self.derive_key(self.salt, password)

        # encode it using Base64 and return it
        return base64.urlsafe_b64encode(derived_key)

    
    def progressbar(self):
        self.percom = 100//(self.count+1)
        self.progressBar.setValue(self.percom*self.num)
    
    
    def encryptor(self, path):
        frnet = Fernet(self.key)

        with open(path, "rb") as file:
            # read all file data
            file_data = file.read()

        # encrypt data
        try:
            encrypted_data = frnet.encrypt(file_data)
        except cryptography.fernet.InvalidToken:
            self.status_text.append("[!] Invalid token, most likely the password is incorrect")
            return False
            
        fname = path.split("/")[-1]
        self.status_text.append(f"[+] Encrypting {fname}")

        # write the encrypted file
        with open(path, "wb") as file:
            file.write(encrypted_data)
        
        self.progressbar()


    def encrypt(self, files):
        try:
            for file in files:
                self.num+=1
                self.encryptor(file)
            self.status_text.append("\n[+] Encryption successfully completed!\n")
        except Exception as e:
            print(e)
            self.status_text.append("\n[+] Encryption failed!\n")            
        self.status_text.append("-"*44)
    
    
    def decryptor(self, path):
        f = Fernet(self.key)
    
        with open(path, "rb") as file:
            # read all file data
            file_data = file.read()

        # decrypt data
        try:
            decrypted_data = f.decrypt(file_data)
        except cryptography.fernet.InvalidToken:
            self.status_text.append("[!] Invalid token, most likely the password is incorrect")
            return False
            
        fname = path.split("/")[-1]
        self.status_text.append(f"[+] Decrypting {fname}")
    
        # write the decrypted file
        with open(path, "wb") as file:
            file.write(decrypted_data)
        
        self.progressbar()


    def decrypt(self, files):
        try:
            for file in files:
                self.num+=1
                self.decryptor(file)        
            self.status_text.append("\n[+] Decryption successfully completed!")
        except:
            self.status_text.append("\n[+] Decryption failed!")
            
        self.status_text.append("-"*44)
        
    
if __name__ == "__main__":
    app = QtWidgets.QApplication(sys.argv)
    pysafe = QtWidgets.QMainWindow()
    pysafe.setStyleSheet(qdarkstyle.load_stylesheet_pyqt5())
    ui = Ui_MainWindow()
    ui.setupUi(pysafe)
    pysafe.show()
    sys.exit(app.exec_())