diff --git a/malsim/sims/mal_simulator.py b/malsim/sims/mal_simulator.py index 54438b84..28c5877e 100644 --- a/malsim/sims/mal_simulator.py +++ b/malsim/sims/mal_simulator.py @@ -32,13 +32,14 @@ class MalSimAgent: name: str type: AgentType - # Contains reward from last step of MalSimulator + # Contains current agent reward in the simulation + # Attackers get positive rewards, defenders negative reward: int = 0 # Contains possible actions for the agent in the next step action_surface: list[AttackGraphNode] = field(default_factory=list) - # Fields that tell if agent is 'dead' + # Fields that tell if agent is 'dead' / disabled truncated: bool = False terminated: bool = False @@ -59,7 +60,7 @@ def __init__(self, name: str): super().__init__(name, AgentType.DEFENDER) class MalSimulator(): - """A simple MAL Simulator that works on the AttackGraph + """A MAL Simulator that works on the AttackGraph Allows user to register agents (defender and attacker) and lets the agents perform actions step by step and updates @@ -172,8 +173,10 @@ def _init_agent_action_surfaces(self): """Set agent action surfaces according to current state""" for agent in self.agents_dict.values(): if agent.type == AgentType.ATTACKER: - attacker_id = agent.attacker_id - attacker = self.attack_graph.get_attacker_by_id(attacker_id) + # Get the Attacker object + attacker = \ + self.attack_graph.get_attacker_by_id(agent.attacker_id) + # Get current action surface agent.action_surface = \ query.get_attack_surface(attacker) @@ -395,7 +398,7 @@ def step( # Peform agent actions # Note: by design, defenders perform actions - # before attackers (see register_agent) + # before attackers (see _register_agent) for agent_name in self.agents: agent = self.get_agent(agent_name) agent_actions = actions.get(agent.name, [])