oxbrbtx/PrimAITE
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

#2869 - Commit before changing branches. Addition of properties to Agent classes and removal of if/else chain in game.py

Browse Source
This commit is contained in:
Charlie Crane
2024-11-21 14:45:35 +00:00
parent 75d4ef2dfd
commit 7435a4dee8
3 changed files with 45 additions and 121 deletions
Download Patch File Download Diff File Expand all files Collapse all files

63
src/primaite/game/agent/interface.py
View File

@@ -96,11 +96,6 @@ class AbstractAgent(BaseModel, ABC, identifier="Abstract_Agent"):
"""Base class for scripted and RL agents."""
_registry: ClassVar[Dict[str, Type[AbstractAgent]]] = {}
action_manager: Optional[ActionManager]
observation_manager: Optional[ObservationManager]
reward_function: Optional[RewardFunction]
config: "AbstractAgent.ConfigSchema"
class ConfigSchema(BaseModel):
@@ -121,39 +116,12 @@ class AbstractAgent(BaseModel, ABC, identifier="Abstract_Agent"):
type: str
agent_name: ClassVar[str]
agent_settings = Optional[AgentSettings] = None
history: List[AgentHistoryItem] = []
logger: AgentLog = AgentLog(agent_name)
# def __init__(
# self,
# agent_name: Optional[str],
# action_space: Optional[ActionManager],
# observation_space: Optional[ObservationManager],
# reward_function: Optional[RewardFunction],
# agent_settings: Optional[AgentSettings] = None,
# ) -> None:
# """
# Initialize an agent.
# :param agent_name: Unique string identifier for the agent, for reporting and multi-agent purposes.
# :type agent_name: Optional[str]
# :param action_space: Action space for the agent.
# :type action_space: Optional[ActionManager]
# :param observation_space: Observation space for the agent.
# :type observation_space: Optional[ObservationSpace]
# :param reward_function: Reward function for the agent.
# :type reward_function: Optional[RewardFunction]
# :param agent_settings: Configurable Options for Abstracted Agents
# :type agent_settings: Optional[AgentSettings]
# """
# self.agent_name: str = agent_name or "unnamed_agent"
# self.action_manager: Optional[ActionManager] = action_space
# self.observation_manager: Optional[ObservationManager] = observation_space
# self.reward_function: Optional[RewardFunction] = reward_function
# self.agent_settings = agent_settings or AgentSettings()
# self.history: List[AgentHistoryItem] = []
# self.logger = AgentLog(agent_name)
history: List[AgentHistoryItem] = []
action_manager: Optional[ActionManager] = None
observation_manager: Optional[ObservationManager] = None
reward_function: Optional[RewardFunction] = None
agent_settings = Optional[AgentSettings] = None
def __init_subclass__(cls, identifier: str, **kwargs: Any) -> None:
super().__init_subclass__(**kwargs)
@@ -161,16 +129,25 @@ class AbstractAgent(BaseModel, ABC, identifier="Abstract_Agent"):
raise ValueError(f"Cannot create a new agent under reserved name {identifier}")
cls._registry[identifier] = cls
@property
def observation_manager(self) -> ObservationManager:
"""Returns the agents observation manager."""
return self.config.observation_manager
@property
def action_manager(self) -> ActionManager:
"""Returns the agents action manager."""
return self.config.action_manager
@property
def reward_function(self) -> RewardFunction:
"""Returns the agents reward function."""
return self.config.reward_function
@classmethod
def from_config(cls, config: Dict) -> "AbstractAgent":
"""Creates an agent component from a configuration dictionary."""
obj = cls(config=cls.ConfigSchema(**config))
# Pull managers out of config section for ease of use (?)
obj.observation_manager = obj.config.observation_manager
obj.action_manager = obj.config.action_manager
obj.reward_function = obj.config.reward_function
return obj
def update_observation(self, state: Dict) -> ObsType:

23
src/primaite/game/agent/scripted_agents/data_manipulation_bot.py
View File

@@ -10,9 +10,6 @@ from primaite.game.agent.interface import AbstractScriptedAgent
class DataManipulationAgent(AbstractScriptedAgent, identifier="Data_Manipulation_Agent"):
"""Agent that uses a DataManipulationBot to perform an SQL injection attack."""
next_execution_timestep: int = 0
starting_node_idx: int = 0
config: "DataManipulationAgent.ConfigSchema"
class ConfigSchema(AbstractScriptedAgent.ConfigSchema):
@@ -27,13 +24,23 @@ class DataManipulationAgent(AbstractScriptedAgent, identifier="Data_Manipulation
super().__init__(*args, **kwargs)
self.setup_agent()
@property
def next_execution_timestep(self):
"""Returns the agents next execution timestep."""
return self.config.next_execution_timestep
@property
def starting_node_name(self):
"""Returns the agents starting node name."""
return self.config.starting_node_name
def _set_next_execution_timestep(self, timestep: int) -> None:
"""Set the next execution timestep with a configured random variance.
:param timestep: The timestep to add variance to.
"""
random_timestep_increment = random.randint(
-self.agent_settings.start_settings.variance, self.agent_settings.start_settings.variance
-self.config.agent_settings.start_settings.variance, self.config.agent_settings.start_settings.variance
)
self.next_execution_timestep = timestep + random_timestep_increment
@@ -48,11 +55,11 @@ class DataManipulationAgent(AbstractScriptedAgent, identifier="Data_Manipulation
:rtype: Tuple[str, Dict]
"""
if timestep < self.next_execution_timestep:
self.logger.debug(msg="Performing do nothing action")
self.config.logger.debug(msg="Performing do nothing action")
return "do_nothing", {}
self._set_next_execution_timestep(timestep + self.agent_settings.start_settings.frequency)
self.logger.info(msg="Performing a data manipulation attack!")
self._set_next_execution_timestep(timestep + self.config.agent_settings.start_settings.frequency)
self.config.logger.info(msg="Performing a data manipulation attack!")
return "node_application_execute", {
"node_name": self.config.starting_node_name,
"application_name": self.config.starting_application_name,
@@ -68,4 +75,4 @@ class DataManipulationAgent(AbstractScriptedAgent, identifier="Data_Manipulation
# we are assuming that every node in the node manager has a data manipulation application at idx 0
num_nodes = len(self.action_manager.node_names)
self.starting_node_idx = random.randint(0, num_nodes - 1)
self.logger.debug(msg=f"Select Start Node ID: {self.starting_node_idx}")
self.config.logger.debug(msg=f"Select Start Node ID: {self.starting_node_idx}")