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

15
src/primaite/agents/agent.py
View File

@@ -21,7 +21,8 @@ _LOGGER = getLogger(__name__)
def get_session_path(session_timestamp: datetime) -> Path:
"""Get the directory path the session will output to.
"""
Get the directory path the session will output to.
This is set in the format of:
~/primaite/sessions/<yyyy-mm-dd>/<yyyy-mm-dd>_<hh-mm-ss>.
@@ -194,7 +195,8 @@ class AgentSessionABC(ABC):
self,
**kwargs,
):
"""Train the agent.
"""
Train the agent.
:param kwargs: Any agent-specific key-word args to be passed.
"""
@@ -211,7 +213,8 @@ class AgentSessionABC(ABC):
self,
**kwargs,
):
"""Evaluate the agent.
"""
Evaluate the agent.
:param kwargs: Any agent-specific key-word args to be passed.
"""
@@ -340,7 +343,8 @@ class HardCodedAgentSessionABC(AgentSessionABC):
self,
**kwargs,
):
"""Train the agent.
"""
Train the agent.
:param kwargs: Any agent-specific key-word args to be passed.
"""
@@ -354,7 +358,8 @@ class HardCodedAgentSessionABC(AgentSessionABC):
self,
**kwargs,
):
"""Evaluate the agent.
"""
Evaluate the agent.
:param kwargs: Any agent-specific key-word args to be passed.
"""

24
src/primaite/agents/hardcoded_acl.py
View File

@@ -46,7 +46,8 @@ class HardCodedACLAgent(HardCodedAgentSessionABC):
return blocked_green_iers
def get_matching_acl_rules_for_ier(self, ier, acl, nodes):
"""Get matching ACL rules for an IER.
"""
Get matching ACL rules for an IER.
TODO: Add params and return in docstring.
TODO: Typehint params and return.
@@ -62,7 +63,8 @@ class HardCodedACLAgent(HardCodedAgentSessionABC):
return matching_rules
def get_blocking_acl_rules_for_ier(self, ier, acl, nodes):
"""Get blocking ACL rules for an IER.
"""
Get blocking ACL rules for an IER.
.. warning::
Can return empty dict but IER can still be blocked by default
@@ -81,7 +83,8 @@ class HardCodedACLAgent(HardCodedAgentSessionABC):
return blocked_rules
def get_allow_acl_rules_for_ier(self, ier, acl, nodes):
"""Get all allowing ACL rules for an IER.
"""
Get all allowing ACL rules for an IER.
TODO: Add params and return in docstring.
TODO: Typehint params and return.
@@ -105,7 +108,8 @@ class HardCodedACLAgent(HardCodedAgentSessionABC):
nodes,
services_list,
):
"""Get matching ACL rules.
"""
Get matching ACL rules.
TODO: Add params and return in docstring.
TODO: Typehint params and return.
@@ -136,7 +140,8 @@ class HardCodedACLAgent(HardCodedAgentSessionABC):
nodes,
services_list,
):
"""Get the ALLOW ACL rules.
"""
Get the ALLOW ACL rules.
TODO: Add params and return in docstring.
TODO: Typehint params and return.
@@ -168,7 +173,8 @@ class HardCodedACLAgent(HardCodedAgentSessionABC):
nodes,
services_list,
):
"""Get the DENY ACL rules.
"""
Get the DENY ACL rules.
TODO: Add params and return in docstring.
TODO: Typehint params and return.
@@ -191,7 +197,8 @@ class HardCodedACLAgent(HardCodedAgentSessionABC):
return allowed_rules
def _calculate_action_full_view(self, obs):
"""Calculate a good acl-based action for the blue agent to take.
"""
Calculate a good acl-based action for the blue agent to take.
Knowledge of just the observation space is insufficient for a perfect solution, as we need to know:
@@ -355,7 +362,8 @@ class HardCodedACLAgent(HardCodedAgentSessionABC):
return action
def _calculate_action_basic_view(self, obs):
"""Calculate a good acl-based action for the blue agent to take.
"""
Calculate a good acl-based action for the blue agent to take.
Uses ONLY information from the current observation with NO knowledge
of previous actions taken and NO reward feedback.

3
src/primaite/agents/hardcoded_node.py
View File

@@ -6,7 +6,8 @@ class HardCodedNodeAgent(HardCodedAgentSessionABC):
"""An Agent Session class that implements a deterministic Node agent."""
def _calculate_action(self, obs):
"""Calculate a good node-based action for the blue agent to take.
"""
Calculate a good node-based action for the blue agent to take.
TODO: Add params and return in docstring.
TODO: Typehint params and return.

6
src/primaite/agents/rllib.py
View File

@@ -140,7 +140,8 @@ class RLlibAgent(AgentSessionABC):
self,
**kwargs,
):
"""Evaluate the agent.
"""
Evaluate the agent.
:param kwargs: Any agent-specific key-word args to be passed.
"""
@@ -158,7 +159,8 @@ class RLlibAgent(AgentSessionABC):
self,
**kwargs,
):
"""Evaluate the agent.
"""
Evaluate the agent.
:param kwargs: Any agent-specific key-word args to be passed.
"""

6
src/primaite/agents/sb3.py
View File

@@ -89,7 +89,8 @@ class SB3Agent(AgentSessionABC):
self,
**kwargs,
):
"""Train the agent.
"""
Train the agent.
:param kwargs: Any agent-specific key-word args to be passed.
"""
@@ -109,7 +110,8 @@ class SB3Agent(AgentSessionABC):
deterministic: bool = True,
**kwargs,
):
"""Evaluate the agent.
"""
Evaluate the agent.
:param deterministic: Whether the evaluation is deterministic.
:param kwargs: Any agent-specific key-word args to be passed.

63
src/primaite/agents/utils.py
View File

@@ -11,7 +11,8 @@ from primaite.common.enums import (
def transform_action_node_readable(action):
"""Convert a node action from enumerated format to readable format.
"""
Convert a node action from enumerated format to readable format.
example:
[1, 3, 1, 0] -> [1, 'SERVICE', 'PATCHING', 0]
@@ -33,7 +34,8 @@ def transform_action_node_readable(action):
def transform_action_acl_readable(action):
"""Transform an ACL action to a more readable format.
"""
Transform an ACL action to a more readable format.
example:
[0, 1, 2, 5, 0, 1] -> ['NONE', 'ALLOW', 2, 5, 'ANY', 1]
@@ -57,7 +59,8 @@ def transform_action_acl_readable(action):
def is_valid_node_action(action):
"""Is the node action an actual valid action.
"""
Is the node action an actual valid action.
Only uses information about the action to determine if the action has an effect
@@ -92,7 +95,8 @@ def is_valid_node_action(action):
def is_valid_acl_action(action):
"""Is the ACL action an actual valid action.
"""
Is the ACL action an actual valid action.
Only uses information about the action to determine if the action has an effect.
@@ -124,7 +128,8 @@ def is_valid_acl_action(action):
def is_valid_acl_action_extra(action):
"""Harsher version of valid acl actions, does not allow action.
"""
Harsher version of valid acl actions, does not allow action.
TODO: Add params and return in docstring.
TODO: Typehint params and return.
@@ -147,7 +152,8 @@ def is_valid_acl_action_extra(action):
def transform_change_obs_readable(obs):
"""Transform list of transactions to readable list of each observation property.
"""
Transform list of transactions to readable list of each observation property.
example:
np.array([[1,2,1,3],[2,1,1,1]]) -> [[1, 2], ['OFF', 'ON'], ['GOOD', 'GOOD'], ['COMPROMISED', 'GOOD']]
@@ -169,7 +175,8 @@ def transform_change_obs_readable(obs):
def transform_obs_readable(obs):
"""Transform observation to readable format.
"""
Transform observation to readable format.
np.array([[1,2,1,3],[2,1,1,1]]) -> [[1, 'OFF', 'GOOD', 'COMPROMISED'], [2, 'ON', 'GOOD', 'GOOD']]
@@ -185,7 +192,8 @@ def transform_obs_readable(obs):
def convert_to_new_obs(obs, num_nodes=10):
"""Convert original gym Box observation space to new multiDiscrete observation space.
"""
Convert original gym Box observation space to new multiDiscrete observation space.
TODO: Add params and return in docstring.
TODO: Typehint params and return.
@@ -196,7 +204,8 @@ def convert_to_new_obs(obs, num_nodes=10):
def convert_to_old_obs(obs, num_nodes=10, num_links=10, num_services=1):
"""Convert to old observation.
"""
Convert to old observation.
Links filled with 0's as no information is included in new observation space.
@@ -232,7 +241,8 @@ def convert_to_old_obs(obs, num_nodes=10, num_links=10, num_services=1):
def describe_obs_change(obs1, obs2, num_nodes=10, num_links=10, num_services=1):
"""Return string describing change between two observations.
"""
Return string describing change between two observations.
example:
obs_1 = array([[1, 1, 1, 1, 3], [2, 1, 1, 1, 1]])
@@ -260,7 +270,8 @@ def describe_obs_change(obs1, obs2, num_nodes=10, num_links=10, num_services=1):
def _describe_obs_change_helper(obs_change, is_link):
"""Helper funcion to describe what has changed.
"""
Helper funcion to describe what has changed.
example:
[ 1 -1 -1 -1 1] -> "ID 1: Service 1 changed to GOOD"
@@ -295,7 +306,8 @@ def _describe_obs_change_helper(obs_change, is_link):
def transform_action_node_enum(action):
"""Convert a node action from readable string format, to enumerated format.
"""
Convert a node action from readable string format, to enumerated format.
example:
[1, 'SERVICE', 'PATCHING', 0] -> [1, 3, 1, 0]
@@ -326,7 +338,8 @@ def transform_action_node_enum(action):
def transform_action_node_readable(action):
"""Convert a node action from enumerated format to readable format.
"""
Convert a node action from enumerated format to readable format.
example:
[1, 3, 1, 0] -> [1, 'SERVICE', 'PATCHING', 0]
@@ -348,7 +361,8 @@ def transform_action_node_readable(action):
def node_action_description(action):
"""Generate string describing a node-based action.
"""
Generate string describing a node-based action.
TODO: Add params and return in docstring.
TODO: Typehint params and return.
@@ -375,7 +389,8 @@ def node_action_description(action):
def transform_action_acl_enum(action):
"""Convert acl action from readable str format, to enumerated format.
"""
Convert acl action from readable str format, to enumerated format.
TODO: Add params and return in docstring.
TODO: Typehint params and return.
@@ -397,7 +412,8 @@ def transform_action_acl_enum(action):
def acl_action_description(action):
"""Generate string describing an acl-based action.
"""
Generate string describing an acl-based action.
TODO: Add params and return in docstring.
TODO: Typehint params and return.
@@ -417,7 +433,8 @@ def acl_action_description(action):
def get_node_of_ip(ip, node_dict):
"""Get the node ID of an IP address.
"""
Get the node ID of an IP address.
node_dict: dictionary of nodes where key is ID, and value is the node (can be ontained from env.nodes)
@@ -431,7 +448,8 @@ def get_node_of_ip(ip, node_dict):
def is_valid_node_action(action):
"""Is the node action an actual valid action.
"""
Is the node action an actual valid action.
Only uses information about the action to determine if the action has an effect
@@ -464,7 +482,8 @@ def is_valid_node_action(action):
def is_valid_acl_action(action):
"""Is the ACL action an actual valid action.
"""
Is the ACL action an actual valid action.
Only uses information about the action to determine if the action has an effect
@@ -496,7 +515,8 @@ def is_valid_acl_action(action):
def is_valid_acl_action_extra(action):
"""Harsher version of valid acl actions, does not allow action.
"""
Harsher version of valid acl actions, does not allow action.
TODO: Add params and return in docstring.
TODO: Typehint params and return.
@@ -519,7 +539,8 @@ def is_valid_acl_action_extra(action):
def get_new_action(old_action, action_dict):
"""Get new action (e.g. 32) from old action e.g. [1,1,1,0].
"""
Get new action (e.g. 32) from old action e.g. [1,1,1,0].
Old_action can be either node or acl action type