Update gym to gymnasium

.gitignore

@@ -31,4 +31,6 @@ CybORG/Emulator/AWS/Config/awsconfig.ini
 *.swp
 *.pstats
 
-.ipynb_checkpoints/
+.ipynb_checkpoints/
+
+CybORG/Evaluation/

CybORG/Agents/Wrappers/BaseWrapper.py

@@ -18,8 +18,8 @@ def step(self, agent=None, action=None) -> Results:
         result.action_space = self.action_space_change(result.action_space)
         return result
 
-    def reset(self, agent=None):
-        result = self.env.reset(agent)
+    def reset(self, agent=None, **kwargs):
+        result = self.env.reset(agent, **kwargs)
         result.action_space = self.action_space_change(result.action_space)
         result.observation = self.observation_change(result.observation)
         return result

CybORG/Agents/Wrappers/BlueTableWrapper.py

@@ -6,51 +6,51 @@
 from CybORG.Agents.Wrappers.BaseWrapper import BaseWrapper
 from CybORG.Agents.Wrappers.TrueTableWrapper import TrueTableWrapper
 
+
 class BlueTableWrapper(BaseWrapper):
-    def __init__(self,env=None,agent=None,output_mode='table'):
-        super().__init__(env,agent)
+    def __init__(self, env=None, agent=None, output_mode='table'):
+        super().__init__(env, agent)
         self.env = TrueTableWrapper(env=env, agent=agent)
         self.agent = agent
 
         self.baseline = None
         self.output_mode = output_mode
         self.blue_info = {}
 
-    def reset(self, agent='Blue'):        
-        result = self.env.reset(agent)
+    def reset(self, agent='Blue', **kwargs):
+        result = self.env.reset(agent, **kwargs)
         obs = result.observation
         if agent == 'Blue':
             self._process_initial_obs(obs)
-            obs = self.observation_change(obs,baseline=True)
+            obs = self.observation_change(obs, baseline=True)
         result.observation = obs
         return result
 
     def step(self, agent=None, action=None) -> Results:
         result = self.env.step(agent, action)
-        obs = result.observation
+        obs = result.observation  # NOTE The original Blue agent observation. Should be equivalent to self.get_observation("Blue")
         if agent == 'Blue':
             obs = self.observation_change(obs)
         result.observation = obs
         result.action_space = self.action_space_change(result.action_space)
         return result
 
-    def get_table(self,output_mode='blue_table'):
+    def get_table(self, output_mode='blue_table'):
         if output_mode == 'blue_table':
             return self._create_blue_table(success=None)
         elif output_mode == 'true_table':
             return self.env.get_table()
 
-    def observation_change(self,observation,baseline=False):
+    def observation_change(self, observation, baseline=False):
         obs = observation if type(observation) == dict else observation.data
         obs = deepcopy(observation)
         success = obs['success']
-
         self._process_last_action()
         anomaly_obs = self._detect_anomalies(obs) if not baseline else obs
         del obs['success']
         # TODO check what info is for baseline
         info = self._process_anomalies(anomaly_obs)
-        if baseline:
+        if baseline:  # NOTE: baseline refers to the first observation by the blue agent
             for host in info:
                 info[host][-2] = 'None'
                 info[host][-1] = 'No'
@@ -82,14 +82,14 @@ def _process_initial_obs(self, obs):
             subnet = interface['Subnet']
             ip = str(interface['IP Address'])
             hostname = host['System info']['Hostname']
-            self.blue_info[hostname] = [str(subnet),str(ip),hostname, 'None','No']
+            self.blue_info[hostname] = [str(subnet), str(ip), hostname, 'None', 'No']
         return self.blue_info
 
     def _process_last_action(self):
         action = self.get_last_action(agent='Blue')
         if action is not None:
             name = action.__class__.__name__
-            hostname = action.get_params()['hostname'] if name in ('Restore','Remove') else None
+            hostname = action.get_params()['hostname'] if name in ('Restore', 'Remove') else None
 
             if name == 'Restore':
                 self.blue_info[hostname][-1] = 'No'
@@ -98,13 +98,14 @@ def _process_last_action(self):
                 if compromised != 'No':
                     self.blue_info[hostname][-1] = 'Unknown'
 
-    def _detect_anomalies(self,obs):
+    def _detect_anomalies(self, obs):
         if self.baseline is None:
-            raise TypeError('BlueTableWrapper was unable to establish baseline. This usually means the environment was not reset before calling the step method.')
+            raise TypeError(
+                'BlueTableWrapper was unable to establish baseline. This usually means the environment was not reset before calling the step method.')
 
         anomaly_dict = {}
 
-        for hostid,host in obs.items():
+        for hostid, host in obs.items():
             if hostid == 'success':
                 continue
 
@@ -114,7 +115,7 @@ def _detect_anomalies(self,obs):
 
             host_anomalies = {}
             if 'Files' in host:
-                baseline_files = host_baseline.get('Files',[])
+                baseline_files = host_baseline.get('Files', [])
                 anomalous_files = []
                 for f in host['Files']:
                     if f not in baseline_files:
@@ -123,7 +124,7 @@ def _detect_anomalies(self,obs):
                     host_anomalies['Files'] = anomalous_files
 
             if 'Processes' in host:
-                baseline_processes = host_baseline.get('Processes',[])
+                baseline_processes = host_baseline.get('Processes', [])
                 anomalous_processes = []
                 for p in host['Processes']:
                     if p not in baseline_processes:
@@ -136,16 +137,18 @@ def _detect_anomalies(self,obs):
 
         return anomaly_dict
 
-    def _process_anomalies(self,anomaly_dict):
+    def _process_anomalies(self, anomaly_dict):
         info = deepcopy(self.blue_info)
         for hostid, host_anomalies in anomaly_dict.items():
             assert len(host_anomalies) > 0
             if 'Processes' in host_anomalies:
-                connection_type = self._interpret_connections(host_anomalies['Processes'])
-                info[hostid][-2] = connection_type
-                if connection_type == 'Exploit':
-                    info[hostid][-1] = 'User'
-                    self.blue_info[hostid][-1] = 'User'
+                # NOTE John Cardiff fix
+                if "Connections" in host_anomalies['Processes'][-1]:
+                    connection_type = self._interpret_connections(host_anomalies['Processes'])
+                    info[hostid][-2] = connection_type
+                    if connection_type == 'Exploit':
+                        info[hostid][-1] = 'User'
+                        self.blue_info[hostid][-1] = 'User'
             if 'Files' in host_anomalies:
                 malware = [f['Density'] >= 0.9 for f in host_anomalies['Files']]
                 if any(malware):
@@ -154,47 +157,43 @@ def _process_anomalies(self,anomaly_dict):
 
         return info
 
-    def _interpret_connections(self,activity:list):                
+    def _interpret_connections(self, activity: list):
         num_connections = len(activity)
-
         ports = set([item['Connections'][0]['local_port'] \
-            for item in activity if 'Connections' in item])
+                     for item in activity if 'Connections' in item])
         port_focus = len(ports)
 
         remote_ports = set([item['Connections'][0].get('remote_port') \
-            for item in activity if 'Connections' in item])
+                            for item in activity if 'Connections' in item])
         if None in remote_ports:
             remote_ports.remove(None)
 
-        if num_connections >= 3 and port_focus >=3:
+        if num_connections >= 3 and port_focus >= 3:
             anomaly = 'Scan'
         elif 4444 in remote_ports:
             anomaly = 'Exploit'
         elif num_connections >= 3 and port_focus == 1:
             anomaly = 'Exploit'
-        elif 'Service Name' in activity[0]:
-            anomaly = 'None'
         else:
             anomaly = 'Scan'
 
         return anomaly
 
     # def _malware_analysis(self,obs,hostname):
-        # anomaly_dict = {hostname: {'Files': []}}
-        # if hostname in obs:
-            # if 'Files' in obs[hostname]:
-                # files = obs[hostname]['Files']
-            # else:
-                # return anomaly_dict
-        # else:
-            # return anomaly_dict
+    # anomaly_dict = {hostname: {'Files': []}}
+    # if hostname in obs:
+    # if 'Files' in obs[hostname]:
+    # files = obs[hostname]['Files']
+    # else:
+    # return anomaly_dict
+    # else:
+    # return anomaly_dict
 
-        # for f in files:
-            # if f['Density'] >= 0.9:
-                # anomaly_dict[hostname]['Files'].append(f)
-
-        # return anomaly_dict
+    # for f in files:
+    # if f['Density'] >= 0.9:
+    # anomaly_dict[hostname]['Files'].append(f)
 
+    # return anomaly_dict
 
     def _create_blue_table(self, success):
         table = PrettyTable([
@@ -203,10 +202,10 @@ def _create_blue_table(self, success):
             'Hostname',
             'Activity',
             'Compromised'
-            ])
-        for hostid in self.info:
+        ])
+        for hostid in self.info:  # NOTE extracted from the anomalies detection method
             table.add_row(self.info[hostid])
-        
+
         table.sortby = 'Hostname'
         table.success = success
         return table
@@ -219,11 +218,11 @@ def _create_vector(self, success):
             # Activity
             activity = row[3]
             if activity == 'None':
-                value = [0,0]
+                value = [0, 0]
             elif activity == 'Scan':
-                value = [1,0]
+                value = [1, 0]
             elif activity == 'Exploit':
-                value = [1,1]
+                value = [1, 1]
             else:
                 raise ValueError('Table had invalid Access Level')
             proto_vector.extend(value)
@@ -235,16 +234,16 @@ def _create_vector(self, success):
             elif compromised == 'Unknown':
                 value = [1, 0]
             elif compromised == 'User':
-                value = [0,1]
+                value = [0, 1]
             elif compromised == 'Privileged':
-                value = [1,1]
+                value = [1, 1]
             else:
                 raise ValueError('Table had invalid Access Level')
             proto_vector.extend(value)
 
         return np.array(proto_vector)
 
-    def get_attr(self,attribute:str):
+    def get_attr(self, attribute: str):
         return self.env.get_attr(attribute)
 
     def get_observation(self, agent: str):
@@ -255,13 +254,13 @@ def get_observation(self, agent: str):
 
         return output
 
-    def get_agent_state(self,agent:str):
+    def get_agent_state(self, agent: str):
         return self.get_attr('get_agent_state')(agent)
 
-    def get_action_space(self,agent):
+    def get_action_space(self, agent):
         return self.env.get_action_space(agent)
 
-    def get_last_action(self,agent):
+    def get_last_action(self, agent):
         return self.get_attr('get_last_action')(agent)
 
     def get_ip_map(self):

CybORG/Agents/Wrappers/ChallengeWrapper.py

@@ -1,4 +1,4 @@
-from gym import Env
+from gymnasium import Env
 from CybORG.Agents.Wrappers import BaseWrapper, OpenAIGymWrapper, BlueTableWrapper,RedTableWrapper,EnumActionWrapper
 
 
@@ -26,17 +26,17 @@ def __init__(self, agent_name: str, env, agent=None,
         self.step_counter = None
 
     def step(self,action=None):
-        obs, reward, done, info = self.env.step(action=action)
+        obs, reward, terminated, truncated, info = self.env.step(action=action)
 
         self.step_counter += 1
         if self.max_steps is not None and self.step_counter >= self.max_steps:
-            done = True
+            truncated = True
 
-        return obs, reward, done, info
+        return obs, reward, terminated, truncated, info
 
-    def reset(self):
+    def reset(self, **kwargs):
         self.step_counter = 0
-        return self.env.reset()
+        return self.env.reset(**kwargs)
 
     def get_attr(self,attribute:str):
         return self.env.get_attr(attribute)

CybORG/Agents/Wrappers/IntListToAction.py

@@ -30,8 +30,8 @@ def step(self, agent=None, action: list = None) -> Results:
         result.action_name = str(action_obj)
         return result
 
-    def reset(self, agent=None):
-        result = self.env.reset(agent)
+    def reset(self, agent=None, **kwargs):
+        result = self.env.reset(agent, **kwargs)
         result.action_space, result.selection_masks = self.action_space_change(result.action_space)
         self.selection_mask = result.selection_masks
         result.observation = self.observation_change(result.observation)

CybORG/Agents/Wrappers/OpenAIGymWrapper.py

@@ -1,5 +1,5 @@
 import numpy as np
-from gym import spaces, Env
+from gymnasium import spaces, Env
 from typing import Union, List
 from prettytable import PrettyTable
 
@@ -22,19 +22,24 @@ def __init__(self, agent_name: str, env: BaseWrapper = None, agent: BaseAgent =
         self.metadata = {}
         self.action = None
 
-    def step(self, action: Union[int, List[int]] = None) -> (object, float, bool, dict):
+    # 'gym' done = 'gymnasium' terminated or truncated
+    # https://gymnasium.farama.org/content/migration-guide/
+    def step(self, action: Union[int, List[int]] = None) -> tuple[object, float, bool, dict]:
         self.action = action
         result = self.env.step(self.agent_name, action)
         result.observation = self.observation_change(result.observation)
         result.action_space = self.action_space_change(result.action_space)
+        terminated = result.done
+        truncated = False
         info = vars(result)
-        return np.array(result.observation), result.reward, result.done, info
+        return np.array(result.observation), result.reward, terminated, truncated, info
 
-    def reset(self, agent=None):
-        result = self.env.reset(self.agent_name)
+    def reset(self, agent=None, **kwargs):
+        result = self.env.reset(self.agent_name, **kwargs)
         result.action_space = self.action_space_change(result.action_space)
         result.observation = self.observation_change(result.observation)
-        return np.array(result.observation)
+        # gymnasium API: observation, info = env.reset(seed=None, options={})
+        return np.array(result.observation), vars(result)
 
     def render(self):
         # TODO: If FixedFlatWrapper it will error out!

CybORG/Agents/Wrappers/RedTableWrapper.py

@@ -20,13 +20,13 @@ def __init__(self, env=None, agent=None, output_mode='table'):
         self.output_mode = output_mode
         self.success = None
 
-    def reset(self, agent=None):
+    def reset(self, agent=None, **kwargs):
         self.red_info = {}
         self.known_subnets = set()
         self.step_counter = -1
         self.id_tracker = -1
         self.success = None
-        result = self.env.reset(agent)
+        result = self.env.reset(agent, **kwargs)
         if agent =='Red':
             obs = self.observation_change(result.observation)
             result.observation = obs

CybORG/Agents/Wrappers/ReduceActionSpaceWrapper.py

@@ -50,6 +50,6 @@ def action_space_change(self, action_space: dict) -> dict:
     def get_attr(self,attribute:str):
         return self.env.get_attr(attribute)
 
-    def reset(self, agent=None):
+    def reset(self, agent=None, **kwargs):
         self.fixed_size = {}
-        return self.env.reset(agent)
+        return self.env.reset(agent, **kwargs)