diff --git a/src/primaite/game/agent/interface.py b/src/primaite/game/agent/interface.py index 88848479..e641fabb 100644 --- a/src/primaite/game/agent/interface.py +++ b/src/primaite/game/agent/interface.py @@ -6,7 +6,7 @@ from gymnasium.core import ActType, ObsType from pydantic import BaseModel, model_validator from primaite.game.agent.actions import ActionManager -from primaite.game.agent.observations import ObservationManager +from primaite.game.agent.observations.observation_manager import ObservationManager from primaite.game.agent.rewards import RewardFunction if TYPE_CHECKING: @@ -146,23 +146,10 @@ class AbstractAgent(ABC): class AbstractScriptedAgent(AbstractAgent): """Base class for actors which generate their own behaviour.""" - pass - - -class RandomAgent(AbstractScriptedAgent): - """Agent that ignores its observation and acts completely at random.""" - + @abstractmethod def get_action(self, obs: ObsType, timestep: int = 0) -> Tuple[str, Dict]: - """Sample the action space randomly. - - :param obs: Current observation for this agent, not used in RandomAgent - :type obs: ObsType - :param timestep: The current simulation timestep, not used in RandomAgent - :type timestep: int - :return: Action formatted in CAOS format - :rtype: Tuple[str, Dict] - """ - return self.action_manager.get_action(self.action_manager.space.sample()) + """Return an action to be taken in the environment.""" + return super().get_action(obs=obs, timestep=timestep) class ProxyAgent(AbstractAgent): diff --git a/src/primaite/game/agent/observations/__init__.py b/src/primaite/game/agent/observations/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/primaite/game/agent/observations/agent_observations.py b/src/primaite/game/agent/observations/agent_observations.py new file mode 100644 index 00000000..522cdb59 --- /dev/null +++ b/src/primaite/game/agent/observations/agent_observations.py @@ -0,0 +1,188 @@ +from typing import Dict, List, Optional, Tuple, TYPE_CHECKING + +from gymnasium import spaces + +from primaite.game.agent.observations.observations import ( + AbstractObservation, + AclObservation, + ICSObservation, + LinkObservation, + NodeObservation, + NullObservation, +) + +if TYPE_CHECKING: + from primaite.game.game import PrimaiteGame + + +class UC2BlueObservation(AbstractObservation): + """Container for all observations used by the blue agent in UC2. + + TODO: there's no real need for a UC2 blue container class, we should be able to simply use the observation handler + for the purpose of compiling several observation components. + """ + + def __init__( + self, + nodes: List[NodeObservation], + links: List[LinkObservation], + acl: AclObservation, + ics: ICSObservation, + where: Optional[List[str]] = None, + ) -> None: + """Initialise UC2 blue observation. + + :param nodes: List of node observations + :type nodes: List[NodeObservation] + :param links: List of link observations + :type links: List[LinkObservation] + :param acl: The Access Control List observation + :type acl: AclObservation + :param ics: The ICS observation + :type ics: ICSObservation + :param where: Where in the simulation state dict to find information. Not used in this particular observation + because it only compiles other observations and doesn't contribute any new information, defaults to None + :type where: Optional[List[str]], optional + """ + super().__init__() + self.where: Optional[Tuple[str]] = where + + self.nodes: List[NodeObservation] = nodes + self.links: List[LinkObservation] = links + self.acl: AclObservation = acl + self.ics: ICSObservation = ics + + self.default_observation: Dict = { + "NODES": {i + 1: n.default_observation for i, n in enumerate(self.nodes)}, + "LINKS": {i + 1: l.default_observation for i, l in enumerate(self.links)}, + "ACL": self.acl.default_observation, + "ICS": self.ics.default_observation, + } + + def observe(self, state: Dict) -> Dict: + """Generate observation based on the current state of the simulation. + + :param state: Simulation state dictionary + :type state: Dict + :return: Observation + :rtype: Dict + """ + if self.where is None: + return self.default_observation + + obs = {} + obs["NODES"] = {i + 1: node.observe(state) for i, node in enumerate(self.nodes)} + obs["LINKS"] = {i + 1: link.observe(state) for i, link in enumerate(self.links)} + obs["ACL"] = self.acl.observe(state) + obs["ICS"] = self.ics.observe(state) + + return obs + + @property + def space(self) -> spaces.Space: + """ + Gymnasium space object describing the observation space shape. + + :return: Space + :rtype: spaces.Space + """ + return spaces.Dict( + { + "NODES": spaces.Dict({i + 1: node.space for i, node in enumerate(self.nodes)}), + "LINKS": spaces.Dict({i + 1: link.space for i, link in enumerate(self.links)}), + "ACL": self.acl.space, + "ICS": self.ics.space, + } + ) + + @classmethod + def from_config(cls, config: Dict, game: "PrimaiteGame") -> "UC2BlueObservation": + """Create UC2 blue observation from a config. + + :param config: Dictionary containing the configuration for this UC2 blue observation. This includes the nodes, + links, ACL and ICS observations. + :type config: Dict + :param game: Reference to the PrimaiteGame object that spawned this observation. + :type game: PrimaiteGame + :return: Constructed UC2 blue observation + :rtype: UC2BlueObservation + """ + node_configs = config["nodes"] + + num_services_per_node = config["num_services_per_node"] + num_folders_per_node = config["num_folders_per_node"] + num_files_per_folder = config["num_files_per_folder"] + num_nics_per_node = config["num_nics_per_node"] + nodes = [ + NodeObservation.from_config( + config=n, + game=game, + num_services_per_node=num_services_per_node, + num_folders_per_node=num_folders_per_node, + num_files_per_folder=num_files_per_folder, + num_nics_per_node=num_nics_per_node, + ) + for n in node_configs + ] + + link_configs = config["links"] + links = [LinkObservation.from_config(config=link, game=game) for link in link_configs] + + acl_config = config["acl"] + acl = AclObservation.from_config(config=acl_config, game=game) + + ics_config = config["ics"] + ics = ICSObservation.from_config(config=ics_config, game=game) + new = cls(nodes=nodes, links=links, acl=acl, ics=ics, where=["network"]) + return new + + +class UC2RedObservation(AbstractObservation): + """Container for all observations used by the red agent in UC2.""" + + def __init__(self, nodes: List[NodeObservation], where: Optional[List[str]] = None) -> None: + super().__init__() + self.where: Optional[List[str]] = where + self.nodes: List[NodeObservation] = nodes + + self.default_observation: Dict = { + "NODES": {i + 1: n.default_observation for i, n in enumerate(self.nodes)}, + } + + def observe(self, state: Dict) -> Dict: + """Generate observation based on the current state of the simulation.""" + if self.where is None: + return self.default_observation + + obs = {} + obs["NODES"] = {i + 1: node.observe(state) for i, node in enumerate(self.nodes)} + return obs + + @property + def space(self) -> spaces.Space: + """Gymnasium space object describing the observation space shape.""" + return spaces.Dict( + { + "NODES": spaces.Dict({i + 1: node.space for i, node in enumerate(self.nodes)}), + } + ) + + @classmethod + def from_config(cls, config: Dict, game: "PrimaiteGame") -> "UC2RedObservation": + """ + Create UC2 red observation from a config. + + :param config: Dictionary containing the configuration for this UC2 red observation. + :type config: Dict + :param game: Reference to the PrimaiteGame object that spawned this observation. + :type game: PrimaiteGame + """ + node_configs = config["nodes"] + nodes = [NodeObservation.from_config(config=cfg, game=game) for cfg in node_configs] + return cls(nodes=nodes, where=["network"]) + + +class UC2GreenObservation(NullObservation): + """Green agent observation. As the green agent's actions don't depend on the observation, this is empty.""" + + pass diff --git a/src/primaite/game/agent/observations/observation_manager.py b/src/primaite/game/agent/observations/observation_manager.py new file mode 100644 index 00000000..400345fa --- /dev/null +++ b/src/primaite/game/agent/observations/observation_manager.py @@ -0,0 +1,73 @@ +from typing import Dict, TYPE_CHECKING + +from gymnasium.core import ObsType + +from primaite.game.agent.observations.agent_observations import ( + UC2BlueObservation, + UC2GreenObservation, + UC2RedObservation, +) +from primaite.game.agent.observations.observations import AbstractObservation + +if TYPE_CHECKING: + from primaite.game.game import PrimaiteGame + + +class ObservationManager: + """ + Manage the observations of an Agent. + + The observation space has the purpose of: + 1. Reading the outputted state from the PrimAITE Simulation. + 2. Selecting parts of the simulation state that are requested by the simulation config + 3. Formatting this information so an agent can use it to make decisions. + """ + + # TODO: Dear code reader: This class currently doesn't do much except hold an observation object. It will be changed + # to have more of it's own behaviour, and it will replace UC2BlueObservation and UC2RedObservation during the next + # refactor. + + def __init__(self, observation: AbstractObservation) -> None: + """Initialise observation space. + + :param observation: Observation object + :type observation: AbstractObservation + """ + self.obs: AbstractObservation = observation + self.current_observation: ObsType + + def update(self, state: Dict) -> Dict: + """ + Generate observation based on the current state of the simulation. + + :param state: Simulation state dictionary + :type state: Dict + """ + self.current_observation = self.obs.observe(state) + return self.current_observation + + @property + def space(self) -> None: + """Gymnasium space object describing the observation space shape.""" + return self.obs.space + + @classmethod + def from_config(cls, config: Dict, game: "PrimaiteGame") -> "ObservationManager": + """Create observation space from a config. + + :param config: Dictionary containing the configuration for this observation space. + It should contain the key 'type' which selects which observation class to use (from a choice of: + UC2BlueObservation, UC2RedObservation, UC2GreenObservation) + The other key is 'options' which are passed to the constructor of the selected observation class. + :type config: Dict + :param game: Reference to the PrimaiteGame object that spawned this observation. + :type game: PrimaiteGame + """ + if config["type"] == "UC2BlueObservation": + return cls(UC2BlueObservation.from_config(config.get("options", {}), game=game)) + elif config["type"] == "UC2RedObservation": + return cls(UC2RedObservation.from_config(config.get("options", {}), game=game)) + elif config["type"] == "UC2GreenObservation": + return cls(UC2GreenObservation.from_config(config.get("options", {}), game=game)) + else: + raise ValueError("Observation space type invalid") diff --git a/src/primaite/game/agent/observations.py b/src/primaite/game/agent/observations/observations.py similarity index 79% rename from src/primaite/game/agent/observations.py rename to src/primaite/game/agent/observations/observations.py index 82e11fe0..6d6614f4 100644 --- a/src/primaite/game/agent/observations.py +++ b/src/primaite/game/agent/observations/observations.py @@ -4,7 +4,6 @@ from ipaddress import IPv4Address from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING from gymnasium import spaces -from gymnasium.core import ObsType from primaite import getLogger from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE @@ -822,236 +821,3 @@ class ICSObservation(NullObservation): """ICS observation placeholder, currently not implemented so always returns a single 0.""" pass - - -class UC2BlueObservation(AbstractObservation): - """Container for all observations used by the blue agent in UC2. - - TODO: there's no real need for a UC2 blue container class, we should be able to simply use the observation handler - for the purpose of compiling several observation components. - """ - - def __init__( - self, - nodes: List[NodeObservation], - links: List[LinkObservation], - acl: AclObservation, - ics: ICSObservation, - where: Optional[List[str]] = None, - ) -> None: - """Initialise UC2 blue observation. - - :param nodes: List of node observations - :type nodes: List[NodeObservation] - :param links: List of link observations - :type links: List[LinkObservation] - :param acl: The Access Control List observation - :type acl: AclObservation - :param ics: The ICS observation - :type ics: ICSObservation - :param where: Where in the simulation state dict to find information. Not used in this particular observation - because it only compiles other observations and doesn't contribute any new information, defaults to None - :type where: Optional[List[str]], optional - """ - super().__init__() - self.where: Optional[Tuple[str]] = where - - self.nodes: List[NodeObservation] = nodes - self.links: List[LinkObservation] = links - self.acl: AclObservation = acl - self.ics: ICSObservation = ics - - self.default_observation: Dict = { - "NODES": {i + 1: n.default_observation for i, n in enumerate(self.nodes)}, - "LINKS": {i + 1: l.default_observation for i, l in enumerate(self.links)}, - "ACL": self.acl.default_observation, - "ICS": self.ics.default_observation, - } - - def observe(self, state: Dict) -> Dict: - """Generate observation based on the current state of the simulation. - - :param state: Simulation state dictionary - :type state: Dict - :return: Observation - :rtype: Dict - """ - if self.where is None: - return self.default_observation - - obs = {} - obs["NODES"] = {i + 1: node.observe(state) for i, node in enumerate(self.nodes)} - obs["LINKS"] = {i + 1: link.observe(state) for i, link in enumerate(self.links)} - obs["ACL"] = self.acl.observe(state) - obs["ICS"] = self.ics.observe(state) - - return obs - - @property - def space(self) -> spaces.Space: - """ - Gymnasium space object describing the observation space shape. - - :return: Space - :rtype: spaces.Space - """ - return spaces.Dict( - { - "NODES": spaces.Dict({i + 1: node.space for i, node in enumerate(self.nodes)}), - "LINKS": spaces.Dict({i + 1: link.space for i, link in enumerate(self.links)}), - "ACL": self.acl.space, - "ICS": self.ics.space, - } - ) - - @classmethod - def from_config(cls, config: Dict, game: "PrimaiteGame") -> "UC2BlueObservation": - """Create UC2 blue observation from a config. - - :param config: Dictionary containing the configuration for this UC2 blue observation. This includes the nodes, - links, ACL and ICS observations. - :type config: Dict - :param game: Reference to the PrimaiteGame object that spawned this observation. - :type game: PrimaiteGame - :return: Constructed UC2 blue observation - :rtype: UC2BlueObservation - """ - node_configs = config["nodes"] - - num_services_per_node = config["num_services_per_node"] - num_folders_per_node = config["num_folders_per_node"] - num_files_per_folder = config["num_files_per_folder"] - num_nics_per_node = config["num_nics_per_node"] - nodes = [ - NodeObservation.from_config( - config=n, - game=game, - num_services_per_node=num_services_per_node, - num_folders_per_node=num_folders_per_node, - num_files_per_folder=num_files_per_folder, - num_nics_per_node=num_nics_per_node, - ) - for n in node_configs - ] - - link_configs = config["links"] - links = [LinkObservation.from_config(config=link, game=game) for link in link_configs] - - acl_config = config["acl"] - acl = AclObservation.from_config(config=acl_config, game=game) - - ics_config = config["ics"] - ics = ICSObservation.from_config(config=ics_config, game=game) - new = cls(nodes=nodes, links=links, acl=acl, ics=ics, where=["network"]) - return new - - -class UC2RedObservation(AbstractObservation): - """Container for all observations used by the red agent in UC2.""" - - def __init__(self, nodes: List[NodeObservation], where: Optional[List[str]] = None) -> None: - super().__init__() - self.where: Optional[List[str]] = where - self.nodes: List[NodeObservation] = nodes - - self.default_observation: Dict = { - "NODES": {i + 1: n.default_observation for i, n in enumerate(self.nodes)}, - } - - def observe(self, state: Dict) -> Dict: - """Generate observation based on the current state of the simulation.""" - if self.where is None: - return self.default_observation - - obs = {} - obs["NODES"] = {i + 1: node.observe(state) for i, node in enumerate(self.nodes)} - return obs - - @property - def space(self) -> spaces.Space: - """Gymnasium space object describing the observation space shape.""" - return spaces.Dict( - { - "NODES": spaces.Dict({i + 1: node.space for i, node in enumerate(self.nodes)}), - } - ) - - @classmethod - def from_config(cls, config: Dict, game: "PrimaiteGame") -> "UC2RedObservation": - """ - Create UC2 red observation from a config. - - :param config: Dictionary containing the configuration for this UC2 red observation. - :type config: Dict - :param game: Reference to the PrimaiteGame object that spawned this observation. - :type game: PrimaiteGame - """ - node_configs = config["nodes"] - nodes = [NodeObservation.from_config(config=cfg, game=game) for cfg in node_configs] - return cls(nodes=nodes, where=["network"]) - - -class UC2GreenObservation(NullObservation): - """Green agent observation. As the green agent's actions don't depend on the observation, this is empty.""" - - pass - - -class ObservationManager: - """ - Manage the observations of an Agent. - - The observation space has the purpose of: - 1. Reading the outputted state from the PrimAITE Simulation. - 2. Selecting parts of the simulation state that are requested by the simulation config - 3. Formatting this information so an agent can use it to make decisions. - """ - - # TODO: Dear code reader: This class currently doesn't do much except hold an observation object. It will be changed - # to have more of it's own behaviour, and it will replace UC2BlueObservation and UC2RedObservation during the next - # refactor. - - def __init__(self, observation: AbstractObservation) -> None: - """Initialise observation space. - - :param observation: Observation object - :type observation: AbstractObservation - """ - self.obs: AbstractObservation = observation - self.current_observation: ObsType - - def update(self, state: Dict) -> Dict: - """ - Generate observation based on the current state of the simulation. - - :param state: Simulation state dictionary - :type state: Dict - """ - self.current_observation = self.obs.observe(state) - return self.current_observation - - @property - def space(self) -> None: - """Gymnasium space object describing the observation space shape.""" - return self.obs.space - - @classmethod - def from_config(cls, config: Dict, game: "PrimaiteGame") -> "ObservationManager": - """Create observation space from a config. - - :param config: Dictionary containing the configuration for this observation space. - It should contain the key 'type' which selects which observation class to use (from a choice of: - UC2BlueObservation, UC2RedObservation, UC2GreenObservation) - The other key is 'options' which are passed to the constructor of the selected observation class. - :type config: Dict - :param game: Reference to the PrimaiteGame object that spawned this observation. - :type game: PrimaiteGame - """ - if config["type"] == "UC2BlueObservation": - return cls(UC2BlueObservation.from_config(config.get("options", {}), game=game)) - elif config["type"] == "UC2RedObservation": - return cls(UC2RedObservation.from_config(config.get("options", {}), game=game)) - elif config["type"] == "UC2GreenObservation": - return cls(UC2GreenObservation.from_config(config.get("options", {}), game=game)) - else: - raise ValueError("Observation space type invalid") diff --git a/src/primaite/game/agent/scripted_agents/__init__.py b/src/primaite/game/agent/scripted_agents/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/primaite/game/agent/data_manipulation_bot.py b/src/primaite/game/agent/scripted_agents/data_manipulation_bot.py similarity index 100% rename from src/primaite/game/agent/data_manipulation_bot.py rename to src/primaite/game/agent/scripted_agents/data_manipulation_bot.py diff --git a/src/primaite/game/agent/scripted_agents.py b/src/primaite/game/agent/scripted_agents/probabilistic_agent.py similarity index 97% rename from src/primaite/game/agent/scripted_agents.py rename to src/primaite/game/agent/scripted_agents/probabilistic_agent.py index 5111df32..9cddc978 100644 --- a/src/primaite/game/agent/scripted_agents.py +++ b/src/primaite/game/agent/scripted_agents/probabilistic_agent.py @@ -7,7 +7,7 @@ from gymnasium.core import ObsType from primaite.game.agent.actions import ActionManager from primaite.game.agent.interface import AbstractScriptedAgent -from primaite.game.agent.observations import ObservationManager +from primaite.game.agent.observations.observation_manager import ObservationManager from primaite.game.agent.rewards import RewardFunction diff --git a/src/primaite/game/agent/scripted_agents/random_agent.py b/src/primaite/game/agent/scripted_agents/random_agent.py new file mode 100644 index 00000000..34a4b5ac --- /dev/null +++ b/src/primaite/game/agent/scripted_agents/random_agent.py @@ -0,0 +1,21 @@ +from typing import Dict, Tuple + +from gymnasium.core import ObsType + +from primaite.game.agent.interface import AbstractScriptedAgent + + +class RandomAgent(AbstractScriptedAgent): + """Agent that ignores its observation and acts completely at random.""" + + def get_action(self, obs: ObsType, timestep: int = 0) -> Tuple[str, Dict]: + """Sample the action space randomly. + + :param obs: Current observation for this agent, not used in RandomAgent + :type obs: ObsType + :param timestep: The current simulation timestep, not used in RandomAgent + :type timestep: int + :return: Action formatted in CAOS format + :rtype: Tuple[str, Dict] + """ + return self.action_manager.get_action(self.action_manager.space.sample()) diff --git a/src/primaite/game/game.py b/src/primaite/game/game.py index 394a8154..33f9186b 100644 --- a/src/primaite/game/game.py +++ b/src/primaite/game/game.py @@ -6,11 +6,11 @@ from pydantic import BaseModel, ConfigDict from primaite import getLogger from primaite.game.agent.actions import ActionManager -from primaite.game.agent.data_manipulation_bot import DataManipulationAgent from primaite.game.agent.interface import AbstractAgent, AgentSettings, ProxyAgent -from primaite.game.agent.observations import ObservationManager +from primaite.game.agent.observations.observation_manager import ObservationManager from primaite.game.agent.rewards import RewardFunction -from primaite.game.agent.scripted_agents import ProbabilisticAgent +from primaite.game.agent.scripted_agents.data_manipulation_bot import DataManipulationAgent +from primaite.game.agent.scripted_agents.probabilistic_agent import ProbabilisticAgent from primaite.simulator.network.hardware.base import NodeOperatingState from primaite.simulator.network.hardware.nodes.host.computer import Computer from primaite.simulator.network.hardware.nodes.host.host_node import NIC diff --git a/tests/conftest.py b/tests/conftest.py index a117a1ef..20600e73 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -10,7 +10,8 @@ from _pytest.monkeypatch import MonkeyPatch from primaite import getLogger, PRIMAITE_PATHS from primaite.game.agent.actions import ActionManager from primaite.game.agent.interface import AbstractAgent -from primaite.game.agent.observations import ICSObservation, ObservationManager +from primaite.game.agent.observations.observation_manager import ObservationManager +from primaite.game.agent.observations.observations import ICSObservation from primaite.game.agent.rewards import RewardFunction from primaite.game.game import PrimaiteGame from primaite.session.session import PrimaiteSession diff --git a/tests/integration_tests/configuration_file_parsing/software_installation_and_configuration.py b/tests/integration_tests/configuration_file_parsing/software_installation_and_configuration.py index 3aff59af..f993af5f 100644 --- a/tests/integration_tests/configuration_file_parsing/software_installation_and_configuration.py +++ b/tests/integration_tests/configuration_file_parsing/software_installation_and_configuration.py @@ -5,8 +5,8 @@ from typing import Union import yaml from primaite.config.load import data_manipulation_config_path -from primaite.game.agent.data_manipulation_bot import DataManipulationAgent from primaite.game.agent.interface import ProxyAgent, RandomAgent +from primaite.game.agent.scripted_agents.data_manipulation_bot import DataManipulationAgent from primaite.game.game import APPLICATION_TYPES_MAPPING, PrimaiteGame, SERVICE_TYPES_MAPPING from primaite.simulator.network.container import Network from primaite.simulator.network.hardware.nodes.host.computer import Computer diff --git a/tests/integration_tests/game_layer/test_actions.py b/tests/integration_tests/game_layer/test_actions.py index 8911632c..740fb491 100644 --- a/tests/integration_tests/game_layer/test_actions.py +++ b/tests/integration_tests/game_layer/test_actions.py @@ -10,28 +10,14 @@ # 4. Check that the simulation has changed in the way that I expect. # 5. Repeat for all actions. -from typing import Dict, Tuple +from typing import Tuple import pytest -from primaite.game.agent.actions import ActionManager -from primaite.game.agent.interface import AbstractAgent, ProxyAgent -from primaite.game.agent.observations import ICSObservation, ObservationManager -from primaite.game.agent.rewards import RewardFunction +from primaite.game.agent.interface import ProxyAgent from primaite.game.game import PrimaiteGame from primaite.simulator.file_system.file_system_item_abc import FileSystemItemHealthStatus -from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState -from primaite.simulator.network.hardware.nodes.host.computer import Computer -from primaite.simulator.network.hardware.nodes.host.server import Server -from primaite.simulator.network.hardware.nodes.network.router import ACLAction, Router -from primaite.simulator.network.hardware.nodes.network.switch import Switch -from primaite.simulator.network.transmission.network_layer import IPProtocol -from primaite.simulator.network.transmission.transport_layer import Port -from primaite.simulator.sim_container import Simulation from primaite.simulator.system.applications.web_browser import WebBrowser -from primaite.simulator.system.services.dns.dns_client import DNSClient -from primaite.simulator.system.services.dns.dns_server import DNSServer -from primaite.simulator.system.services.web_server.web_server import WebServer from primaite.simulator.system.software import SoftwareHealthState diff --git a/tests/integration_tests/game_layer/test_observations.py b/tests/integration_tests/game_layer/test_observations.py index d1301759..b6aed30b 100644 --- a/tests/integration_tests/game_layer/test_observations.py +++ b/tests/integration_tests/game_layer/test_observations.py @@ -1,6 +1,6 @@ from gymnasium import spaces -from primaite.game.agent.observations import FileObservation +from primaite.game.agent.observations.observations import FileObservation from primaite.simulator.network.hardware.nodes.host.computer import Computer from primaite.simulator.sim_container import Simulation diff --git a/tests/integration_tests/network/test_capture_nmne.py b/tests/integration_tests/network/test_capture_nmne.py index 698bfc72..4bbde32f 100644 --- a/tests/integration_tests/network/test_capture_nmne.py +++ b/tests/integration_tests/network/test_capture_nmne.py @@ -1,4 +1,4 @@ -from primaite.game.agent.observations import NicObservation +from primaite.game.agent.observations.observations import NicObservation from primaite.simulator.network.hardware.nodes.host.server import Server from primaite.simulator.network.nmne import set_nmne_config from primaite.simulator.sim_container import Simulation diff --git a/tests/unit_tests/_primaite/_game/_agent/test_probabilistic_agent.py b/tests/unit_tests/_primaite/_game/_agent/test_probabilistic_agent.py index 73228e36..c556cfad 100644 --- a/tests/unit_tests/_primaite/_game/_agent/test_probabilistic_agent.py +++ b/tests/unit_tests/_primaite/_game/_agent/test_probabilistic_agent.py @@ -1,7 +1,8 @@ from primaite.game.agent.actions import ActionManager -from primaite.game.agent.observations import ICSObservation, ObservationManager +from primaite.game.agent.observations.observation_manager import ObservationManager +from primaite.game.agent.observations.observations import ICSObservation from primaite.game.agent.rewards import RewardFunction -from primaite.game.agent.scripted_agents import ProbabilisticAgent +from primaite.game.agent.scripted_agents.probabilistic_agent import ProbabilisticAgent def test_probabilistic_agent():