From 0356596a41fff50fba180c7b04f3f6bfb6cc8ec9 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 11 Jul 2022 11:05:47 -0400 Subject: [PATCH] Island: Add ILockableEncryptor.reset_key() --- .../encryption/i_lockable_encryptor.py | 6 ++++++ .../encryption/repository_encryptor.py | 7 +++++-- .../encryption/test_repository_encryptor.py | 16 ++++++++++++++++ 3 files changed, 27 insertions(+), 2 deletions(-) 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 b1b52f1d8..f0f5ac7d3 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 @@ -43,6 +43,12 @@ class ILockableEncryptor(IEncryptor): Lock the encryptor, making it unusable """ + @abstractmethod + def reset_key(self): + """ + Reset the encryptor's key + """ + @abstractmethod def encrypt(self, plaintext: bytes) -> bytes: """ 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 0e0310147..b4d0722f4 100644 --- a/monkey/monkey_island/cc/server_utils/encryption/repository_encryptor.py +++ b/monkey/monkey_island/cc/server_utils/encryption/repository_encryptor.py @@ -1,4 +1,3 @@ -import os import secrets from pathlib import Path @@ -22,7 +21,7 @@ class RepositoryEncryptor(ILockableEncryptor): self._key_based_encryptor = self._initialize_key_based_encryptor() def _initialize_key_based_encryptor(self): - if os.path.exists(self._key_file): + if self._key_file.is_file(): return self._load_key() return self._create_key() @@ -46,6 +45,10 @@ class RepositoryEncryptor(ILockableEncryptor): def lock(self): self._key_based_encryptor = None + def reset_key(self): + if self._key_file.is_file(): + self._key_file.unlink() + def encrypt(self, plaintext: bytes) -> bytes: if self._key_based_encryptor is None: raise LockedKeyError("Cannot encrypt while the encryptor is locked)") 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 74b02fb9e..6b99ff36d 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 @@ -68,3 +68,19 @@ def test_lock(encryptor): with pytest.raises(LockedKeyError): encryptor.decrypt(encrypted_data) + + +def test_reset(encryptor, key_file): + encryptor.unlock(SECRET) + key_file_hash_1 = get_file_sha256_hash(key_file) + + encryptor.reset_key() + encryptor.unlock(SECRET) + key_file_hash_2 = get_file_sha256_hash(key_file) + + assert key_file_hash_1 != key_file_hash_2 + + +def test_reset_before_unlock(encryptor): + # Test will fail if an exception is raised + encryptor.reset_key()