From 4cbed6dce93dc299a47b32d5167a322b494bab37 Mon Sep 17 00:00:00 2001 From: VakarisZ Date: Fri, 1 Oct 2021 12:34:21 +0300 Subject: [PATCH] Fix typos and rename files/classes related to data store encryptor. Change PasswordBasedBytesEncryptor interface to use bytes instead of io.BytesIO --- .../monkey_island/cc/resources/auth/auth.py | 4 +-- .../cc/resources/auth/credential_utils.py | 4 +-- .../cc/server_utils/encryption/__init__.py | 4 +-- .../encryption/data_store_encryptor.py | 25 ++++++++----------- ....py => password_based_bytes_encryption.py} | 19 +++++++------- ...py => password_based_string_encryptior.py} | 15 +++++------ 6 files changed, 32 insertions(+), 39 deletions(-) rename monkey/monkey_island/cc/server_utils/encryption/{password_based_byte_encryption.py => password_based_bytes_encryption.py} (77%) rename monkey/monkey_island/cc/server_utils/encryption/{password_based_string_encryption.py => password_based_string_encryptior.py} (58%) diff --git a/monkey/monkey_island/cc/resources/auth/auth.py b/monkey/monkey_island/cc/resources/auth/auth.py index b6c5adb60..06fcb9e07 100644 --- a/monkey/monkey_island/cc/resources/auth/auth.py +++ b/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) diff --git a/monkey/monkey_island/cc/resources/auth/credential_utils.py b/monkey/monkey_island/cc/resources/auth/credential_utils.py index 689d4cc0b..874906edb 100644 --- a/monkey/monkey_island/cc/resources/auth/credential_utils.py +++ b/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", "") diff --git a/monkey/monkey_island/cc/server_utils/encryption/__init__.py b/monkey/monkey_island/cc/server_utils/encryption/__init__.py index 531659a9e..e9c19d75e 100644 --- a/monkey/monkey_island/cc/server_utils/encryption/__init__.py +++ b/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, 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 c08a9cb70..624e663b4 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 @@ -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 ( - PasswordBasedByteEncryptor, +from monkey_island.cc.server_utils.encryption.password_based_bytes_encryption import ( + 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 = ( - PasswordBasedByteEncryptor(secret).decrypt(io.BytesIO(encrypted_key)).getvalue() - ) - self._key_base_encryptor = KeyBasedEncryptor(cipher_key) + cipher_key = PasswordBasedBytesEncryptor(secret).decrypt(encrypted_key) + self._key_based_encryptor = KeyBasedEncryptor(cipher_key) def _create_new_key(self, secret: str): cipher_key = Random.new().read(self._BLOCK_SIZE) - encrypted_key = ( - PasswordBasedByteEncryptor(secret).encrypt(io.BytesIO(cipher_key)).getvalue() - ) + encrypted_key = PasswordBasedBytesEncryptor(secret).encrypt(cipher_key) 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): diff --git a/monkey/monkey_island/cc/server_utils/encryption/password_based_byte_encryption.py b/monkey/monkey_island/cc/server_utils/encryption/password_based_bytes_encryption.py similarity index 77% rename from monkey/monkey_island/cc/server_utils/encryption/password_based_byte_encryption.py rename to monkey/monkey_island/cc/server_utils/encryption/password_based_bytes_encryption.py index 569dc0d12..2e7c05819 100644 --- a/monkey/monkey_island/cc/server_utils/encryption/password_based_byte_encryption.py +++ b/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): diff --git a/monkey/monkey_island/cc/server_utils/encryption/password_based_string_encryption.py b/monkey/monkey_island/cc/server_utils/encryption/password_based_string_encryptior.py similarity index 58% rename from monkey/monkey_island/cc/server_utils/encryption/password_based_string_encryption.py rename to monkey/monkey_island/cc/server_utils/encryption/password_based_string_encryptior.py index ea796a441..716455a5b 100644 --- a/monkey/monkey_island/cc/server_utils/encryption/password_based_string_encryption.py +++ b/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 ( - PasswordBasedByteEncryptor, +from monkey_island.cc.server_utils.encryption.password_based_bytes_encryption import ( + 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 = PasswordBasedByteEncryptor(self.password).encrypt(plaintext_stream) + ciphertext = PasswordBasedBytesEncryptor(self.password).encrypt(plaintext.encode()) - 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) - return plaintext_stream.getvalue().decode("utf-8") + plaintext_stream = PasswordBasedBytesEncryptor(self.password).decrypt(ciphertext) + return plaintext_stream.decode() def is_encrypted(ciphertext: str) -> bool: