UT: Use encryptor fixture in agent encryption tests

This commit is contained in:
Kekoa Kaaikala 2022-09-20 14:03:10 +00:00
parent c0869aebba
commit dea7b4f74a
1 changed files with 9 additions and 17 deletions

View File

@ -6,7 +6,6 @@ import pytest
from common.agent_event_serializers import PydanticAgentEventSerializer from common.agent_event_serializers import PydanticAgentEventSerializer
from common.agent_events import AbstractAgentEvent from common.agent_events import AbstractAgentEvent
from monkey_island.cc.repository.agent_event_encryption import decrypt_event, encrypt_event from monkey_island.cc.repository.agent_event_encryption import decrypt_event, encrypt_event
from monkey_island.cc.server_utils.encryption import RepositoryEncryptor
class FakeAgentEvent(AbstractAgentEvent): class FakeAgentEvent(AbstractAgentEvent):
@ -25,21 +24,14 @@ def key_file(tmp_path):
return tmp_path / "test_key.bin" return tmp_path / "test_key.bin"
@pytest.fixture
def encryptor(key_file):
encryptor = RepositoryEncryptor(key_file)
encryptor.unlock(b"password")
return encryptor
@pytest.fixture @pytest.fixture
def serializer(): def serializer():
return PydanticAgentEventSerializer(FakeAgentEvent) return PydanticAgentEventSerializer(FakeAgentEvent)
def test_agent_event_encryption__encrypts(encryptor, serializer): def test_agent_event_encryption__encrypts(repository_encryptor, serializer):
data = serializer.serialize(EVENT) data = serializer.serialize(EVENT)
encrypted_data = encrypt_event(encryptor.encrypt, data) encrypted_data = encrypt_event(repository_encryptor.encrypt, data)
# Encrypted fields have the "encrypted_" prefix # Encrypted fields have the "encrypted_" prefix
assert "encrypted_data" in encrypted_data assert "encrypted_data" in encrypted_data
@ -48,25 +40,25 @@ def test_agent_event_encryption__encrypts(encryptor, serializer):
assert encrypted_data["encrypted_dict_data"] is not EVENT.dict_data assert encrypted_data["encrypted_dict_data"] is not EVENT.dict_data
def test_agent_event_encryption__decrypts(encryptor, serializer): def test_agent_event_encryption__decrypts(repository_encryptor, serializer):
data = serializer.serialize(EVENT) data = serializer.serialize(EVENT)
encrypted_data = encrypt_event(encryptor.encrypt, data) encrypted_data = encrypt_event(repository_encryptor.encrypt, data)
decrypted_data = decrypt_event(encryptor.decrypt, encrypted_data) decrypted_data = decrypt_event(repository_encryptor.decrypt, encrypted_data)
deserialized_event = serializer.deserialize(decrypted_data) deserialized_event = serializer.deserialize(decrypted_data)
assert deserialized_event == EVENT assert deserialized_event == EVENT
def test_agent_event_encryption__encryption_throws(encryptor): def test_agent_event_encryption__encryption_throws(repository_encryptor):
data = "Not a dict." data = "Not a dict."
with pytest.raises(TypeError): with pytest.raises(TypeError):
encrypt_event(encryptor.encrypt, data, fields=[]) encrypt_event(repository_encryptor.encrypt, data, fields=[])
def test_agent_event_encryption__decryption_throws(encryptor): def test_agent_event_encryption__decryption_throws(repository_encryptor):
data = "Not a dict." data = "Not a dict."
with pytest.raises(TypeError): with pytest.raises(TypeError):
decrypt_event(encryptor.decrypt, data) decrypt_event(repository_encryptor.decrypt, data)