Ran pre-commit hook on all files and performed changes to fix flake8 failures

.flake8

@@ -0,0 +1,12 @@
+[flake8]
+max-line-length=120
+extend-ignore =
+    D105
+    D107
+    D100
+    D104
+    E203
+    E712
+    D401
+exclude =
+    docs/source/*

docs/conf.py

@@ -3,12 +3,15 @@
 # For the full list of built-in configuration values, see the documentation:
 # https://www.sphinx-doc.org/en/master/usage/configuration.html
 
+import datetime
+
 # -- Project information -----------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information
 import os
 import sys
-import datetime
+
-import furo
+import furo  # noqa
+
 sys.path.insert(0, os.path.abspath("../"))
 
 
@@ -33,7 +36,6 @@ templates_path = ["_templates"]
 exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"]
 
 
-
 # -- Options for HTML output -------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output
 

docs/source/dependencies.rst

@@ -22,5 +22,3 @@ The latest release of PrimAITE has been tested against the following versions of
 * gym 0.21.0
 * matplotlib 3.6.2
 * stable_baselines_3 1.6.2
-
-

docs/source/session.rst

@@ -84,5 +84,3 @@ In order to execute a session, carry out the following steps:
 2. Start a console window (type “CMD” in path window, or start a console window first and navigate to “[Install Directory]\\Primaite\\Primaite\\”)
 3. Type “python main.py”
 4. The session will start with an output indicating the current episode, and average reward value for the episode
-
- 

setup.py

@@ -10,6 +10,7 @@ class bdist_wheel(_bdist_wheel):  # noqa
         # Source: https://stackoverflow.com/a/45150383
         self.root_is_pure = False
 
+
 setup(
     cmdclass={
         "bdist_wheel": bdist_wheel,

src/primaite/acl/access_control_list.py

@@ -1,25 +1,19 @@
 # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
-"""
+"""A class that implements the access control list implementation for the network."""
-A class that implements the access control list implementation for the network
-"""
 
 from primaite.acl.acl_rule import ACLRule
 
-class AccessControlList():
+
-    """
+class AccessControlList:
-    Access Control List class
+    """Access Control List class."""
-    """
 
     def __init__(self):
-        """
+        """Init."""
-        Init
-        """
-
         self.acl = {}  # A dictionary of ACL Rules
 
     def check_address_match(self, _rule, _source_ip_address, _dest_ip_address):
         """
-        Checks for IP address matches
+        Checks for IP address matches.
 
         Args:
             _rule: The rule being checked
@@ -29,18 +23,28 @@ class AccessControlList():
         Returns:
              True if match; False otherwise.
         """
-
+        if (
-        if ((_rule.get_source_ip() == _source_ip_address and _rule.get_dest_ip() == _dest_ip_address) or
+            (
-            (_rule.get_source_ip() == "ANY" and _rule.get_dest_ip() == _dest_ip_address) or
+                _rule.get_source_ip() == _source_ip_address
-            (_rule.get_source_ip() == _source_ip_address and _rule.get_dest_ip() == "ANY") or
+                and _rule.get_dest_ip() == _dest_ip_address
-            (_rule.get_source_ip() == "ANY" and _rule.get_dest_ip() == "ANY")):
+            )
+            or (
+                _rule.get_source_ip() == "ANY"
+                and _rule.get_dest_ip() == _dest_ip_address
+            )
+            or (
+                _rule.get_source_ip() == _source_ip_address
+                and _rule.get_dest_ip() == "ANY"
+            )
+            or (_rule.get_source_ip() == "ANY" and _rule.get_dest_ip() == "ANY")
+        ):
             return True
         else:
             return False
 
     def is_blocked(self, _source_ip_address, _dest_ip_address, _protocol, _port):
         """
-        Checks for rules that block a protocol / port
+        Checks for rules that block a protocol / port.
 
         Args:
             _source_ip_address: the source IP address to check
@@ -51,11 +55,17 @@ class AccessControlList():
         Returns:
              Indicates block if all conditions are satisfied.
         """
-        
         for rule_key, rule_value in self.acl.items():
-            if self.check_address_match(rule_value, _source_ip_address, _dest_ip_address):
+            if self.check_address_match(
-                if ((rule_value.get_protocol() == _protocol or rule_value.get_protocol() == "ANY") and
+                rule_value, _source_ip_address, _dest_ip_address
-                    (str(rule_value.get_port()) == str(_port) or rule_value.get_port() == "ANY")):
+            ):
+                if (
+                    rule_value.get_protocol() == _protocol
+                    or rule_value.get_protocol() == "ANY"
+                ) and (
+                    str(rule_value.get_port()) == str(_port)
+                    or rule_value.get_port() == "ANY"
+                ):
                     # There's a matching rule. Get the permission
                     if rule_value.get_permission() == "DENY":
                         return True
@@ -67,7 +77,7 @@ class AccessControlList():
 
     def add_rule(self, _permission, _source_ip, _dest_ip, _protocol, _port):
         """
-        Adds a new rule
+        Adds a new rule.
 
         Args:
             _permission: the permission value (e.g. "ALLOW" or "DENY")
@@ -76,14 +86,13 @@ class AccessControlList():
             _protocol: the protocol
             _port: the port
         """
-
         new_rule = ACLRule(_permission, _source_ip, _dest_ip, _protocol, str(_port))
         hash_value = hash(new_rule)
         self.acl[hash_value] = new_rule
 
     def remove_rule(self, _permission, _source_ip, _dest_ip, _protocol, _port):
         """
-        Removes a rule
+        Removes a rule.
 
         Args:
             _permission: the permission value (e.g. "ALLOW" or "DENY")
@@ -92,25 +101,21 @@ class AccessControlList():
             _protocol: the protocol
             _port: the port
         """
-
         rule = ACLRule(_permission, _source_ip, _dest_ip, _protocol, str(_port))
         hash_value = hash(rule)
         # There will not always be something 'popable' since the agent will be trying random things
         try:
             self.acl.pop(hash_value)
-        except:
+        except Exception:
             return
 
     def remove_all_rules(self):
-        """
+        """Removes all rules."""
-        Removes all rules
-        """
-
         self.acl.clear()
 
     def get_dictionary_hash(self, _permission, _source_ip, _dest_ip, _protocol, _port):
         """
-        Produces a hash value for a rule
+        Produces a hash value for a rule.
 
         Args:
             _permission: the permission value (e.g. "ALLOW" or "DENY")
@@ -122,13 +127,6 @@ class AccessControlList():
         Returns:
              Hash value based on rule parameters.
         """
-
         rule = ACLRule(_permission, _source_ip, _dest_ip, _protocol, str(_port))
         hash_value = hash(rule)
         return hash_value
-
-
-
-
-
-

src/primaite/acl/acl_rule.py

@@ -1,16 +1,13 @@
 # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
-"""
+"""A class that implements an access control list rule."""
-A class that implements an access control list rule
-"""
 
-class ACLRule():
+
-    """
+class ACLRule:
-    Access Control List Rule class
+    """Access Control List Rule class."""
-    """
 
     def __init__(self, _permission, _source_ip, _dest_ip, _protocol, _port):
         """
-        Init
+        Init.
 
         Args:
             _permission: The permission (ALLOW or DENY)
@@ -19,7 +16,6 @@ class ACLRule():
             _protocol: The rule protocol
             _port: The rule port
         """
-
         self.permission = _permission
         self.source_ip = _source_ip
         self.dest_ip = _dest_ip
@@ -28,47 +24,45 @@ class ACLRule():
 
     def __hash__(self):
         """
-        Override the hash function
+        Override the hash function.
 
         Returns:
              Returns hash of core parameters.
         """
-
+        return hash(
-        return hash((self.permission, self.source_ip, self.dest_ip, self.protocol, self.port))
+            (self.permission, self.source_ip, self.dest_ip, self.protocol, self.port)
+        )
 
     def get_permission(self):
         """
-        Gets the permission attribute
+        Gets the permission attribute.
 
         Returns:
              Returns permission attribute
         """
-
         return self.permission
 
     def get_source_ip(self):
         """
-        Gets the source IP address attribute
+        Gets the source IP address attribute.
 
         Returns:
              Returns source IP address attribute
         """
-
         return self.source_ip
 
     def get_dest_ip(self):
         """
-        Gets the desintation IP address attribute
+        Gets the desintation IP address attribute.
 
         Returns:
              Returns destination IP address attribute
         """
-
         return self.dest_ip
 
     def get_protocol(self):
         """
-        Gets the protocol attribute
+        Gets the protocol attribute.
 
         Returns:
              Returns protocol attribute
@@ -77,12 +71,9 @@ class ACLRule():
 
     def get_port(self):
         """
-        Gets the port attribute
+        Gets the port attribute.
 
         Returns:
              Returns port attribute
         """
-
         return self.port
-        
-

src/primaite/common/__init__.py

@@ -1,2 +1 @@
-
 # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.

src/primaite/common/config_values_main.py

@@ -1,18 +1,12 @@
 # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
-"""
+"""The config class."""
-The config class
+
-"""
 
 class config_values_main(object):
-    """
+    """Class to hold main config values."""
-    Class to hold main config values
-    """
 
     def __init__(self):
-        """
+        """Init."""
-        Init
-        """
-
         # Generic
         self.agent_identifier = ""  # the agent in use
         self.num_episodes = 0  # number of episodes to train over
@@ -22,7 +16,9 @@ class config_values_main(object):
         self.session_type = ""  # the session type to run (TRAINING or EVALUATION)
 
         # Environment
-        self.observation_space_high_value = 0       # The high value for the observation space
+        self.observation_space_high_value = (
+            0  # The high value for the observation space
+        )
 
         # Reward values
         # Generic
@@ -59,11 +55,15 @@ class config_values_main(object):
         self.repairing_should_be_good = 0
         self.repairing_should_be_restoring = 0
         self.repairing_should_be_corrupt = 0
-        self.repairing_should_be_destroyed = 0   # Repairing does not fix destroyed state - you need to restore
+        self.repairing_should_be_destroyed = (
+            0  # Repairing does not fix destroyed state - you need to restore
+        )
         self.repairing = 0
         self.restoring_should_be_good = 0
         self.restoring_should_be_repairing = 0
-        self.restoring_should_be_corrupt = 0     # Not the optimal method (as repair will fix corruption)
+        self.restoring_should_be_corrupt = (
+            0  # Not the optimal method (as repair will fix corruption)
+        )
         self.restoring_should_be_destroyed = 0
         self.restoring = 0
         self.corrupt_should_be_good = 0
@@ -88,4 +88,3 @@ class config_values_main(object):
         self.file_system_repairing_limit = 0  # The time take to repair a file
         self.file_system_restoring_limit = 0  # The time take to restore a file
         self.file_system_scanning_limit = 0  # The time taken to scan the file system
-

src/primaite/common/enums.py

@@ -1,14 +1,11 @@
 # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
-"""
+"""Enumerations for APE."""
-Enumerations for APE
-"""
 
 from enum import Enum
 
+
 class TYPE(Enum):
-    """
+    """Node type enumeration."""
-    Node type enumeration
-    """
 
     CCTV = 1
     SWITCH = 2
@@ -21,10 +18,9 @@ class TYPE(Enum):
     ACTUATOR = 9
     SERVER = 10
 
+
 class PRIORITY(Enum):
-    """
+    """Node priority enumeration."""
-    Node priority enumeration
-    """
 
     P1 = 1
     P2 = 2
@@ -32,48 +28,43 @@ class PRIORITY(Enum):
     P4 = 4
     P5 = 5
 
+
 class HARDWARE_STATE(Enum):
-    """
+    """Node hardware state enumeration."""
-    Node hardware state enumeration
-    """
 
     ON = 1
     OFF = 2
     RESETTING = 3
 
+
 class SOFTWARE_STATE(Enum):
-    """
+    """O/S or Service state enumeration."""
-    O/S or Service state enumeration
-    """
 
     GOOD = 1
     PATCHING = 2
     COMPROMISED = 3
     OVERWHELMED = 4
 
+
 class NODE_POL_TYPE(Enum):
-    """
+    """Node Pattern of Life type enumeration."""
-    Node Pattern of Life type enumeration
-    """
 
     OPERATING = 1
     OS = 2
     SERVICE = 3
     FILE = 4
 
+
 class NODE_POL_INITIATOR(Enum):
-    """
+    """Node Pattern of Life initiator enumeration."""
-    Node Pattern of Life initiator enumeration
-    """
 
     DIRECT = 1
     IER = 2
     SERVICE = 3
 
+
 class PROTOCOL(Enum):
-    """
+    """Service protocol enumeration."""
-    Service protocol enumeration
-    """
 
     LDAP = 0
     FTP = 1
@@ -84,18 +75,16 @@ class PROTOCOL(Enum):
     TCP = 6
     NONE = 7
 
+
 class ACTION_TYPE(Enum):
-    """
+    """Action type enumeration."""
-    Action type enumeration
-    """
 
     NODE = 0
     ACL = 1
 
+
 class FILE_SYSTEM_STATE(Enum):
-    """
+    """File System State."""
-    File System State
-    """
 
     GOOD = 1
     CORRUPT = 2

src/primaite/common/protocol.py

@@ -1,59 +1,47 @@
 # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
-"""
+"""The protocol class."""
-The protocol class
+
-"""
 
 class Protocol(object):
-    """
+    """Protocol class."""
-    Protocol class
-    """
 
     def __init__(self, _name):
         """
-        Init
+        Init.
 
         Args:
             _name: The protocol name
         """
-
         self.name = _name
         self.load = 0  # bps
 
     def get_name(self):
         """
-        Gets the protocol name
+        Gets the protocol name.
 
         Returns:
              The protocol name
         """
-
         return self.name
 
     def get_load(self):
         """
-        Gets the protocol load
+        Gets the protocol load.
 
         Returns:
              The protocol load (bps)
         """
-
         return self.load
 
     def add_load(self, _load):
         """
-        Adds load to the protocol
+        Adds load to the protocol.
 
         Args:
             _load: The load to add
         """
-
         self.load += _load
 
     def clear_load(self):
-        """
+        """Clears the load on this protocol."""
-        Clears the load on this protocol
-        """
-
         self.load = 0
-
-

src/primaite/common/service.py

@@ -1,25 +1,21 @@
 # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
-"""
+"""The Service class."""
-The Service class
-"""
 
 from primaite.common.enums import SOFTWARE_STATE
 
+
 class Service(object):
-    """
+    """Service class."""
-    Service class
-    """
 
     def __init__(self, _name, _port, _state):
         """
-        Init
+        Init.
 
         Args:
             _name: The service name
             _port: The service port
             _state: The service state
         """
-
         self.name = _name
         self.port = _port
         self.state = _state
@@ -27,74 +23,61 @@ class Service(object):
 
     def set_name(self, _name):
         """
-        Sets the service name
+        Sets the service name.
 
         Args:
             _name: The service name
         """
-
         self.name = _name
 
     def get_name(self):
         """
-        Gets the service name
+        Gets the service name.
 
         Returns:
              The service name
         """
-
         return self.name
 
     def set_port(self, _port):
         """
-        Sets the service port
+        Sets the service port.
 
         Args:
             _port: The service port
         """
-
         self.port = _port
 
     def get_port(self):
         """
-        Gets the service port
+        Gets the service port.
 
         Returns:
              The service port
         """
-
         return self.port
 
     def set_state(self, _state):
         """
-        Sets the service state
+        Sets the service state.
 
         Args:
             _state: The service state
         """
-
         self.state = _state
 
     def get_state(self):
         """
-        Gets the service state
+        Gets the service state.
 
         Returns:
              The service state
         """
-
         return self.state
 
     def reduce_patching_count(self):
-        """
+        """Reduces the patching count for the service."""
-        Reduces the patching count for the service
-        """
-
         self.patching_count -= 1
         if self.patching_count <= 0:
             self.patching_count = 0
             self.state = SOFTWARE_STATE.GOOD
-
-
-
-

src/primaite/environment/__init__.py

@@ -1,2 +1 @@
-
 # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.

src/primaite/environment/primaite_env.py

@@ -1,39 +1,45 @@
 # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
-"""
+"""Main environment module containing the PRIMmary AI Training Evironment (Primaite) class."""
-Main environment module containing the PRIMmary AI Training Evironment (Primaite) class
-"""
 
-import numpy as np
-import networkx as nx
 import copy
 import csv
-import yaml
-import os.path
 import logging
-
+import os.path
-from gym import Env, spaces
-from matplotlib import pyplot as plt
 from datetime import datetime
 
-from primaite.common.enums import *
+import networkx as nx
+import numpy as np
+import yaml
+from gym import Env, spaces
+from matplotlib import pyplot as plt
+
+from primaite.acl.access_control_list import AccessControlList
+from primaite.common.enums import (
+    ACTION_TYPE,
+    FILE_SYSTEM_STATE,
+    HARDWARE_STATE,
+    NODE_POL_INITIATOR,
+    NODE_POL_TYPE,
+    PRIORITY,
+    SOFTWARE_STATE,
+    TYPE,
+)
+from primaite.common.service import Service
+from primaite.environment.reward import calculate_reward_function
 from primaite.links.link import Link
-from primaite.pol.ier import IER
+from primaite.nodes.active_node import ActiveNode
 from primaite.nodes.node_state_instruction_green import NodeStateInstructionGreen
 from primaite.nodes.node_state_instruction_red import NodeStateInstructionRed
-from primaite.pol.green_pol import apply_iers, apply_node_pol
-from primaite.pol.red_agent_pol import apply_red_agent_iers, apply_red_agent_node_pol
-from primaite.nodes.active_node import ActiveNode
 from primaite.nodes.passive_node import PassiveNode
 from primaite.nodes.service_node import ServiceNode
-from primaite.common.service import Service
+from primaite.pol.green_pol import apply_iers, apply_node_pol
-from primaite.acl.access_control_list import AccessControlList
+from primaite.pol.ier import IER
-from primaite.environment.reward import calculate_reward_function
+from primaite.pol.red_agent_pol import apply_red_agent_iers, apply_red_agent_node_pol
 from primaite.transactions.transaction import Transaction
 
+
 class Primaite(Env):
-    """
+    """PRIMmary AI Training Evironment (Primaite) class."""
-    PRIMmary AI Training Evironment (Primaite) class
-    """
 
     # Observation / Action Space contants
     OBSERVATION_SPACE_FIXED_PARAMETERS = 4
@@ -46,7 +52,7 @@ class Primaite(Env):
 
     def __init__(self, _config_values, _transaction_list):
         """
-        Init
+        Init.
 
         Args:
             _episode_steps: The number of steps for the episode
@@ -54,7 +60,6 @@ class Primaite(Env):
             _transaction_list: The list of transactions to populate
             _agent_identifier: Identifier for the agent
         """
-
         super(Primaite, self).__init__()
 
         # Take a copy of the config values
@@ -140,10 +145,12 @@ class Primaite(Env):
 
         # Open the config file and build the environment laydown
         try:
-            self.config_file = open("config/" + self.config_values.config_filename_use_case, "r")
+            self.config_file = open(
+                "config/" + self.config_values.config_filename_use_case, "r"
+            )
             self.config_data = yaml.safe_load(self.config_file)
             self.load_config()
-        except Exception as e:
+        except Exception:
             logging.error("Could not load the environment configuration")
             logging.error("Exception occured", exc_info=True)
 
@@ -165,14 +172,14 @@ class Primaite(Env):
             now = datetime.now()  # current date and time
             time = now.strftime("%Y%m%d_%H%M%S")
 
-            path = 'outputs/diagrams'
+            path = "outputs/diagrams"
             is_dir = os.path.isdir(path)
             if not is_dir:
                 os.makedirs(path)
             filename = "outputs/diagrams/network_" + time + ".png"
             plt.savefig(filename, format="PNG")
             plt.clf()
-        except Exception as a:
+        except Exception:
             logging.error("Could not save network diagram")
             logging.error("Exception occured", exc_info=True)
             print("Could not save network diagram")
@@ -194,16 +201,22 @@ class Primaite(Env):
         # - service F state          |      service F loading
         # - service G state          |      service G loading
 
-        # Calculate the number of items that need to be included in the observation space
+        # Calculate the number of items that need to be included in the
+        # observation space
         num_items = self.num_links + self.num_nodes
-        # Set the number of observation parameters, being # of services plus id, operating state, file system state and O/S state (i.e. 4)
+        # Set the number of observation parameters, being # of services plus id,
-        self.num_observation_parameters = self.num_services + self.OBSERVATION_SPACE_FIXED_PARAMETERS
+        # operating state, file system state and O/S state (i.e. 4)
+        self.num_observation_parameters = (
+            self.num_services + self.OBSERVATION_SPACE_FIXED_PARAMETERS
+        )
         # Define the observation shape
         self.observation_shape = (num_items, self.num_observation_parameters)
-        self.observation_space = spaces.Box(low=0, 
+        self.observation_space = spaces.Box(
+            low=0,
             high=self.config_values.observation_space_high_value,
             shape=self.observation_shape,
-                                            dtype=np.int64)
+            dtype=np.int64,
+        )
 
         # This is the observation that is sent back via the rest and step functions
         self.env_obs = np.zeros(self.observation_shape, dtype=np.int64)
@@ -216,7 +229,14 @@ class Primaite(Env):
             # [0, 4] - what property it's acting on (0 = nothing, state, o/s state, service state, file system state)
             # [0, 3] - action on property (0 = nothing, On / Scan, Off / Repair, Reset / Patch / Restore)
             # [0, num services] - resolves to service ID (0 = nothing, resolves to service)
-            self.action_space = spaces.MultiDiscrete([self.num_nodes, self.ACTION_SPACE_NODE_PROPERTY_VALUES, self.ACTION_SPACE_NODE_ACTION_VALUES, self.num_services])
+            self.action_space = spaces.MultiDiscrete(
+                [
+                    self.num_nodes,
+                    self.ACTION_SPACE_NODE_PROPERTY_VALUES,
+                    self.ACTION_SPACE_NODE_ACTION_VALUES,
+                    self.num_services,
+                ]
+            )
         else:
             logging.info("Action space type ACL selected")
             # Terms (for ACL action space):
@@ -226,35 +246,45 @@ class Primaite(Env):
             # [0, num nodes] - Dest IP (0 = any, then 1 -> x resolving to IP addresses)
             # [0, num services] - Protocol (0 = any, then 1 -> x resolving to protocol)
             # [0, num ports] - Port (0 = any, then 1 -> x resolving to port)
-            self.action_space = spaces.MultiDiscrete([self.ACTION_SPACE_ACL_ACTION_VALUES, self.ACTION_SPACE_ACL_PERMISSION_VALUES, self.num_nodes + 1, self.num_nodes + 1, self.num_services + 1, self.num_ports + 1])
+            self.action_space = spaces.MultiDiscrete(
+                [
+                    self.ACTION_SPACE_ACL_ACTION_VALUES,
+                    self.ACTION_SPACE_ACL_PERMISSION_VALUES,
+                    self.num_nodes + 1,
+                    self.num_nodes + 1,
+                    self.num_services + 1,
+                    self.num_ports + 1,
+                ]
+            )
 
         # Set up a csv to store the results of the training
         try:
             now = datetime.now()  # current date and time
             time = now.strftime("%Y%m%d_%H%M%S")
-            header = ['Episode', 'Average Reward']
+            header = ["Episode", "Average Reward"]
 
             # Check whether the output/rerults folder exists (doesn't exist by default install)
-            path = 'outputs/results/'
+            path = "outputs/results/"
             is_dir = os.path.isdir(path)
             if not is_dir:
                 os.makedirs(path)
             filename = "outputs/results/average_reward_per_episode_" + time + ".csv"
-            self.csv_file = open(filename, 'w', encoding='UTF8', newline='')
+            self.csv_file = open(filename, "w", encoding="UTF8", newline="")
             self.csv_writer = csv.writer(self.csv_file)
             self.csv_writer.writerow(header)
-        except Exception as e:
+        except Exception:
-            logging.error("Could not create csv file to hold average reward per episode")
+            logging.error(
+                "Could not create csv file to hold average reward per episode"
+            )
             logging.error("Exception occured", exc_info=True)
 
     def reset(self):
         """
-        AI Gym Reset function
+        AI Gym Reset function.
 
         Returns:
              Environment observation space (reset)
         """
-
         csv_data = self.episode_count, self.average_reward
         self.csv_writer.writerow(csv_data)
 
@@ -280,7 +310,7 @@ class Primaite(Env):
 
     def step(self, action):
         """
-        AI Gym Step function
+        AI Gym Step function.
 
         Args:
             action: Action space from agent
@@ -291,7 +321,6 @@ class Primaite(Env):
              done: Indicates episode is complete if True
              step_info: Additional information relating to this step
         """
-
         if self.step_count == 0:
             print("Episode: " + str(self.episode_count) + " running")
 
@@ -306,7 +335,9 @@ class Primaite(Env):
             link_value.clear_traffic()
 
         # Create a Transaction (metric) object for this step
-        transaction = Transaction(datetime.now(), self.agent_identifier, self.episode_count, self.step_count)
+        transaction = Transaction(
+            datetime.now(), self.agent_identifier, self.episode_count, self.step_count
+        )
         # Load the initial observation space into the transaction
         transaction.set_obs_space_pre(copy.deepcopy(self.env_obs))
         # Load the action space into the transaction
@@ -317,17 +348,40 @@ class Primaite(Env):
 
         # 2. Apply PoL
         apply_node_pol(self.nodes, self.node_pol, self.step_count)  # Node PoL
-        apply_iers(self.network, self.nodes, self.links, self.green_iers, self.acl, self.step_count)        # Network PoL  
+        apply_iers(
+            self.network,
+            self.nodes,
+            self.links,
+            self.green_iers,
+            self.acl,
+            self.step_count,
+        )  # Network PoL
         # Take snapshots of nodes and links
         self.nodes_post_pol = copy.deepcopy(self.nodes)
         self.links_post_pol = copy.deepcopy(self.links)
         # Reference
         apply_node_pol(self.nodes_reference, self.node_pol, self.step_count)  # Node PoL
-        apply_iers(self.network_reference, self.nodes_reference, self.links_reference, self.green_iers, self.acl, self.step_count)      # Network PoL  
+        apply_iers(
+            self.network_reference,
+            self.nodes_reference,
+            self.links_reference,
+            self.green_iers,
+            self.acl,
+            self.step_count,
+        )  # Network PoL
 
         # 3. Implement Red Action
-        apply_red_agent_iers(self.network, self.nodes, self.links, self.red_iers, self.acl, self.step_count)
+        apply_red_agent_iers(
-        apply_red_agent_node_pol(self.nodes, self.red_iers, self.red_node_pol, self.step_count)
+            self.network,
+            self.nodes,
+            self.links,
+            self.red_iers,
+            self.acl,
+            self.step_count,
+        )
+        apply_red_agent_node_pol(
+            self.nodes, self.red_iers, self.red_node_pol, self.step_count
+        )
         # Take snapshots of nodes and links
         self.nodes_post_red = copy.deepcopy(self.nodes)
         self.links_post_red = copy.deepcopy(self.links)
@@ -335,24 +389,48 @@ class Primaite(Env):
         # 4. Implement Blue Action
         self.interpret_action_and_apply(action)
 
-        # 5. Reapply normal and Red agent IER PoL, as we need to see what effect the blue agent action has had (if any) on link status
+        # 5. Reapply normal and Red agent IER PoL, as we need to see what
+        # effect the blue agent action has had (if any) on link status
         # Need to clear traffic on all links first
         for link_key, link_value in self.links.items():
             link_value.clear_traffic()
-        apply_iers(self.network, self.nodes, self.links, self.green_iers, self.acl, self.step_count)
+        apply_iers(
-        apply_red_agent_iers(self.network, self.nodes, self.links, self.red_iers, self.acl, self.step_count)
+            self.network,
+            self.nodes,
+            self.links,
+            self.green_iers,
+            self.acl,
+            self.step_count,
+        )
+        apply_red_agent_iers(
+            self.network,
+            self.nodes,
+            self.links,
+            self.red_iers,
+            self.acl,
+            self.step_count,
+        )
         # Take snapshots of nodes and links
         self.nodes_post_blue = copy.deepcopy(self.nodes)
         self.links_post_blue = copy.deepcopy(self.links)
 
         # 6. Calculate reward signal (for RL)
-        reward = calculate_reward_function(self.nodes_post_pol, self.nodes_post_blue, self.nodes_reference, self.green_iers, self.red_iers, self.step_count, self.config_values)
+        reward = calculate_reward_function(
+            self.nodes_post_pol,
+            self.nodes_post_blue,
+            self.nodes_reference,
+            self.green_iers,
+            self.red_iers,
+            self.step_count,
+            self.config_values,
+        )
         # print("Step reward: " + str(reward))
         self.total_reward += reward
         if self.step_count == self.episode_steps:
             self.average_reward = self.total_reward / self.step_count
             if self.config_values.session_type == "EVALUATION":
-                # For evaluation, need to trigger the done value = True when step count is reached in order to prevent neverending episode
+                # For evaluation, need to trigger the done value = True when
+                # step count is reached in order to prevent neverending episode
                 done = True
             print("Average reward: " + str(self.average_reward))
         # Load the reward into the transaction
@@ -373,38 +451,33 @@ class Primaite(Env):
         return self.env_obs, reward, done, self.step_info
 
     def __close__(self):
-        """
+        """Override close function."""
-        Override close function
-        """
-
         self.csv_file.close()
         self.config_file.close()
 
     def init_acl(self):
-        """
+        """Initialise the Access Control List."""
-        Initialise the Access Control List
-        """
-
         self.acl.remove_all_rules()
 
     def output_link_status(self):
-        """
+        """Output the link status of all links to the console."""
-        Output the link status of all links to the console
-        """
-
         for link_key, link_value in self.links.items():
             print("Link ID: " + link_value.get_id())
             for protocol in link_value.get_protocol_list():
-                print("    Protocol: " + protocol.get_name().name + ", Load: " + str(protocol.get_load()))
+                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
         """
-
         # At the moment, actions are only affecting nodes
         if self.action_type == ACTION_TYPE.NODE:
             self.apply_actions_to_nodes(_action)
@@ -413,12 +486,11 @@ class Primaite(Env):
 
     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
         """
-
         node_id = _action[0]
         node_property = _action[1]
         property_action = _action[2]
@@ -427,7 +499,7 @@ class Primaite(Env):
         # Check that the action is requesting a valid node
         try:
             node = self.nodes[str(node_id)]
-        except:
+        except Exception:
             return
 
         if node_property == 0:
@@ -472,7 +544,9 @@ class Primaite(Env):
                     return
                 elif property_action == 1:
                     # Patch (valid action if it's good or compromised)
-                    node.set_service_state(self.services_list[service_index], SOFTWARE_STATE.PATCHING)
+                    node.set_service_state(
+                        self.services_list[service_index], SOFTWARE_STATE.PATCHING
+                    )
             else:
                 # Node is not of Service Type
                 return
@@ -488,7 +562,10 @@ class Primaite(Env):
                 elif property_action == 2:
                     # Repair
                     # You cannot repair a destroyed file system - it needs restoring
-                    if node.get_file_system_state_actual() != FILE_SYSTEM_STATE.DESTROYED:
+                    if (
+                        node.get_file_system_state_actual()
+                        != FILE_SYSTEM_STATE.DESTROYED
+                    ):
                         node.set_file_system_state(FILE_SYSTEM_STATE.REPAIRING)
                 elif property_action == 3:
                     # Restore
@@ -501,12 +578,11 @@ class Primaite(Env):
 
     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
         """
-
         action_decision = _action[0]
         action_permission = _action[1]
         action_source_ip = _action[2]
@@ -556,18 +632,31 @@ class Primaite(Env):
             # Now add or remove
             if action_decision == 1:
                 # Add the rule
-                self.acl.add_rule(acl_rule_permission, acl_rule_source, acl_rule_destination, acl_rule_protocol, acl_rule_port)
+                self.acl.add_rule(
+                    acl_rule_permission,
+                    acl_rule_source,
+                    acl_rule_destination,
+                    acl_rule_protocol,
+                    acl_rule_port,
+                )
             elif action_decision == 2:
                 # Remove the rule
-                self.acl.remove_rule(acl_rule_permission, acl_rule_source, acl_rule_destination, acl_rule_protocol, acl_rule_port)
+                self.acl.remove_rule(
+                    acl_rule_permission,
+                    acl_rule_source,
+                    acl_rule_destination,
+                    acl_rule_protocol,
+                    acl_rule_port,
+                )
             else:
                 return
 
     def apply_time_based_updates(self):
         """
-        Updates anything that needs to count down and then change state (e.g. reset / patching status)
+        Updates anything that needs to count down and then change state.
-        """
 
+        e.g. reset / patching status
+        """
         for node_key, node in self.nodes.items():
             if node.get_state() == HARDWARE_STATE.RESETTING:
                 node.update_resetting_status()
@@ -605,10 +694,7 @@ class Primaite(Env):
                 pass
 
     def update_environent_obs(self):
-        """
+        """Updates the observation space based on the node and link status."""
-        # Updates the observation space based on the node and link status
-        """  
-        
         item_index = 0
 
         # Do nodes first
@@ -617,7 +703,9 @@ class Primaite(Env):
             self.env_obs[item_index][1] = node.get_state().value
             if isinstance(node, ActiveNode) or isinstance(node, ServiceNode):
                 self.env_obs[item_index][2] = node.get_os_state().value
-                self.env_obs[item_index][3] = node.get_file_system_state_observed().value
+                self.env_obs[item_index][
+                    3
+                ] = node.get_file_system_state_observed().value
             else:
                 self.env_obs[item_index][2] = 0
                 self.env_obs[item_index][3] = 0
@@ -625,7 +713,9 @@ class Primaite(Env):
             if isinstance(node, ServiceNode):
                 for service in self.services_list:
                     if node.has_service(service):
-                        self.env_obs[item_index][service_index] = node.get_service_state(service).value
+                        self.env_obs[item_index][
+                            service_index
+                        ] = node.get_service_state(service).value
                     else:
                         self.env_obs[item_index][service_index] = 0
                     service_index += 1
@@ -650,10 +740,7 @@ class Primaite(Env):
             item_index += 1
 
     def load_config(self):
-        """
+        """Loads config data in order to build the environment configuration."""
-        # Loads config data in order to build the environment configuration
-        """ 
-
         for item in self.config_data:
             if item["itemType"] == "NODE":
                 # Create a node
@@ -697,12 +784,11 @@ class Primaite(Env):
 
     def create_node(self, item):
         """
-        Creates a node from config data
+        Creates a node from config data.
 
         Args:
             item: A config data item
         """
-
         # All nodes have these parameters
         node_id = item["id"]
         node_name = item["name"]
@@ -712,19 +798,46 @@ class Primaite(Env):
         node_hardware_state = HARDWARE_STATE[item["hardwareState"]]
 
         if node_base_type == "PASSIVE":
-            node = PassiveNode(node_id, node_name, node_type, node_priority, node_hardware_state, self.config_values)
+            node = PassiveNode(
+                node_id,
+                node_name,
+                node_type,
+                node_priority,
+                node_hardware_state,
+                self.config_values,
+            )
         elif node_base_type == "ACTIVE":
             # Active nodes have IP address, operating system state and file system state
             node_ip_address = item["ipAddress"]
             node_software_state = SOFTWARE_STATE[item["softwareState"]]
             node_file_system_state = FILE_SYSTEM_STATE[item["fileSystemState"]]
-            node = ActiveNode(node_id, node_name, node_type, node_priority, node_hardware_state, node_ip_address, node_software_state, node_file_system_state, self.config_values)
+            node = ActiveNode(
+                node_id,
+                node_name,
+                node_type,
+                node_priority,
+                node_hardware_state,
+                node_ip_address,
+                node_software_state,
+                node_file_system_state,
+                self.config_values,
+            )
         elif node_base_type == "SERVICE":
             # Service nodes have IP address, operating system state, file system state and list of services
             node_ip_address = item["ipAddress"]
             node_software_state = SOFTWARE_STATE[item["softwareState"]]
             node_file_system_state = FILE_SYSTEM_STATE[item["fileSystemState"]]
-            node = ServiceNode(node_id, node_name, node_type, node_priority, node_hardware_state, node_ip_address, node_software_state, node_file_system_state, self.config_values)
+            node = ServiceNode(
+                node_id,
+                node_name,
+                node_type,
+                node_priority,
+                node_hardware_state,
+                node_ip_address,
+                node_software_state,
+                node_file_system_state,
+                self.config_values,
+            )
             node_services = item["services"]
             for service in node_services:
                 service_protocol = service["name"]
@@ -752,12 +865,11 @@ class Primaite(Env):
 
     def create_link(self, item):
         """
-        Creates a link from config data
+        Creates a link from config data.
 
         Args:
             item: A config data item
         """
-
         link_id = item["id"]
         link_name = item["name"]
         link_bandwidth = item["bandwidth"]
@@ -771,7 +883,13 @@ class Primaite(Env):
         self.network.add_edge(source_node, dest_node, id=link_name)
 
         # Add link to link dictionary
-        self.links[link_name] = Link(link_id, link_bandwidth, source_node.get_name(), dest_node.get_name(), self.services_list)
+        self.links[link_name] = Link(
+            link_id,
+            link_bandwidth,
+            source_node.get_name(),
+            dest_node.get_name(),
+            self.services_list,
+        )
 
         # Reference
         source_node_ref = self.nodes_reference[link_source]
@@ -781,16 +899,21 @@ class Primaite(Env):
         self.network_reference.add_edge(source_node_ref, dest_node_ref, id=link_name)
 
         # Add link to link dictionary (reference)
-        self.links_reference[link_name] = Link(link_id, link_bandwidth, source_node_ref.get_name(), dest_node_ref.get_name(), self.services_list)
+        self.links_reference[link_name] = Link(
+            link_id,
+            link_bandwidth,
+            source_node_ref.get_name(),
+            dest_node_ref.get_name(),
+            self.services_list,
+        )
 
     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
         """
-
         ier_id = item["id"]
         ier_start_step = item["startStep"]
         ier_end_step = item["endStep"]
@@ -802,16 +925,25 @@ class Primaite(Env):
         ier_mission_criticality = item["missionCriticality"]
 
         # Create IER and add to green IER dictionary
-        self.green_iers[ier_id] = IER(ier_id, ier_start_step, ier_end_step, ier_load, ier_protocol, ier_port, ier_source, ier_destination, ier_mission_criticality)
+        self.green_iers[ier_id] = IER(
+            ier_id,
+            ier_start_step,
+            ier_end_step,
+            ier_load,
+            ier_protocol,
+            ier_port,
+            ier_source,
+            ier_destination,
+            ier_mission_criticality,
+        )
 
     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
         """
-
         ier_id = item["id"]
         ier_start_step = item["startStep"]
         ier_end_step = item["endStep"]
@@ -823,16 +955,25 @@ class Primaite(Env):
         ier_mission_criticality = item["missionCriticality"]
 
         # Create IER and add to red IER dictionary
-        self.red_iers[ier_id] = IER(ier_id, ier_start_step, ier_end_step, ier_load, ier_protocol, ier_port, ier_source, ier_destination, ier_mission_criticality)
+        self.red_iers[ier_id] = IER(
+            ier_id,
+            ier_start_step,
+            ier_end_step,
+            ier_load,
+            ier_protocol,
+            ier_port,
+            ier_source,
+            ier_destination,
+            ier_mission_criticality,
+        )
 
     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
         """
-
         pol_id = item["id"]
         pol_start_step = item["startStep"]
         pol_end_step = item["endStep"]
@@ -850,16 +991,23 @@ class Primaite(Env):
             pol_protocol = item["protocol"]
             pol_state = SOFTWARE_STATE[item["state"]]
 
-        self.node_pol[pol_id] = NodeStateInstructionGreen(pol_id, pol_start_step, pol_end_step, pol_node, pol_type, pol_protocol, pol_state)
+        self.node_pol[pol_id] = NodeStateInstructionGreen(
+            pol_id,
+            pol_start_step,
+            pol_end_step,
+            pol_node,
+            pol_type,
+            pol_protocol,
+            pol_state,
+        )
 
     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
         """
-
         pol_id = item["id"]
         pol_start_step = item["startStep"]
         pol_end_step = item["endStep"]
@@ -880,32 +1028,48 @@ class Primaite(Env):
         pol_source_node_service = item["sourceNodeService"]
         pol_source_node_service_state = item["sourceNodeServiceState"]
 
-        self.red_node_pol[pol_id] = NodeStateInstructionRed(pol_id, pol_start_step, pol_end_step, pol_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[pol_id] = NodeStateInstructionRed(
+            pol_id,
+            pol_start_step,
+            pol_end_step,
+            pol_target_node_id,
+            pol_initiator,
+            pol_type,
+            pol_protocol,
+            pol_state,
+            pol_source_node_id,
+            pol_source_node_service,
+            pol_source_node_service_state,
+        )
 
     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
         """
-
         acl_rule_permission = item["permission"]
         acl_rule_source = item["source"]
         acl_rule_destination = item["destination"]
         acl_rule_protocol = item["protocol"]
         acl_rule_port = item["port"]
 
-        self.acl.add_rule(acl_rule_permission, acl_rule_source, acl_rule_destination, acl_rule_protocol, acl_rule_port)
+        self.acl.add_rule(
+            acl_rule_permission,
+            acl_rule_source,
+            acl_rule_destination,
+            acl_rule_protocol,
+            acl_rule_port,
+        )
 
     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
         """
-
         service_list = services["serviceList"]
 
         for service in service_list:
@@ -917,12 +1081,11 @@ class Primaite(Env):
 
     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
         """
-
         ports_list = ports["portsList"]
 
         for port in ports_list:
@@ -934,31 +1097,30 @@ class Primaite(Env):
 
     def get_action_info(self, action_info):
         """
-        Extracts action_info
+        Extracts action_info.
 
         Args:
             item: A config data item representing action info
         """
-
         self.action_type = ACTION_TYPE[action_info["type"]]
 
-
     def get_steps_info(self, steps_info):
         """
-        Extracts steps_info
+        Extracts steps_info.
 
         Args:
             item: A config data item representing steps info
         """
-
         self.episode_steps = int(steps_info["steps"])
         logging.info("Training episodes have " + str(self.episode_steps) + " steps")
 
     def reset_environment(self):
         """
-        # Resets environment using config data config data in order to build the environment configuration
+        # Resets environment.
-        """ 
 
+        Uses config data config data in order to build the environment
+        configuration.
+        """
         for item in self.config_data:
             if item["itemType"] == "NODE":
                 # Reset a node's state (normal and reference)
@@ -970,7 +1132,6 @@ class Primaite(Env):
                 # Do nothing (bad formatting or not relevant to reset)
                 pass
 
-
         # Reset the IER status so they are not running initially
         # Green IERs
         for ier_key, ier_value in self.green_iers.items():
@@ -981,12 +1142,11 @@ class Primaite(Env):
 
     def reset_node(self, item):
         """
-        Resets the statuses of a node
+        Resets the statuses of a node.
 
         Args:
             item: A config data item
         """
-
         # All nodes have these parameters
         node_id = item["id"]
         node_base_type = item["baseType"]
@@ -1027,10 +1187,3 @@ class Primaite(Env):
         else:
             # Bad formatting
             pass
-
-
-
-
-
-
-

src/primaite/environment/reward.py

@@ -1,15 +1,21 @@
 # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
-"""
+"""Implements reward function."""
-Implements reward function
+from primaite.common.enums import FILE_SYSTEM_STATE, HARDWARE_STATE, SOFTWARE_STATE
-"""
-
-from primaite.common.enums import *
 from primaite.nodes.active_node import ActiveNode
 from primaite.nodes.service_node import ServiceNode
 
-def calculate_reward_function(initial_nodes, final_nodes, reference_nodes, green_iers, red_iers, step_count, config_values):
+
+def calculate_reward_function(
+    initial_nodes,
+    final_nodes,
+    reference_nodes,
+    green_iers,
+    red_iers,
+    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
@@ -20,7 +26,6 @@ def calculate_reward_function(initial_nodes, final_nodes, reference_nodes, green
         step_count: current step
         config_values: Config values
     """
-
     reward_value = 0
 
     # For each node, compare operating state, o/s operating state, service states
@@ -29,19 +34,27 @@ def calculate_reward_function(initial_nodes, final_nodes, reference_nodes, green
         reference_node = reference_nodes[node_key]
 
         # Operating State
-        reward_value += score_node_operating_state(final_node, initial_node, reference_node, config_values)
+        reward_value += score_node_operating_state(
+            final_node, initial_node, reference_node, config_values
+        )
 
         # Operating System State
-        if (isinstance(final_node, ActiveNode) or isinstance(final_node, ServiceNode)):
+        if isinstance(final_node, ActiveNode) or isinstance(final_node, ServiceNode):
-            reward_value += score_node_os_state(final_node, initial_node, reference_node, config_values)
+            reward_value += score_node_os_state(
+                final_node, initial_node, reference_node, config_values
+            )
 
         # Service State
-        if (isinstance(final_node, ServiceNode)):
+        if isinstance(final_node, ServiceNode):
-            reward_value += score_node_service_state(final_node, initial_node, reference_node, config_values)
+            reward_value += score_node_service_state(
+                final_node, initial_node, reference_node, config_values
+            )
 
         # File System State
         if isinstance(final_node, ActiveNode):
-            reward_value += score_node_file_system(final_node, initial_node, reference_node, config_values)
+            reward_value += score_node_file_system(
+                final_node, initial_node, reference_node, config_values
+            )
 
     # Go through each red IER - penalise if it is running
     for ier_key, ier_value in red_iers.items():
@@ -57,14 +70,17 @@ def calculate_reward_function(initial_nodes, final_nodes, reference_nodes, green
         stop_step = ier_value.get_end_step()
         if step_count >= start_step and step_count <= stop_step:
             if not ier_value.get_is_running():
-                reward_value += config_values.green_ier_blocked * ier_value.get_mission_criticality()
+                reward_value += (
+                    config_values.green_ier_blocked
+                    * ier_value.get_mission_criticality()
+                )
 
     return reward_value
 
 
 def score_node_operating_state(final_node, initial_node, reference_node, config_values):
     """
-    Calculates score relating to the operating state of a node
+    Calculates score relating to the operating state of a node.
 
     Args:
         final_node: The node after red and blue agents take effect
@@ -72,7 +88,6 @@ def score_node_operating_state(final_node, initial_node, reference_node, config_
         reference_node: The node if there had been no red or blue effect
         config_values: Config values
     """
-    
     score = 0
     final_node_operating_state = final_node.get_state()
     initial_node_operating_state = initial_node.get_state()
@@ -112,9 +127,10 @@ def score_node_operating_state(final_node, initial_node, reference_node, config_
 
     return score
 
+
 def score_node_os_state(final_node, initial_node, reference_node, config_values):
     """
-    Calculates score relating to the operating system state of a node
+    Calculates score relating to the operating system state of a node.
 
     Args:
         final_node: The node after red and blue agents take effect
@@ -122,7 +138,6 @@ def score_node_os_state(final_node, initial_node, reference_node, config_values)
         reference_node: The node if there had been no red or blue effect
         config_values: Config values
     """
-
     score = 0
     final_node_os_state = final_node.get_os_state()
     initial_node_os_state = initial_node.get_os_state()
@@ -164,9 +179,10 @@ def score_node_os_state(final_node, initial_node, reference_node, config_values)
 
     return score
 
+
 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
@@ -174,7 +190,6 @@ def score_node_service_state(final_node, initial_node, reference_node, config_va
         reference_node: The node if there had been no red or blue effect
         config_values: Config values
     """
-
     score = 0
     final_node_services = final_node.get_services()
     initial_node_services = initial_node.get_services()
@@ -237,16 +252,16 @@ def score_node_service_state(final_node, initial_node, reference_node, config_va
 
     return score
 
+
 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
         initial_node: The node before red and blue agents take effect
         reference_node: The node if there had been no red or blue effect
     """
-
     score = 0
     final_node_file_system_state = final_node.get_file_system_state_actual()
     initial_node_file_system_state = initial_node.get_file_system_state_actual()

src/primaite/links/link.py

@@ -1,19 +1,15 @@
 # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
-"""
+"""The link class."""
-The link class
-"""
 
 from primaite.common.protocol import Protocol
-from primaite.common.enums import *
+
 
 class Link(object):
-    """
+    """Link class."""
-    Link class
-    """
 
     def __init__(self, _id, _bandwidth, _source_node_name, _dest_node_name, _services):
         """
-        Init
+        Init.
 
         Args:
             _id: The IER id
@@ -22,7 +18,6 @@ class Link(object):
             _dest_node_name: The name of the destination node
             _protocols: The protocols to add to the link
         """
-
         self.id = _id
         self.bandwidth = _bandwidth
         self.source_node_name = _source_node_name
@@ -35,72 +30,65 @@ class Link(object):
 
     def add_protocol(self, _protocol):
         """
-        Adds a new protocol to the list of protocols on this link
+        Adds a new protocol to the list of protocols on this link.
 
         Args:
             _protocol: The protocol to be added (enum)
         """
-
         self.protocol_list.append(Protocol(_protocol))
 
     def get_id(self):
         """
-        Gets link ID
+        Gets link ID.
 
         Returns:
              Link ID
         """
-
         return self.id
 
     def get_source_node_name(self):
         """
-        Gets source node name
+        Gets source node name.
 
         Returns:
              Source node name
         """
-
         return self.source_node_name
 
     def get_dest_node_name(self):
         """
-        Gets destination node name
+        Gets destination node name.
 
         Returns:
              Destination node name
         """
-
         return self.dest_node_name
 
     def get_bandwidth(self):
         """
-        Gets bandwidth of link
+        Gets bandwidth of link.
 
         Returns:
              Link bandwidth (bps)
         """
-
         return self.bandwidth
 
     def get_protocol_list(self):
         """
-        Gets list of protocols on this link
+        Gets list of protocols on this link.
 
         Returns:
              List of protocols on this link
         """
-
         return self.protocol_list
 
     def get_current_load(self):
         """
-        Gets current total load on this link
+        Gets current total load on this link.
 
         Returns:
              Total load on this link (bps)
         """
-
         total_load = 0
         for protocol in self.protocol_list:
             total_load += protocol.get_load()
@@ -108,13 +96,12 @@ class Link(object):
 
     def add_protocol_load(self, _protocol, _load):
         """
-        Adds a loading to a protocol on this link
+        Adds a loading to a protocol on this link.
 
         Args:
             _protocol: The protocol to load
             _load: The amount to load (bps)
         """
-
         for protocol in self.protocol_list:
             if protocol.get_name() == _protocol:
                 protocol.add_load(_load)
@@ -122,11 +109,6 @@ class Link(object):
                 pass
 
     def clear_traffic(self):
-        """
+        """Clears all traffic on this link."""
-        Clears all traffic on this link
-        """
-
         for protocol in self.protocol_list:
             protocol.clear_load()
-
-

src/primaite/main.py

@@ -1,37 +1,31 @@
 # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
 """
-Primaite - main (harness) module
+Primaite - main (harness) module.
 
 Coding Standards: PEP 8
 """
 
-from sys import exc_info
-import time
-import yaml
-import os.path
 import logging
+import os.path
+import time
 from datetime import datetime
 
+import yaml
+from stable_baselines3 import A2C, PPO
+from stable_baselines3.common.evaluation import evaluate_policy
+from stable_baselines3.ppo import MlpPolicy as PPOMlp
+
+from primaite.common.config_values_main import config_values_main
 from primaite.environment.primaite_env import Primaite
 from primaite.transactions.transactions_to_file import write_transaction_to_file
-from primaite.common.config_values_main import config_values_main
 
-from stable_baselines3 import PPO
+# FUNCTIONS #
-from stable_baselines3.ppo import MlpPolicy as PPOMlp
-from stable_baselines3 import A2C
-from stable_baselines3.common.env_checker import check_env
-from stable_baselines3.common.evaluation import evaluate_policy
 
-################################# FUNCTIONS ######################################
 
 def run_generic():
-    """
+    """Run against a generic agent."""
-    Run against a generic agent
-    """
-
     for episode in range(0, config_values.num_episodes):
         for step in range(0, config_values.num_steps):
-
             # Send the observation space to the agent to get an action
             # TEMP - random action for now
             # action = env.blue_agent_action(obs)
@@ -54,15 +48,20 @@ def run_generic():
 
 
 def run_stable_baselines3_ppo():
-    """
+    """Run against a stable_baselines3 PPO agent."""
-    Run against a stable_baselines3 PPO agent
-    """ 
-
     if config_values.load_agent == True:
         try:
-            agent = PPO.load(config_values.agent_load_file, env, verbose=0, n_steps=config_values.num_steps)
+            agent = PPO.load(
-        except:
+                config_values.agent_load_file,
-            print("ERROR: Could not load agent at location: " + config_values.agent_load_file)
+                env,
+                verbose=0,
+                n_steps=config_values.num_steps,
+            )
+        except Exception:
+            print(
+                "ERROR: Could not load agent at location: "
+                + config_values.agent_load_file
+            )
             logging.error("Could not load agent")
             logging.error("Exception occured", exc_info=True)
     else:
@@ -83,16 +82,22 @@ def run_stable_baselines3_ppo():
 
     env.close()
 
-def run_stable_baselines3_a2c():
-    """
-    Run against a stable_baselines3 A2C agent
-    """
 
+def run_stable_baselines3_a2c():
+    """Run against a stable_baselines3 A2C agent."""
     if config_values.load_agent == True:
         try:
-            agent = A2C.load(config_values.agent_load_file, env, verbose=0, n_steps=config_values.num_steps)
+            agent = A2C.load(
-        except:
+                config_values.agent_load_file,
-            print("ERROR: Could not load agent at location: " + config_values.agent_load_file)
+                env,
+                verbose=0,
+                n_steps=config_values.num_steps,
+            )
+        except Exception:
+            print(
+                "ERROR: Could not load agent at location: "
+                + config_values.agent_load_file
+            )
             logging.error("Could not load agent")
             logging.error("Exception occured", exc_info=True)
     else:
@@ -113,133 +118,203 @@ def run_stable_baselines3_a2c():
 
     env.close()
 
-def save_agent(_agent):
-    """
-    Persist an agent (only works for stable baselines3 agents at present)
-    """
 
+def save_agent(_agent):
+    """Persist an agent (only works for stable baselines3 agents at present)."""
     now = datetime.now()  # current date and time
     time = now.strftime("%Y%m%d_%H%M%S")
 
     try:
-        path = 'outputs/agents/'
+        path = "outputs/agents/"
         is_dir = os.path.isdir(path)
         if not is_dir:
             os.makedirs(path)
         filename = "outputs/agents/agent_saved_" + time
         _agent.save(filename)
         logging.info("Trained agent saved as " + filename)
-    except Exception as e:
+    except Exception:
         logging.error("Could not save agent")
         logging.error("Exception occured", exc_info=True)
 
-def configure_logging():
-    """
-    Configures logging
-    """
 
+def configure_logging():
+    """Configures logging."""
     try:
         now = datetime.now()  # current date and time
         time = now.strftime("%Y%m%d_%H%M%S")
         filename = "logs/app_" + time + ".log"
-        path = 'logs/'
+        path = "logs/"
         is_dir = os.path.isdir(path)
         if not is_dir:
             os.makedirs(path)
-        logging.basicConfig(filename=filename, filemode='w', format='%(asctime)s - %(levelname)s - %(message)s', datefmt='%d-%b-%y %H:%M:%S', level=logging.INFO)
+        logging.basicConfig(
-    except:
+            filename=filename,
+            filemode="w",
+            format="%(asctime)s - %(levelname)s - %(message)s",
+            datefmt="%d-%b-%y %H:%M:%S",
+            level=logging.INFO,
+        )
+    except Exception:
         print("ERROR: Could not start logging")
 
-def load_config_values():
-    """
-    Loads the config values from the main config file into a config object
-    """
 
+def load_config_values():
+    """Loads the config values from the main config file into a config object."""
     try:
         # Generic
-        config_values.agent_identifier = config_data['agentIdentifier']           
+        config_values.agent_identifier = config_data["agentIdentifier"]
-        config_values.num_episodes = int(config_data['numEpisodes'])                                
+        config_values.num_episodes = int(config_data["numEpisodes"])
-        config_values.time_delay = int(config_data['timeDelay'])                  
+        config_values.time_delay = int(config_data["timeDelay"])
-        config_values.config_filename_use_case = config_data['configFilename']   
+        config_values.config_filename_use_case = config_data["configFilename"]
-        config_values.session_type = config_data['sessionType']
+        config_values.session_type = config_data["sessionType"]
-        config_values.load_agent = bool(config_data['loadAgent'])
+        config_values.load_agent = bool(config_data["loadAgent"])
-        config_values.agent_load_file = config_data['agentLoadFile']
+        config_values.agent_load_file = config_data["agentLoadFile"]
         # Environment
-        config_values.observation_space_high_value = int(config_data['observationSpaceHighValue'])
+        config_values.observation_space_high_value = int(
+            config_data["observationSpaceHighValue"]
+        )
         # Reward values
         # Generic
-        config_values.all_ok = int(config_data['allOk']) 
+        config_values.all_ok = int(config_data["allOk"])
         # Node Operating State
-        config_values.off_should_be_on = int(config_data['offShouldBeOn'])
+        config_values.off_should_be_on = int(config_data["offShouldBeOn"])
-        config_values.off_should_be_resetting = int(config_data['offShouldBeResetting'])
+        config_values.off_should_be_resetting = int(config_data["offShouldBeResetting"])
-        config_values.on_should_be_off = int(config_data['onShouldBeOff'])
+        config_values.on_should_be_off = int(config_data["onShouldBeOff"])
-        config_values.on_should_be_resetting = int(config_data['onShouldBeResetting'])
+        config_values.on_should_be_resetting = int(config_data["onShouldBeResetting"])
-        config_values.resetting_should_be_on = int(config_data['resettingShouldBeOn'])
+        config_values.resetting_should_be_on = int(config_data["resettingShouldBeOn"])
-        config_values.resetting_should_be_off = int(config_data['resettingShouldBeOff']) 
+        config_values.resetting_should_be_off = int(config_data["resettingShouldBeOff"])
-        config_values.resetting = int(config_data['resetting'])
+        config_values.resetting = int(config_data["resetting"])
         # Node O/S or Service State
-        config_values.good_should_be_patching = int(config_data['goodShouldBePatching'])
+        config_values.good_should_be_patching = int(config_data["goodShouldBePatching"])
-        config_values.good_should_be_compromised = int(config_data['goodShouldBeCompromised'])
+        config_values.good_should_be_compromised = int(
-        config_values.good_should_be_overwhelmed = int(config_data['goodShouldBeOverwhelmed'])
+            config_data["goodShouldBeCompromised"]
-        config_values.patching_should_be_good = int(config_data['patchingShouldBeGood'])
+        )
-        config_values.patching_should_be_compromised = int(config_data['patchingShouldBeCompromised'])
+        config_values.good_should_be_overwhelmed = int(
-        config_values.patching_should_be_overwhelmed = int(config_data['patchingShouldBeOverwhelmed'])
+            config_data["goodShouldBeOverwhelmed"]
-        config_values.patching = int(config_data['patching'])
+        )
-        config_values.compromised_should_be_good = int(config_data['compromisedShouldBeGood'])
+        config_values.patching_should_be_good = int(config_data["patchingShouldBeGood"])
-        config_values.compromised_should_be_patching = int(config_data['compromisedShouldBePatching'])
+        config_values.patching_should_be_compromised = int(
-        config_values.compromised_should_be_overwhelmed = int(config_data['compromisedShouldBeOverwhelmed'])
+            config_data["patchingShouldBeCompromised"]
-        config_values.compromised = int(config_data['compromised'])     
+        )
-        config_values.overwhelmed_should_be_good = int(config_data['overwhelmedShouldBeGood'])
+        config_values.patching_should_be_overwhelmed = int(
-        config_values.overwhelmed_should_be_patching = int(config_data['overwhelmedShouldBePatching'])
+            config_data["patchingShouldBeOverwhelmed"]
-        config_values.overwhelmed_should_be_compromised = int(config_data['overwhelmedShouldBeCompromised'])
+        )
-        config_values.overwhelmed = int(config_data['overwhelmed'])
+        config_values.patching = int(config_data["patching"])
+        config_values.compromised_should_be_good = int(
+            config_data["compromisedShouldBeGood"]
+        )
+        config_values.compromised_should_be_patching = int(
+            config_data["compromisedShouldBePatching"]
+        )
+        config_values.compromised_should_be_overwhelmed = int(
+            config_data["compromisedShouldBeOverwhelmed"]
+        )
+        config_values.compromised = int(config_data["compromised"])
+        config_values.overwhelmed_should_be_good = int(
+            config_data["overwhelmedShouldBeGood"]
+        )
+        config_values.overwhelmed_should_be_patching = int(
+            config_data["overwhelmedShouldBePatching"]
+        )
+        config_values.overwhelmed_should_be_compromised = int(
+            config_data["overwhelmedShouldBeCompromised"]
+        )
+        config_values.overwhelmed = int(config_data["overwhelmed"])
         # Node File System State
-        config_values.good_should_be_repairing = int(config_data['goodShouldBeRepairing'])
+        config_values.good_should_be_repairing = int(
-        config_values.good_should_be_restoring = int(config_data['goodShouldBeRestoring'])
+            config_data["goodShouldBeRepairing"]
-        config_values.good_should_be_corrupt = int(config_data['goodShouldBeCorrupt'])
+        )
-        config_values.good_should_be_destroyed = int(config_data['goodShouldBeDestroyed'])
+        config_values.good_should_be_restoring = int(
-        config_values.repairing_should_be_good = int(config_data['repairingShouldBeGood'])
+            config_data["goodShouldBeRestoring"]
-        config_values.repairing_should_be_restoring = int(config_data['repairingShouldBeRestoring'])
+        )
-        config_values.repairing_should_be_corrupt = int(config_data['repairingShouldBeCorrupt'])
+        config_values.good_should_be_corrupt = int(config_data["goodShouldBeCorrupt"])
-        config_values.repairing_should_be_destroyed = int(config_data['repairingShouldBeDestroyed'])   
+        config_values.good_should_be_destroyed = int(
-        config_values.repairing = int(config_data['repairing'])
+            config_data["goodShouldBeDestroyed"]
-        config_values.restoring_should_be_good = int(config_data['restoringShouldBeGood'])
+        )
-        config_values.restoring_should_be_repairing = int(config_data['restoringShouldBeRepairing'])
+        config_values.repairing_should_be_good = int(
-        config_values.restoring_should_be_corrupt = int(config_data['restoringShouldBeCorrupt'])     
+            config_data["repairingShouldBeGood"]
-        config_values.restoring_should_be_destroyed = int(config_data['restoringShouldBeDestroyed'])
+        )
-        config_values.restoring = int(config_data['restoring'])
+        config_values.repairing_should_be_restoring = int(
-        config_values.corrupt_should_be_good = int(config_data['corruptShouldBeGood'])
+            config_data["repairingShouldBeRestoring"]
-        config_values.corrupt_should_be_repairing = int(config_data['corruptShouldBeRepairing'])
+        )
-        config_values.corrupt_should_be_restoring = int(config_data['corruptShouldBeRestoring'])
+        config_values.repairing_should_be_corrupt = int(
-        config_values.corrupt_should_be_destroyed = int(config_data['corruptShouldBeDestroyed'])
+            config_data["repairingShouldBeCorrupt"]
-        config_values.corrupt = int(config_data['corrupt'])
+        )
-        config_values.destroyed_should_be_good = int(config_data['destroyedShouldBeGood'])
+        config_values.repairing_should_be_destroyed = int(
-        config_values.destroyed_should_be_repairing = int(config_data['destroyedShouldBeRepairing'])
+            config_data["repairingShouldBeDestroyed"]
-        config_values.destroyed_should_be_restoring = int(config_data['destroyedShouldBeRestoring'])
+        )
-        config_values.destroyed_should_be_corrupt = int(config_data['destroyedShouldBeCorrupt'])
+        config_values.repairing = int(config_data["repairing"])
-        config_values.destroyed = int(config_data['destroyed'])
+        config_values.restoring_should_be_good = int(
-        config_values.scanning = int(config_data['scanning'])
+            config_data["restoringShouldBeGood"]
+        )
+        config_values.restoring_should_be_repairing = int(
+            config_data["restoringShouldBeRepairing"]
+        )
+        config_values.restoring_should_be_corrupt = int(
+            config_data["restoringShouldBeCorrupt"]
+        )
+        config_values.restoring_should_be_destroyed = int(
+            config_data["restoringShouldBeDestroyed"]
+        )
+        config_values.restoring = int(config_data["restoring"])
+        config_values.corrupt_should_be_good = int(config_data["corruptShouldBeGood"])
+        config_values.corrupt_should_be_repairing = int(
+            config_data["corruptShouldBeRepairing"]
+        )
+        config_values.corrupt_should_be_restoring = int(
+            config_data["corruptShouldBeRestoring"]
+        )
+        config_values.corrupt_should_be_destroyed = int(
+            config_data["corruptShouldBeDestroyed"]
+        )
+        config_values.corrupt = int(config_data["corrupt"])
+        config_values.destroyed_should_be_good = int(
+            config_data["destroyedShouldBeGood"]
+        )
+        config_values.destroyed_should_be_repairing = int(
+            config_data["destroyedShouldBeRepairing"]
+        )
+        config_values.destroyed_should_be_restoring = int(
+            config_data["destroyedShouldBeRestoring"]
+        )
+        config_values.destroyed_should_be_corrupt = int(
+            config_data["destroyedShouldBeCorrupt"]
+        )
+        config_values.destroyed = int(config_data["destroyed"])
+        config_values.scanning = int(config_data["scanning"])
         # IER status
-        config_values.red_ier_running = int(config_data['redIerRunning'])
+        config_values.red_ier_running = int(config_data["redIerRunning"])
-        config_values.green_ier_blocked = int(config_data['greenIerBlocked'])
+        config_values.green_ier_blocked = int(config_data["greenIerBlocked"])
         # Patching / Reset durations
-        config_values.os_patching_duration = int(config_data['osPatchingDuration'])                         
+        config_values.os_patching_duration = int(config_data["osPatchingDuration"])
-        config_values.node_reset_duration = int(config_data['nodeResetDuration'])                           
+        config_values.node_reset_duration = int(config_data["nodeResetDuration"])
-        config_values.service_patching_duration = int(config_data['servicePatchingDuration']) 
+        config_values.service_patching_duration = int(
-        config_values.file_system_repairing_limit = int(config_data['fileSystemRepairingLimit']) 
+            config_data["servicePatchingDuration"]
-        config_values.file_system_restoring_limit = int(config_data['fileSystemRestoringLimit']) 
+        )
-        config_values.file_system_scanning_limit = int(config_data['fileSystemScanningLimit']) 
+        config_values.file_system_repairing_limit = int(
+            config_data["fileSystemRepairingLimit"]
+        )
+        config_values.file_system_restoring_limit = int(
+            config_data["fileSystemRestoringLimit"]
+        )
+        config_values.file_system_scanning_limit = int(
+            config_data["fileSystemScanningLimit"]
+        )
 
         logging.info("Training agent: " + config_values.agent_identifier)
-        logging.info("Training environment config: " + config_values.config_filename_use_case)
+        logging.info(
-        logging.info("Training cycle has " + str(config_values.num_episodes) + " episodes")
+            "Training environment config: " + config_values.config_filename_use_case
+        )
+        logging.info(
+            "Training cycle has " + str(config_values.num_episodes) + " episodes"
+        )
 
-    except Exception as e:
+    except Exception:
         logging.error("Could not save load config data")
         logging.error("Exception occured", exc_info=True)
 
 
-################################# MAIN PROCESS ############################################
+# MAIN PROCESS #
 
 # Starting point
 
@@ -257,7 +332,7 @@ try:
     config_values = config_values_main()
     # Load in config data
     load_config_values()
-except Exception as e:
+except Exception:
     logging.error("Could not load main config")
     logging.error("Exception occured", exc_info=True)
 
@@ -275,7 +350,7 @@ transaction_list = []
 try:
     env = Primaite(config_values, transaction_list)
     logging.info("PrimAITE environment created")
-except Exception as e:
+except Exception:
     logging.error("Could not create PrimAITE environment")
     logging.error("Exception occured", exc_info=True)
 
@@ -302,11 +377,3 @@ config_file_main.close
 
 print("Finished")
 logging.info("Finished")
-
-
-
-
-
-
-
-

src/primaite/nodes/active_node.py

@@ -1,19 +1,26 @@
 # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
-"""
+"""An Active Node (i.e. not an actuator)."""
-An Active Node (i.e. not an actuator)
+from primaite.common.enums import FILE_SYSTEM_STATE, SOFTWARE_STATE
-"""
-
 from primaite.nodes.node import Node
-from primaite.common.enums import *
+
 
 class ActiveNode(Node):
-    """
+    """Active Node class."""
-    Active Node class
-    """
 
-    def __init__(self, _id, _name, _type, _priority, _state, _ip_address, _os_state, _file_system_state, _config_values):
+    def __init__(
+        self,
+        _id,
+        _name,
+        _type,
+        _priority,
+        _state,
+        _ip_address,
+        _os_state,
+        _file_system_state,
+        _config_values,
+    ):
         """
-        Init
+        Init.
 
         Args:
             _id: The node ID
@@ -26,7 +33,6 @@ class ActiveNode(Node):
             _file_system_state: The node file system state
             _config_values: The config values
         """
-
         super().__init__(_id, _name, _type, _priority, _state, _config_values)
         self.ip_address = _ip_address
         # Related to O/S
@@ -39,20 +45,18 @@ class ActiveNode(Node):
         self.file_system_scanning_count = 0
         self.file_system_action_count = 0
 
-
     def set_ip_address(self, _ip_address):
         """
-        Sets IP address
+        Sets IP address.
 
         Args:
             _ip_address: IP address
         """
-
         self.ip_address = _ip_address
 
     def get_ip_address(self):
         """
-        Gets IP address
+        Gets IP address.
 
         Returns:
              IP address
@@ -61,24 +65,22 @@ class ActiveNode(Node):
 
     def set_os_state(self, _os_state):
         """
-        Sets operating system state
+        Sets operating system state.
 
         Args:
             _os_state: Operating system state
         """
-
         self.os_state = _os_state
         if _os_state == SOFTWARE_STATE.PATCHING:
             self.patching_count = self.config_values.os_patching_duration
 
     def set_os_state_if_not_compromised(self, _os_state):
         """
-        Sets operating system state if the node is not compromised
+        Sets operating system state if the node is not compromised.
 
         Args:
             _os_state: Operating system state
         """
-
         if self.os_state != SOFTWARE_STATE.COMPROMISED:
             self.os_state = _os_state
             if _os_state == SOFTWARE_STATE.PATCHING:
@@ -86,19 +88,15 @@ class ActiveNode(Node):
 
     def get_os_state(self):
         """
-        Gets operating system state
+        Gets operating system state.
 
         Returns:
              Operating system state
         """
-
         return self.os_state
 
     def update_os_patching_status(self):
-        """
+        """Updates operating system status based on patching cycle."""
-        Updates operating system status based on patching cycle
-        """
-
         self.patching_count -= 1
         if self.patching_count <= 0:
             self.patching_count = 0
@@ -106,87 +104,88 @@ class ActiveNode(Node):
 
     def set_file_system_state(self, _file_system_state):
         """
-        Sets the file system state (actual and observed)
+        Sets the file system state (actual and observed).
 
         Args:
             _file_system_state: File system state
         """
-
         self.file_system_state_actual = _file_system_state
 
         if _file_system_state == FILE_SYSTEM_STATE.REPAIRING:
-            self.file_system_action_count = self.config_values.file_system_repairing_limit
+            self.file_system_action_count = (
+                self.config_values.file_system_repairing_limit
+            )
             self.file_system_state_observed = FILE_SYSTEM_STATE.REPAIRING
         elif _file_system_state == FILE_SYSTEM_STATE.RESTORING:
-            self.file_system_action_count = self.config_values.file_system_restoring_limit
+            self.file_system_action_count = (
+                self.config_values.file_system_restoring_limit
+            )
             self.file_system_state_observed = FILE_SYSTEM_STATE.RESTORING
         elif _file_system_state == FILE_SYSTEM_STATE.GOOD:
             self.file_system_state_observed = FILE_SYSTEM_STATE.GOOD
 
     def set_file_system_state_if_not_compromised(self, _file_system_state):
         """
-        Sets the file system state (actual and observed) if not in a compromised state
+        Sets the file system state (actual and observed) if not in a compromised state.
+
         Use for green PoL to prevent it overturning a compromised state
 
         Args:
             _file_system_state: File system state
         """
-
+        if (
-        if self.file_system_state_actual != FILE_SYSTEM_STATE.CORRUPT and self.file_system_state_actual != FILE_SYSTEM_STATE.DESTROYED:
+            self.file_system_state_actual != FILE_SYSTEM_STATE.CORRUPT
+            and self.file_system_state_actual != FILE_SYSTEM_STATE.DESTROYED
+        ):
             self.file_system_state_actual = _file_system_state
 
             if _file_system_state == FILE_SYSTEM_STATE.REPAIRING:
-                self.file_system_action_count = self.config_values.file_system_repairing_limit
+                self.file_system_action_count = (
+                    self.config_values.file_system_repairing_limit
+                )
                 self.file_system_state_observed = FILE_SYSTEM_STATE.REPAIRING
             elif _file_system_state == FILE_SYSTEM_STATE.RESTORING:
-                self.file_system_action_count = self.config_values.file_system_restoring_limit
+                self.file_system_action_count = (
+                    self.config_values.file_system_restoring_limit
+                )
                 self.file_system_state_observed = FILE_SYSTEM_STATE.RESTORING
             elif _file_system_state == FILE_SYSTEM_STATE.GOOD:
                 self.file_system_state_observed = FILE_SYSTEM_STATE.GOOD
 
     def get_file_system_state_actual(self):
         """
-        Gets file system state (actual)
+        Gets file system state (actual).
 
         Returns:
              File system state (actual)
         """
-
         return self.file_system_state_actual
 
     def get_file_system_state_observed(self):
         """
-        Gets file system state (observed)
+        Gets file system state (observed).
 
         Returns:
              File system state (observed)
         """
-
         return self.file_system_state_observed
 
     def start_file_system_scan(self):
-        """
+        """Starts a file system scan."""
-        Starts a file system scan
-        """
-
         self.file_system_scanning = True
         self.file_system_scanning_count = self.config_values.file_system_scanning_limit
 
     def is_scanning_file_system(self):
         """
-        Gets true/false on whether file system is being scanned
+        Gets true/false on whether file system is being scanned.
 
         Returns:
              True if file system is being scanned
         """
-
         return self.file_system_scanning
 
     def update_file_system_state(self):
-        """
+        """Updates file system status based on scanning / restore / repair cycle."""
-        Updates file system status based on scanning / restore / repair cycle
-        """
-
         # Deprecate both the action count (for restoring or reparing) and the scanning count
         self.file_system_action_count -= 1
         self.file_system_scanning_count -= 1
@@ -194,7 +193,10 @@ class ActiveNode(Node):
         # Reparing / Restoring updates
         if self.file_system_action_count <= 0:
             self.file_system_action_count = 0
-            if self.file_system_state_actual == FILE_SYSTEM_STATE.REPAIRING or self.file_system_state_actual == FILE_SYSTEM_STATE.RESTORING:
+            if (
+                self.file_system_state_actual == FILE_SYSTEM_STATE.REPAIRING
+                or self.file_system_state_actual == FILE_SYSTEM_STATE.RESTORING
+            ):
                 self.file_system_state_actual = FILE_SYSTEM_STATE.GOOD
                 self.file_system_state_observed = FILE_SYSTEM_STATE.GOOD
 

src/primaite/nodes/node.py

@@ -1,18 +1,14 @@
 # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
-"""
+"""The base Node class."""
-The base Node class
+from primaite.common.enums import HARDWARE_STATE
-"""
 
-from primaite.common.enums import *
 
 class Node:
-    """
+    """Node class."""
-    Node class
-    """
 
     def __init__(self, _id, _name, _type, _priority, _state, _config_values):
         """
-        Init
+        Init.
 
         Args:
             _id: The node id
@@ -21,7 +17,6 @@ class Node:
             _priority: The priority of the node
             _state: The state of the node
         """
-
         self.id = _id
         self.name = _name
         self.type = _type
@@ -31,146 +26,115 @@ class Node:
         self.config_values = _config_values
 
     def __repr__(self):
-        """
+        """Returns the name of the node."""
-        Returns the name of the node
-        """
-
         return self.name
 
     def set_id(self, _id):
         """
-        Sets the node ID
+        Sets the node ID.
 
         Args:
             _id: The node ID
         """
-
         self.id = _id
 
     def get_id(self):
         """
-        Gets the node ID
+        Gets the node ID.
 
         Returns:
              The node ID
         """
-
         return self.id
 
     def set_name(self, _name):
         """
-        Sets the node name
+        Sets the node name.
 
         Args:
             _name: The node name
         """
-
         self.name = _name
 
     def get_name(self):
         """
-        Gets the node name
+        Gets the node name.
 
         Returns:
              The node name
         """
-
         return self.name
 
     def set_type(self, _type):
         """
-        Sets the node type
+        Sets the node type.
 
         Args:
             _type: The node type
         """
-
         self.type = _type
 
     def get_type(self):
         """
-        Gets the node type
+        Gets the node type.
 
         Returns:
              The node type
         """
-
         return self.type
 
     def set_priority(self, _priority):
         """
-        Sets the node priority
+        Sets the node priority.
 
         Args:
             _priority: The node priority
         """
-
         self.priority = _priority
 
     def get_priority(self):
         """
-        Gets the node priority
+        Gets the node priority.
 
         Returns:
              The node priority
         """
-
         return self.priority
 
     def set_state(self, _state):
         """
-        Sets the node state
+        Sets the node state.
 
         Args:
             _state: The node state
         """
-
         self.operating_state = _state
 
     def get_state(self):
         """
-        Gets the node operating state
+        Gets the node operating state.
 
         Returns:
              The node operating state
         """
-
         return self.operating_state
 
     def turn_on(self):
-        """
+        """Sets the node state to ON."""
-        Sets the node state to ON
-        """
-
         self.operating_state = HARDWARE_STATE.ON
 
     def turn_off(self):
-        """
+        """Sets the node state to OFF."""
-        Sets the node state to OFF
-        """
-
         self.operating_state = HARDWARE_STATE.OFF
 
     def reset(self):
-        """
+        """Sets the node state to Resetting and starts the reset count."""
-        Sets the node state to Resetting and starts the reset count
-        """
-
         self.operating_state = HARDWARE_STATE.RESETTING
         self.resetting_count = self.config_values.node_reset_duration
 
     def update_resetting_status(self):
-        """
+        """Updates the resetting count."""
-        Updates the resetting count
-        """
-
         self.resetting_count -= 1
         if self.resetting_count <= 0:
             self.resetting_count = 0
             self.operating_state = HARDWARE_STATE.ON
-
-
-
-
-
-

src/primaite/nodes/node_state_instruction_green.py

@@ -1,16 +1,22 @@
 # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
-"""
+"""Defines node behaviour for Green PoL."""
-Defines node behaviour for Green PoL
+
-"""
 
 class NodeStateInstructionGreen(object):
-    """
+    """The Node State Instruction class."""
-    The Node State Instruction class
-    """
 
-    def __init__(self, _id, _start_step, _end_step, _node_id, _node_pol_type, _service_name, _state):
+    def __init__(
+        self,
+        _id,
+        _start_step,
+        _end_step,
+        _node_id,
+        _node_pol_type,
+        _service_name,
+        _state,
+    ):
         """
-        Init
+        Init.
 
         Args:
             _id: The node state instruction id
@@ -21,7 +27,6 @@ class NodeStateInstructionGreen(object):
             _service_name: The service name
             _state: The state (node or service)
         """
-
         self.id = _id
         self.start_step = _start_step
         self.end_step = _end_step
@@ -32,61 +37,54 @@ class NodeStateInstructionGreen(object):
 
     def get_start_step(self):
         """
-        Gets the start step
+        Gets the start step.
 
         Returns:
              The start step
         """
-
         return self.start_step
 
     def get_end_step(self):
         """
-        Gets the end step
+        Gets the end step.
 
         Returns:
              The end step
         """
-
         return self.end_step
 
     def get_node_id(self):
         """
-        Gets the node ID
+        Gets the node ID.
 
         Returns:
              The node ID
         """
-
         return self.node_id
 
     def get_node_pol_type(self):
         """
-        Gets the node pattern of life type (enum)
+        Gets the node pattern of life type (enum).
 
         Returns:
              The node pattern of life type (enum)
         """
-
         return self.node_pol_type
 
     def get_service_name(self):
         """
-        Gets the service name
+        Gets the service name.
 
         Returns:
              The service name
         """
-
         return self.service_name
 
     def get_state(self):
         """
-        Gets the state (node or service)
+        Gets the state (node or service).
 
         Returns:
              The state (node or service)
         """
-
         return self.state
-        

src/primaite/nodes/node_state_instruction_red.py

@@ -1,16 +1,26 @@
 # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
-"""
+"""Defines node behaviour for Green PoL."""
-Defines node behaviour for Green PoL
+
-"""
 
 class NodeStateInstructionRed(object):
-    """
+    """The Node State Instruction class."""
-    The Node State Instruction class
-    """
 
-    def __init__(self, _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):
+    def __init__(
+        self,
+        _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,
+    ):
         """
-        Init
+        Init.
 
         Args:
             _id: The node state instruction id
@@ -25,7 +35,6 @@ class NodeStateInstructionRed(object):
             _pol_source_node_service: The source node service (used for initiator type SERVICE)
             _pol_source_node_service_state: The source node service state (used for initiator type SERVICE)
         """
-
         self.id = _id
         self.start_step = _start_step
         self.end_step = _end_step
@@ -40,101 +49,90 @@ class NodeStateInstructionRed(object):
 
     def get_start_step(self):
         """
-        Gets the start step
+        Gets the start step.
 
         Returns:
              The start step
         """
-
         return self.start_step
 
     def get_end_step(self):
         """
-        Gets the end step
+        Gets the end step.
 
         Returns:
              The end step
         """
-
         return self.end_step
 
     def get_target_node_id(self):
         """
-        Gets the node ID
+        Gets the node ID.
 
         Returns:
              The node ID
         """
-
         return self.target_node_id
 
     def get_initiator(self):
         """
-        Gets the initiator
+        Gets the initiator.
 
         Returns:
              The initiator
         """
-
         return self.initiator
 
     def get_pol_type(self):
         """
-        Gets the node pattern of life type (enum)
+        Gets the node pattern of life type (enum).
 
         Returns:
              The node pattern of life type (enum)
         """
-
         return self.pol_type
 
     def get_service_name(self):
         """
-        Gets the service name
+        Gets the service name.
 
         Returns:
              The service name
         """
-
         return self.service_name
 
     def get_state(self):
         """
-        Gets the state (node or service)
+        Gets the state (node or service).
 
         Returns:
              The state (node or service)
         """
-
         return self.state
 
     def get_source_node_id(self):
         """
-        Gets the source node id (used for initiator type SERVICE)
+        Gets the source node id (used for initiator type SERVICE).
 
         Returns:
              The source node id
         """
-
         return self.source_node_id
 
     def get_source_node_service(self):
         """
-        Gets the source node service (used for initiator type SERVICE)
+        Gets the source node service (used for initiator type SERVICE).
 
         Returns:
              The source node service
         """
-
         return self.source_node_service
 
     def get_source_node_service_state(self):
         """
-        Gets the source node service state (used for initiator type SERVICE)
+        Gets the source node service state (used for initiator type SERVICE).
 
         Returns:
              The source node service state
         """
-
         return self.source_node_service_state
-        

src/primaite/nodes/passive_node.py

@@ -1,18 +1,15 @@
 # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
-"""
+"""The Passive Node class (i.e. an actuator)."""
-The Passive Node class (i.e. an actuator)
-"""
 
 from primaite.nodes.node import Node
 
+
 class PassiveNode(Node):
-    """
+    """The Passive Node class."""
-    The Passive Node class
-    """
 
     def __init__(self, _id, _name, _type, _priority, _state, _config_values):
         """
-        Init
+        Init.
 
         Args:
             _id: The node id
@@ -21,17 +18,15 @@ class PassiveNode(Node):
             _priority: The priority of the node
             _state: The state of the node
         """
-
         # Pass through to Super for now
         super().__init__(_id, _name, _type, _priority, _state, _config_values)
 
     def get_ip_address(self):
         """
-        Gets the node IP address
+        Gets the node IP address.
 
         Returns:
              The node IP address
         """
-
         # No concept of IP address for passive nodes for now
         return ""

src/primaite/nodes/service_node.py

@@ -1,19 +1,26 @@
 # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
-"""
+"""A Service Node (i.e. not an actuator)."""
-A Service Node (i.e. not an actuator)
+from primaite.common.enums import SOFTWARE_STATE
-"""
-
 from primaite.nodes.active_node import ActiveNode
-from primaite.common.enums import *
+
 
 class ServiceNode(ActiveNode):
-    """
+    """ServiceNode class."""
-    ServiceNode class
-    """
 
-    def __init__(self, _id, _name, _type, _priority, _state, _ip_address, _os_state, _file_system_state, _config_values):
+    def __init__(
+        self,
+        _id,
+        _name,
+        _type,
+        _priority,
+        _state,
+        _ip_address,
+        _os_state,
+        _file_system_state,
+        _config_values,
+    ):
         """
-        Init
+        Init.
 
         Args:
             _id: The node id
@@ -25,38 +32,44 @@ class ServiceNode(ActiveNode):
             _osState: The operating system state of the node
             _file_system_state: The file system state of the node
         """
-
+        super().__init__(
-        super().__init__(_id, _name, _type, _priority, _state, _ip_address, _os_state, _file_system_state, _config_values)
+            _id,
+            _name,
+            _type,
+            _priority,
+            _state,
+            _ip_address,
+            _os_state,
+            _file_system_state,
+            _config_values,
+        )
         self.services = {}
 
     def add_service(self, _service):
         """
-        Adds a service to the node
+        Adds a service to the node.
 
         Args:
             _service: The service to add
         """
-
         self.services[_service.get_name()] = _service
 
     def get_services(self):
         """
-        Gets the dictionary of services on this node
+        Gets the dictionary of services on this node.
 
         Returns:
              Dictionary of services on this node
         """
-
         return self.services
 
     def has_service(self, _protocol):
         """
-        Indicates whether a service is on a node
+        Indicates whether a service is on a node.
 
         Returns:
              True if service (protocol) is on the node
         """
-
         for service_key, service_value in self.services.items():
             if service_key == _protocol:
                 return True
@@ -66,12 +79,11 @@ class ServiceNode(ActiveNode):
 
     def service_running(self, _protocol):
         """
-        Indicates whether a service is in a running state on the node
+        Indicates whether a service is in a running state on the node.
 
         Returns:
              True if service (protocol) is in a running state on the node
         """
-
         for service_key, service_value in self.services.items():
             if service_key == _protocol:
                 if service_value.get_state() != SOFTWARE_STATE.PATCHING:
@@ -84,12 +96,11 @@ class ServiceNode(ActiveNode):
 
     def service_is_overwhelmed(self, _protocol):
         """
-        Indicates whether a service is in an overwhelmed state on the node
+        Indicates whether a service is in an overwhelmed state on the node.
 
         Returns:
              True if service (protocol) is in an overwhelmed state on the node
         """
-
         for service_key, service_value in self.services.items():
             if service_key == _protocol:
                 if service_value.get_state() == SOFTWARE_STATE.OVERWHELMED:
@@ -102,61 +113,61 @@ class ServiceNode(ActiveNode):
 
     def set_service_state(self, _protocol, _state):
         """
-        Sets the state of a service (protocol) on the node
+        Sets the state of a service (protocol) on the node.
 
         Args:
             _protocol: The service (protocol)
             _state: The state value
         """
-
         for service_key, service_value in self.services.items():
             if service_key == _protocol:
                 # Can't set to compromised if you're in a patching state
-                if (_state == SOFTWARE_STATE.COMPROMISED and service_value.get_state() != SOFTWARE_STATE.PATCHING) or _state != SOFTWARE_STATE.COMPROMISED:
+                if (
+                    _state == SOFTWARE_STATE.COMPROMISED
+                    and service_value.get_state() != SOFTWARE_STATE.PATCHING
+                ) or _state != SOFTWARE_STATE.COMPROMISED:
                     service_value.set_state(_state)
                 else:
                     # Do nothing
                     pass
                 if _state == SOFTWARE_STATE.PATCHING:
-                    service_value.patching_count = self.config_values.service_patching_duration
+                    service_value.patching_count = (
+                        self.config_values.service_patching_duration
+                    )
                 else:
                     # Do nothing
                     pass
 
     def set_service_state_if_not_compromised(self, _protocol, _state):
         """
-        Sets the state of a service (protocol) on the node if the operating state is not "compromised"
+        Sets the state of a service (protocol) on the node if the operating state is not "compromised".
 
         Args:
             _protocol: The service (protocol)
             _state: The state value
         """
-
         for service_key, service_value in self.services.items():
             if service_key == _protocol:
                 if service_value.get_state() != SOFTWARE_STATE.COMPROMISED:
                     service_value.set_state(_state)
                     if _state == SOFTWARE_STATE.PATCHING:
-                        service_value.patching_count = self.config_values.service_patching_duration
+                        service_value.patching_count = (
+                            self.config_values.service_patching_duration
+                        )
 
     def get_service_state(self, _protocol):
         """
-        Gets the state of a service
+        Gets the state of a service.
 
         Returns:
              The state of the service
         """
-
         for service_key, service_value in self.services.items():
             if service_key == _protocol:
                 return service_value.get_state()
 
     def update_services_patching_status(self):
-        """
+        """Updates the patching counter for any service that are patching."""
-        Updates the patching counter for any service that are patching
-        """
-
         for service_key, service_value in self.services.items():
             if service_value.get_state() == SOFTWARE_STATE.PATCHING:
                 service_value.reduce_patching_count()
-

src/primaite/pol/green_pol.py

@@ -1,19 +1,18 @@
 # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
-"""
+"""Implements Pattern of Life on the network (nodes and links)."""
-Implements Pattern of Life on the network (nodes and links)
-"""
 
 from networkx import shortest_path
 
-from primaite.common.enums import *
+from primaite.common.enums import HARDWARE_STATE, NODE_POL_TYPE, SOFTWARE_STATE, TYPE
 from primaite.nodes.active_node import ActiveNode
 from primaite.nodes.service_node import ServiceNode
 
 _VERBOSE = False
 
+
 def apply_iers(network, nodes, links, iers, acl, step):
     """
-    Applies IERs to the links (link pattern of life)
+    Applies IERs to the links (link pattern of life).
 
     Args:
         network: The network modelled in the environment
@@ -21,9 +20,8 @@ def apply_iers(network, nodes, links, iers, acl, step):
         links: The links within the environment
         iers: The IERs to apply to the links
         acl: The Access Control List
-        step: The step number
+        step: The step number.
     """
-
     if _VERBOSE:
         print("Applying IERs")
 
@@ -55,7 +53,10 @@ def apply_iers(network, nodes, links, iers, acl, step):
             # 1. Check the source node situation
             if source_node.get_type() == TYPE.SWITCH:
                 # It's a switch
-                if source_node.get_state() == HARDWARE_STATE.ON and source_node.get_os_state() != SOFTWARE_STATE.PATCHING:
+                if (
+                    source_node.get_state() == HARDWARE_STATE.ON
+                    and source_node.get_os_state() != SOFTWARE_STATE.PATCHING
+                ):
                     source_valid = True
                 else:
                     # IER no longer valid
@@ -66,9 +67,14 @@ def apply_iers(network, nodes, links, iers, acl, step):
                 pass
             else:
                 # It's not a switch or an actuator (so active node)
-                if source_node.get_state() == HARDWARE_STATE.ON and source_node.get_os_state() != SOFTWARE_STATE.PATCHING:
+                if (
+                    source_node.get_state() == HARDWARE_STATE.ON
+                    and source_node.get_os_state() != SOFTWARE_STATE.PATCHING
+                ):
                     if source_node.has_service(protocol):
-                        if source_node.service_running(protocol) and not source_node.service_is_overwhelmed(protocol):
+                        if source_node.service_running(
+                            protocol
+                        ) and not source_node.service_is_overwhelmed(protocol):
                             source_valid = True
                         else:
                             source_valid = False
@@ -80,11 +86,13 @@ def apply_iers(network, nodes, links, iers, acl, step):
                     # Do nothing - IER no longer valid
                     source_valid = False
 
-
             # 2. Check the dest node situation
             if dest_node.get_type() == TYPE.SWITCH:
                 # It's a switch
-                if dest_node.get_state() == HARDWARE_STATE.ON and dest_node.get_os_state() != SOFTWARE_STATE.PATCHING:
+                if (
+                    dest_node.get_state() == HARDWARE_STATE.ON
+                    and dest_node.get_os_state() != SOFTWARE_STATE.PATCHING
+                ):
                     dest_valid = True
                 else:
                     # IER no longer valid
@@ -94,9 +102,14 @@ def apply_iers(network, nodes, links, iers, acl, step):
                 pass
             else:
                 # It's not a switch or an actuator (so active node)
-                if dest_node.get_state() == HARDWARE_STATE.ON and dest_node.get_os_state() != SOFTWARE_STATE.PATCHING:
+                if (
+                    dest_node.get_state() == HARDWARE_STATE.ON
+                    and dest_node.get_os_state() != SOFTWARE_STATE.PATCHING
+                ):
                     if dest_node.has_service(protocol):
-                        if dest_node.service_running(protocol) and not dest_node.service_is_overwhelmed(protocol):
+                        if dest_node.service_running(
+                            protocol
+                        ) and not dest_node.service_is_overwhelmed(protocol):
                             dest_valid = True
                         else:
                             dest_valid = False
@@ -109,10 +122,21 @@ def apply_iers(network, nodes, links, iers, acl, step):
                     dest_valid = False
 
             # 3. Check that the ACL doesn't block it
-            acl_block = acl.is_blocked(source_node.get_ip_address(), dest_node.get_ip_address(), protocol, port)
+            acl_block = acl.is_blocked(
+                source_node.get_ip_address(), dest_node.get_ip_address(), protocol, port
+            )
             if acl_block:
                 if _VERBOSE:
-                    print("ACL block on source: " + source_node.get_ip_address() + ", dest: " + dest_node.get_ip_address() + ", protocol: " + protocol + ", port: " + port)
+                    print(
+                        "ACL block on source: "
+                        + source_node.get_ip_address()
+                        + ", dest: "
+                        + dest_node.get_ip_address()
+                        + ", protocol: "
+                        + protocol
+                        + ", port: "
+                        + port
+                    )
             else:
                 if _VERBOSE:
                     print("No ACL block")
@@ -131,7 +155,10 @@ def apply_iers(network, nodes, links, iers, acl, step):
 
                 # We might have a switch in the path, so check all nodes are operational
                 for node in path_node_list:
-                    if node.get_state() != HARDWARE_STATE.ON or node.get_os_state() == SOFTWARE_STATE.PATCHING:
+                    if (
+                        node.get_state() != HARDWARE_STATE.ON
+                        or node.get_os_state() == SOFTWARE_STATE.PATCHING
+                    ):
                         path_valid = False
 
                 if path_valid:
@@ -143,8 +170,10 @@ def apply_iers(network, nodes, links, iers, acl, step):
                     # Check that the link capacity is not exceeded by the new load
                     while count < path_node_list_length - 1:
                         # Get the link between the next two nodes
-                        edge_dict = network.get_edge_data(path_node_list[count], path_node_list[count+1])
+                        edge_dict = network.get_edge_data(
-                        link_id = edge_dict[0].get('id')
+                            path_node_list[count], path_node_list[count + 1]
+                        )
+                        link_id = edge_dict[0].get("id")
                         link = links[link_id]
                         # Check whether the new load exceeds the bandwidth
                         if (link.get_current_load() + load) > link.get_bandwidth():
@@ -160,8 +189,10 @@ def apply_iers(network, nodes, links, iers, acl, step):
                         count = 0
                         while count < path_node_list_length - 1:
                             # Get the link between the next two nodes
-                            edge_dict = network.get_edge_data(path_node_list[count], path_node_list[count+1])
+                            edge_dict = network.get_edge_data(
-                            link_id = edge_dict[0].get('id')
+                                path_node_list[count], path_node_list[count + 1]
+                            )
+                            link_id = edge_dict[0].get("id")
                             link = links[link_id]
                             # Add the load from this IER
                             link.add_protocol_load(protocol, load)
@@ -183,16 +214,16 @@ def apply_iers(network, nodes, links, iers, acl, step):
             # Do nothing - IER no longer valid
             pass
 
+
 def apply_node_pol(nodes, node_pol, step):
     """
-    Applies node pattern of life
+    Applies node pattern of life.
 
     Args:
         nodes: The nodes within the environment
         node_pol: The node pattern of life to apply
-        step: The step number
+        step: The step number.
     """
-
     if _VERBOSE:
         print("Applying Node PoL")
 

src/primaite/pol/ier.py

@@ -1,17 +1,29 @@
 # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
 """
-Information Exchange Requirements for APE
+Information Exchange Requirements for APE.
-Used to represent an information flow from source to destination
+
+Used to represent an information flow from source to destination.
 """
 
+
 class IER(object):
-    """
+    """Information Exchange Requirement class."""
-    Information Exchange Requirement class
-    """
 
-    def __init__(self, _id, _start_step, _end_step, _load, _protocol, _port, _source_node_id, _dest_node_id, _mission_criticality, _running=False):
+    def __init__(
+        self,
+        _id,
+        _start_step,
+        _end_step,
+        _load,
+        _protocol,
+        _port,
+        _source_node_id,
+        _dest_node_id,
+        _mission_criticality,
+        _running=False,
+    ):
         """
-        Init
+        Init.
 
         Args:
             _id: The IER id
@@ -25,7 +37,6 @@ class IER(object):
             _mission_criticality: Criticality of this IER to the mission (0 none, 5 mission critical)
             _running: Indicates whether the IER is currently running
         """
-
         self.id = _id
         self.start_step = _start_step
         self.end_step = _end_step
@@ -39,97 +50,88 @@ class IER(object):
 
     def get_id(self):
         """
-        Gets IER ID
+        Gets IER ID.
 
         Returns:
              IER ID
         """
-
         return self.id
 
     def get_start_step(self):
         """
-        Gets IER start step
+        Gets IER start step.
 
         Returns:
              IER start step
         """
-
         return self.start_step
 
     def get_end_step(self):
         """
-        Gets IER end step
+        Gets IER end step.
 
         Returns:
              IER end step
         """
-
         return self.end_step
 
     def get_load(self):
         """
-        Gets IER load
+        Gets IER load.
 
         Returns:
              IER load
         """
-
         return self.load
 
     def get_protocol(self):
         """
-        Gets IER protocol
+        Gets IER protocol.
 
         Returns:
              IER protocol
         """
-
         return self.protocol
 
     def get_port(self):
         """
-        Gets IER port
+        Gets IER port.
 
         Returns:
              IER port
         """
-
         return self.port
 
     def get_source_node_id(self):
         """
-        Gets IER source node ID
+        Gets IER source node ID.
 
         Returns:
              IER source node ID
         """
-
         return self.source_node_id
 
     def get_dest_node_id(self):
         """
-        Gets IER destination node ID
+        Gets IER destination node ID.
 
         Returns:
              IER destination node ID
         """
-
         return self.dest_node_id
 
     def get_is_running(self):
         """
-        Informs whether the IER is currently running
+        Informs whether the IER is currently running.
 
         Returns:
              True if running
         """
-
         return self.running
 
     def set_is_running(self, _value):
         """
-        Sets the running state of the IER
+        Sets the running state of the IER.
 
         Args:
              _value: running status
@@ -138,10 +140,9 @@ class IER(object):
 
     def get_mission_criticality(self):
         """
-        Gets the IER mission criticality (used in the reward function)
+        Gets the IER mission criticality (used in the reward function).
 
         Returns:
              Mission criticality value (0 lowest to 5 highest)
         """
-
         return self.mission_criticality

src/primaite/pol/red_agent_pol.py

@@ -1,19 +1,24 @@
 # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
-"""
+"""Implements POL on the network (nodes and links) resulting from the red agent attack."""
-Implements Pattern of Life on the network (nodes and links) resulting from the red agent attack
-"""
 
 from networkx import shortest_path
 
-from primaite.common.enums import *
+from primaite.common.enums import (
+    HARDWARE_STATE,
+    NODE_POL_INITIATOR,
+    NODE_POL_TYPE,
+    SOFTWARE_STATE,
+    TYPE,
+)
 from primaite.nodes.active_node import ActiveNode
 from primaite.nodes.service_node import ServiceNode
 
 _VERBOSE = False
 
+
 def apply_red_agent_iers(network, nodes, links, iers, acl, step):
     """
-    Applies IERs to the links (link pattern of life) resulting from red agent attack
+    Applies IERs to the links (link POL) resulting from red agent attack.
 
     Args:
         network: The network modelled in the environment
@@ -21,9 +26,8 @@ def apply_red_agent_iers(network, nodes, links, iers, acl, step):
         links: The links within the environment
         iers: The red agent IERs to apply to the links
         acl: The Access Control List
-        step: The step number
+        step: The step number.
     """
-
     # Go through each IER and check the conditions for it being applied
     # If everything is in place, apply the IER protocol load to the relevant links
     for ier_key, ier_value in iers.items():
@@ -66,7 +70,10 @@ def apply_red_agent_iers(network, nodes, links, iers, acl, step):
                 if source_node.get_state() == HARDWARE_STATE.ON:
                     if source_node.has_service(protocol):
                         # Red agents IERs can only be valid if the source service is in a compromised state
-                        if source_node.get_service_state(protocol) == SOFTWARE_STATE.COMPROMISED:
+                        if (
+                            source_node.get_service_state(protocol)
+                            == SOFTWARE_STATE.COMPROMISED
+                        ):
                             source_valid = True
                         else:
                             source_valid = False
@@ -78,7 +85,6 @@ def apply_red_agent_iers(network, nodes, links, iers, acl, step):
                     # Do nothing - IER no longer valid
                     source_valid = False
 
-
             # 2. Check the dest node situation
             if dest_node.get_type() == TYPE.SWITCH:
                 # It's a switch
@@ -105,10 +111,21 @@ def apply_red_agent_iers(network, nodes, links, iers, acl, step):
                     dest_valid = False
 
             # 3. Check that the ACL doesn't block it
-            acl_block = acl.is_blocked(source_node.get_ip_address(), dest_node.get_ip_address(), protocol, port)
+            acl_block = acl.is_blocked(
+                source_node.get_ip_address(), dest_node.get_ip_address(), protocol, port
+            )
             if acl_block:
                 if _VERBOSE:
-                    print("ACL block on source: " + source_node.get_ip_address() + ", dest: " + dest_node.get_ip_address() + ", protocol: " + protocol + ", port: " + port)
+                    print(
+                        "ACL block on source: "
+                        + source_node.get_ip_address()
+                        + ", dest: "
+                        + dest_node.get_ip_address()
+                        + ", protocol: "
+                        + protocol
+                        + ", port: "
+                        + port
+                    )
             else:
                 if _VERBOSE:
                     print("No ACL block")
@@ -140,8 +157,10 @@ def apply_red_agent_iers(network, nodes, links, iers, acl, step):
                     # Check that the link capacity is not exceeded by the new load
                     while count < path_node_list_length - 1:
                         # Get the link between the next two nodes
-                        edge_dict = network.get_edge_data(path_node_list[count], path_node_list[count+1])
+                        edge_dict = network.get_edge_data(
-                        link_id = edge_dict[0].get('id')
+                            path_node_list[count], path_node_list[count + 1]
+                        )
+                        link_id = edge_dict[0].get("id")
                         link = links[link_id]
                         # Check whether the new load exceeds the bandwidth
                         if (link.get_current_load() + load) > link.get_bandwidth():
@@ -157,8 +176,10 @@ def apply_red_agent_iers(network, nodes, links, iers, acl, step):
                         count = 0
                         while count < path_node_list_length - 1:
                             # Get the link between the next two nodes
-                            edge_dict = network.get_edge_data(path_node_list[count], path_node_list[count+1])
+                            edge_dict = network.get_edge_data(
-                            link_id = edge_dict[0].get('id')
+                                path_node_list[count], path_node_list[count + 1]
+                            )
+                            link_id = edge_dict[0].get("id")
                             link = links[link_id]
                             # Add the load from this IER
                             link.add_protocol_load(protocol, load)
@@ -185,17 +206,17 @@ def apply_red_agent_iers(network, nodes, links, iers, acl, step):
 
     pass
 
+
 def apply_red_agent_node_pol(nodes, iers, node_pol, step):
     """
-    Applies node pattern of life
+    Applies node pattern of life.
 
     Args:
         nodes: The nodes within the environment
         iers: The red agent IERs
         node_pol: The red agent node pattern of life to apply
-        step: The step number
+        step: The step number.
     """
-
     if _VERBOSE:
         print("Applying Node Red Agent PoL")
 
@@ -209,7 +230,9 @@ def apply_red_agent_node_pol(nodes, iers, node_pol, step):
         state = node_instruction.get_state()
         source_node_id = node_instruction.get_source_node_id()
         source_node_service_name = node_instruction.get_source_node_service()
-        source_node_service_state_value = node_instruction.get_source_node_service_state()
+        source_node_service_state_value = (
+            node_instruction.get_source_node_service_state()
+        )
 
         passed_checks = False
 
@@ -228,7 +251,10 @@ def apply_red_agent_node_pol(nodes, iers, node_pol, step):
                 # Need to check the condition of a service on another node
                 source_node = nodes[source_node_id]
                 if source_node.has_service(source_node_service_name):
-                    if source_node.get_service_state(source_node_service_name) == SOFTWARE_STATE[source_node_service_state_value]:
+                    if (
+                        source_node.get_service_state(source_node_service_name)
+                        == SOFTWARE_STATE[source_node_service_state_value]
+                    ):
                         passed_checks = True
                     else:
                         # Do nothing, no matching state value
@@ -248,7 +274,9 @@ def apply_red_agent_node_pol(nodes, iers, node_pol, step):
                     target_node.set_state(state)
                 elif pol_type == NODE_POL_TYPE.OS:
                     # Change OS state
-                    if isinstance(target_node, ActiveNode) or isinstance(target_node, ServiceNode):
+                    if isinstance(target_node, ActiveNode) or isinstance(
+                        target_node, ServiceNode
+                    ):
                         target_node.set_os_state(state)
                 elif pol_type == NODE_POL_TYPE.SERVICE:
                     # Change a service state
@@ -256,7 +284,9 @@ def apply_red_agent_node_pol(nodes, iers, node_pol, step):
                         target_node.set_service_state(service_name, state)
                 else:
                     # Change the file system status
-                    if isinstance(target_node, ActiveNode) or isinstance(target_node, ServiceNode):
+                    if isinstance(target_node, ActiveNode) or isinstance(
+                        target_node, ServiceNode
+                    ):
                         target_node.set_file_system_state(state)
             else:
                 if _VERBOSE:
@@ -265,13 +295,22 @@ def apply_red_agent_node_pol(nodes, iers, node_pol, step):
             # PoL is not valid in this time step
             pass
 
-def is_red_ier_incoming(node, iers, node_pol_type):
 
+def is_red_ier_incoming(node, iers, node_pol_type):
+    """
+    Checks if the RED IER is incoming.
+
+    TODO: Write more descriptive docstring with params and returns.
+    """
     node_id = node.get_id()
 
     for ier_key, ier_value in iers.items():
         if ier_value.get_is_running() and ier_value.get_dest_node_id() == node_id:
-            if node_pol_type == NODE_POL_TYPE.OPERATING or node_pol_type == NODE_POL_TYPE.OS or node_pol_type == NODE_POL_TYPE.FILE:
+            if (
+                node_pol_type == NODE_POL_TYPE.OPERATING
+                or node_pol_type == NODE_POL_TYPE.OS
+                or node_pol_type == NODE_POL_TYPE.FILE
+            ):
                 # It's looking to change operating state, file system or O/S state, so valid
                 return True
             elif node_pol_type == NODE_POL_TYPE.SERVICE:
@@ -297,5 +336,3 @@ def is_red_ier_incoming(node, iers, node_pol_type):
         else:
             # The IER destination is not this node, or the IER is not running
             return False
-            
-

src/primaite/transactions/transaction.py

@@ -1,15 +1,13 @@
 # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
-"""
+"""The Transaction class."""
-The Transaction class
+
-"""
 
 class Transaction(object):
-    """
+    """Transaction class."""
-    Transaction class
+
-    """
     def __init__(self, _timestamp, _agent_identifier, _episode_number, _step_number):
         """
-        Init
+        Init.
 
         Args:
             _timestamp: The time this object was created
@@ -17,7 +15,6 @@ class Transaction(object):
             _episode_number: The episode number
             _step_number: The step number
         """
-
         self.timestamp = _timestamp
         self.agent_identifier = _agent_identifier
         self.episode_number = _episode_number
@@ -25,45 +22,36 @@ class Transaction(object):
 
     def set_obs_space_pre(self, _obs_space_pre):
         """
-        Sets the observation space (pre)
+        Sets the observation space (pre).
 
         Args:
             _obs_space_pre: The observation space before any actions are taken
         """
-
         self.obs_space_pre = _obs_space_pre
 
     def set_obs_space_post(self, _obs_space_post):
         """
-        Sets the observation space (post)
+        Sets the observation space (post).
 
         Args:
             _obs_space_post: The observation space after any actions are taken
         """
-
         self.obs_space_post = _obs_space_post
 
     def set_reward(self, _reward):
         """
-        Sets the reward
+        Sets the reward.
 
         Args:
             _reward: The reward value
         """
-
         self.reward = _reward
 
     def set_action_space(self, _action_space):
         """
-        Sets the action space
+        Sets the action space.
 
         Args:
             _action_space: The action space invoked by the agent
         """
-
         self.action_space = _action_space
-
-
-
-
-

src/primaite/transactions/transactions_to_file.py

@@ -1,40 +1,35 @@
 # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
-"""
+"""Writes the Transaction log list out to file for evaluation to utilse."""
-Writes the Transaction log list out to file for evaluation to utilse
-"""
 
 import csv
 import logging
 import os.path
-
 from datetime import datetime
 
-from primaite.transactions.transaction import Transaction
 
 def turn_action_space_to_array(_action_space):
     """
-    Turns action space into a string array so it can be saved to csv
+    Turns action space into a string array so it can be saved to csv.
 
     Args:
-        _action_space: The action space
+        _action_space: The action space.
     """
-
     return_array = []
     for x in range(len(_action_space)):
         return_array.append(str(_action_space[x]))
 
     return return_array
 
+
 def turn_obs_space_to_array(_obs_space, _obs_assets, _obs_features):
     """
-    Turns observation space into a string array so it can be saved to csv
+    Turns observation space into a string array so it can be saved to csv.
 
     Args:
         _obs_space: The observation space
         _obs_assets: The number of assets (i.e. nodes or links) in the observation space
         _obs_features: The number of features associated with the asset
     """
-
     return_array = []
     for x in range(_obs_assets):
         for y in range(_obs_features):
@@ -42,15 +37,15 @@ def turn_obs_space_to_array(_obs_space, _obs_assets, _obs_features):
 
     return return_array
 
+
 def write_transaction_to_file(_transaction_list):
     """
-    Writes transaction logs to file to support training evaluation
+    Writes transaction logs to file to support training evaluation.
 
     Args:
         _transaction_list: The list of transactions from all steps and all episodes
-        _num_episodes: The number of episodes that were conducted
+        _num_episodes: The number of episodes that were conducted.
     """
-
     # Get the first transaction and use it to determine the makeup of the observation space and action space
     # Label the obs space fields in csv as "OSI_1_1", "OSN_1_1" and action space as "AS_1"
     # This will be tied into the PrimAITE Use Case so that they make sense
@@ -62,43 +57,53 @@ def write_transaction_to_file(_transaction_list):
     # Create the action space headers array
     action_header = []
     for x in range(action_length):
-        action_header.append('AS_' + str(x))
+        action_header.append("AS_" + str(x))
 
     # Create the observation space headers array
     obs_header_initial = []
     obs_header_new = []
     for x in range(obs_assets):
         for y in range(obs_features):
-            obs_header_initial.append('OSI_' + str(x) + '_' + str(y))
+            obs_header_initial.append("OSI_" + str(x) + "_" + str(y))
-            obs_header_new.append('OSN_' + str(x) + '_' + str(y))
+            obs_header_new.append("OSN_" + str(x) + "_" + str(y))
 
     # Open up a csv file
-    header = ['Timestamp', 'Episode', 'Step', 'Reward']
+    header = ["Timestamp", "Episode", "Step", "Reward"]
     header = header + action_header + obs_header_initial + obs_header_new
     now = datetime.now()  # current date and time
     time = now.strftime("%Y%m%d_%H%M%S")
 
     try:
-        path = 'outputs/results/'
+        path = "outputs/results/"
         is_dir = os.path.isdir(path)
         if not is_dir:
             os.makedirs(path)
 
         filename = "outputs/results/all_transactions_" + time + ".csv"
-        csv_file = open(filename, 'w', encoding='UTF8', newline='')
+        csv_file = open(filename, "w", encoding="UTF8", newline="")
         csv_writer = csv.writer(csv_file)
         csv_writer.writerow(header)
 
         for transaction in _transaction_list:
-            csv_data = [str(transaction.timestamp),  str(transaction.episode_number), str(transaction.step_number), str(transaction.reward)]
+            csv_data = [
-            csv_data = csv_data + turn_action_space_to_array(transaction.action_space) + \
+                str(transaction.timestamp),
-                turn_obs_space_to_array(transaction.obs_space_pre, obs_assets, obs_features) + \
+                str(transaction.episode_number),
-                turn_obs_space_to_array(transaction.obs_space_post, obs_assets, obs_features)
+                str(transaction.step_number),
+                str(transaction.reward),
+            ]
+            csv_data = (
+                csv_data
+                + turn_action_space_to_array(transaction.action_space)
+                + turn_obs_space_to_array(
+                    transaction.obs_space_pre, obs_assets, obs_features
+                )
+                + turn_obs_space_to_array(
+                    transaction.obs_space_post, obs_assets, obs_features
+                )
+            )
             csv_writer.writerow(csv_data)
 
         csv_file.close()
-    except Exception as e:
+    except Exception:
         logging.error("Could not save the transaction file")
         logging.error("Exception occured", exc_info=True)
-
-

tests/test_acl.py

@@ -1,61 +1,48 @@
-
 # Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
-"""
+"""Used to tes the ACL functions."""
-Used to tes the ACL functions
-"""
 
-from primaite.acl.acl_rule import ACLRule
 from primaite.acl.access_control_list import AccessControlList
+from primaite.acl.acl_rule import ACLRule
+
 
 def test_acl_address_match_1():
-    """
+    """Test that matching IP addresses produce True."""
-    Test that matching IP addresses produce True
-    """
-
     acl = AccessControlList()
 
     rule = ACLRule("ALLOW", "192.168.1.1", "192.168.1.2", "TCP", "80")
 
     assert acl.check_address_match(rule, "192.168.1.1", "192.168.1.2") == True
 
-def test_acl_address_match_2():
-    """
-    Test that mismatching IP addresses produce False
-    """
 
+def test_acl_address_match_2():
+    """Test that mismatching IP addresses produce False."""
     acl = AccessControlList()
 
     rule = ACLRule("ALLOW", "192.168.1.1", "192.168.1.2", "TCP", "80")
 
     assert acl.check_address_match(rule, "192.168.1.1", "192.168.1.3") == False
 
-def test_acl_address_match_3():
-    """
-    Test the ANY condition for source IP addresses produce True
-    """
 
+def test_acl_address_match_3():
+    """Test the ANY condition for source IP addresses produce True."""
     acl = AccessControlList()
 
     rule = ACLRule("ALLOW", "ANY", "192.168.1.2", "TCP", "80")
 
     assert acl.check_address_match(rule, "192.168.1.1", "192.168.1.2") == True
 
-def test_acl_address_match_4():
-    """
-    Test the ANY condition for dest IP addresses produce True
-    """
 
+def test_acl_address_match_4():
+    """Test the ANY condition for dest IP addresses produce True."""
     acl = AccessControlList()
 
     rule = ACLRule("ALLOW", "192.168.1.1", "ANY", "TCP", "80")
 
     assert acl.check_address_match(rule, "192.168.1.1", "192.168.1.2") == True
 
-def test_check_acl_block_affirmative():
-    """
-    Test the block function (affirmative)
-    """
 
+def test_check_acl_block_affirmative():
+    """Test the block function (affirmative)."""
     # Create the Access Control List
     acl = AccessControlList()
 
@@ -66,15 +53,19 @@ def test_check_acl_block_affirmative():
     acl_rule_protocol = "TCP"
     acl_rule_port = "80"
 
-    acl.add_rule(acl_rule_permission, acl_rule_source, acl_rule_destination, acl_rule_protocol, acl_rule_port)
+    acl.add_rule(
+        acl_rule_permission,
+        acl_rule_source,
+        acl_rule_destination,
+        acl_rule_protocol,
+        acl_rule_port,
+    )
 
     assert acl.is_blocked("192.168.1.1", "192.168.1.2", "TCP", "80") == False
 
-def test_check_acl_block_negative():
-    """
-    Test the block function (negative)
-    """
 
+def test_check_acl_block_negative():
+    """Test the block function (negative)."""
     # Create the Access Control List
     acl = AccessControlList()
 
@@ -85,21 +76,27 @@ def test_check_acl_block_negative():
     acl_rule_protocol = "TCP"
     acl_rule_port = "80"
 
-    acl.add_rule(acl_rule_permission, acl_rule_source, acl_rule_destination, acl_rule_protocol, acl_rule_port)
+    acl.add_rule(
+        acl_rule_permission,
+        acl_rule_source,
+        acl_rule_destination,
+        acl_rule_protocol,
+        acl_rule_port,
+    )
 
     assert acl.is_blocked("192.168.1.1", "192.168.1.2", "TCP", "80") == True
 
-def test_rule_hash():
-    """
-    Test the rule hash
-    """
 
+def test_rule_hash():
+    """Test the rule hash."""
     # Create the Access Control List
     acl = AccessControlList()
 
     rule = ACLRule("DENY", "192.168.1.1", "192.168.1.2", "TCP", "80")
     hash_value_local = hash(rule)
 
-    hash_value_remote = acl.get_dictionary_hash("DENY", "192.168.1.1", "192.168.1.2", "TCP", "80")
+    hash_value_remote = acl.get_dictionary_hash(
+        "DENY", "192.168.1.1", "192.168.1.2", "TCP", "80"
+    )
 
     assert hash_value_local == hash_value_remote