Modify config encryption logic: don't save the file in the backend, just encrypt it and send it back to the frontend

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()

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

monkey/monkey_island/cc/services/utils/file_handler.py

@@ -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