You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

138 lines
4.6 KiB
Python

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
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"])
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"])
return new_data
def create_default_config(self):
"""
Create a default config, overwrites all settings and generates a new key
:return:
"""
self.data = {
"db": {
"host": "localhost",
"port": 5432,
"database": "artnet",
"user": "your_user",
"password": "enter_password_via_gui"
},
}
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)