diff --git a/monkey/monkey_island/cc/resources/configuration_export.py b/monkey/monkey_island/cc/resources/configuration_export.py index 4cecce634..00dc00fd0 100644 --- a/monkey/monkey_island/cc/resources/configuration_export.py +++ b/monkey/monkey_island/cc/resources/configuration_export.py @@ -1,42 +1,23 @@ import json -import os import flask_restful from flask import jsonify, request from monkey_island.cc.resources.auth.auth import jwt_required -from monkey_island.cc.server_utils.consts import DEFAULT_DATA_DIR from monkey_island.cc.services.config import ConfigService -from monkey_island.cc.services.utils.file_handler import encrypt_file_with_password +from monkey_island.cc.services.utils.config_encryption import encrypt_config class ConfigurationExport(flask_restful.Resource): @jwt_required def get(self): - return jsonify( - config_encrypted=self.file_encryption_successful, - plaintext_removed=self.plaintext_file_removal_successful, - ) + return jsonify(encrypted_config=self.encrypted_config) @jwt_required def post(self): - data = json.loads(request.data) + password = json.loads(request.data)["password"] + plaintext_config = ConfigService.get_config() - config = ConfigService.get_config() - - config_filename = "monkey.conf" - plaintext_config_path = os.path.join(DEFAULT_DATA_DIR, config_filename) - with open(plaintext_config_path) as file: - json.dump(config, file) - - self.file_encryption_successful = self.plaintext_file_removal_successful = False - if "password" in data: - encrypted_config_path = os.path.join(DEFAULT_DATA_DIR, f"encrypted_{config_filename}") - ( - self.file_encryption_successful, - self.plaintext_file_removal_successful, - ) = encrypt_file_with_password( - plaintext_config_path, encrypted_config_path, data["password"] - ) + self.encrypted_config = encrypt_config(plaintext_config, password) return self.get() diff --git a/monkey/monkey_island/cc/services/utils/config_encryption.py b/monkey/monkey_island/cc/services/utils/config_encryption.py new file mode 100644 index 000000000..3c3488b47 --- /dev/null +++ b/monkey/monkey_island/cc/services/utils/config_encryption.py @@ -0,0 +1,37 @@ +import io +import json +from typing import Dict + +import pyAesCrypt + +BUFFER_SIZE = 64 * 1024 + + +def encrypt_config(config: Dict, password: str) -> bytes: + plaintext_config_stream = io.BytesIO(json.dumps(config).encode()) + ciphertext_config_stream = io.BytesIO() + + pyAesCrypt.encryptStream( + plaintext_config_stream, ciphertext_config_stream, password, BUFFER_SIZE + ) + + ciphertext_config_bytes = ciphertext_config_stream.getvalue() + return ciphertext_config_bytes + + +def decrypt_config(enc_config: bytes, password: str) -> Dict: + ciphertext_config_stream = io.BytesIO(enc_config) + dec_plaintext_config_stream = io.BytesIO() + + len_ciphertext_config_stream = len(ciphertext_config_stream.getvalue()) + + pyAesCrypt.decryptStream( + ciphertext_config_stream, + dec_plaintext_config_stream, + password, + BUFFER_SIZE, + len_ciphertext_config_stream, + ) + + plaintext_config = json.loads(dec_plaintext_config_stream.getvalue().decode("utf-8")) + return plaintext_config diff --git a/monkey/monkey_island/cc/services/utils/file_handler.py b/monkey/monkey_island/cc/services/utils/file_handler.py deleted file mode 100644 index db71ea08a..000000000 --- a/monkey/monkey_island/cc/services/utils/file_handler.py +++ /dev/null @@ -1,39 +0,0 @@ -import logging -import os -from typing import Optional, Tuple - -import pyAesCrypt - -logger = logging.getLogger(__name__) - - -def encrypt_file_with_password( - plaintext_file_path: str, - encrypted_file_path: str, - password: str, - should_remove_plaintext_file: bool = True, -) -> Tuple[bool, Optional[bool]]: - - file_encryption_successful = False - try: - pyAesCrypt.encryptFile(plaintext_file_path, encrypted_file_path, password) - file_encryption_successful = True - except Exception as ex: - logger.error(f"Could not encrypt config file: {str(ex)}") - - plaintext_file_removal_successful = False - if file_encryption_successful and should_remove_plaintext_file: - plaintext_file_removal_successful = remove_file(plaintext_file_path) - - return file_encryption_successful, plaintext_file_removal_successful - - -def remove_file(path: str) -> bool: - file_removal_successful = False - try: - os.remove_file(path) - file_removal_successful = True - except Exception as ex: - logger.error(f"Could not remove plaintext file: {str(ex)}") - - return file_removal_successful