import os import base64 import copy import yaml from cryptography.exceptions import AlreadyFinalized from cryptography.exceptions import InvalidTag from cryptography.exceptions import UnsupportedAlgorithm from cryptography.hazmat.backends import default_backend from cryptography.hazmat.primitives import hashes from cryptography.hazmat.primitives.ciphers.aead import AESGCM from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC # TODO make the config actually safe, hard-coded encryption password is BAD (but a bit better than plaintext) class ConfigReader: kdf = None nonce = None CONFIG_VERSION = 3 def __init__(self, root_folder: str, password: str): if root_folder is None: raise Exception("No root folder was defined!") self.__key = None self.__aesgcm = None self.data = {} self.__password = password self.__config_location = os.path.join(root_folder, ".artnet", "artnet.config") ConfigReader.__create_kdf() self.__key = ConfigReader.kdf.derive(bytes(password.encode('utf-8'))) self.__aesgcm = AESGCM(self.__key) if os.path.exists(self.__config_location) and os.path.isfile(self.__config_location): self.read_config() else: if not os.path.exists(os.path.join(root_folder, ".artnet")): os.mkdir(os.path.join(root_folder, ".artnet")) self.create_default_config() self.read_config() def update_config(self): """ Update the written config with the local settings :return: """ file = open(self.__config_location, "w") yaml.dump(stream=file, data=self.encrypt_sensitive_data(copy.deepcopy(self.data))) def read_config(self): """ Read the config from file and overwrite local settings :return: """ print(f"Config file: {os.path.join(os.getcwd(), self.__config_location)}") file = open(self.__config_location, "r") data = yaml.safe_load(stream=file) self.data = self.decrypt_sensitive_data(data) def encrypt_sensitive_data(self, data: dict) -> dict: """ Encrypts the sensitive portions of the data :return: """ new_data = data new_data["db"]["password"] = self.encrypt_text(data["db"]["password"]) new_data["db"]["user"] = self.encrypt_text(data["db"]["user"]) new_data["file_root"] = self.encrypt_text(data["file_root"]) return new_data def decrypt_sensitive_data(self, data: dict) -> dict: """ Decrypts the sensitive portions of the data :return: """ new_data = data new_data["db"]["password"] = self.decrypt_text(data["db"]["password"]) new_data["db"]["user"] = self.decrypt_text(data["db"]["user"]) new_data["file_root"] = self.decrypt_text(data["file_root"]) return new_data def create_default_config(self): """ Create a default config, overwrites all settings and generates a new key :return: """ self.data = { "version": ConfigReader.CONFIG_VERSION, "db": { "host": "localhost", "port": 5432, "database": "artnet", "user": "your_user", "password": "enter_password_via_gui" }, "file_root": "", } self.update_config() @staticmethod def __create_kdf(): ConfigReader.kdf = PBKDF2HMAC(algorithm=hashes.SHA512(), length=32, salt=bytes("ArtN3t.WhatElse?".encode('utf-8')), iterations=10000, backend=default_backend() ) ConfigReader.nonce = bytes("qt34nvßn".encode('utf-8')) def encrypt_text(self, text: str) -> bytes: cipher_text_bytes = self.__aesgcm.encrypt(data=text.encode('utf-8'), associated_data=None, nonce=ConfigReader.nonce ) return base64.urlsafe_b64encode(cipher_text_bytes) def decrypt_text(self, cipher: str) -> str: if ConfigReader.kdf is None: ConfigReader.__create_kdf() try: decrypted_cipher_text_bytes = self.__aesgcm.decrypt( nonce=ConfigReader.nonce, data=base64.urlsafe_b64decode(cipher), associated_data=None ) except InvalidTag: raise Exception("Could not decrypt Text! Wrong Password?") return decrypted_cipher_text_bytes.decode('utf-8') if __name__ == "__main__": cr = ConfigReader(password="MySuperDuperKey", root_folder=".") print(cr.data)