Use Agent.Role

Author: ondrej-lukas

agents/action_list_base_agent.py

@@ -2,7 +2,7 @@
 import numpy as np
 import json
 from os import path
-from netsecgame import Action, Observation, GameState, BaseAgent, generate_valid_actions
+from netsecgame import Action, Observation, GameState, BaseAgent, generate_valid_actions, AgentRole
 
 
 class ActionListAgent(BaseAgent):
@@ -14,7 +14,7 @@ class ActionListAgent(BaseAgent):
     Compatible with the WhiteBoxNSGCoordinator.
     """
 
-    def __init__(self, host, port, role: str):
+    def __init__(self, host, port, role: AgentRole):
         super().__init__(host, port, role)
         self._action_list = []
         self._action_to_idx = {}
@@ -70,7 +70,7 @@ def get_valid_action_mask(self , state: GameState) -> np.ndarray:
 
     args = parser.parse_args()
     log_filename = path.dirname(path.abspath(__file__)) + '/action_List_base_agent.log'
-    agent = ActionListAgent(args.host, args.port, "Attacker")
+    agent = ActionListAgent(args.host, args.port, AgentRole.Attacker)
     observation = agent.register()
     print(f"Total actions: {len(agent._action_list)}")
     print(f"Valid action mask: {agent.get_valid_action_mask(observation.state)}")

agents/attackers/markov_chain_agent/markov_chain_agent.py

@@ -9,7 +9,7 @@
 from os import path, makedirs
 
 
-from AIDojoCoordinator.game_components import Action, Observation, AgentStatus, ActionType
+from AIDojoCoordinator.game_components import Action, Observation, AgentStatus, ActionType, AgentRole
 
 sys.path.append(path.dirname(path.dirname(path.dirname(path.abspath(__file__) ))))
 
@@ -312,7 +312,7 @@ def save_solutions_json(self, filepath: str):
         level=logging.INFO
     )
 
-    agent = MarkovChainAgent(args.host, args.port, "Attacker", args.episodes)
+    agent = MarkovChainAgent(args.host, args.port, AgentRole.Attacker, args.episodes)
     observation = agent.register()
 
     if not args.evaluate:

agents/benign/random/benign_random_agent.py

@@ -6,13 +6,11 @@
 import numpy as np
 import time
 from random import uniform, choice
-from AIDojoCoordinator.game_components import Action, Observation, ActionType
-from NetSecGameAgents.agents.base_agent import BaseAgent
-from NetSecGameAgents.agents.agent_utils import generate_valid_actions
+from netsecgame import Action, Observation, ActionType, BaseAgent, generate_valid_actions, AgentRole
 
 class RandomBenignAgent(BaseAgent):
 
-    def __init__(self, host:str, port:int, role:str, allowed_actions:list, apm_limit:int=None) -> None:
+    def __init__(self, host:str, port:int, role:AgentRole, allowed_actions:list, apm_limit:int=None) -> None:
         super().__init__(host, port, role)
         self._allowed_actions = allowed_actions
         self._apm_limit = apm_limit
@@ -28,6 +26,10 @@ def play_game(self, num_episodes=1):
         """
 
         observation = self.register()
+        if observation is None:
+            self._logger.error("Registration failed. Terminating.")
+            return
+
         returns = []
         for episode in range(num_episodes):
             episodic_returns = []
@@ -79,5 +81,5 @@ def select_action(self, observation:Observation)->Action:
     logging.basicConfig(filename=os.path.join(args.logdir, "benign_random_agent.log"), filemode='w', format='%(asctime)s %(name)s %(levelname)s %(message)s', datefmt='%H:%M:%S',level=logging.DEBUG)
 
     # Create agent
-    agent = RandomBenignAgent(args.host, args.port, "Benign", allowed_actions=[ActionType.FindData, ActionType.ExfiltrateData, ActionType.FindServices], apm_limit=args.apm)
+    agent = RandomBenignAgent(args.host, args.port, AgentRole.Benign, allowed_actions=[ActionType.FindData, ActionType.ExfiltrateData, ActionType.FindServices], apm_limit=args.apm)
     agent.play_game(args.episodes)

agents/defenders/random/random_agent.py

@@ -8,7 +8,7 @@
 import mlflow
 
 from random import uniform, choice
-from AIDojoCoordinator.game_components import Action, Observation, ActionType
+from AIDojoCoordinator.game_components import Action, Observation, ActionType, AgentRole
 from NetSecGameAgents.agents.base_agent import BaseAgent
 from NetSecGameAgents.agents.agent_utils import generate_valid_actions
 
@@ -88,7 +88,7 @@ def select_action(self, observation:Observation)->Action:
     logging.basicConfig(filename=os.path.join(args.logdir, "defender_random_agent.log"), filemode='w', format='%(asctime)s %(name)s %(levelname)s %(message)s', datefmt='%H:%M:%S',level=logging.DEBUG)
 
     # Create agent
-    agent = RandomDefenderAgent(args.host, args.port, "Defender", allowed_actions=[ActionType.FindData, ActionType.ExfiltrateData, ActionType.FindServices, ActionType.BlockIP], apm_limit=args.apm)
+    agent = RandomDefenderAgent(args.host, args.port, AgentRole.Defender, allowed_actions=[ActionType.FindData, ActionType.ExfiltrateData, ActionType.FindServices, ActionType.BlockIP], apm_limit=args.apm)
 
     print("Agent created. Starting interaction with the game server...")