901 - fixed test_single_action_space.py test

This commit is contained in:
SunilSamra
2023-07-13 11:45:23 +01:00
parent f8cb18c654
commit 0ab4dab72a
6 changed files with 4 additions and 10 deletions

View File

@@ -80,11 +80,8 @@ class AccessControlList:
Indicates block if all conditions are satisfied.
"""
for rule in self.acl:
print("loops through rule", rule, isinstance(rule, ACLRule))
if isinstance(rule, ACLRule):
print("finds rule")
if self.check_address_match(rule, _source_ip_address, _dest_ip_address):
print("source and dest ip match")
if (rule.get_protocol() == _protocol or rule.get_protocol() == "ANY") and (
str(rule.get_port()) == str(_port) or rule.get_port() == "ANY"
):

View File

@@ -267,7 +267,6 @@ def load(file_path: Union[str, Path], legacy_file: bool = False) -> TrainingConf
:raises TypeError: When the TrainingConfig object cannot be created
using the values from the config file read from ``file_path``.
"""
print("FILE PATH", file_path)
if not isinstance(file_path, Path):
file_path = Path(file_path)
if file_path.exists():

View File

@@ -444,12 +444,11 @@ class Primaite(Env):
_action: The action space from the agent
"""
# At the moment, actions are only affecting nodes
if self.training_config.action_type == ActionType.NODE:
self.apply_actions_to_nodes(_action)
elif self.training_config.action_type == ActionType.ACL:
self.apply_actions_to_acl(_action)
elif len(self.action_dict[_action]) == 7: # ACL actions in multidiscrete form have len 6
elif len(self.action_dict[_action]) == 7: # ACL actions in multidiscrete form have len 7
self.apply_actions_to_acl(_action)
elif len(self.action_dict[_action]) == 4: # Node actions in multdiscrete (array) from have len 4
self.apply_actions_to_nodes(_action)
@@ -1248,7 +1247,6 @@ class Primaite(Env):
# Combine the Node dict and ACL dict
combined_action_dict = {**acl_action_dict, **new_node_action_dict}
print("combined dict", combined_action_dict.items())
return combined_action_dict
def _create_random_red_agent(self):

View File

@@ -42,7 +42,7 @@ apply_implicit_rule: True
# Implicit ACL firewall rule at end of lists to be default action or no rule can be selected (ALLOW or DENY)
implicit_acl_rule: DENY
# Total number of ACL rules allowed in the environment
max_number_acl_rules: 3
max_number_acl_rules: 10
observation_space:
components:

View File

@@ -58,7 +58,6 @@ class TempPrimaiteSession(PrimaiteSession):
def __exit__(self, type, value, tb):
shutil.rmtree(self.session_path)
# shutil.rmtree(self.session_path.parent)
_LOGGER.debug(f"Deleted temp session directory: {self.session_path}")

View File

@@ -2,6 +2,7 @@ import time
import pytest
from primaite.acl.acl_rule import ACLRule
from primaite.common.enums import HardwareState
from primaite.environment.primaite_env import Primaite
from tests import TEST_CONFIG_ROOT
@@ -112,7 +113,7 @@ def test_agent_is_executing_actions_from_both_spaces(temp_primaite_session):
# This 1 rule added to the implicit deny means there should be 2 rules in total.
rules_count = 0
for rule in acl_rules_list:
if rule != -1:
if isinstance(rule, ACLRule):
rules_count += 1
# Therefore these statements below MUST be true
assert computer_node_hardware_state == HardwareState.OFF