Island: Rename IslandModeEnum -> IslandMode

This commit is contained in:
Mike Salvatore 2022-07-01 12:43:13 -04:00
parent d88fc86d90
commit f78fa73563
10 changed files with 35 additions and 35 deletions

View File

@ -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

View File

@ -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):

View File

@ -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)

View File

@ -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

View File

@ -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

View File

@ -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
)

View File

@ -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)

View File

@ -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():

View File

@ -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

View File

@ -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(