diff --git a/monkey/monkey_island/cc/models/__init__.py b/monkey/monkey_island/cc/models/__init__.py index 2fcf42fd1..7f7cde2d7 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 +from .simulation import Simulation, SimulationSchema, IslandModeEnum diff --git a/monkey/monkey_island/cc/models/simulation.py b/monkey/monkey_island/cc/models/simulation.py index 4f64b6a1a..693287160 100644 --- a/monkey/monkey_island/cc/models/simulation.py +++ b/monkey/monkey_island/cc/models/simulation.py @@ -1,11 +1,16 @@ from __future__ import annotations from dataclasses import dataclass +from enum import Enum from marshmallow import Schema, post_load from marshmallow_enum import EnumField -from monkey_island.cc.services.mode.mode_enum import IslandModeEnum + +class IslandModeEnum(Enum): + UNSET = "unset" + RANSOMWARE = "ransomware" + ADVANCED = "advanced" @dataclass(frozen=True) diff --git a/monkey/monkey_island/cc/repository/file_simulation_repository.py b/monkey/monkey_island/cc/repository/file_simulation_repository.py index 021a7083e..66db53a7f 100644 --- a/monkey/monkey_island/cc/repository/file_simulation_repository.py +++ b/monkey/monkey_island/cc/repository/file_simulation_repository.py @@ -2,9 +2,8 @@ import dataclasses import io from monkey_island.cc import repository -from monkey_island.cc.models import Simulation, SimulationSchema +from monkey_island.cc.models import IslandModeEnum, Simulation, SimulationSchema from monkey_island.cc.repository import IFileRepository, ISimulationRepository, RetrievalError -from monkey_island.cc.services.mode.mode_enum import IslandModeEnum SIMULATION_STATE_FILE_NAME = "simulation_state.json" diff --git a/monkey/monkey_island/cc/repository/i_simulation_repository.py b/monkey/monkey_island/cc/repository/i_simulation_repository.py index 816b691e1..80dda8076 100644 --- a/monkey/monkey_island/cc/repository/i_simulation_repository.py +++ b/monkey/monkey_island/cc/repository/i_simulation_repository.py @@ -1,7 +1,6 @@ from abc import ABC, abstractmethod -from monkey_island.cc.models import Simulation -from monkey_island.cc.services.mode.mode_enum import IslandModeEnum +from monkey_island.cc.models import IslandModeEnum, Simulation class ISimulationRepository(ABC): diff --git a/monkey/monkey_island/cc/resources/island_mode.py b/monkey/monkey_island/cc/resources/island_mode.py index 1a9483eec..43ca46d2d 100644 --- a/monkey/monkey_island/cc/resources/island_mode.py +++ b/monkey/monkey_island/cc/resources/island_mode.py @@ -3,10 +3,10 @@ import logging from flask import make_response, request +from monkey_island.cc.models import 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 -from monkey_island.cc.services.mode.mode_enum import IslandModeEnum logger = logging.getLogger(__name__) diff --git a/monkey/monkey_island/cc/services/island_mode_service.py b/monkey/monkey_island/cc/services/island_mode_service.py index 5183bdb8c..7bb08aa60 100644 --- a/monkey/monkey_island/cc/services/island_mode_service.py +++ b/monkey/monkey_island/cc/services/island_mode_service.py @@ -1,6 +1,6 @@ from common.configuration import AgentConfiguration +from monkey_island.cc.models import IslandModeEnum from monkey_island.cc.repository import IAgentConfigurationRepository, ISimulationRepository -from monkey_island.cc.services.mode.mode_enum import IslandModeEnum class IslandModeService: diff --git a/monkey/tests/monkey_island/in_memory_simulation_configuration.py b/monkey/tests/monkey_island/in_memory_simulation_configuration.py index e6e037fcd..7044889d2 100644 --- a/monkey/tests/monkey_island/in_memory_simulation_configuration.py +++ b/monkey/tests/monkey_island/in_memory_simulation_configuration.py @@ -1,8 +1,7 @@ import dataclasses -from monkey_island.cc.models import Simulation +from monkey_island.cc.models import IslandModeEnum, Simulation from monkey_island.cc.repository import ISimulationRepository -from monkey_island.cc.services.mode.mode_enum import IslandModeEnum class InMemorySimulationRepository(ISimulationRepository): 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 ca4bc6b8c..3ed23b83a 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,9 +1,8 @@ import pytest from tests.monkey_island import OpenErrorFileRepository, SingleFileRepository -from monkey_island.cc.models import Simulation +from monkey_island.cc.models import IslandModeEnum, Simulation from monkey_island.cc.repository import FileSimulationRepository, RetrievalError -from monkey_island.cc.services.mode.mode_enum import IslandModeEnum @pytest.fixture 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 ae3905923..d875fddb8 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,10 +5,10 @@ import pytest from tests.common import StubDIContainer from tests.monkey_island import InMemorySimulationRepository +from monkey_island.cc.models import IslandModeEnum 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 -from monkey_island.cc.services.mode.mode_enum import IslandModeEnum class MockIslandModeService(IslandModeService): 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 4d5f5f4bc..9e1ce30bf 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,8 +2,8 @@ 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.services import IslandModeService -from monkey_island.cc.services.mode.mode_enum import IslandModeEnum @pytest.fixture