diff --git a/monkey/monkey_island/cc/server_utils/encryption/encryptors/key_based_encryptor.py b/monkey/monkey_island/cc/server_utils/encryption/encryptors/key_based_encryptor.py index b5fe92d96..78aaacf9c 100644 --- a/monkey/monkey_island/cc/server_utils/encryption/encryptors/key_based_encryptor.py +++ b/monkey/monkey_island/cc/server_utils/encryption/encryptors/key_based_encryptor.py @@ -5,6 +5,7 @@ import logging # is maintained. from Crypto import Random # noqa: DUO133 # nosec: B413 from Crypto.Cipher import AES # noqa: DUO133 # nosec: B413 +from Crypto.Util import Padding # noqa: DUO133 from monkey_island.cc.server_utils.encryption import IEncryptor @@ -37,11 +38,8 @@ class KeyBasedEncryptor(IEncryptor): cipher = AES.new(self._key, AES.MODE_CBC, cipher_iv) return self._unpad(cipher.decrypt(enc_message[AES.block_size :]).decode()) - # TODO: Review and evaluate the security of the padding function - def _pad(self, message): - return message + (self._BLOCK_SIZE - (len(message) % self._BLOCK_SIZE)) * chr( - self._BLOCK_SIZE - (len(message) % self._BLOCK_SIZE) - ) + def _pad(self, message: str) -> str: + return Padding.pad(message.encode(), self._BLOCK_SIZE).decode() - def _unpad(self, message: str): - return message[0 : -ord(message[len(message) - 1])] + def _unpad(self, message: str) -> str: + return Padding.unpad(message.encode(), self._BLOCK_SIZE).decode()