feature\1522:

Create random red agent behaviour.

src/primaite/config/training_config.py

@@ -1,7 +1,7 @@
 # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Any, Dict, Final, Union, Optional
+from typing import Any, Dict, Final, Optional, Union
 
 import yaml
 
@@ -21,6 +21,9 @@ class TrainingConfig:
     agent_identifier: str = "STABLE_BASELINES3_A2C"
     "The Red Agent algo/class to be used."
 
+    red_agent_identifier: str = "RANDOM"
+    "Creates Random Red Agent Attacks"
+
     action_type: ActionType = ActionType.ANY
     "The ActionType to use."
 
@@ -167,8 +170,7 @@ def main_training_config_path() -> Path:
     return path
 
 
-def load(file_path: Union[str, Path],
-         legacy_file: bool = False) -> TrainingConfig:
+def load(file_path: Union[str, Path], legacy_file: bool = False) -> TrainingConfig:
     """
     Read in a training config yaml file.
 
@@ -213,9 +215,7 @@ def load(file_path: Union[str, Path],
 
 
 def convert_legacy_training_config_dict(
-        legacy_config_dict: Dict[str, Any],
-        num_steps: int = 256,
-        action_type: str = "ANY"
+    legacy_config_dict: Dict[str, Any], num_steps: int = 256, action_type: str = "ANY"
 ) -> Dict[str, Any]:
     """
     Convert a legacy training config dict to the new format.
@@ -227,10 +227,7 @@ def convert_legacy_training_config_dict(
         don't have action_type values.
     :return: The converted training config dict.
     """
-    config_dict = {
-        "num_steps": num_steps,
-        "action_type": action_type
-    }
+    config_dict = {"num_steps": num_steps, "action_type": action_type}
     for legacy_key, value in legacy_config_dict.items():
         new_key = _get_new_key_from_legacy(legacy_key)
         if new_key:

src/primaite/environment/primaite_env.py

@@ -14,8 +14,7 @@ from gym import Env, spaces
 from matplotlib import pyplot as plt
 
 from primaite.acl.access_control_list import AccessControlList
-from primaite.agents.utils import is_valid_acl_action_extra, \
-    is_valid_node_action
+from primaite.agents.utils import is_valid_acl_action_extra, is_valid_node_action
 from primaite.common.custom_typing import NodeUnion
 from primaite.common.enums import (
     ActionType,
@@ -24,8 +23,9 @@ from primaite.common.enums import (
     NodePOLInitiator,
     NodePOLType,
     NodeType,
+    ObservationType,
     Priority,
-    SoftwareState, ObservationType,
+    SoftwareState,
 )
 from primaite.common.service import Service
 from primaite.config import training_config
@@ -35,15 +35,13 @@ from primaite.environment.reward import calculate_reward_function
 from primaite.links.link import Link
 from primaite.nodes.active_node import ActiveNode
 from primaite.nodes.node import Node
-from primaite.nodes.node_state_instruction_green import \
-    NodeStateInstructionGreen
+from primaite.nodes.node_state_instruction_green import NodeStateInstructionGreen
 from primaite.nodes.node_state_instruction_red import NodeStateInstructionRed
 from primaite.nodes.passive_node import PassiveNode
 from primaite.nodes.service_node import ServiceNode
 from primaite.pol.green_pol import apply_iers, apply_node_pol
 from primaite.pol.ier import IER
-from primaite.pol.red_agent_pol import apply_red_agent_iers, \
-    apply_red_agent_node_pol
+from primaite.pol.red_agent_pol import apply_red_agent_iers, apply_red_agent_node_pol
 from primaite.transactions.transaction import Transaction
 
 _LOGGER = logging.getLogger(__name__)
@@ -177,7 +175,6 @@ class Primaite(Env):
         # It will be initialised later.
         self.obs_handler: ObservationsHandler
 
-
         # Open the config file and build the environment laydown
         with open(self._lay_down_config_path, "r") as file:
             # Open the config file and build the environment laydown
@@ -238,7 +235,9 @@ class Primaite(Env):
             self.action_dict = self.create_node_and_acl_action_dict()
             self.action_space = spaces.Discrete(len(self.action_dict))
         else:
-            _LOGGER.info(f"Invalid action type selected: {self.training_config.action_type}")
+            _LOGGER.info(
+                f"Invalid action type selected: {self.training_config.action_type}"
+            )
         # Set up a csv to store the results of the training
         try:
             header = ["Episode", "Average Reward"]
@@ -275,6 +274,10 @@ class Primaite(Env):
         # Does this for both live and reference nodes
         self.reset_environment()
 
+        # Create a random red agent to use for this episode
+        if self.training_config.red_agent_identifier == "RANDOM":
+            self.create_random_red_agent()
+
         # Reset counters and totals
         self.total_reward = 0
         self.step_count = 0
@@ -379,7 +382,7 @@ class Primaite(Env):
             self.step_count,
             self.training_config,
         )
-        #print(f"    Step {self.step_count} Reward: {str(reward)}")
+        print(f"    Step {self.step_count} Reward: {str(reward)}")
         self.total_reward += reward
         if self.step_count == self.episode_steps:
             self.average_reward = self.total_reward / self.step_count
@@ -1033,7 +1036,6 @@ class Primaite(Env):
         """
         self.observation_type = ObservationType[observation_info["type"]]
 
-
     def get_action_info(self, action_info):
         """
         Extracts action_info.
@@ -1216,3 +1218,152 @@ class Primaite(Env):
         # Combine the Node dict and ACL dict
         combined_action_dict = {**acl_action_dict, **new_node_action_dict}
         return combined_action_dict
+
+    def create_random_red_agent(self):
+        """Decide on random red agent for the episode to be called in env.reset()."""
+
+        # Reset the current red iers and red node pol
+        self.red_iers = {}
+        self.red_node_pol = {}
+
+        # Decide how many nodes become compromised
+        node_list = list(self.nodes.values())
+        computers = [node for node in node_list if node.node_type == NodeType.COMPUTER]
+        max_num_nodes_compromised = len(
+            computers
+        )  # only computers can become compromised
+        # random select between 1 and max_num_nodes_compromised
+        num_nodes_to_compromise = np.random.randint(1, max_num_nodes_compromised + 1)
+
+        # Decide which of the nodes to compromise
+        nodes_to_be_compromised = np.random.choice(computers, num_nodes_to_compromise)
+
+        # For each of the nodes to be compromised decide which step they become compromised
+        max_step_compromised = (
+            self.episode_steps // 2
+        )  # always compromise in first half of episode
+
+        # Bandwidth for all links
+        bandwidths = [i.get_bandwidth() for i in list(self.links.values())]
+        servers = [node for node in node_list if node.node_type == NodeType.SERVER]
+
+        for n, node in enumerate(nodes_to_be_compromised):
+            # 1: Use Node PoL to set node to compromised
+
+            _id = str(1000 + n)  # doesn't really matter, make sure it doesn't duplicate
+            _start_step = np.random.randint(
+                2, max_step_compromised + 1
+            )  # step compromised
+            _end_step = _start_step  # Become compromised on 1 step
+            _target_node_id = node.node_id
+            _pol_initiator = "DIRECT"
+            _pol_type = NodePOLType["SERVICE"]  # All computers are service nodes
+            pol_service_name = np.random.choice(
+                list(node.get_services().keys())
+            )  # Random service may wish to change this, currently always TCP)
+            pol_protocol = pol_protocol
+            _pol_state = SoftwareState.COMPROMISED
+            is_entry_node = True  # Assumes all computers in network are entry nodes
+            _pol_source_node_id = _pol_source_node_id
+            _pol_source_node_service = _pol_source_node_service
+            _pol_source_node_service_state = _pol_source_node_service_state
+            red_pol = NodeStateInstructionRed(
+                _id,
+                _start_step,
+                _end_step,
+                _target_node_id,
+                _pol_initiator,
+                _pol_type,
+                pol_protocol,
+                _pol_state,
+                _pol_source_node_id,
+                _pol_source_node_service,
+                _pol_source_node_service_state,
+            )
+
+            self.red_node_pol[_id] = red_pol
+
+            # 2: Launch the attack from compromised node - set the IER
+
+            ier_id = str(2000 + n)
+            # Launch the attack after node is compromised, and not right at the end of the episode
+            ier_start_step = np.random.randint(
+                _start_step + 2, int(self.episode_steps * 0.8)
+            )
+            ier_end_step = self.episode_steps
+            ier_source_node_id = node.get_id()
+            # Randomise the load, as a percentage of a random link bandwith
+            ier_load = np.random.uniform(low=0.4, high=0.8) * np.random.choice(
+                bandwidths
+            )
+            ier_protocol = pol_service_name  # Same protocol as compromised node
+            ier_service = node.get_services()[
+                pol_service_name
+            ]  # same service as defined in the pol
+            ier_port = ier_service.get_port()
+            ier_mission_criticality = (
+                0  # Red IER will never be important to green agent success
+            )
+            # We choose a node to attack based on the first that applies:
+            # a. Green IERs, select dest node of the red ier based on dest node of green IER
+            # b. Attack a random server that doesn't have a DENY acl rule in default config
+            # c. Attack a random server
+            possible_ier_destinations = [
+                ier.get_dest_node_id()
+                for ier in list(self.green_iers.values())
+                if ier.get_source_node_id() == node.get_id()
+            ]
+            if len(possible_ier_destinations) < 1:
+                for server in servers:
+                    if not self.acl.is_blocked(
+                        node.get_ip_address(),
+                        server.ip_address,
+                        ier_service,
+                        ier_port,
+                    ):
+                        possible_ier_destinations.append(server.node_id)
+            if len(possible_ier_destinations) < 1:
+                # If still none found choose from all servers
+                possible_ier_destinations = [server.node_id for server in servers]
+            ier_dest = np.random.choice(possible_ier_destinations)
+            self.red_iers[ier_id] = IER(
+                ier_id,
+                ier_start_step,
+                ier_end_step,
+                ier_load,
+                ier_protocol,
+                ier_port,
+                ier_source_node_id,
+                ier_dest,
+                ier_mission_criticality,
+            )
+
+            # 3: Make sure the targetted node can be set to overwhelmed - with node pol
+            # TODO remove duplicate red pol for same targetted service - must take into account start step
+
+            o_pol_id = str(3000 + n)
+            o_pol_start_step = ier_start_step  # Can become compromised the same step attack is launched
+            o_pol_end_step = (
+                self.episode_steps
+            )  # Can become compromised at any timestep after start
+            o_pol_node_id = ier_dest  # Node effected is the one targetted by the IER
+            o_pol_node_type = NodePOLType["SERVICE"]  # Always targets service nodes
+            o_pol_service_name = (
+                ier_protocol  # Same protocol/service as the IER uses to attack
+            )
+            o_pol_new_state = SoftwareState["OVERWHELMED"]
+            o_pol_entry_node = False  # Assumes servers are not entry nodes
+            o_red_pol = NodeStateInstructionRed(
+                _id,
+                _start_step,
+                _end_step,
+                _target_node_id,
+                _pol_initiator,
+                _pol_type,
+                pol_protocol,
+                _pol_state,
+                _pol_source_node_id,
+                _pol_source_node_service,
+                _pol_source_node_service_state,
+            )
+            self.red_node_pol[o_pol_id] = o_red_pol

src/primaite/nodes/node_state_instruction_red.py

@@ -137,3 +137,20 @@ class NodeStateInstructionRed(object):
              The source node service state
         """
         return self.source_node_service_state
+
+    def __repr__(self):
+        return (
+            f"{self.__class__.__name__}("
+            f"id={self.id}, "
+            f"start_step={self.start_step}, "
+            f"end_step={self.end_step}, "
+            f"target_node_id={self.target_node_id}, "
+            f"initiator={self.initiator}, "
+            f"pol_type={self.pol_type}, "
+            f"service_name={self.service_name}, "
+            f"state={self.state}, "
+            f"source_node_id={self.source_node_id}, "
+            f"source_node_service={self.source_node_service}, "
+            f"source_node_service_state={self.source_node_service_state}"
+            f")"
+        )