oxbrbtx/PrimAITE
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Standardise docstring summary line placement.

Browse Source
This commit is contained in:
Marek Wolan
2023-07-07 10:28:00 +01:00
parent 86725064ec
commit f4b98542b6
36 changed files with 350 additions and 175 deletions
Download Patch File Download Diff File Expand all files Collapse all files

36
src/primaite/environment/observations.py
View File

@@ -50,7 +50,8 @@ class AbstractObservationComponent(ABC):
class NodeLinkTable(AbstractObservationComponent):
"""Table with nodes and links as rows and hardware/software status as cols.
"""
Table with nodes and links as rows and hardware/software status as cols.
This will create the observation space formatted as a table of integers.
There is one row per node, followed by one row per link.
@@ -74,7 +75,8 @@ class NodeLinkTable(AbstractObservationComponent):
_DATA_TYPE: type = np.int64
def __init__(self, env: "Primaite"):
"""Initialise a NodeLinkTable observation space component.
"""
Initialise a NodeLinkTable observation space component.
:param env: Training environment.
:type env: Primaite
@@ -100,7 +102,8 @@ class NodeLinkTable(AbstractObservationComponent):
self.structure = self.generate_structure()
def update(self):
"""Update the observation based on current environment state.
"""
Update the observation based on current environment state.
The structure of the observation space is described in :class:`.NodeLinkTable`
"""
@@ -181,7 +184,8 @@ class NodeLinkTable(AbstractObservationComponent):
class NodeStatuses(AbstractObservationComponent):
"""Flat list of nodes' hardware, OS, file system, and service states.
"""
Flat list of nodes' hardware, OS, file system, and service states.
The MultiDiscrete observation space can be though of as a one-dimensional vector of discrete states, represented by
integers.
@@ -234,7 +238,8 @@ class NodeStatuses(AbstractObservationComponent):
self.structure = self.generate_structure()
def update(self):
"""Update the observation based on current environment state.
"""
Update the observation based on current environment state.
The structure of the observation space is described in :class:`.NodeStatuses`
"""
@@ -287,7 +292,8 @@ class NodeStatuses(AbstractObservationComponent):
class LinkTrafficLevels(AbstractObservationComponent):
"""Flat list of traffic levels encoded into banded categories.
"""
Flat list of traffic levels encoded into banded categories.
For each link, total traffic or traffic per service is encoded into a categorical value.
For example, if ``quantisation_levels=5``, the traffic levels represent these values:
@@ -354,7 +360,8 @@ class LinkTrafficLevels(AbstractObservationComponent):
self.structure = self.generate_structure()
def update(self):
"""Update the observation based on current environment state.
"""
Update the observation based on current environment state.
The structure of the observation space is described in :class:`.LinkTrafficLevels`
"""
@@ -395,7 +402,8 @@ class LinkTrafficLevels(AbstractObservationComponent):
class ObservationsHandler:
"""Component-based observation space handler.
"""
Component-based observation space handler.
This allows users to configure observation spaces by mixing and matching components. Each component can also define
further parameters to make them more flexible.
@@ -436,7 +444,8 @@ class ObservationsHandler:
self._flat_observation = spaces.flatten(self._space, self._observation)
def register(self, obs_component: AbstractObservationComponent):
"""Add a component for this handler to track.
"""
Add a component for this handler to track.
:param obs_component: The component to add.
:type obs_component: AbstractObservationComponent
@@ -445,7 +454,8 @@ class ObservationsHandler:
self.update_space()
def deregister(self, obs_component: AbstractObservationComponent):
"""Remove a component from this handler.
"""
Remove a component from this handler.
:param obs_component: Which component to remove. It must exist within this object's
``registered_obs_components`` attribute.
@@ -488,7 +498,8 @@ class ObservationsHandler:
@classmethod
def from_config(cls, env: "Primaite", obs_space_config: dict):
"""Parse a config dictinary, return a new observation handler populated with new observation component objects.
"""
Parse a config dictinary, return a new observation handler populated with new observation component objects.
The expected format for the config dictionary is:
@@ -533,7 +544,8 @@ class ObservationsHandler:
return handler
def describe_structure(self):
"""Create a list of names for the features of the obs space.
"""
Create a list of names for the features of the obs space.
The order of labels follows the flattened version of the space.
"""

69
src/primaite/environment/primaite_env.py
View File

@@ -255,7 +255,8 @@ class Primaite(Env):
self.total_step_count = 0
def reset(self):
"""AI Gym Reset function.
"""
AI Gym Reset function.
Returns:
Environment observation space (reset)
@@ -291,7 +292,8 @@ class Primaite(Env):
return self.env_obs
def step(self, action):
"""AI Gym Step function.
"""
AI Gym Step function.
Args:
action: Action space from agent
@@ -429,7 +431,8 @@ class Primaite(Env):
print(" Protocol: " + protocol.get_name().name + ", Load: " + str(protocol.get_load()))
def interpret_action_and_apply(self, _action):
"""Applies agent actions to the nodes and Access Control List.
"""
Applies agent actions to the nodes and Access Control List.
Args:
_action: The action space from the agent
@@ -448,7 +451,8 @@ class Primaite(Env):
logging.error("Invalid action type found")
def apply_actions_to_nodes(self, _action):
"""Applies agent actions to the nodes.
"""
Applies agent actions to the nodes.
Args:
_action: The action space from the agent
@@ -535,7 +539,8 @@ class Primaite(Env):
return
def apply_actions_to_acl(self, _action):
"""Applies agent actions to the Access Control List [TO DO].
"""
Applies agent actions to the Access Control List [TO DO].
Args:
_action: The action space from the agent
@@ -612,7 +617,8 @@ class Primaite(Env):
return
def apply_time_based_updates(self):
"""Updates anything that needs to count down and then change state.
"""
Updates anything that needs to count down and then change state.
e.g. reset / patching status
"""
@@ -653,7 +659,8 @@ class Primaite(Env):
pass
def init_observations(self) -> Tuple[spaces.Space, np.ndarray]:
"""Create the environment's observation handler.
"""
Create the environment's observation handler.
:return: The observation space, initial observation (zeroed out array with the correct shape)
:rtype: Tuple[spaces.Space, np.ndarray]
@@ -709,7 +716,8 @@ class Primaite(Env):
print("Environment configuration loaded")
def create_node(self, item):
"""Creates a node from config data.
"""
Creates a node from config data.
Args:
item: A config data item
@@ -789,7 +797,8 @@ class Primaite(Env):
self.network_reference.add_nodes_from([node_ref])
def create_link(self, item: Dict):
"""Creates a link from config data.
"""
Creates a link from config data.
Args:
item: A config data item
@@ -832,7 +841,8 @@ class Primaite(Env):
)
def create_green_ier(self, item):
"""Creates a green IER from config data.
"""
Creates a green IER from config data.
Args:
item: A config data item
@@ -872,7 +882,8 @@ class Primaite(Env):
)
def create_red_ier(self, item):
"""Creates a red IER from config data.
"""
Creates a red IER from config data.
Args:
item: A config data item
@@ -901,7 +912,8 @@ class Primaite(Env):
)
def create_green_pol(self, item):
"""Creates a green PoL object from config data.
"""
Creates a green PoL object from config data.
Args:
item: A config data item
@@ -934,7 +946,8 @@ class Primaite(Env):
)
def create_red_pol(self, item):
"""Creates a red PoL object from config data.
"""
Creates a red PoL object from config data.
Args:
item: A config data item
@@ -974,7 +987,8 @@ class Primaite(Env):
)
def create_acl_rule(self, item):
"""Creates an ACL rule from config data.
"""
Creates an ACL rule from config data.
Args:
item: A config data item
@@ -994,7 +1008,8 @@ class Primaite(Env):
)
def create_services_list(self, services):
"""Creates a list of services (enum) from config data.
"""
Creates a list of services (enum) from config data.
Args:
item: A config data item representing the services
@@ -1009,7 +1024,8 @@ class Primaite(Env):
self.num_services = len(self.services_list)
def create_ports_list(self, ports):
"""Creates a list of ports from config data.
"""
Creates a list of ports from config data.
Args:
item: A config data item representing the ports
@@ -1024,7 +1040,8 @@ class Primaite(Env):
self.num_ports = len(self.ports_list)
def get_observation_info(self, observation_info):
"""Extracts observation_info.
"""
Extracts observation_info.
:param observation_info: Config item that defines which type of observation space to use
:type observation_info: str
@@ -1032,7 +1049,8 @@ class Primaite(Env):
self.observation_type = ObservationType[observation_info["type"]]
def get_action_info(self, action_info):
"""Extracts action_info.
"""
Extracts action_info.
Args:
item: A config data item representing action info
@@ -1040,7 +1058,8 @@ class Primaite(Env):
self.action_type = ActionType[action_info["type"]]
def save_obs_config(self, obs_config: dict):
"""Cache the config for the observation space.
"""
Cache the config for the observation space.
This is necessary as the observation space can't be built while reading the config,
it must be done after all the nodes, links, and services have been initialised.
@@ -1052,7 +1071,8 @@ class Primaite(Env):
self.obs_config = obs_config
def reset_environment(self):
"""# Resets environment.
"""
Resets environment.
Uses config data config data in order to build the environment configuration.
"""
@@ -1076,7 +1096,8 @@ class Primaite(Env):
ier_value.set_is_running(False)
def reset_node(self, item):
"""Resets the statuses of a node.
"""
Resets the statuses of a node.
Args:
item: A config data item
@@ -1123,7 +1144,8 @@ class Primaite(Env):
pass
def create_node_action_dict(self):
"""Creates a dictionary mapping each possible discrete action to more readable multidiscrete action.
"""
Creates a dictionary mapping each possible discrete action to more readable multidiscrete action.
Note: Only actions that have the potential to change the state exist in the mapping (except for key 0)
@@ -1187,7 +1209,8 @@ class Primaite(Env):
return actions
def create_node_and_acl_action_dict(self):
"""Create a dictionary mapping each possible discrete action to a more readable mutlidiscrete action.
"""
Create a dictionary mapping each possible discrete action to a more readable mutlidiscrete action.
The dictionary contains actions of both Node and ACL action types.
"""

15
src/primaite/environment/reward.py
View File

@@ -21,7 +21,8 @@ def calculate_reward_function(
step_count,
config_values,
):
"""Compares the states of the initial and final nodes/links to get a reward.
"""
Compares the states of the initial and final nodes/links to get a reward.
Args:
initial_nodes: The nodes before red and blue agents take effect
@@ -94,7 +95,8 @@ def calculate_reward_function(
def score_node_operating_state(final_node, initial_node, reference_node, config_values):
"""Calculates score relating to the hardware state of a node.
"""
Calculates score relating to the hardware state of a node.
Args:
final_node: The node after red and blue agents take effect
@@ -142,7 +144,8 @@ def score_node_operating_state(final_node, initial_node, reference_node, config_
def score_node_os_state(final_node, initial_node, reference_node, config_values):
"""Calculates score relating to the Software State of a node.
"""
Calculates score relating to the Software State of a node.
Args:
final_node: The node after red and blue agents take effect
@@ -192,7 +195,8 @@ def score_node_os_state(final_node, initial_node, reference_node, config_values)
def score_node_service_state(final_node, initial_node, reference_node, config_values):
"""Calculates score relating to the service state(s) of a node.
"""
Calculates score relating to the service state(s) of a node.
Args:
final_node: The node after red and blue agents take effect
@@ -263,7 +267,8 @@ def score_node_service_state(final_node, initial_node, reference_node, config_va
def score_node_file_system(final_node, initial_node, reference_node, config_values):
"""Calculates score relating to the file system state of a node.
"""
Calculates score relating to the file system state of a node.
Args:
final_node: The node after red and blue agents take effect