From 78fd3f90121d31825c07bad35fbdc768c4b5d3ef Mon Sep 17 00:00:00 2001 From: Agnibho Mondal Date: Tue, 12 Sep 2023 20:42:25 +0530 Subject: [PATCH] Support for password protected private key --- filehandler.py | 4 ++-- signature.py | 23 ++++++++++----------- window.py | 55 +++++++++++++++++++++++++++++++++----------------- 3 files changed, 50 insertions(+), 32 deletions(-) diff --git a/filehandler.py b/filehandler.py index dbd4ce1..943b976 100644 --- a/filehandler.py +++ b/filehandler.py @@ -61,10 +61,10 @@ class FileHandler(): with ZipFile(self.file, "r", strict_timestamps=False) as source: source.extractall(self.directory.name) - def sign(self): + def sign(self, password=""): with open(os.path.join(self.directory.name, "prescription.json"), "r") as file: data=file.read() - signature=Signature.sign(data, certificate=config["certificate"], privkey=config["private_key"]) + signature=Signature.sign(data, certificate=config["certificate"], privkey=config["private_key"], password=password) with open(os.path.join(self.directory.name, "signature.p7m"), "w") as file: file.write(signature) shutil.copyfile(config["certificate"], os.path.join(self.directory.name, "certificate.pem")) diff --git a/signature.py b/signature.py index 7587411..604a2a3 100644 --- a/signature.py +++ b/signature.py @@ -11,18 +11,17 @@ from hashlib import sha256 class Signature(): - def sign(data, certificate, privkey): - try: - hash=sha256(data.encode()).hexdigest() - smime=SMIME.SMIME() - smime.load_key(privkey, certificate) - p7=smime.sign(BIO.MemoryBuffer(hash.encode()), SMIME.PKCS7_DETACHED) - out=BIO.MemoryBuffer() - smime.write(out, p7) - return(out.read().decode()) - except Exception as e: - print(e) - return None + def sign(data, certificate, privkey, password=""): + def get_password(*args): + print(password) + return bytes(password, "ascii") + hash=sha256(data.encode()).hexdigest() + smime=SMIME.SMIME() + smime.load_key(privkey, certificate, get_password) + p7=smime.sign(BIO.MemoryBuffer(hash.encode()), SMIME.PKCS7_DETACHED) + out=BIO.MemoryBuffer() + smime.write(out, p7) + return(out.read().decode()) def verify(data, certificate, signature): hash=sha256(data.encode()).hexdigest() diff --git a/window.py b/window.py index 0a68b93..6bb25c8 100644 --- a/window.py +++ b/window.py @@ -7,10 +7,13 @@ import os, sys, datetime, dateutil.parser, shutil from PyQt6.QtCore import Qt, QDateTime, QSize, pyqtSignal -from PyQt6.QtWidgets import QWidget, QMainWindow, QMessageBox, QLabel, QPushButton, QLineEdit, QTextEdit, QDateTimeEdit, QListWidget, QComboBox, QCheckBox, QVBoxLayout, QHBoxLayout, QFormLayout, QToolBar, QTabWidget, QStatusBar, QFileDialog, QCompleter, QSizePolicy +from PyQt6.QtWidgets import QWidget, QMainWindow, QMessageBox, QLabel, QPushButton, QLineEdit, QTextEdit, QDateTimeEdit, QListWidget, QComboBox, QCheckBox, QVBoxLayout, QHBoxLayout, QFormLayout, QToolBar, QTabWidget, QStatusBar, QFileDialog, QInputDialog, QCompleter, QSizePolicy from PyQt6.QtGui import QAction, QIcon from pathlib import Path from hashlib import md5 +from M2Crypto.EVP import EVPError +from M2Crypto.BIO import BIOError +from M2Crypto.SMIME import SMIME_Error from config import config from prescription import Prescription from renderer import Renderer @@ -88,29 +91,45 @@ class MainWindow(QMainWindow): sys.exit() def cmd_render(self): - self.update_instance() + self.refresh() if(self.save_state==md5(self.prescription.get_json().encode()).hexdigest()): - target=self.renderer.render(self.current_file.directory.name) - self.signal_view.emit(target) - self.renderbox.showMaximized() + try: + target=self.renderer.render(self.current_file.directory.name) + self.signal_view.emit(target) + self.renderbox.showMaximized() + except FileNotFoundError as e: + print(e) + QMessageBox.information(self, "Save first", "Please save the file before rendering.") + else: QMessageBox.information(self, "Save first", "Please save the file before rendering.") def cmd_sign(self): - self.update_instance() + self.refresh() if(self.save_state==md5(self.prescription.get_json().encode()).hexdigest()): - try: - self.current_file.sign() - self.cmd_save() - except FileNotFoundError as e: - print(e) - QMessageBox.information(self, "Save first", "Please save the file before signing.") - except TypeError as e: - print(e) - QMessageBox.information(self, "Configure", "Please add valid key and certificate to the config file.") - except Exception as e: - print(e) - QMessageBox.information(self, "Failed", "Failed to sign.") + password, ok=QInputDialog.getText(self, "Enter password", "Private key password", QLineEdit.EchoMode.Password) + if(ok): + try: + self.current_file.sign(password) + self.cmd_save() + except FileNotFoundError as e: + print(e) + QMessageBox.information(self, "Save first", "Please save the file before signing.") + except TypeError as e: + print(e) + QMessageBox.information(self, "Configure", "Please add valid key and certificate to the config file.") + except EVPError as e: + print(e) + QMessageBox.information(self, "Check password", "Failed to load key. Please check if password is correct.") + except BIOError as e: + print(e) + QMessageBox.information(self, "Not found", "Certifcate and/or key not found.") + except SMIME_Error as e: + print(e) + QMessageBox.information(self, "Failed to load", "Failed to sign. Please check if certificate and key match.") + except Exception as e: + print(e) + QMessageBox.information(self, "Failed", "Failed to sign.") else: QMessageBox.information(self, "Save first", "Please save the file before signing.") -- 2.39.2