#2417 Add categorisation and updated new configs from merge

src/primaite/game/agent/observations/file_system_observations.py

@@ -43,6 +43,26 @@ class FileObservation(AbstractObservation, identifier="FILE"):
         if self.include_num_access:
             self.default_observation["num_access"] = 0
 
+        # TODO: allow these to be configured in yaml
+        self.high_threshold = 10
+        self.med_threshold = 5
+        self.low_threshold = 0
+
+    def _categorise_num_access(self, num_access: int) -> int:
+        """
+        Represent number of file accesses as a categorical variable.
+
+        :param num_access: Number of file accesses.
+        :return: Bin number corresponding to the number of accesses.
+        """
+        if num_access > self.high_threshold:
+            return 3
+        elif num_access > self.med_threshold:
+            return 2
+        elif num_access > self.low_threshold:
+            return 1
+        return 0
+
     def observe(self, state: Dict) -> ObsType:
         """
         Generate observation based on the current state of the simulation.
@@ -57,8 +77,7 @@ class FileObservation(AbstractObservation, identifier="FILE"):
             return self.default_observation
         obs = {"health_status": file_state["visible_status"]}
         if self.include_num_access:
-            obs["num_access"] = file_state["num_access"]
+            obs["num_access"] = self._categorise_num_access(file_state["num_access"])
-            # raise NotImplementedError("TODO: need to fix num_access to use thresholds instead of raw value.")
         return obs
 
     @property

src/primaite/game/agent/observations/firewall_observation.py

@@ -214,9 +214,8 @@ class FirewallObservation(AbstractObservation, identifier="FIREWALL"):
         :return: Constructed firewall observation instance.
         :rtype: FirewallObservation
         """
-        where = parent_where + ["nodes", config.hostname]
         return cls(
-            where=where,
+            where=parent_where + ["nodes", config.hostname],
             ip_list=config.ip_list,
             wildcard_list=config.wildcard_list,
             port_list=config.port_list,

src/primaite/game/agent/observations/observation_manager.py

@@ -185,9 +185,11 @@ class ObservationManager:
         Create observation space from a config.
 
         :param config: Dictionary containing the configuration for this observation space.
-            It should contain the key 'type' which selects which observation class to use (from a choice of:
+            If None, a blank observation space is created.
-            UC2BlueObservation, UC2RedObservation, UC2GreenObservation)
+            Otherwise, this must be a Dict with a type field and options field.
-            The other key is 'options' which are passed to the constructor of the selected observation class.
+            type: string that corresponds to one of the observation identifiers that are provided when subclassing
+            AbstractObservation
+            options: this must adhere to the chosen observation type's ConfigSchema nested class.
         :type config: Dict
         :param game: Reference to the PrimaiteGame object that spawned this observation.
         :type game: PrimaiteGame

src/primaite/game/agent/observations/software_observation.py

@@ -98,6 +98,26 @@ class ApplicationObservation(AbstractObservation, identifier="APPLICATION"):
         self.where = where
         self.default_observation = {"operating_status": 0, "health_status": 0, "num_executions": 0}
 
+        # TODO: allow these to be configured in yaml
+        self.high_threshold = 10
+        self.med_threshold = 5
+        self.low_threshold = 0
+
+    def _categorise_num_executions(self, num_executions: int) -> int:
+        """
+        Represent number of file accesses as a categorical variable.
+
+        :param num_access: Number of file accesses.
+        :return: Bin number corresponding to the number of accesses.
+        """
+        if num_executions > self.high_threshold:
+            return 3
+        elif num_executions > self.med_threshold:
+            return 2
+        elif num_executions > self.low_threshold:
+            return 1
+        return 0
+
     def observe(self, state: Dict) -> ObsType:
         """
         Generate observation based on the current state of the simulation.
@@ -113,7 +133,7 @@ class ApplicationObservation(AbstractObservation, identifier="APPLICATION"):
         return {
             "operating_status": application_state["operating_state"],
             "health_status": application_state["health_state_visible"],
-            "num_executions": application_state["num_executions"],
+            "num_executions": self._categorise_num_executions(application_state["num_executions"]),
         }
 
     @property

tests/assets/configs/firewall_actions_network.yaml

@@ -64,25 +64,67 @@ agents:
   - ref: defender
     team: BLUE
     type: ProxyAgent
+
     observation_space:
-      type: UC2BlueObservation
+      type: CUSTOM
       options:
-        num_services_per_node: 1
+        components:
-        num_folders_per_node: 1
+          - type: NODES
-        num_files_per_folder: 1
+            label: NODES
-        num_nics_per_node: 2
+            options:
-        nodes:
+              hosts:
-        - node_hostname: client_1
+                - hostname: client_1
-        links:
+              num_services: 1
-        - link_ref: client_1___switch_1
+              num_applications: 0
-        acl:
+              num_folders: 1
-          options:
+              num_files: 1
-            max_acl_rules: 10
+              num_nics: 2
-          router_hostname: router_1
+              include_num_access: false
-          ip_address_order:
+              include_nmne: true
-            - node_hostname: client_1
+              routers:
-              nic_num: 1
+                - hostname: router_1
-        ics: null
+              num_ports: 0
+              ip_list:
+                - 192.168.0.10
+              wildcard_list:
+                - 0.0.0.1
+              port_list:
+                - 80
+                - 5432
+              protocol_list:
+                - ICMP
+                - TCP
+                - UDP
+              num_rules: 10
+
+          - type: LINKS
+            label: LINKS
+            options:
+              link_references:
+                - client_1___switch_1
+          - type: "NONE"
+            label: ICS
+            options: {}
+
+    # observation_space:
+    #   type: UC2BlueObservation
+    #   options:
+    #     num_services_per_node: 1
+    #     num_folders_per_node: 1
+    #     num_files_per_folder: 1
+    #     num_nics_per_node: 2
+    #     nodes:
+    #     - node_hostname: client_1
+    #     links:
+    #     - link_ref: client_1___switch_1
+    #     acl:
+    #       options:
+    #         max_acl_rules: 10
+    #       router_hostname: router_1
+    #       ip_address_order:
+    #         - node_hostname: client_1
+    #           nic_num: 1
+    #     ics: null
     action_space:
       action_list:
         - type: DONOTHING

tests/assets/configs/test_application_install.yaml

@@ -41,8 +41,7 @@ agents:
         0: 0.3
         1: 0.6
         2: 0.1
-    observation_space:
+    observation_space: null
-      type: UC2GreenObservation
     action_space:
       action_list:
         - type: DONOTHING
@@ -91,8 +90,7 @@ agents:
         0: 0.3
         1: 0.6
         2: 0.1
-    observation_space:
+    observation_space: null
-      type: UC2GreenObservation
     action_space:
       action_list:
         - type: DONOTHING
@@ -141,10 +139,7 @@ agents:
     team: RED
     type: RedDatabaseCorruptingAgent
 
-    observation_space:
+    observation_space: null
-      type: UC2RedObservation
-      options:
-        nodes: {}
 
     action_space:
       action_list:
@@ -177,61 +172,73 @@ agents:
     type: ProxyAgent
 
     observation_space:
-      type: UC2BlueObservation
+      type: CUSTOM
       options:
-        num_services_per_node: 1
+        components:
-        num_folders_per_node: 1
+          - type: NODES
-        num_files_per_folder: 1
+            label: NODES
-        num_nics_per_node: 2
+            options:
-        nodes:
+              hosts:
-        - node_hostname: domain_controller
+                - hostname: domain_controller
-          services:
+                - hostname: web_server
-          - service_name: DNSServer
+                  services:
-        - node_hostname: web_server
+                    - service_name: WebServer
-          services:
+                - hostname: database_server
-          - service_name: WebServer
+                  folders:
-        - node_hostname: database_server
+                    - folder_name: database
-          folders:
+                      files:
-          - folder_name: database
+                      - file_name: database.db
-            files:
+                - hostname: backup_server
-            - file_name: database.db
+                - hostname: security_suite
-        - node_hostname: backup_server
+                - hostname: client_1
-        - node_hostname: security_suite
+                - hostname: client_2
-        - node_hostname: client_1
+              num_services: 1
-        - node_hostname: client_2
+              num_applications: 0
-        links:
+              num_folders: 1
-        - link_ref: router_1___switch_1
+              num_files: 1
-        - link_ref: router_1___switch_2
+              num_nics: 2
-        - link_ref: switch_1___domain_controller
+              include_num_access: false
-        - link_ref: switch_1___web_server
+              include_nmne: true
-        - link_ref: switch_1___database_server
+              routers:
-        - link_ref: switch_1___backup_server
+                - hostname: router_1
-        - link_ref: switch_1___security_suite
+              num_ports: 0
-        - link_ref: switch_2___client_1
+              ip_list:
-        - link_ref: switch_2___client_2
+                - 192.168.1.10
-        - link_ref: switch_2___security_suite
+                - 192.168.1.12
-        acl:
+                - 192.168.1.14
-          options:
+                - 192.168.1.16
-            max_acl_rules: 10
+                - 192.168.1.110
-          router_hostname: router_1
+                - 192.168.10.21
-          ip_address_order:
+                - 192.168.10.22
-            - node_hostname: domain_controller
+                - 192.168.10.110
-              nic_num: 1
+              wildcard_list:
-            - node_hostname: web_server
+                - 0.0.0.1
-              nic_num: 1
+              port_list:
-            - node_hostname: database_server
+                - 80
-              nic_num: 1
+                - 5432
-            - node_hostname: backup_server
+              protocol_list:
-              nic_num: 1
+                - ICMP
-            - node_hostname: security_suite
+                - TCP
-              nic_num: 1
+                - UDP
-            - node_hostname: client_1
+              num_rules: 10
-              nic_num: 1
+
-            - node_hostname: client_2
+          - type: LINKS
-              nic_num: 1
+            label: LINKS
-            - node_hostname: security_suite
+            options:
-              nic_num: 2
+              link_references:
-        ics: null
+                - router_1___switch_1
+                - router_1___switch_2
+                - switch_1___domain_controller
+                - switch_1___web_server
+                - switch_1___database_server
+                - switch_1___backup_server
+                - switch_1___security_suite
+                - switch_2___client_1
+                - switch_2___client_2
+                - switch_2___security_suite
+          - type: "NONE"
+            label: ICS
+            options: {}
 
     action_space:
       action_list:

tests/integration_tests/game_layer/observations/test_file_system_observations.py

@@ -72,3 +72,6 @@ def test_folder_observation(simulation):
 
     observation_state = root_folder_obs.observe(simulation.describe_state())
     assert observation_state.get("health_status") == 3  # file is corrupt therefore folder is corrupted too
+
+
+# TODO: Add tests to check num access is correct.

tests/integration_tests/game_layer/observations/test_link_observations.py

@@ -51,31 +51,8 @@ def simulation() -> Simulation:
     return sim
 
 
-def test_link_observation(simulation):
+def test_link_observation():
-    """Test the link observation."""
+    """Check the shape and contents of the link observation."""
-    # get a link
-    link: Link = next(iter(simulation.network.links.values()))
-
-    computer: Computer = simulation.network.get_node_by_hostname("computer")
-    server: Server = simulation.network.get_node_by_hostname("server")
-
-    simulation.apply_timestep(0)  # some pings when network was made - reset with apply timestep
-
-    link_obs = LinkObservation(where=["network", "links", link.uuid])
-
-    assert link_obs.space["PROTOCOLS"]["ALL"] == spaces.Discrete(11)  # test that the spaces are 0-10 including 0 and 10
-
-    observation_state = link_obs.observe(simulation.describe_state())
-    assert observation_state.get("PROTOCOLS") is not None
-    assert observation_state["PROTOCOLS"]["ALL"] == 0
-
-    computer.ping(server.network_interface.get(1).ip_address)
-
-    observation_state = link_obs.observe(simulation.describe_state())
-    assert observation_state["PROTOCOLS"]["ALL"] == 1
-
-
-def test_link_observation_again():
     net = Network()
     sim = Simulation(network=net)
     switch = Switch(hostname="switch", num_ports=5, operating_state=NodeOperatingState.ON)
@@ -102,6 +79,8 @@ def test_link_observation_again():
     assert "PROTOCOLS" in link_2_obs
     assert "ALL" in link_1_obs["PROTOCOLS"]
     assert "ALL" in link_2_obs["PROTOCOLS"]
+    assert link_1_observation.space["PROTOCOLS"]["ALL"] == spaces.Discrete(11)
+    assert link_2_observation.space["PROTOCOLS"]["ALL"] == spaces.Discrete(11)
     assert link_1_obs["PROTOCOLS"]["ALL"] == 0
     assert link_2_obs["PROTOCOLS"]["ALL"] == 0
 

tests/integration_tests/game_layer/test_observations.py

@@ -19,3 +19,8 @@ def test_file_observation():
     )
     assert dog_file_obs.observe(state) == {"health_status": 1}
     assert dog_file_obs.space == spaces.Dict({"health_status": spaces.Discrete(6)})
+
+
+# TODO:
+# def test_file_num_access():
+#     ...