forked from p15670423/monkey
Island: Rename IslandModeEnum -> IslandMode
This commit is contained in:
parent
d88fc86d90
commit
f78fa73563
|
@ -8,4 +8,4 @@ from .monkey_ttl import MonkeyTtl
|
||||||
from .pba_results import PbaResults
|
from .pba_results import PbaResults
|
||||||
from monkey_island.cc.models.report.report import Report
|
from monkey_island.cc.models.report.report import Report
|
||||||
from .stolen_credentials import StolenCredentials
|
from .stolen_credentials import StolenCredentials
|
||||||
from .simulation import Simulation, SimulationSchema, IslandModeEnum
|
from .simulation import Simulation, SimulationSchema, IslandMode
|
||||||
|
|
|
@ -7,7 +7,7 @@ from marshmallow import Schema, post_load
|
||||||
from marshmallow_enum import EnumField
|
from marshmallow_enum import EnumField
|
||||||
|
|
||||||
|
|
||||||
class IslandModeEnum(Enum):
|
class IslandMode(Enum):
|
||||||
UNSET = "unset"
|
UNSET = "unset"
|
||||||
RANSOMWARE = "ransomware"
|
RANSOMWARE = "ransomware"
|
||||||
ADVANCED = "advanced"
|
ADVANCED = "advanced"
|
||||||
|
@ -15,11 +15,11 @@ class IslandModeEnum(Enum):
|
||||||
|
|
||||||
@dataclass(frozen=True)
|
@dataclass(frozen=True)
|
||||||
class Simulation:
|
class Simulation:
|
||||||
mode: IslandModeEnum = IslandModeEnum.UNSET
|
mode: IslandMode = IslandMode.UNSET
|
||||||
|
|
||||||
|
|
||||||
class SimulationSchema(Schema):
|
class SimulationSchema(Schema):
|
||||||
mode = EnumField(IslandModeEnum)
|
mode = EnumField(IslandMode)
|
||||||
|
|
||||||
@post_load
|
@post_load
|
||||||
def _make_simulation(self, data, **kwargs):
|
def _make_simulation(self, data, **kwargs):
|
||||||
|
|
|
@ -2,7 +2,7 @@ import dataclasses
|
||||||
import io
|
import io
|
||||||
|
|
||||||
from monkey_island.cc import repository
|
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
|
from monkey_island.cc.repository import IFileRepository, ISimulationRepository, RetrievalError
|
||||||
|
|
||||||
SIMULATION_STATE_FILE_NAME = "simulation_state.json"
|
SIMULATION_STATE_FILE_NAME = "simulation_state.json"
|
||||||
|
@ -31,10 +31,10 @@ class FileSimulationRepository(ISimulationRepository):
|
||||||
SIMULATION_STATE_FILE_NAME, io.BytesIO(simulation_json.encode())
|
SIMULATION_STATE_FILE_NAME, io.BytesIO(simulation_json.encode())
|
||||||
)
|
)
|
||||||
|
|
||||||
def get_mode(self) -> IslandModeEnum:
|
def get_mode(self) -> IslandMode:
|
||||||
return self.get_simulation().mode
|
return self.get_simulation().mode
|
||||||
|
|
||||||
def set_mode(self, mode: IslandModeEnum):
|
def set_mode(self, mode: IslandMode):
|
||||||
old_simulation = self.get_simulation()
|
old_simulation = self.get_simulation()
|
||||||
new_simulation = dataclasses.replace(old_simulation, mode=mode)
|
new_simulation = dataclasses.replace(old_simulation, mode=mode)
|
||||||
self.save_simulation(new_simulation)
|
self.save_simulation(new_simulation)
|
||||||
|
|
|
@ -1,6 +1,6 @@
|
||||||
from abc import ABC, abstractmethod
|
from abc import ABC, abstractmethod
|
||||||
|
|
||||||
from monkey_island.cc.models import IslandModeEnum, Simulation
|
from monkey_island.cc.models import IslandMode, Simulation
|
||||||
|
|
||||||
|
|
||||||
class ISimulationRepository(ABC):
|
class ISimulationRepository(ABC):
|
||||||
|
@ -25,7 +25,7 @@ class ISimulationRepository(ABC):
|
||||||
pass
|
pass
|
||||||
|
|
||||||
@abstractmethod
|
@abstractmethod
|
||||||
def get_mode(self) -> IslandModeEnum:
|
def get_mode(self) -> IslandMode:
|
||||||
"""
|
"""
|
||||||
Get's the island's current mode
|
Get's the island's current mode
|
||||||
|
|
||||||
|
@ -36,7 +36,7 @@ class ISimulationRepository(ABC):
|
||||||
pass
|
pass
|
||||||
|
|
||||||
@abstractmethod
|
@abstractmethod
|
||||||
def set_mode(self, mode: IslandModeEnum):
|
def set_mode(self, mode: IslandMode):
|
||||||
"""
|
"""
|
||||||
Set the island's mode
|
Set the island's mode
|
||||||
|
|
||||||
|
|
|
@ -3,7 +3,7 @@ import logging
|
||||||
|
|
||||||
from flask import make_response, request
|
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.AbstractResource import AbstractResource
|
||||||
from monkey_island.cc.resources.request_authentication import jwt_required
|
from monkey_island.cc.resources.request_authentication import jwt_required
|
||||||
from monkey_island.cc.services import IslandModeService
|
from monkey_island.cc.services import IslandModeService
|
||||||
|
|
|
@ -1,5 +1,5 @@
|
||||||
from common.configuration import AgentConfiguration
|
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
|
from monkey_island.cc.repository import IAgentConfigurationRepository, ISimulationRepository
|
||||||
|
|
||||||
|
|
||||||
|
@ -25,7 +25,7 @@ class IslandModeService:
|
||||||
"""
|
"""
|
||||||
return self._simulation_repository.get_mode()
|
return self._simulation_repository.get_mode()
|
||||||
|
|
||||||
def set_mode(self, mode: IslandModeEnum):
|
def set_mode(self, mode: IslandMode):
|
||||||
"""
|
"""
|
||||||
Set the island's mode
|
Set the island's mode
|
||||||
|
|
||||||
|
@ -35,8 +35,8 @@ class IslandModeService:
|
||||||
self._simulation_repository.set_mode(mode)
|
self._simulation_repository.set_mode(mode)
|
||||||
self._set_configuration(mode)
|
self._set_configuration(mode)
|
||||||
|
|
||||||
def _set_configuration(self, mode: IslandModeEnum):
|
def _set_configuration(self, mode: IslandMode):
|
||||||
if mode == IslandModeEnum.RANSOMWARE:
|
if mode == IslandMode.RANSOMWARE:
|
||||||
self._agent_configuration_repository.store_configuration(
|
self._agent_configuration_repository.store_configuration(
|
||||||
self._default_ransomware_agent_configuration
|
self._default_ransomware_agent_configuration
|
||||||
)
|
)
|
||||||
|
|
|
@ -1,6 +1,6 @@
|
||||||
import dataclasses
|
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
|
from monkey_island.cc.repository import ISimulationRepository
|
||||||
|
|
||||||
|
|
||||||
|
@ -14,8 +14,8 @@ class InMemorySimulationRepository(ISimulationRepository):
|
||||||
def save_simulation(self, simulation: Simulation):
|
def save_simulation(self, simulation: Simulation):
|
||||||
self._simulation = simulation
|
self._simulation = simulation
|
||||||
|
|
||||||
def get_mode(self) -> IslandModeEnum:
|
def get_mode(self) -> IslandMode:
|
||||||
return self._simulation.mode
|
return self._simulation.mode
|
||||||
|
|
||||||
def set_mode(self, mode: IslandModeEnum):
|
def set_mode(self, mode: IslandMode):
|
||||||
self._simulation = dataclasses.replace(self._simulation, mode=mode)
|
self._simulation = dataclasses.replace(self._simulation, mode=mode)
|
||||||
|
|
|
@ -1,7 +1,7 @@
|
||||||
import pytest
|
import pytest
|
||||||
from tests.monkey_island import OpenErrorFileRepository, SingleFileRepository
|
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
|
from monkey_island.cc.repository import FileSimulationRepository, RetrievalError
|
||||||
|
|
||||||
|
|
||||||
|
@ -10,7 +10,7 @@ def simulation_repository():
|
||||||
return FileSimulationRepository(SingleFileRepository())
|
return FileSimulationRepository(SingleFileRepository())
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.parametrize("mode", list(IslandModeEnum))
|
@pytest.mark.parametrize("mode", list(IslandMode))
|
||||||
def test_save_simulation(simulation_repository, mode):
|
def test_save_simulation(simulation_repository, mode):
|
||||||
simulation = Simulation(mode)
|
simulation = Simulation(mode)
|
||||||
simulation_repository.save_simulation(simulation)
|
simulation_repository.save_simulation(simulation)
|
||||||
|
@ -25,13 +25,13 @@ def test_get_default_simulation(simulation_repository):
|
||||||
|
|
||||||
|
|
||||||
def test_set_mode(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):
|
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():
|
def test_get_simulation_retrieval_error():
|
||||||
|
|
|
@ -5,7 +5,7 @@ import pytest
|
||||||
from tests.common import StubDIContainer
|
from tests.common import StubDIContainer
|
||||||
from tests.monkey_island import InMemorySimulationRepository
|
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.repository import RetrievalError
|
||||||
from monkey_island.cc.resources.island_mode import IslandMode as IslandModeResource
|
from monkey_island.cc.resources.island_mode import IslandMode as IslandModeResource
|
||||||
from monkey_island.cc.services import IslandModeService
|
from monkey_island.cc.services import IslandModeService
|
||||||
|
@ -15,10 +15,10 @@ class MockIslandModeService(IslandModeService):
|
||||||
def __init__(self):
|
def __init__(self):
|
||||||
self._simulation_repository = InMemorySimulationRepository()
|
self._simulation_repository = InMemorySimulationRepository()
|
||||||
|
|
||||||
def get_mode(self) -> IslandModeEnum:
|
def get_mode(self) -> IslandMode:
|
||||||
return self._simulation_repository.get_mode()
|
return self._simulation_repository.get_mode()
|
||||||
|
|
||||||
def set_mode(self, mode: IslandModeEnum):
|
def set_mode(self, mode: IslandMode):
|
||||||
self._simulation_repository.set_mode(mode)
|
self._simulation_repository.set_mode(mode)
|
||||||
|
|
||||||
|
|
||||||
|
@ -33,7 +33,7 @@ def flask_client(build_flask_client):
|
||||||
|
|
||||||
@pytest.mark.parametrize(
|
@pytest.mark.parametrize(
|
||||||
"mode",
|
"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):
|
def test_island_mode_post(flask_client, mode):
|
||||||
resp = flask_client.post(
|
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:
|
with build_flask_client(container) as flask_client:
|
||||||
resp = flask_client.post(
|
resp = flask_client.post(
|
||||||
IslandModeResource.urls[0],
|
IslandModeResource.urls[0],
|
||||||
data=json.dumps({"mode": IslandModeEnum.RANSOMWARE.value}),
|
data=json.dumps({"mode": IslandMode.RANSOMWARE.value}),
|
||||||
follow_redirects=True,
|
follow_redirects=True,
|
||||||
)
|
)
|
||||||
|
|
||||||
assert resp.status_code == 500
|
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):
|
def test_island_mode_endpoint(flask_client, mode):
|
||||||
flask_client.post(
|
flask_client.post(
|
||||||
IslandModeResource.urls[0], data=json.dumps({"mode": mode}), follow_redirects=True
|
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)
|
resp_get = flask_client.get(IslandModeResource.urls[0], follow_redirects=True)
|
||||||
assert resp_post.status_code == 422
|
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
|
||||||
|
|
|
@ -2,7 +2,7 @@ import pytest
|
||||||
from tests.monkey_island import InMemoryAgentConfigurationRepository, InMemorySimulationRepository
|
from tests.monkey_island import InMemoryAgentConfigurationRepository, InMemorySimulationRepository
|
||||||
|
|
||||||
from common.configuration import DEFAULT_AGENT_CONFIGURATION, DEFAULT_RANSOMWARE_AGENT_CONFIGURATION
|
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
|
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):
|
def test_set_mode(island_mode_service, mode):
|
||||||
island_mode_service.set_mode(mode)
|
island_mode_service.set_mode(mode)
|
||||||
assert island_mode_service.get_mode() == mode
|
assert island_mode_service.get_mode() == mode
|
||||||
|
@ -30,9 +30,9 @@ def test_set_mode(island_mode_service, mode):
|
||||||
@pytest.mark.parametrize(
|
@pytest.mark.parametrize(
|
||||||
"mode, expected_config",
|
"mode, expected_config",
|
||||||
[
|
[
|
||||||
(IslandModeEnum.UNSET, DEFAULT_AGENT_CONFIGURATION),
|
(IslandMode.UNSET, DEFAULT_AGENT_CONFIGURATION),
|
||||||
(IslandModeEnum.ADVANCED, DEFAULT_AGENT_CONFIGURATION),
|
(IslandMode.ADVANCED, DEFAULT_AGENT_CONFIGURATION),
|
||||||
(IslandModeEnum.RANSOMWARE, DEFAULT_RANSOMWARE_AGENT_CONFIGURATION),
|
(IslandMode.RANSOMWARE, DEFAULT_RANSOMWARE_AGENT_CONFIGURATION),
|
||||||
],
|
],
|
||||||
)
|
)
|
||||||
def test_set_mode_sets_config(
|
def test_set_mode_sets_config(
|
||||||
|
|
Loading…
Reference in New Issue