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.
142 lines
4.8 KiB
Python
142 lines
4.8 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
|
|
|
|
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)
|