Fix typos and rename files/classes related to data store encryptor. Change PasswordBasedBytesEncryptor interface to use bytes instead of io.BytesIO

monkey/monkey_island/cc/resources/auth/auth.py

@@ -10,7 +10,7 @@ from jwt import PyJWTError
 import monkey_island.cc.environment.environment_singleton as env_singleton
 import monkey_island.cc.resources.auth.user_store as user_store
 from monkey_island.cc.resources.auth.credential_utils import (
-    get_creds_from_request,
+    get_credentials_from_request,
     password_matches_hash,
 )
 from monkey_island.cc.server_utils.encryption.data_store_encryptor import setup_datastore_key
@@ -41,7 +41,7 @@ class Authenticate(flask_restful.Resource):
             "password": "my_password"
         }
         """
-        username, password = get_creds_from_request(request)
+        username, password = get_credentials_from_request(request)
 
         if _credentials_match_registered_user(username, password):
             setup_datastore_key(username, password)

monkey/monkey_island/cc/resources/auth/credential_utils.py

@@ -19,13 +19,13 @@ def password_matches_hash(plaintext_password, password_hash):
 
 
 def get_user_credentials_from_request(_request) -> UserCreds:
-    username, password = get_creds_from_request(_request)
+    username, password = get_credentials_from_request(_request)
     password_hash = hash_password(password)
 
     return UserCreds(username, password_hash)
 
 
-def get_creds_from_request(_request: Request) -> Tuple[str, str]:
+def get_credentials_from_request(_request: Request) -> Tuple[str, str]:
     cred_dict = json.loads(request.data)
     username = cred_dict.get("username", "")
     password = cred_dict.get("password", "")

monkey/monkey_island/cc/server_utils/encryption/__init__.py

@@ -1,10 +1,10 @@
 from monkey_island.cc.server_utils.encryption.i_encryptor import IEncryptor
 from monkey_island.cc.server_utils.encryption.key_based_encryptor import KeyBasedEncryptor
-from monkey_island.cc.server_utils.encryption.password_based_string_encryption import (
+from monkey_island.cc.server_utils.encryption.password_based_string_encryptior import (
     PasswordBasedStringEncryptor,
     is_encrypted,
 )
-from .password_based_byte_encryption import InvalidCredentialsError, InvalidCiphertextError
+from .password_based_bytes_encryption import InvalidCredentialsError, InvalidCiphertextError
 from monkey_island.cc.server_utils.encryption.data_store_encryptor import (
     DataStoreEncryptor,
     get_datastore_encryptor,

monkey/monkey_island/cc/server_utils/encryption/data_store_encryptor.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import io
 import os
 
 # PyCrypto is deprecated, but we use pycryptodome, which uses the exact same imports but
@@ -10,8 +9,8 @@ from typing import Union
 from Crypto import Random  # noqa: DUO133  # nosec: B413
 
 from monkey_island.cc.server_utils.encryption import KeyBasedEncryptor
-from monkey_island.cc.server_utils.encryption.password_based_byte_encryption import (
+from monkey_island.cc.server_utils.encryption.password_based_bytes_encryption import (
-    PasswordBasedByteEncryptor,
+    PasswordBasedBytesEncryptor,
 )
 from monkey_island.cc.server_utils.file_utils import open_new_securely_permissioned_file
 
@@ -24,7 +23,7 @@ class DataStoreEncryptor:
 
     def __init__(self, key_file_dir: str):
         self.key_file_path = os.path.join(key_file_dir, self._KEY_FILENAME)
-        self._key_base_encryptor = None
+        self._key_based_encryptor = None
 
     def init_key(self, secret: str):
         if os.path.exists(self.key_file_path):
@@ -35,28 +34,24 @@ class DataStoreEncryptor:
     def _load_existing_key(self, secret: str):
         with open(self.key_file_path, "rb") as f:
             encrypted_key = f.read()
-        cipher_key = (
+        cipher_key = PasswordBasedBytesEncryptor(secret).decrypt(encrypted_key)
-            PasswordBasedByteEncryptor(secret).decrypt(io.BytesIO(encrypted_key)).getvalue()
+        self._key_based_encryptor = KeyBasedEncryptor(cipher_key)
-        )
-        self._key_base_encryptor = KeyBasedEncryptor(cipher_key)
 
     def _create_new_key(self, secret: str):
         cipher_key = Random.new().read(self._BLOCK_SIZE)
-        encrypted_key = (
+        encrypted_key = PasswordBasedBytesEncryptor(secret).encrypt(cipher_key)
-            PasswordBasedByteEncryptor(secret).encrypt(io.BytesIO(cipher_key)).getvalue()
-        )
         with open_new_securely_permissioned_file(self.key_file_path, "wb") as f:
             f.write(encrypted_key)
-        self._key_base_encryptor = KeyBasedEncryptor(cipher_key)
+        self._key_based_encryptor = KeyBasedEncryptor(cipher_key)
 
     def is_key_setup(self) -> bool:
-        return self._key_base_encryptor is not None
+        return self._key_based_encryptor is not None
 
     def enc(self, message: str):
-        return self._key_base_encryptor.encrypt(message)
+        return self._key_based_encryptor.encrypt(message)
 
     def dec(self, enc_message: str):
-        return self._key_base_encryptor.decrypt(enc_message)
+        return self._key_based_encryptor.decrypt(enc_message)
 
 
 def initialize_datastore_encryptor(key_file_dir: str):

monkey/monkey_island/cc/server_utils/encryption/password_based_bytes_encryption.py

@@ -1,6 +1,5 @@
 import io
 import logging
-from io import BytesIO
 
 import pyAesCrypt
 
@@ -17,28 +16,30 @@ logger = logging.getLogger(__name__)
 # Note: password != key
 
 
-class PasswordBasedByteEncryptor(IEncryptor):
+class PasswordBasedBytesEncryptor(IEncryptor):
 
     _BUFFER_SIZE = pyAesCrypt.crypto.bufferSizeDef
 
     def __init__(self, password: str):
         self.password = password
 
-    def encrypt(self, plaintext: BytesIO) -> BytesIO:
+    def encrypt(self, plaintext: bytes) -> bytes:
         ciphertext_stream = io.BytesIO()
 
-        pyAesCrypt.encryptStream(plaintext, ciphertext_stream, self.password, self._BUFFER_SIZE)
+        pyAesCrypt.encryptStream(
+            io.BytesIO(plaintext), ciphertext_stream, self.password, self._BUFFER_SIZE
+        )
 
-        return ciphertext_stream
+        return ciphertext_stream.getvalue()
 
-    def decrypt(self, ciphertext: BytesIO) -> BytesIO:
+    def decrypt(self, ciphertext: bytes) -> bytes:
         plaintext_stream = io.BytesIO()
 
-        ciphertext_stream_len = len(ciphertext.getvalue())
+        ciphertext_stream_len = len(ciphertext)
 
         try:
             pyAesCrypt.decryptStream(
-                ciphertext,
+                io.BytesIO(ciphertext),
                 plaintext_stream,
                 self.password,
                 self._BUFFER_SIZE,
@@ -51,7 +52,7 @@ class PasswordBasedByteEncryptor(IEncryptor):
             else:
                 logger.info("The corrupt ciphertext provided.")
                 raise InvalidCiphertextError
-        return plaintext_stream
+        return plaintext_stream.getvalue()
 
 
 class InvalidCredentialsError(Exception):

monkey/monkey_island/cc/server_utils/encryption/password_based_string_encryptior.py

@@ -1,12 +1,11 @@
 import base64
-import io
 import logging
 
 import pyAesCrypt
 
 from monkey_island.cc.server_utils.encryption import IEncryptor
-from monkey_island.cc.server_utils.encryption.password_based_byte_encryption import (
+from monkey_island.cc.server_utils.encryption.password_based_bytes_encryption import (
-    PasswordBasedByteEncryptor,
+    PasswordBasedBytesEncryptor,
 )
 
 logger = logging.getLogger(__name__)
@@ -20,17 +19,15 @@ class PasswordBasedStringEncryptor(IEncryptor):
         self.password = password
 
     def encrypt(self, plaintext: str) -> str:
-        plaintext_stream = io.BytesIO(plaintext.encode())
+        ciphertext = PasswordBasedBytesEncryptor(self.password).encrypt(plaintext.encode())
-        ciphertext = PasswordBasedByteEncryptor(self.password).encrypt(plaintext_stream)
 
-        return base64.b64encode(ciphertext.getvalue()).decode()
+        return base64.b64encode(ciphertext).decode()
 
     def decrypt(self, ciphertext: str) -> str:
         ciphertext = base64.b64decode(ciphertext)
-        ciphertext_stream = io.BytesIO(ciphertext)
 
-        plaintext_stream = PasswordBasedByteEncryptor(self.password).decrypt(ciphertext_stream)
+        plaintext_stream = PasswordBasedBytesEncryptor(self.password).decrypt(ciphertext)
-        return plaintext_stream.getvalue().decode("utf-8")
+        return plaintext_stream.decode()
 
 
 def is_encrypted(ciphertext: str) -> bool: