diff --git a/monkey/monkey_island/cc/server_utils/encryption/data_store_encryptor.py b/monkey/monkey_island/cc/server_utils/encryption/data_store_encryptor.py index 3e5415f0d..fb4308a33 100644 --- a/monkey/monkey_island/cc/server_utils/encryption/data_store_encryptor.py +++ b/monkey/monkey_island/cc/server_utils/encryption/data_store_encryptor.py @@ -10,25 +10,27 @@ from .i_encryptor import IEncryptor from .key_based_encryptor import KeyBasedEncryptor from .password_based_bytes_encryptor import PasswordBasedBytesEncryptor +_KEY_FILE_NAME = "mongo_key.bin" + _encryptor: Union[None, IEncryptor] = None class DataStoreEncryptor(IEncryptor): _KEY_LENGTH_BYTES = 32 - def __init__(self, secret: str, key_file_path: Path): - self._key_file_path = key_file_path + def __init__(self, secret: str, key_file: Path): + self._key_file = key_file self._password_based_encryptor = PasswordBasedBytesEncryptor(secret) self._key_based_encryptor = self._initialize_key_based_encryptor() def _initialize_key_based_encryptor(self): - if os.path.exists(self._key_file_path): + if os.path.exists(self._key_file): return self._load_existing_key() return self._create_new_key() def _load_existing_key(self) -> KeyBasedEncryptor: - with open(self._key_file_path, "rb") as f: + with open(self._key_file, "rb") as f: encrypted_key = f.read() plaintext_key = self._password_based_encryptor.decrypt(encrypted_key) @@ -38,7 +40,7 @@ class DataStoreEncryptor(IEncryptor): plaintext_key = Random.new().read(DataStoreEncryptor._KEY_LENGTH_BYTES) encrypted_key = self._password_based_encryptor.encrypt(plaintext_key) - with open_new_securely_permissioned_file(self._key_file_path, "wb") as f: + with open_new_securely_permissioned_file(self._key_file, "wb") as f: f.write(encrypted_key) return KeyBasedEncryptor(plaintext_key) @@ -50,22 +52,24 @@ class DataStoreEncryptor(IEncryptor): return self._key_based_encryptor.decrypt(ciphertext) -def reset_datastore_encryptor(key_file_dir: str, secret: str, key_file_name: str = "mongo_key.bin"): - key_file_path = Path(key_file_dir) / key_file_name +def reset_datastore_encryptor(key_file_dir: str, secret: str, key_file_name: str = _KEY_FILE_NAME): + key_file = Path(key_file_dir) / key_file_name - if key_file_path.is_file(): - key_file_path.unlink() + if key_file.is_file(): + key_file.unlink() - unlock_datastore_encryptor(key_file_dir, secret, key_file_name) + _initialize_datastore_encryptor(key_file, secret) -def unlock_datastore_encryptor( - key_file_dir: str, secret: str, key_file_name: str = "mongo_key.bin" -): +def unlock_datastore_encryptor(key_file_dir: str, secret: str, key_file_name: str = _KEY_FILE_NAME): + key_file = Path(key_file_dir) / key_file_name + _initialize_datastore_encryptor(key_file, secret) + + +def _initialize_datastore_encryptor(key_file: Path, secret: str): global _encryptor - key_file_path = Path(key_file_dir) / key_file_name - _encryptor = DataStoreEncryptor(secret, key_file_path) + _encryptor = DataStoreEncryptor(secret, key_file) def get_datastore_encryptor() -> IEncryptor: