ArtNet-Desktop/ArtNet/file/config_reader.py

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(os.path.isfile(self.__config_location))
        print(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) -> str:
        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)