diff --git a/monkey/monkey_island/cc/server_utils/encryption/__init__.py b/monkey/monkey_island/cc/server_utils/encryption/__init__.py index d90468cc4..592d42381 100644 --- a/monkey/monkey_island/cc/server_utils/encryption/__init__.py +++ b/monkey/monkey_island/cc/server_utils/encryption/__init__.py @@ -7,7 +7,7 @@ from .password_based_bytes_encryptor import ( InvalidCredentialsError, InvalidCiphertextError, ) -from .i_lockable_encryptor import ILockableEncryptor, LockedKeyError +from .i_lockable_encryptor import ILockableEncryptor, LockedKeyError, UnlockError from .repository_encryptor import RepositoryEncryptor from .data_store_encryptor import ( get_datastore_encryptor, diff --git a/monkey/monkey_island/cc/server_utils/encryption/i_lockable_encryptor.py b/monkey/monkey_island/cc/server_utils/encryption/i_lockable_encryptor.py index f0f5ac7d3..231af0a16 100644 --- a/monkey/monkey_island/cc/server_utils/encryption/i_lockable_encryptor.py +++ b/monkey/monkey_island/cc/server_utils/encryption/i_lockable_encryptor.py @@ -21,6 +21,12 @@ class LockedKeyError(Exception): """ +class UnlockError(Exception): + """ + Raised if an error occurs while attempting to unlock an ILockableEncryptor + """ + + class ILockableEncryptor(IEncryptor): """ An encryptor that can be locked or unlocked. @@ -35,6 +41,7 @@ class ILockableEncryptor(IEncryptor): Unlock the encryptor :param secret: A secret that must be used to access the ILockableEncryptor's key material. + :raises UnlockError: If the ILockableEncryptor could not be unlocked """ @abstractmethod diff --git a/monkey/monkey_island/cc/server_utils/encryption/repository_encryptor.py b/monkey/monkey_island/cc/server_utils/encryption/repository_encryptor.py index b4d0722f4..2c845fe43 100644 --- a/monkey/monkey_island/cc/server_utils/encryption/repository_encryptor.py +++ b/monkey/monkey_island/cc/server_utils/encryption/repository_encryptor.py @@ -3,7 +3,7 @@ from pathlib import Path from monkey_island.cc.server_utils.file_utils import open_new_securely_permissioned_file -from . import ILockableEncryptor, LockedKeyError +from . import ILockableEncryptor, LockedKeyError, UnlockError from .key_based_encryptor import KeyBasedEncryptor from .password_based_bytes_encryptor import PasswordBasedBytesEncryptor @@ -17,8 +17,11 @@ class RepositoryEncryptor(ILockableEncryptor): self._key_based_encryptor = None def unlock(self, secret: bytes): - self._password_based_encryptor = PasswordBasedBytesEncryptor(secret.decode()) - self._key_based_encryptor = self._initialize_key_based_encryptor() + try: + self._password_based_encryptor = PasswordBasedBytesEncryptor(secret.decode()) + self._key_based_encryptor = self._initialize_key_based_encryptor() + except Exception as err: + raise UnlockError(err) def _initialize_key_based_encryptor(self): if self._key_file.is_file(): diff --git a/monkey/tests/unit_tests/monkey_island/cc/server_utils/encryption/test_repository_encryptor.py b/monkey/tests/unit_tests/monkey_island/cc/server_utils/encryption/test_repository_encryptor.py index 6b99ff36d..d1592219a 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/server_utils/encryption/test_repository_encryptor.py +++ b/monkey/tests/unit_tests/monkey_island/cc/server_utils/encryption/test_repository_encryptor.py @@ -4,7 +4,11 @@ import string import pytest from common.utils.file_utils import get_file_sha256_hash -from monkey_island.cc.server_utils.encryption import LockedKeyError, RepositoryEncryptor +from monkey_island.cc.server_utils.encryption import ( + LockedKeyError, + RepositoryEncryptor, + UnlockError, +) PLAINTEXT = b"Hello, Monkey!" SECRET = b"53CR31" @@ -51,6 +55,20 @@ def test_existing_key_reused(encryptor, key_file): assert key_file_hash_1 == key_file_hash_2 +def test_unlock_os_error(encryptor, key_file): + key_file.mkdir() + + with pytest.raises(UnlockError): + encryptor.unlock(SECRET) + + +def test_unlock_wrong_password(encryptor): + encryptor.unlock(SECRET) + + with pytest.raises(UnlockError): + encryptor.unlock(b"WRONG_PASSWORD") + + def test_use_locked_encryptor__encrypt(encryptor): with pytest.raises(LockedKeyError): encryptor.encrypt(PLAINTEXT)