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

#1355 - Carried out full renaming in node.py, active_node.py, passive_node.py, and service_node.py to make params and variable names explicit.

Browse Source 
- Made the same renaming in the yaml laydown config files.
- Added Type hints wherever I've been.
- Added a custom NodeType in custom_typing.py to encompass the Union of ActiveNode, PassiveNode, ServiceNode.
This commit is contained in:
Chris McCarthy
2023-05-25 21:03:11 +01:00
parent 00e9d1f88d
commit 32a4d9e459
31 changed files with 965 additions and 1081 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -6,6 +6,7 @@ import csv
import logging
import os.path
from datetime import datetime
from typing import Dict
import networkx as nx
import numpy as np
@@ -14,20 +15,22 @@ from gym import Env, spaces
from matplotlib import pyplot as plt
from primaite.acl.access_control_list import AccessControlList
from primaite.common.custom_typing import NodeType
from primaite.common.enums import (
ACTION_TYPE,
FILE_SYSTEM_STATE,
HARDWARE_STATE,
NODE_POL_INITIATOR,
NODE_POL_TYPE,
PRIORITY,
SOFTWARE_STATE,
TYPE,
ActionType,
FileSystemState,
HardwareState,
NodePOLInitiator,
NodePOLType,
NodeType,
Priority,
SoftwareState,
)
from primaite.common.service import Service
from primaite.environment.reward import calculate_reward_function
from primaite.links.link import Link
from primaite.nodes.active_node import ActiveNode
from primaite.nodes.node import Node
from primaite.nodes.node_state_instruction_green import NodeStateInstructionGreen
from primaite.nodes.node_state_instruction_red import NodeStateInstructionRed
from primaite.nodes.passive_node import PassiveNode
@@ -77,19 +80,19 @@ class Primaite(Env):
self.agent_identifier = self.config_values.agent_identifier
# Create a dictionary to hold all the nodes
self.nodes = {}
self.nodes: Dict[str, NodeType] = {}
# Create a dictionary to hold a reference set of nodes
self.nodes_reference = {}
self.nodes_reference: Dict[str, NodeType] = {}
# Create a dictionary to hold all the links
self.links = {}
self.links: Dict[str, Link] = {}
# Create a dictionary to hold a reference set of links
self.links_reference = {}
self.links_reference: Dict[str, Link] = {}
# Create a dictionary to hold all the green IERs (this will come from an external source)
self.green_iers = {}
self.green_iers: Dict[str, IER] = {}
# Create a dictionary to hold all the node PoLs (this will come from an external source)
self.node_pol = {}
@@ -190,8 +193,8 @@ class Primaite(Env):
# For each item, we send:
# - [For Nodes] | [For Links]
# - node ID | link ID
# - operating state | N/A
# - operating system state | N/A
# - hardware state | N/A
# - Software State | N/A
# - file system state | N/A
# - service A state | service A loading
# - service B state | service B loading
@@ -205,7 +208,7 @@ class Primaite(Env):
# 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)
# hardware state, file system state and SoftwareState (i.e. 4)
self.num_observation_parameters = (
self.num_services + self.OBSERVATION_SPACE_FIXED_PARAMETERS
)
@@ -222,13 +225,13 @@ class Primaite(Env):
self.env_obs = np.zeros(self.observation_shape, dtype=np.int64)
# Define Action Space - depends on action space type (Node or ACL)
if self.action_type == ACTION_TYPE.NODE:
if self.action_type == ActionType.NODE:
_LOGGER.info("Action space type NODE selected")
# Terms (for node action space):
# [0, num nodes] - node ID (0 = nothing, node ID)
# [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)
# [0, 4] - what property it's acting on (0 = nothing, state, SoftwareState, service state, file system state) # noqa
# [0, 3] - action on property (0 = nothing, On / Scan, Off / Repair, Reset / Patch / Restore) # noqa
# [0, num services] - resolves to service ID (0 = nothing, resolves to service) # noqa
self.action_space = spaces.MultiDiscrete(
[
self.num_nodes,
@@ -402,7 +405,7 @@ class Primaite(Env):
self.step_count,
self.config_values,
)
print(f" Step {self.step_count} Reward: {str(reward)}")
# print(f" Step {self.step_count} Reward: {str(reward)}")
self.total_reward += reward
if self.step_count == self.episode_steps:
self.average_reward = self.total_reward / self.step_count
@@ -441,7 +444,7 @@ class Primaite(Env):
"""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():
for protocol in link_value.protocol_list:
print(
" Protocol: "
+ protocol.get_name().name
@@ -457,7 +460,7 @@ class Primaite(Env):
_action: The action space from the agent
"""
# At the moment, actions are only affecting nodes
if self.action_type == ACTION_TYPE.NODE:
if self.action_type == ActionType.NODE:
self.apply_actions_to_nodes(_action)
else:
self.apply_actions_to_acl(_action)
@@ -484,32 +487,32 @@ class Primaite(Env):
# This is the do nothing action
return
elif node_property == 1:
# This is an action on the node Operating State
# This is an action on the node Hardware State
if property_action == 0:
# Do nothing
return
elif property_action == 1:
# Turn on (only applicable if it's OFF, not if it's patching)
if node.get_state() == HARDWARE_STATE.OFF:
if node.hardware_state == HardwareState.OFF:
node.turn_on()
elif property_action == 2:
# Turn off
node.turn_off()
elif property_action == 3:
# Reset (only applicable if it's ON)
if node.get_state() == HARDWARE_STATE.ON:
if node.hardware_state == HardwareState.ON:
node.reset()
else:
return
elif node_property == 2:
if isinstance(node, ActiveNode) or isinstance(node, ServiceNode):
# This is an action on the node Operating System State
# This is an action on the node Software State
if property_action == 0:
# Do nothing
return
elif property_action == 1:
# Patch (valid action if it's good or compromised)
node.set_os_state(SOFTWARE_STATE.PATCHING)
node.software_state = SoftwareState.PATCHING
else:
# Node is not of Active or Service Type
return
@@ -523,7 +526,7 @@ class Primaite(Env):
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
self.services_list[service_index], SoftwareState.PATCHING
)
else:
# Node is not of Service Type
@@ -540,14 +543,11 @@ 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
):
node.set_file_system_state(FILE_SYSTEM_STATE.REPAIRING)
if node.file_system_state_actual != FileSystemState.DESTROYED:
node.set_file_system_state(FileSystemState.REPAIRING)
elif property_action == 3:
# Restore
node.set_file_system_state(FILE_SYSTEM_STATE.RESTORING)
node.set_file_system_state(FileSystemState.RESTORING)
else:
# Node is not of Active Type
return
@@ -584,7 +584,7 @@ class Primaite(Env):
else:
node = list(self.nodes.values())[action_source_ip - 1]
if isinstance(node, ServiceNode) or isinstance(node, ActiveNode):
acl_rule_source = node.get_ip_address()
acl_rule_source = node.ip_address
else:
return
# Destination IP value
@@ -593,7 +593,7 @@ class Primaite(Env):
else:
node = list(self.nodes.values())[action_destination_ip - 1]
if isinstance(node, ServiceNode) or isinstance(node, ActiveNode):
acl_rule_destination = node.get_ip_address()
acl_rule_destination = node.ip_address
else:
return
# Protocol value
@@ -636,13 +636,13 @@ class Primaite(Env):
e.g. reset / patching status
"""
for node_key, node in self.nodes.items():
if node.get_state() == HARDWARE_STATE.RESETTING:
if node.hardware_state == HardwareState.RESETTING:
node.update_resetting_status()
else:
pass
if isinstance(node, ActiveNode) or isinstance(node, ServiceNode):
node.update_file_system_state()
if node.get_os_state() == SOFTWARE_STATE.PATCHING:
if node.software_state == SoftwareState.PATCHING:
node.update_os_patching_status()
else:
pass
@@ -654,13 +654,13 @@ class Primaite(Env):
pass
for node_key, node in self.nodes_reference.items():
if node.get_state() == HARDWARE_STATE.RESETTING:
if node.hardware_state == HardwareState.RESETTING:
node.update_resetting_status()
else:
pass
if isinstance(node, ActiveNode) or isinstance(node, ServiceNode):
node.update_file_system_state()
if node.get_os_state() == SOFTWARE_STATE.PATCHING:
if node.software_state == SoftwareState.PATCHING:
node.update_os_patching_status()
else:
pass
@@ -677,13 +677,11 @@ class Primaite(Env):
# Do nodes first
for node_key, node in self.nodes.items():
self.env_obs[item_index][0] = int(node.get_id())
self.env_obs[item_index][1] = node.get_state().value
self.env_obs[item_index][0] = int(node.node_id)
self.env_obs[item_index][1] = node.hardware_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][2] = node.software_state.value
self.env_obs[item_index][3] = node.file_system_state_observed.value
else:
self.env_obs[item_index][2] = 0
self.env_obs[item_index][3] = 0
@@ -768,14 +766,14 @@ class Primaite(Env):
item: A config data item
"""
# All nodes have these parameters
node_id = item["id"]
node_id = item["node_id"]
node_name = item["name"]
node_base_type = item["baseType"]
node_type = TYPE[item["nodeType"]]
node_priority = PRIORITY[item["priority"]]
node_hardware_state = HARDWARE_STATE[item["hardwareState"]]
node_class = item["node_class"]
node_type = NodeType[item["node_type"]]
node_priority = Priority[item["priority"]]
node_hardware_state = HardwareState[item["hardware_state"]]
if node_base_type == "PASSIVE":
if node_class == "PASSIVE":
node = PassiveNode(
node_id,
node_name,
@@ -784,11 +782,11 @@ class Primaite(Env):
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"]]
elif node_class == "ACTIVE":
# Active nodes have IP address, Software State and file system state
node_ip_address = item["ip_address"]
node_software_state = SoftwareState[item["software_state"]]
node_file_system_state = FileSystemState[item["file_system_state"]]
node = ActiveNode(
node_id,
node_name,
@@ -800,11 +798,11 @@ class Primaite(Env):
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"]]
elif node_class == "SERVICE":
# Service nodes have IP address, Software State, file system state and list of services
node_ip_address = item["ip_address"]
node_software_state = SoftwareState[item["software_state"]]
node_file_system_state = FileSystemState[item["file_system_state"]]
node = ServiceNode(
node_id,
node_name,
@@ -820,7 +818,7 @@ class Primaite(Env):
for service in node_services:
service_protocol = service["name"]
service_port = service["port"]
service_state = SOFTWARE_STATE[service["state"]]
service_state = SoftwareState[service["state"]]
node.add_service(Service(service_protocol, service_port, service_state))
else:
# Bad formatting
@@ -841,7 +839,7 @@ class Primaite(Env):
# Add node to network (reference)
self.network_reference.add_nodes_from([node_ref])
def create_link(self, item):
def create_link(self, item: Dict):
"""
Creates a link from config data.
@@ -854,8 +852,8 @@ class Primaite(Env):
link_source = item["source"]
link_destination = item["destination"]
source_node = self.nodes[link_source]
dest_node = self.nodes[link_destination]
source_node: Node = self.nodes[link_source]
dest_node: Node = self.nodes[link_destination]
# Add link to network
self.network.add_edge(source_node, dest_node, id=link_name)
@@ -864,14 +862,14 @@ class Primaite(Env):
self.links[link_name] = Link(
link_id,
link_bandwidth,
source_node.get_name(),
dest_node.get_name(),
source_node.name,
dest_node.name,
self.services_list,
)
# Reference
source_node_ref = self.nodes_reference[link_source]
dest_node_ref = self.nodes_reference[link_destination]
source_node_ref: Node = self.nodes_reference[link_source]
dest_node_ref: Node = self.nodes_reference[link_destination]
# Add link to network (reference)
self.network_reference.add_edge(source_node_ref, dest_node_ref, id=link_name)
@@ -880,8 +878,8 @@ class Primaite(Env):
self.links_reference[link_name] = Link(
link_id,
link_bandwidth,
source_node_ref.get_name(),
dest_node_ref.get_name(),
source_node_ref.name,
dest_node_ref.name,
self.services_list,
)
@@ -956,18 +954,18 @@ class Primaite(Env):
pol_start_step = item["startStep"]
pol_end_step = item["endStep"]
pol_node = item["nodeId"]
pol_type = NODE_POL_TYPE[item["type"]]
pol_type = NodePOLType[item["type"]]
# State depends on whether this is Operating, O/S, file system or Service PoL type
if pol_type == NODE_POL_TYPE.OPERATING:
pol_state = HARDWARE_STATE[item["state"]]
# State depends on whether this is Operating, Software, file system or Service PoL type
if pol_type == NodePOLType.OPERATING:
pol_state = HardwareState[item["state"]]
pol_protocol = ""
elif pol_type == NODE_POL_TYPE.FILE:
pol_state = FILE_SYSTEM_STATE[item["state"]]
elif pol_type == NodePOLType.FILE:
pol_state = FileSystemState[item["state"]]
pol_protocol = ""
else:
pol_protocol = item["protocol"]
pol_state = SOFTWARE_STATE[item["state"]]
pol_state = SoftwareState[item["state"]]
self.node_pol[pol_id] = NodeStateInstructionGreen(
pol_id,
@@ -990,17 +988,17 @@ class Primaite(Env):
pol_start_step = item["startStep"]
pol_end_step = item["endStep"]
pol_target_node_id = item["targetNodeId"]
pol_initiator = NODE_POL_INITIATOR[item["initiator"]]
pol_type = NODE_POL_TYPE[item["type"]]
pol_initiator = NodePOLInitiator[item["initiator"]]
pol_type = NodePOLType[item["type"]]
pol_protocol = item["protocol"]
# State depends on whether this is Operating, O/S, file system or Service PoL type
if pol_type == NODE_POL_TYPE.OPERATING:
pol_state = HARDWARE_STATE[item["state"]]
elif pol_type == NODE_POL_TYPE.FILE:
pol_state = FILE_SYSTEM_STATE[item["state"]]
# State depends on whether this is Operating, Software, file system or Service PoL type
if pol_type == NodePOLType.OPERATING:
pol_state = HardwareState[item["state"]]
elif pol_type == NodePOLType.FILE:
pol_state = FileSystemState[item["state"]]
else:
pol_state = SOFTWARE_STATE[item["state"]]
pol_state = SoftwareState[item["state"]]
pol_source_node_id = item["sourceNodeId"]
pol_source_node_service = item["sourceNodeService"]
@@ -1080,7 +1078,7 @@ class Primaite(Env):
Args:
item: A config data item representing action info
"""
self.action_type = ACTION_TYPE[action_info["type"]]
self.action_type = ActionType[action_info["type"]]
def get_steps_info(self, steps_info):
"""
@@ -1126,38 +1124,38 @@ class Primaite(Env):
item: A config data item
"""
# All nodes have these parameters
node_id = item["id"]
node_base_type = item["baseType"]
node_hardware_state = HARDWARE_STATE[item["hardwareState"]]
node_id = item["node_id"]
node_class = item["node_class"]
node_hardware_state: HardwareState = HardwareState[item["hardware_state"]]
node = self.nodes[node_id]
node: NodeType = self.nodes[node_id]
node_ref = self.nodes_reference[node_id]
# Reset the hardware state (common for all node types)
node.set_state(node_hardware_state)
node_ref.set_state(node_hardware_state)
node.hardware_state = node_hardware_state
node_ref.hardware_state = node_hardware_state
if node_base_type == "ACTIVE":
# Active nodes have operating system state
node_software_state = SOFTWARE_STATE[item["softwareState"]]
node_file_system_state = FILE_SYSTEM_STATE[item["fileSystemState"]]
node.set_os_state(node_software_state)
node_ref.set_os_state(node_software_state)
if node_class == "ACTIVE":
# Active nodes have Software State
node_software_state = SoftwareState[item["software_state"]]
node_file_system_state = FileSystemState[item["file_system_state"]]
node.software_state = node_software_state
node_ref.software_state = node_software_state
node.set_file_system_state(node_file_system_state)
node_ref.set_file_system_state(node_file_system_state)
elif node_base_type == "SERVICE":
# Service nodes have operating system state and list of services
node_software_state = SOFTWARE_STATE[item["softwareState"]]
node_file_system_state = FILE_SYSTEM_STATE[item["fileSystemState"]]
node.set_os_state(node_software_state)
node_ref.set_os_state(node_software_state)
elif node_class == "SERVICE":
# Service nodes have Software State and list of services
node_software_state = SoftwareState[item["software_state"]]
node_file_system_state = FileSystemState[item["file_system_state"]]
node.software_state = node_software_state
node_ref.software_state = node_software_state
node.set_file_system_state(node_file_system_state)
node_ref.set_file_system_state(node_file_system_state)
# Update service states
node_services = item["services"]
for service in node_services:
service_protocol = service["name"]
service_state = SOFTWARE_STATE[service["state"]]
service_state = SoftwareState[service["state"]]
# Update node service state
node.set_service_state(service_protocol, service_state)
# Update reference node service state

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

@@ -1,6 +1,9 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
"""Implements reward function."""
from primaite.common.enums import FILE_SYSTEM_STATE, HARDWARE_STATE, SOFTWARE_STATE
from typing import Dict
from primaite.common.enums import FileSystemState, HardwareState, SoftwareState
from primaite.common.service import Service
from primaite.nodes.active_node import ActiveNode
from primaite.nodes.service_node import ServiceNode
@@ -28,17 +31,17 @@ def calculate_reward_function(
"""
reward_value = 0
# For each node, compare operating state, o/s operating state, service states
# For each node, compare hardware state, SoftwareState, service states
for node_key, final_node in final_nodes.items():
initial_node = initial_nodes[node_key]
reference_node = reference_nodes[node_key]
# Operating State
# Hardware State
reward_value += score_node_operating_state(
final_node, initial_node, reference_node, config_values
)
# Operating System State
# Software State
if isinstance(final_node, ActiveNode) or isinstance(final_node, ServiceNode):
reward_value += score_node_os_state(
final_node, initial_node, reference_node, config_values
@@ -80,7 +83,7 @@ def calculate_reward_function(
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 hardware state of a node.
Args:
final_node: The node after red and blue agents take effect
@@ -89,9 +92,9 @@ def score_node_operating_state(final_node, initial_node, reference_node, config_
config_values: Config values
"""
score = 0
final_node_operating_state = final_node.get_state()
initial_node_operating_state = initial_node.get_state()
reference_node_operating_state = reference_node.get_state()
final_node_operating_state = final_node.hardware_state
initial_node_operating_state = initial_node.hardware_state
reference_node_operating_state = reference_node.hardware_state
if final_node_operating_state == reference_node_operating_state:
# All is well - we're no different from the reference situation
@@ -99,26 +102,26 @@ def score_node_operating_state(final_node, initial_node, reference_node, config_
else:
# We're different from the reference situation
# Need to compare initial and reference (current) state of node (i.e. at every step)
if initial_node_operating_state == HARDWARE_STATE.ON:
if reference_node_operating_state == HARDWARE_STATE.OFF:
if initial_node_operating_state == HardwareState.ON:
if reference_node_operating_state == HardwareState.OFF:
score += config_values.off_should_be_on
elif reference_node_operating_state == HARDWARE_STATE.RESETTING:
elif reference_node_operating_state == HardwareState.RESETTING:
score += config_values.resetting_should_be_on
else:
pass
elif initial_node_operating_state == HARDWARE_STATE.OFF:
if reference_node_operating_state == HARDWARE_STATE.ON:
elif initial_node_operating_state == HardwareState.OFF:
if reference_node_operating_state == HardwareState.ON:
score += config_values.on_should_be_off
elif reference_node_operating_state == HARDWARE_STATE.RESETTING:
elif reference_node_operating_state == HardwareState.RESETTING:
score += config_values.resetting_should_be_off
else:
pass
elif initial_node_operating_state == HARDWARE_STATE.RESETTING:
if reference_node_operating_state == HARDWARE_STATE.ON:
elif initial_node_operating_state == HardwareState.RESETTING:
if reference_node_operating_state == HardwareState.ON:
score += config_values.on_should_be_resetting
elif reference_node_operating_state == HARDWARE_STATE.OFF:
elif reference_node_operating_state == HardwareState.OFF:
score += config_values.off_should_be_resetting
elif reference_node_operating_state == HARDWARE_STATE.RESETTING:
elif reference_node_operating_state == HardwareState.RESETTING:
score += config_values.resetting
else:
pass
@@ -130,7 +133,7 @@ def score_node_operating_state(final_node, initial_node, reference_node, config_
def score_node_os_state(final_node, initial_node, reference_node, config_values):
"""
Calculates score relating to the operating system state of a node.
Calculates score relating to the Software State of a node.
Args:
final_node: The node after red and blue agents take effect
@@ -139,9 +142,9 @@ def score_node_os_state(final_node, initial_node, reference_node, config_values)
config_values: Config values
"""
score = 0
final_node_os_state = final_node.get_os_state()
initial_node_os_state = initial_node.get_os_state()
reference_node_os_state = reference_node.get_os_state()
final_node_os_state = final_node.software_state
initial_node_os_state = initial_node.software_state
reference_node_os_state = reference_node.software_state
if final_node_os_state == reference_node_os_state:
# All is well - we're no different from the reference situation
@@ -149,28 +152,28 @@ def score_node_os_state(final_node, initial_node, reference_node, config_values)
else:
# We're different from the reference situation
# Need to compare initial and reference (current) state of node (i.e. at every step)
if initial_node_os_state == SOFTWARE_STATE.GOOD:
if reference_node_os_state == SOFTWARE_STATE.PATCHING:
if initial_node_os_state == SoftwareState.GOOD:
if reference_node_os_state == SoftwareState.PATCHING:
score += config_values.patching_should_be_good
elif reference_node_os_state == SOFTWARE_STATE.COMPROMISED:
elif reference_node_os_state == SoftwareState.COMPROMISED:
score += config_values.compromised_should_be_good
else:
pass
elif initial_node_os_state == SOFTWARE_STATE.PATCHING:
if reference_node_os_state == SOFTWARE_STATE.GOOD:
elif initial_node_os_state == SoftwareState.PATCHING:
if reference_node_os_state == SoftwareState.GOOD:
score += config_values.good_should_be_patching
elif reference_node_os_state == SOFTWARE_STATE.COMPROMISED:
elif reference_node_os_state == SoftwareState.COMPROMISED:
score += config_values.compromised_should_be_patching
elif reference_node_os_state == SOFTWARE_STATE.PATCHING:
elif reference_node_os_state == SoftwareState.PATCHING:
score += config_values.patching
else:
pass
elif initial_node_os_state == SOFTWARE_STATE.COMPROMISED:
if reference_node_os_state == SOFTWARE_STATE.GOOD:
elif initial_node_os_state == SoftwareState.COMPROMISED:
if reference_node_os_state == SoftwareState.GOOD:
score += config_values.good_should_be_compromised
elif reference_node_os_state == SOFTWARE_STATE.PATCHING:
elif reference_node_os_state == SoftwareState.PATCHING:
score += config_values.patching_should_be_compromised
elif reference_node_os_state == SOFTWARE_STATE.COMPROMISED:
elif reference_node_os_state == SoftwareState.COMPROMISED:
score += config_values.compromised
else:
pass
@@ -191,59 +194,59 @@ def score_node_service_state(final_node, initial_node, reference_node, config_va
config_values: Config values
"""
score = 0
final_node_services = final_node.get_services()
initial_node_services = initial_node.get_services()
reference_node_services = reference_node.get_services()
final_node_services: Dict[str, Service] = final_node.services
initial_node_services: Dict[str, Service] = initial_node.services
reference_node_services: Dict[str, Service] = reference_node.services
for service_key, final_service in final_node_services.items():
reference_service = reference_node_services[service_key]
initial_service = initial_node_services[service_key]
if final_service.get_state() == reference_service.get_state():
if final_service.software_state == reference_service.software_state:
# All is well - we're no different from the reference situation
score += config_values.all_ok
else:
# We're different from the reference situation
# Need to compare initial and reference state of node (i.e. at every step)
if initial_service.get_state() == SOFTWARE_STATE.GOOD:
if reference_service.get_state() == SOFTWARE_STATE.PATCHING:
if initial_service.software_state == SoftwareState.GOOD:
if reference_service.software_state == SoftwareState.PATCHING:
score += config_values.patching_should_be_good
elif reference_service.get_state() == SOFTWARE_STATE.COMPROMISED:
elif reference_service.software_state == SoftwareState.COMPROMISED:
score += config_values.compromised_should_be_good
elif reference_service.get_state() == SOFTWARE_STATE.OVERWHELMED:
elif reference_service.software_state == SoftwareState.OVERWHELMED:
score += config_values.overwhelmed_should_be_good
else:
pass
elif initial_service.get_state() == SOFTWARE_STATE.PATCHING:
if reference_service.get_state() == SOFTWARE_STATE.GOOD:
elif initial_service.software_state == SoftwareState.PATCHING:
if reference_service.software_state == SoftwareState.GOOD:
score += config_values.good_should_be_patching
elif reference_service.get_state() == SOFTWARE_STATE.COMPROMISED:
elif reference_service.software_state == SoftwareState.COMPROMISED:
score += config_values.compromised_should_be_patching
elif reference_service.get_state() == SOFTWARE_STATE.OVERWHELMED:
elif reference_service.software_state == SoftwareState.OVERWHELMED:
score += config_values.overwhelmed_should_be_patching
elif reference_service.get_state() == SOFTWARE_STATE.PATCHING:
elif reference_service.software_state == SoftwareState.PATCHING:
score += config_values.patching
else:
pass
elif initial_service.get_state() == SOFTWARE_STATE.COMPROMISED:
if reference_service.get_state() == SOFTWARE_STATE.GOOD:
elif initial_service.software_state == SoftwareState.COMPROMISED:
if reference_service.software_state == SoftwareState.GOOD:
score += config_values.good_should_be_compromised
elif reference_service.get_state() == SOFTWARE_STATE.PATCHING:
elif reference_service.software_state == SoftwareState.PATCHING:
score += config_values.patching_should_be_compromised
elif reference_service.get_state() == SOFTWARE_STATE.COMPROMISED:
elif reference_service.software_state == SoftwareState.COMPROMISED:
score += config_values.compromised
elif reference_service.get_state() == SOFTWARE_STATE.OVERWHELMED:
elif reference_service.software_state == SoftwareState.OVERWHELMED:
score += config_values.overwhelmed_should_be_compromised
else:
pass
elif initial_service.get_state() == SOFTWARE_STATE.OVERWHELMED:
if reference_service.get_state() == SOFTWARE_STATE.GOOD:
elif initial_service.software_state == SoftwareState.OVERWHELMED:
if reference_service.software_state == SoftwareState.GOOD:
score += config_values.good_should_be_overwhelmed
elif reference_service.get_state() == SOFTWARE_STATE.PATCHING:
elif reference_service.software_state == SoftwareState.PATCHING:
score += config_values.patching_should_be_overwhelmed
elif reference_service.get_state() == SOFTWARE_STATE.COMPROMISED:
elif reference_service.software_state == SoftwareState.COMPROMISED:
score += config_values.compromised_should_be_overwhelmed
elif reference_service.get_state() == SOFTWARE_STATE.OVERWHELMED:
elif reference_service.software_state == SoftwareState.OVERWHELMED:
score += config_values.overwhelmed
else:
pass
@@ -263,12 +266,12 @@ def score_node_file_system(final_node, initial_node, reference_node, config_valu
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()
reference_node_file_system_state = reference_node.get_file_system_state_actual()
final_node_file_system_state = final_node.file_system_state_actual
initial_node_file_system_state = initial_node.file_system_state_actual
reference_node_file_system_state = reference_node.file_system_state_actual
final_node_scanning_state = final_node.is_scanning_file_system()
reference_node_scanning_state = reference_node.is_scanning_file_system()
final_node_scanning_state = final_node.file_system_scanning
reference_node_scanning_state = reference_node.file_system_scanning
# File System State
if final_node_file_system_state == reference_node_file_system_state:
@@ -277,66 +280,66 @@ def score_node_file_system(final_node, initial_node, reference_node, config_valu
else:
# We're different from the reference situation
# Need to compare initial and reference state of node (i.e. at every step)
if initial_node_file_system_state == FILE_SYSTEM_STATE.GOOD:
if reference_node_file_system_state == FILE_SYSTEM_STATE.REPAIRING:
if initial_node_file_system_state == FileSystemState.GOOD:
if reference_node_file_system_state == FileSystemState.REPAIRING:
score += config_values.repairing_should_be_good
elif reference_node_file_system_state == FILE_SYSTEM_STATE.RESTORING:
elif reference_node_file_system_state == FileSystemState.RESTORING:
score += config_values.restoring_should_be_good
elif reference_node_file_system_state == FILE_SYSTEM_STATE.CORRUPT:
elif reference_node_file_system_state == FileSystemState.CORRUPT:
score += config_values.corrupt_should_be_good
elif reference_node_file_system_state == FILE_SYSTEM_STATE.DESTROYED:
elif reference_node_file_system_state == FileSystemState.DESTROYED:
score += config_values.destroyed_should_be_good
else:
pass
elif initial_node_file_system_state == FILE_SYSTEM_STATE.REPAIRING:
if reference_node_file_system_state == FILE_SYSTEM_STATE.GOOD:
elif initial_node_file_system_state == FileSystemState.REPAIRING:
if reference_node_file_system_state == FileSystemState.GOOD:
score += config_values.good_should_be_repairing
elif reference_node_file_system_state == FILE_SYSTEM_STATE.RESTORING:
elif reference_node_file_system_state == FileSystemState.RESTORING:
score += config_values.restoring_should_be_repairing
elif reference_node_file_system_state == FILE_SYSTEM_STATE.CORRUPT:
elif reference_node_file_system_state == FileSystemState.CORRUPT:
score += config_values.corrupt_should_be_repairing
elif reference_node_file_system_state == FILE_SYSTEM_STATE.DESTROYED:
elif reference_node_file_system_state == FileSystemState.DESTROYED:
score += config_values.destroyed_should_be_repairing
elif reference_node_file_system_state == FILE_SYSTEM_STATE.REPAIRING:
elif reference_node_file_system_state == FileSystemState.REPAIRING:
score += config_values.repairing
else:
pass
elif initial_node_file_system_state == FILE_SYSTEM_STATE.RESTORING:
if reference_node_file_system_state == FILE_SYSTEM_STATE.GOOD:
elif initial_node_file_system_state == FileSystemState.RESTORING:
if reference_node_file_system_state == FileSystemState.GOOD:
score += config_values.good_should_be_restoring
elif reference_node_file_system_state == FILE_SYSTEM_STATE.REPAIRING:
elif reference_node_file_system_state == FileSystemState.REPAIRING:
score += config_values.repairing_should_be_restoring
elif reference_node_file_system_state == FILE_SYSTEM_STATE.CORRUPT:
elif reference_node_file_system_state == FileSystemState.CORRUPT:
score += config_values.corrupt_should_be_restoring
elif reference_node_file_system_state == FILE_SYSTEM_STATE.DESTROYED:
elif reference_node_file_system_state == FileSystemState.DESTROYED:
score += config_values.destroyed_should_be_restoring
elif reference_node_file_system_state == FILE_SYSTEM_STATE.RESTORING:
elif reference_node_file_system_state == FileSystemState.RESTORING:
score += config_values.restoring
else:
pass
elif initial_node_file_system_state == FILE_SYSTEM_STATE.CORRUPT:
if reference_node_file_system_state == FILE_SYSTEM_STATE.GOOD:
elif initial_node_file_system_state == FileSystemState.CORRUPT:
if reference_node_file_system_state == FileSystemState.GOOD:
score += config_values.good_should_be_corrupt
elif reference_node_file_system_state == FILE_SYSTEM_STATE.REPAIRING:
elif reference_node_file_system_state == FileSystemState.REPAIRING:
score += config_values.repairing_should_be_corrupt
elif reference_node_file_system_state == FILE_SYSTEM_STATE.RESTORING:
elif reference_node_file_system_state == FileSystemState.RESTORING:
score += config_values.restoring_should_be_corrupt
elif reference_node_file_system_state == FILE_SYSTEM_STATE.DESTROYED:
elif reference_node_file_system_state == FileSystemState.DESTROYED:
score += config_values.destroyed_should_be_corrupt
elif reference_node_file_system_state == FILE_SYSTEM_STATE.CORRUPT:
elif reference_node_file_system_state == FileSystemState.CORRUPT:
score += config_values.corrupt
else:
pass
elif initial_node_file_system_state == FILE_SYSTEM_STATE.DESTROYED:
if reference_node_file_system_state == FILE_SYSTEM_STATE.GOOD:
elif initial_node_file_system_state == FileSystemState.DESTROYED:
if reference_node_file_system_state == FileSystemState.GOOD:
score += config_values.good_should_be_destroyed
elif reference_node_file_system_state == FILE_SYSTEM_STATE.REPAIRING:
elif reference_node_file_system_state == FileSystemState.REPAIRING:
score += config_values.repairing_should_be_destroyed
elif reference_node_file_system_state == FILE_SYSTEM_STATE.RESTORING:
elif reference_node_file_system_state == FileSystemState.RESTORING:
score += config_values.restoring_should_be_destroyed
elif reference_node_file_system_state == FILE_SYSTEM_STATE.CORRUPT:
elif reference_node_file_system_state == FileSystemState.CORRUPT:
score += config_values.corrupt_should_be_destroyed
elif reference_node_file_system_state == FILE_SYSTEM_STATE.DESTROYED:
elif reference_node_file_system_state == FileSystemState.DESTROYED:
score += config_values.destroyed
else:
pass