island: Change KeyBasedEncryptor's padding functions to use Crypto.Util.Padding

This commit is contained in:
Shreya Malviya 2021-10-04 14:02:51 +05:30
parent 3ab660b8fe
commit fc1affc0e7
1 changed files with 5 additions and 7 deletions

View File

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