From f78fa73563c181f9268c09d2d6598d2e10c9fe3b Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 1 Jul 2022 12:43:13 -0400 Subject: [PATCH] Island: Rename IslandModeEnum -> IslandMode --- monkey/monkey_island/cc/models/__init__.py | 2 +- monkey/monkey_island/cc/models/simulation.py | 6 +++--- .../cc/repository/file_simulation_repository.py | 6 +++--- .../cc/repository/i_simulation_repository.py | 6 +++--- monkey/monkey_island/cc/resources/island_mode.py | 2 +- .../cc/services/island_mode_service.py | 8 ++++---- .../in_memory_simulation_configuration.py | 6 +++--- .../repository/test_file_simulation_repository.py | 10 +++++----- .../monkey_island/cc/resources/test_island_mode.py | 14 +++++++------- .../cc/services/test_island_mode_service.py | 10 +++++----- 10 files changed, 35 insertions(+), 35 deletions(-) diff --git a/monkey/monkey_island/cc/models/__init__.py b/monkey/monkey_island/cc/models/__init__.py index 7f7cde2d7..bf2addf63 100644 --- a/monkey/monkey_island/cc/models/__init__.py +++ b/monkey/monkey_island/cc/models/__init__.py @@ -8,4 +8,4 @@ from .monkey_ttl import MonkeyTtl from .pba_results import PbaResults from monkey_island.cc.models.report.report import Report from .stolen_credentials import StolenCredentials -from .simulation import Simulation, SimulationSchema, IslandModeEnum +from .simulation import Simulation, SimulationSchema, IslandMode diff --git a/monkey/monkey_island/cc/models/simulation.py b/monkey/monkey_island/cc/models/simulation.py index 693287160..88b3ca25d 100644 --- a/monkey/monkey_island/cc/models/simulation.py +++ b/monkey/monkey_island/cc/models/simulation.py @@ -7,7 +7,7 @@ from marshmallow import Schema, post_load from marshmallow_enum import EnumField -class IslandModeEnum(Enum): +class IslandMode(Enum): UNSET = "unset" RANSOMWARE = "ransomware" ADVANCED = "advanced" @@ -15,11 +15,11 @@ class IslandModeEnum(Enum): @dataclass(frozen=True) class Simulation: - mode: IslandModeEnum = IslandModeEnum.UNSET + mode: IslandMode = IslandMode.UNSET class SimulationSchema(Schema): - mode = EnumField(IslandModeEnum) + mode = EnumField(IslandMode) @post_load def _make_simulation(self, data, **kwargs): diff --git a/monkey/monkey_island/cc/repository/file_simulation_repository.py b/monkey/monkey_island/cc/repository/file_simulation_repository.py index 66db53a7f..0b4843fb6 100644 --- a/monkey/monkey_island/cc/repository/file_simulation_repository.py +++ b/monkey/monkey_island/cc/repository/file_simulation_repository.py @@ -2,7 +2,7 @@ import dataclasses import io from monkey_island.cc import repository -from monkey_island.cc.models import IslandModeEnum, Simulation, SimulationSchema +from monkey_island.cc.models import IslandMode, Simulation, SimulationSchema from monkey_island.cc.repository import IFileRepository, ISimulationRepository, RetrievalError SIMULATION_STATE_FILE_NAME = "simulation_state.json" @@ -31,10 +31,10 @@ class FileSimulationRepository(ISimulationRepository): SIMULATION_STATE_FILE_NAME, io.BytesIO(simulation_json.encode()) ) - def get_mode(self) -> IslandModeEnum: + def get_mode(self) -> IslandMode: return self.get_simulation().mode - def set_mode(self, mode: IslandModeEnum): + def set_mode(self, mode: IslandMode): old_simulation = self.get_simulation() new_simulation = dataclasses.replace(old_simulation, mode=mode) self.save_simulation(new_simulation) diff --git a/monkey/monkey_island/cc/repository/i_simulation_repository.py b/monkey/monkey_island/cc/repository/i_simulation_repository.py index 80dda8076..e6230c724 100644 --- a/monkey/monkey_island/cc/repository/i_simulation_repository.py +++ b/monkey/monkey_island/cc/repository/i_simulation_repository.py @@ -1,6 +1,6 @@ from abc import ABC, abstractmethod -from monkey_island.cc.models import IslandModeEnum, Simulation +from monkey_island.cc.models import IslandMode, Simulation class ISimulationRepository(ABC): @@ -25,7 +25,7 @@ class ISimulationRepository(ABC): pass @abstractmethod - def get_mode(self) -> IslandModeEnum: + def get_mode(self) -> IslandMode: """ Get's the island's current mode @@ -36,7 +36,7 @@ class ISimulationRepository(ABC): pass @abstractmethod - def set_mode(self, mode: IslandModeEnum): + def set_mode(self, mode: IslandMode): """ Set the island's mode diff --git a/monkey/monkey_island/cc/resources/island_mode.py b/monkey/monkey_island/cc/resources/island_mode.py index 43ca46d2d..3273f150e 100644 --- a/monkey/monkey_island/cc/resources/island_mode.py +++ b/monkey/monkey_island/cc/resources/island_mode.py @@ -3,7 +3,7 @@ import logging from flask import make_response, request -from monkey_island.cc.models import IslandModeEnum +from monkey_island.cc.models import IslandMode as IslandModeEnum from monkey_island.cc.resources.AbstractResource import AbstractResource from monkey_island.cc.resources.request_authentication import jwt_required from monkey_island.cc.services import IslandModeService diff --git a/monkey/monkey_island/cc/services/island_mode_service.py b/monkey/monkey_island/cc/services/island_mode_service.py index 7bb08aa60..2d65d0067 100644 --- a/monkey/monkey_island/cc/services/island_mode_service.py +++ b/monkey/monkey_island/cc/services/island_mode_service.py @@ -1,5 +1,5 @@ from common.configuration import AgentConfiguration -from monkey_island.cc.models import IslandModeEnum +from monkey_island.cc.models import IslandMode from monkey_island.cc.repository import IAgentConfigurationRepository, ISimulationRepository @@ -25,7 +25,7 @@ class IslandModeService: """ return self._simulation_repository.get_mode() - def set_mode(self, mode: IslandModeEnum): + def set_mode(self, mode: IslandMode): """ Set the island's mode @@ -35,8 +35,8 @@ class IslandModeService: self._simulation_repository.set_mode(mode) self._set_configuration(mode) - def _set_configuration(self, mode: IslandModeEnum): - if mode == IslandModeEnum.RANSOMWARE: + def _set_configuration(self, mode: IslandMode): + if mode == IslandMode.RANSOMWARE: self._agent_configuration_repository.store_configuration( self._default_ransomware_agent_configuration ) diff --git a/monkey/tests/monkey_island/in_memory_simulation_configuration.py b/monkey/tests/monkey_island/in_memory_simulation_configuration.py index 7044889d2..b2b90fb40 100644 --- a/monkey/tests/monkey_island/in_memory_simulation_configuration.py +++ b/monkey/tests/monkey_island/in_memory_simulation_configuration.py @@ -1,6 +1,6 @@ import dataclasses -from monkey_island.cc.models import IslandModeEnum, Simulation +from monkey_island.cc.models import IslandMode, Simulation from monkey_island.cc.repository import ISimulationRepository @@ -14,8 +14,8 @@ class InMemorySimulationRepository(ISimulationRepository): def save_simulation(self, simulation: Simulation): self._simulation = simulation - def get_mode(self) -> IslandModeEnum: + def get_mode(self) -> IslandMode: return self._simulation.mode - def set_mode(self, mode: IslandModeEnum): + def set_mode(self, mode: IslandMode): self._simulation = dataclasses.replace(self._simulation, mode=mode) diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_simulation_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_simulation_repository.py index 3ed23b83a..d01775b98 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_simulation_repository.py +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_simulation_repository.py @@ -1,7 +1,7 @@ import pytest from tests.monkey_island import OpenErrorFileRepository, SingleFileRepository -from monkey_island.cc.models import IslandModeEnum, Simulation +from monkey_island.cc.models import IslandMode, Simulation from monkey_island.cc.repository import FileSimulationRepository, RetrievalError @@ -10,7 +10,7 @@ def simulation_repository(): return FileSimulationRepository(SingleFileRepository()) -@pytest.mark.parametrize("mode", list(IslandModeEnum)) +@pytest.mark.parametrize("mode", list(IslandMode)) def test_save_simulation(simulation_repository, mode): simulation = Simulation(mode) simulation_repository.save_simulation(simulation) @@ -25,13 +25,13 @@ def test_get_default_simulation(simulation_repository): def test_set_mode(simulation_repository): - simulation_repository.set_mode(IslandModeEnum.ADVANCED) + simulation_repository.set_mode(IslandMode.ADVANCED) - assert simulation_repository.get_mode() == IslandModeEnum.ADVANCED + assert simulation_repository.get_mode() == IslandMode.ADVANCED def test_get_mode_default(simulation_repository): - assert simulation_repository.get_mode() == IslandModeEnum.UNSET + assert simulation_repository.get_mode() == IslandMode.UNSET def test_get_simulation_retrieval_error(): diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py index d875fddb8..780553eaa 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py @@ -5,7 +5,7 @@ import pytest from tests.common import StubDIContainer from tests.monkey_island import InMemorySimulationRepository -from monkey_island.cc.models import IslandModeEnum +from monkey_island.cc.models import IslandMode from monkey_island.cc.repository import RetrievalError from monkey_island.cc.resources.island_mode import IslandMode as IslandModeResource from monkey_island.cc.services import IslandModeService @@ -15,10 +15,10 @@ class MockIslandModeService(IslandModeService): def __init__(self): self._simulation_repository = InMemorySimulationRepository() - def get_mode(self) -> IslandModeEnum: + def get_mode(self) -> IslandMode: return self._simulation_repository.get_mode() - def set_mode(self, mode: IslandModeEnum): + def set_mode(self, mode: IslandMode): self._simulation_repository.set_mode(mode) @@ -33,7 +33,7 @@ def flask_client(build_flask_client): @pytest.mark.parametrize( "mode", - [IslandModeEnum.RANSOMWARE.value, IslandModeEnum.ADVANCED.value, IslandModeEnum.UNSET.value], + [IslandMode.RANSOMWARE.value, IslandMode.ADVANCED.value, IslandMode.UNSET.value], ) def test_island_mode_post(flask_client, mode): resp = flask_client.post( @@ -65,14 +65,14 @@ def test_island_mode_post__internal_server_error(build_flask_client): with build_flask_client(container) as flask_client: resp = flask_client.post( IslandModeResource.urls[0], - data=json.dumps({"mode": IslandModeEnum.RANSOMWARE.value}), + data=json.dumps({"mode": IslandMode.RANSOMWARE.value}), follow_redirects=True, ) assert resp.status_code == 500 -@pytest.mark.parametrize("mode", [IslandModeEnum.RANSOMWARE.value, IslandModeEnum.ADVANCED.value]) +@pytest.mark.parametrize("mode", [IslandMode.RANSOMWARE.value, IslandMode.ADVANCED.value]) def test_island_mode_endpoint(flask_client, mode): flask_client.post( IslandModeResource.urls[0], data=json.dumps({"mode": mode}), follow_redirects=True @@ -88,4 +88,4 @@ def test_island_mode_endpoint__invalid_mode(flask_client): ) resp_get = flask_client.get(IslandModeResource.urls[0], follow_redirects=True) assert resp_post.status_code == 422 - assert json.loads(resp_get.data)["mode"] == IslandModeEnum.UNSET.value + assert json.loads(resp_get.data)["mode"] == IslandMode.UNSET.value diff --git a/monkey/tests/unit_tests/monkey_island/cc/services/test_island_mode_service.py b/monkey/tests/unit_tests/monkey_island/cc/services/test_island_mode_service.py index 9e1ce30bf..feb125b49 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/services/test_island_mode_service.py +++ b/monkey/tests/unit_tests/monkey_island/cc/services/test_island_mode_service.py @@ -2,7 +2,7 @@ import pytest from tests.monkey_island import InMemoryAgentConfigurationRepository, InMemorySimulationRepository from common.configuration import DEFAULT_AGENT_CONFIGURATION, DEFAULT_RANSOMWARE_AGENT_CONFIGURATION -from monkey_island.cc.models import IslandModeEnum +from monkey_island.cc.models import IslandMode from monkey_island.cc.services import IslandModeService @@ -21,7 +21,7 @@ def island_mode_service(agent_configuration_repository): ) -@pytest.mark.parametrize("mode", list(IslandModeEnum)) +@pytest.mark.parametrize("mode", list(IslandMode)) def test_set_mode(island_mode_service, mode): island_mode_service.set_mode(mode) assert island_mode_service.get_mode() == mode @@ -30,9 +30,9 @@ def test_set_mode(island_mode_service, mode): @pytest.mark.parametrize( "mode, expected_config", [ - (IslandModeEnum.UNSET, DEFAULT_AGENT_CONFIGURATION), - (IslandModeEnum.ADVANCED, DEFAULT_AGENT_CONFIGURATION), - (IslandModeEnum.RANSOMWARE, DEFAULT_RANSOMWARE_AGENT_CONFIGURATION), + (IslandMode.UNSET, DEFAULT_AGENT_CONFIGURATION), + (IslandMode.ADVANCED, DEFAULT_AGENT_CONFIGURATION), + (IslandMode.RANSOMWARE, DEFAULT_RANSOMWARE_AGENT_CONFIGURATION), ], ) def test_set_mode_sets_config(