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 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
|
||||
|
|
|
@ -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):
|
||||
|
|
|
@ -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)
|
||||
|
|
|
@ -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
|
||||
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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
|
||||
)
|
||||
|
|
|
@ -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)
|
||||
|
|
|
@ -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():
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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(
|
||||
|
|
Loading…
Reference in New Issue