diff --git a/src/primaite/game/agent/scripted_agents/abstract_tap.py b/src/primaite/game/agent/scripted_agents/abstract_tap.py index add29b03..d30ba9b1 100644 --- a/src/primaite/game/agent/scripted_agents/abstract_tap.py +++ b/src/primaite/game/agent/scripted_agents/abstract_tap.py @@ -9,6 +9,7 @@ from gymnasium.core import ObsType from primaite.game.agent.scripted_agents.interface import AbstractAgent, AbstractScriptedAgent +__all__ = ("AbstractTAPAgent") class AbstractTAPAgent(AbstractScriptedAgent, identifier="Abstract_TAP"): """Base class for TAP agents to inherit from.""" @@ -22,11 +23,6 @@ class AbstractTAPAgent(AbstractScriptedAgent, identifier="Abstract_TAP"): starting_node_name: Optional[str] = None - # @property - # def next_execution_timestep(self) -> int: - # """Returns the agents next execution timestep.""" - # return self.next_execution_timestep - @abstractmethod def get_action(self, obs: ObsType, timestep: int = 0) -> Tuple[str, Dict]: """Return an action to be taken in the environment.""" diff --git a/src/primaite/game/agent/scripted_agents/data_manipulation_bot.py b/src/primaite/game/agent/scripted_agents/data_manipulation_bot.py index 84cad9f6..594f1b41 100644 --- a/src/primaite/game/agent/scripted_agents/data_manipulation_bot.py +++ b/src/primaite/game/agent/scripted_agents/data_manipulation_bot.py @@ -5,6 +5,7 @@ from gymnasium.core import ObsType from primaite.game.agent.scripted_agents.abstract_tap import AbstractTAPAgent +__all__ = ("DataManipulationAgent") class DataManipulationAgent(AbstractTAPAgent, identifier="RedDatabaseCorruptingAgent"): """Agent that uses a DataManipulationBot to perform an SQL injection attack.""" @@ -17,11 +18,6 @@ class DataManipulationAgent(AbstractTAPAgent, identifier="RedDatabaseCorruptingA starting_application_name: Optional[str] = None - # def __init__(self, **kwargs: Any) -> None: - # """Initialise DataManipulationAgent.""" - # self.setup_agent() - # super().__init_subclass__(**kwargs) - @property def starting_node_name(self) -> str: """Returns the agents starting node name.""" diff --git a/src/primaite/game/agent/scripted_agents/interface.py b/src/primaite/game/agent/scripted_agents/interface.py index 045d6d12..ab78eee0 100644 --- a/src/primaite/game/agent/scripted_agents/interface.py +++ b/src/primaite/game/agent/scripted_agents/interface.py @@ -17,6 +17,11 @@ from primaite.interface.request import RequestFormat, RequestResponse if TYPE_CHECKING: pass +__all__ = ("AgentHistoryItem", + "AgentStartSettings", + "AbstractAgent", + "AbstractScriptedAgent", + "ProxyAgent") class AgentHistoryItem(BaseModel): """One entry of an agent's action log - what the agent did and how the simulator responded in 1 step.""" @@ -116,7 +121,7 @@ class AbstractAgent(BaseModel): """ model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True) - agent_name: Optional[str] = "Abstract_Agent" + agent_name: ClassVar[str] = "Abstract_Agent" # TODO: Make this a ClassVar[str] like verb in actions? history: List[AgentHistoryItem] = [] _logger: AgentLog = AgentLog(agent_name=agent_name) action_manager: ActionManager diff --git a/src/primaite/game/agent/scripted_agents/probabilistic_agent.py b/src/primaite/game/agent/scripted_agents/probabilistic_agent.py index 750a120f..ba6ba850 100644 --- a/src/primaite/game/agent/scripted_agents/probabilistic_agent.py +++ b/src/primaite/game/agent/scripted_agents/probabilistic_agent.py @@ -8,6 +8,7 @@ from gymnasium.core import ObsType from primaite.game.agent.scripted_agents.interface import AbstractScriptedAgent, AgentSettings +__all__ = ("ProbabilisticAgent") class ProbabilisticAgent(AbstractScriptedAgent, identifier="ProbabilisticAgent"): """Scripted agent which randomly samples its action space with prescribed probabilities for each action.""" diff --git a/src/primaite/game/agent/scripted_agents/random_agent.py b/src/primaite/game/agent/scripted_agents/random_agent.py index d28069e6..fecc235f 100644 --- a/src/primaite/game/agent/scripted_agents/random_agent.py +++ b/src/primaite/game/agent/scripted_agents/random_agent.py @@ -6,6 +6,7 @@ from gymnasium.core import ObsType from primaite.game.agent.scripted_agents.interface import AbstractScriptedAgent +__all__ = ("RandomAgent", "PeriodicAgent") class RandomAgent(AbstractScriptedAgent, identifier="Random_Agent"): """Agent that ignores its observation and acts completely at random.""" @@ -38,18 +39,25 @@ class PeriodicAgent(AbstractScriptedAgent, identifier="Periodic_Agent"): agent_name: str = "Periodic_Agent" """Name of the agent.""" + + # TODO: This is available in config.agent_settings.start_settings.start_step start_step: int = 20 "The timestep at which an agent begins performing it's actions." start_variance: int = 5 "Deviation around the start step." + + # TODO: This is available in config.agent_settings.start_settings.frequency frequency: int = 5 "The number of timesteps to wait between performing actions." + + # TODO: This is available in config.agent_settings.start_settings.variance variance: int = 0 "The amount the frequency can randomly change to." max_executions: int = 999999 "Maximum number of times the agent can execute its action." num_executions: int = 0 """Number of times the agent has executed an action.""" + #TODO: Also in abstract_tap - move up and inherit? Add to AgentStartSettings? next_execution_timestep: int = 0 """Timestep of the next action execution by the agent.""" diff --git a/tests/assets/configs/basic_switched_network.yaml b/tests/assets/configs/basic_switched_network.yaml index fed0f52d..00ba381b 100644 --- a/tests/assets/configs/basic_switched_network.yaml +++ b/tests/assets/configs/basic_switched_network.yaml @@ -63,6 +63,10 @@ agents: start_step: 5 frequency: 4 variance: 3 + action_probabilities: + 0: 0.6 + 1: 0.4 + diff --git a/tests/assets/configs/fix_duration_one_item.yaml b/tests/assets/configs/fix_duration_one_item.yaml index bd0fb61f..62579e35 100644 --- a/tests/assets/configs/fix_duration_one_item.yaml +++ b/tests/assets/configs/fix_duration_one_item.yaml @@ -60,7 +60,9 @@ agents: start_step: 5 frequency: 4 variance: 3 - + action_probabilities: + 0: 0.4 + 1: 0.6 - ref: defender diff --git a/tests/assets/configs/software_fix_duration.yaml b/tests/assets/configs/software_fix_duration.yaml index 1a28258b..3e3d6e22 100644 --- a/tests/assets/configs/software_fix_duration.yaml +++ b/tests/assets/configs/software_fix_duration.yaml @@ -60,7 +60,9 @@ agents: start_step: 5 frequency: 4 variance: 3 - + action_probabilities: + 0: 0.4 + 1: 0.6 - ref: defender