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

Format docstrings

Browse Source
This commit is contained in:
Marek Wolan
2023-07-04 11:11:52 +01:00
parent 91273b2f99
commit 9244c160b1
39 changed files with 208 additions and 409 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -21,8 +21,7 @@ _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>.
@@ -39,11 +38,10 @@ def get_session_path(session_timestamp: datetime) -> Path:
class AgentSessionABC(ABC):
"""
An ABC that manages training and/or evaluation of agents in PrimAITE.
"""An ABC that manages training and/or evaluation of agents in PrimAITE.
This class cannot be directly instantiated and must be inherited from
with all implemented abstract methods implemented.
This class cannot be directly instantiated and must be inherited from with all implemented abstract methods
implemented.
"""
@abstractmethod
@@ -186,8 +184,7 @@ class AgentSessionABC(ABC):
self,
**kwargs,
):
"""
Train the agent.
"""Train the agent.
:param kwargs: Any agent-specific key-word args to be passed.
"""
@@ -204,8 +201,7 @@ class AgentSessionABC(ABC):
self,
**kwargs,
):
"""
Evaluate the agent.
"""Evaluate the agent.
:param kwargs: Any agent-specific key-word args to be passed.
"""
@@ -293,11 +289,10 @@ class AgentSessionABC(ABC):
class HardCodedAgentSessionABC(AgentSessionABC):
"""
An Agent Session ABC for evaluation deterministic agents.
"""An Agent Session ABC for evaluation deterministic agents.
This class cannot be directly instantiated and must be inherited from
with all implemented abstract methods implemented.
This class cannot be directly instantiated and must be inherited from with all implemented abstract methods
implemented.
"""
def __init__(self, training_config_path, lay_down_config_path):
@@ -325,8 +320,7 @@ class HardCodedAgentSessionABC(AgentSessionABC):
self,
**kwargs,
):
"""
Train the agent.
"""Train the agent.
:param kwargs: Any agent-specific key-word args to be passed.
"""
@@ -340,8 +334,7 @@ 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

@@ -23,8 +23,7 @@ class HardCodedACLAgent(HardCodedAgentSessionABC):
return self._calculate_action_full_view(obs)
def get_blocked_green_iers(self, green_iers, acl, nodes):
"""
Get blocked green IERs.
"""Get blocked green IERs.
TODO: Add params and return in docstring.
TODO: Typehint params and return.
@@ -46,8 +45,7 @@ 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.
@@ -63,8 +61,7 @@ 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
@@ -83,8 +80,7 @@ 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.
@@ -108,8 +104,7 @@ 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.
@@ -140,8 +135,7 @@ 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.
@@ -173,8 +167,7 @@ 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.
@@ -197,8 +190,7 @@ 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:

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

@@ -6,8 +6,7 @@ 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

@@ -128,8 +128,7 @@ class RLlibAgent(AgentSessionABC):
self,
**kwargs,
):
"""
Evaluate the agent.
"""Evaluate the agent.
:param kwargs: Any agent-specific key-word args to be passed.
"""
@@ -147,8 +146,7 @@ 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

@@ -77,8 +77,7 @@ class SB3Agent(AgentSessionABC):
self,
**kwargs,
):
"""
Train the agent.
"""Train the agent.
:param kwargs: Any agent-specific key-word args to be passed.
"""
@@ -98,8 +97,7 @@ 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.

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

@@ -3,8 +3,7 @@ from primaite.agents.utils import get_new_action, transform_action_acl_enum, tra
class RandomAgent(HardCodedAgentSessionABC):
"""
A Random Agent.
"""A Random Agent.
Get a completely random action from the action space.
"""
@@ -14,11 +13,9 @@ class RandomAgent(HardCodedAgentSessionABC):
class DummyAgent(HardCodedAgentSessionABC):
"""
A Dummy Agent.
"""A Dummy Agent.
All action spaces setup so dummy action is always 0 regardless of action
type used.
All action spaces setup so dummy action is always 0 regardless of action type used.
"""
def _calculate_action(self, obs):
@@ -26,8 +23,7 @@ class DummyAgent(HardCodedAgentSessionABC):
class DoNothingACLAgent(HardCodedAgentSessionABC):
"""
A do nothing ACL agent.
"""A do nothing ACL agent.
A valid ACL action that has no effect; does nothing.
"""
@@ -41,8 +37,7 @@ class DoNothingACLAgent(HardCodedAgentSessionABC):
class DoNothingNodeAgent(HardCodedAgentSessionABC):
"""
A do nothing Node agent.
"""A do nothing Node agent.
A valid Node action that has no effect; does nothing.
"""

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

@@ -11,8 +11,7 @@ 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]
@@ -34,8 +33,7 @@ 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]
@@ -94,8 +92,7 @@ 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.
@@ -127,8 +124,7 @@ 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.
@@ -151,8 +147,7 @@ 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']]
@@ -174,8 +169,7 @@ 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']]
@@ -191,8 +185,7 @@ 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.
@@ -203,8 +196,7 @@ 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.
@@ -240,8 +232,7 @@ 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]])
@@ -269,8 +260,7 @@ 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"
@@ -305,8 +295,7 @@ 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]
@@ -337,8 +326,7 @@ 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]
@@ -360,8 +348,7 @@ 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.
@@ -388,8 +375,7 @@ 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.
@@ -411,8 +397,7 @@ 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.
@@ -432,8 +417,7 @@ 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)
@@ -480,8 +464,7 @@ 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
@@ -513,8 +496,7 @@ 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.
@@ -537,8 +519,7 @@ 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