From c7257cf00068124166c31e50749db98e22f097f9 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 11 Jul 2022 10:51:37 -0400 Subject: [PATCH] Island: Add RepositoryEncryptor --- .../cc/server_utils/encryption/__init__.py | 3 +- .../encryption/repository_encryptor.py | 59 ++++++++++++++++ .../encryption/test_repository_encryptor.py | 70 +++++++++++++++++++ 3 files changed, 131 insertions(+), 1 deletion(-) create mode 100644 monkey/monkey_island/cc/server_utils/encryption/repository_encryptor.py create mode 100644 monkey/tests/unit_tests/monkey_island/cc/server_utils/encryption/test_repository_encryptor.py diff --git a/monkey/monkey_island/cc/server_utils/encryption/__init__.py b/monkey/monkey_island/cc/server_utils/encryption/__init__.py index c3ed8c46d..d90468cc4 100644 --- a/monkey/monkey_island/cc/server_utils/encryption/__init__.py +++ b/monkey/monkey_island/cc/server_utils/encryption/__init__.py @@ -7,7 +7,8 @@ from .password_based_bytes_encryptor import ( InvalidCredentialsError, InvalidCiphertextError, ) -from .i_lockable_encryptor import ILockableEncryptor +from .i_lockable_encryptor import ILockableEncryptor, LockedKeyError +from .repository_encryptor import RepositoryEncryptor from .data_store_encryptor import ( get_datastore_encryptor, unlock_datastore_encryptor, diff --git a/monkey/monkey_island/cc/server_utils/encryption/repository_encryptor.py b/monkey/monkey_island/cc/server_utils/encryption/repository_encryptor.py new file mode 100644 index 000000000..0e0310147 --- /dev/null +++ b/monkey/monkey_island/cc/server_utils/encryption/repository_encryptor.py @@ -0,0 +1,59 @@ +import os +import secrets +from pathlib import Path + +from monkey_island.cc.server_utils.file_utils import open_new_securely_permissioned_file + +from . import ILockableEncryptor, LockedKeyError +from .key_based_encryptor import KeyBasedEncryptor +from .password_based_bytes_encryptor import PasswordBasedBytesEncryptor + + +class RepositoryEncryptor(ILockableEncryptor): + _KEY_LENGTH_BYTES = 32 + + def __init__(self, key_file: Path): + self._key_file = key_file + self._password_based_encryptor = None + 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() + + def _initialize_key_based_encryptor(self): + if os.path.exists(self._key_file): + return self._load_key() + + return self._create_key() + + def _load_key(self) -> KeyBasedEncryptor: + with open(self._key_file, "rb") as f: + encrypted_key = f.read() + + plaintext_key = self._password_based_encryptor.decrypt(encrypted_key) + return KeyBasedEncryptor(plaintext_key) + + def _create_key(self) -> KeyBasedEncryptor: + plaintext_key = secrets.token_bytes(RepositoryEncryptor._KEY_LENGTH_BYTES) + + encrypted_key = self._password_based_encryptor.encrypt(plaintext_key) + with open_new_securely_permissioned_file(str(self._key_file), "wb") as f: + f.write(encrypted_key) + + return KeyBasedEncryptor(plaintext_key) + + def lock(self): + self._key_based_encryptor = None + + def encrypt(self, plaintext: bytes) -> bytes: + if self._key_based_encryptor is None: + raise LockedKeyError("Cannot encrypt while the encryptor is locked)") + + return self._key_based_encryptor.encrypt(plaintext) + + def decrypt(self, ciphertext: bytes) -> bytes: + if self._key_based_encryptor is None: + raise LockedKeyError("Cannot decrypt while the encryptor is locked)") + + return self._key_based_encryptor.decrypt(ciphertext) 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 new file mode 100644 index 000000000..74b02fb9e --- /dev/null +++ b/monkey/tests/unit_tests/monkey_island/cc/server_utils/encryption/test_repository_encryptor.py @@ -0,0 +1,70 @@ +import random +import string + +import pytest + +from common.utils.file_utils import get_file_sha256_hash +from monkey_island.cc.server_utils.encryption import LockedKeyError, RepositoryEncryptor + +PLAINTEXT = b"Hello, Monkey!" +SECRET = b"53CR31" + + +@pytest.fixture +def key_file(tmp_path): + return tmp_path / "test_key.bin" + + +@pytest.fixture +def encryptor(key_file): + return RepositoryEncryptor(key_file) + + +def test_encryption(encryptor): + plaintext = "".join( + random.choice(string.ascii_uppercase + string.digits) for _ in range(128) # noqa: DUO102 + ).encode() + encryptor.unlock(SECRET) + + encrypted_data = encryptor.encrypt(plaintext) + assert encrypted_data != plaintext + + decrypted_data = encryptor.decrypt(encrypted_data) + assert decrypted_data == plaintext + + +def test_key_creation(encryptor, key_file): + assert not key_file.is_file() + encryptor.unlock(SECRET) + assert key_file.is_file() + + +def test_existing_key_reused(encryptor, key_file): + assert not key_file.is_file() + + encryptor.unlock(SECRET) + key_file_hash_1 = get_file_sha256_hash(key_file) + + encryptor.unlock(SECRET) + key_file_hash_2 = get_file_sha256_hash(key_file) + + assert key_file_hash_1 == key_file_hash_2 + + +def test_use_locked_encryptor__encrypt(encryptor): + with pytest.raises(LockedKeyError): + encryptor.encrypt(PLAINTEXT) + + +def test_use_locked_encryptor__decrypt(encryptor): + with pytest.raises(LockedKeyError): + encryptor.decrypt(PLAINTEXT) + + +def test_lock(encryptor): + encryptor.unlock(SECRET) + encrypted_data = encryptor.encrypt(PLAINTEXT) + encryptor.lock() + + with pytest.raises(LockedKeyError): + encryptor.decrypt(encrypted_data)