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

#901 - Replaced "ALLOW" with RulePermissionType.ALLOW

Browse Source 
- Added Explicit ALLOW to test_configs in order for tests to work
- Added typing to access_control_list.py and acl_rule.py
This commit is contained in:
SunilSamra
2023-07-17 20:40:00 +01:00
parent a2f43b5abc
commit 9520cfea24
7 changed files with 55 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
src/primaite/acl/access_control_list.py
View File

@@ -4,6 +4,7 @@ import logging
from typing import Final, List, Union from typing import Final, List, Union
from primaite.acl.acl_rule import ACLRule from primaite.acl.acl_rule import ACLRule
from primaite.common.enums import RulePermissionType
_LOGGER: Final[logging.Logger] = logging.getLogger(__name__) _LOGGER: Final[logging.Logger] = logging.getLogger(__name__)
@@ -15,12 +16,11 @@ class AccessControlList:
"""Init.""" """Init."""
# Implicit ALLOW or DENY firewall spec # Implicit ALLOW or DENY firewall spec
self.acl_implicit_permission = implicit_permission self.acl_implicit_permission = implicit_permission
print(self.acl_implicit_permission, "ACL IMPLICIT PERMISSION")
# Implicit rule in ACL list # Implicit rule in ACL list
if self.acl_implicit_permission == "DENY": if self.acl_implicit_permission == RulePermissionType.DENY:
self.acl_implicit_rule = ACLRule("DENY", "ANY", "ANY", "ANY", "ANY") self.acl_implicit_rule = ACLRule(RulePermissionType.DENY, "ANY", "ANY", "ANY", "ANY")
elif self.acl_implicit_permission == "ALLOW": elif self.acl_implicit_permission == RulePermissionType.ALLOW:
self.acl_implicit_rule = ACLRule("ALLOW", "ANY", "ANY", "ANY", "ANY") self.acl_implicit_rule = ACLRule(RulePermissionType.ALLOW, "ANY", "ANY", "ANY", "ANY")
else: else:
raise ValueError(f"implicit permission must be ALLOW or DENY, got {self.acl_implicit_permission}") raise ValueError(f"implicit permission must be ALLOW or DENY, got {self.acl_implicit_permission}")
@@ -76,9 +76,9 @@ class AccessControlList:
str(rule.get_port()) == str(_port) or rule.get_port() == "ANY" str(rule.get_port()) == str(_port) or rule.get_port() == "ANY"
): ):
# There's a matching rule. Get the permission # There's a matching rule. Get the permission
if rule.get_permission() == "DENY": if rule.get_permission() == RulePermissionType.DENY:
return True return True
elif rule.get_permission() == "ALLOW": elif rule.get_permission() == RulePermissionType.ALLOW:
return False return False
# If there has been no rule to allow the IER through, it will return a blocked signal by default # If there has been no rule to allow the IER through, it will return a blocked signal by default
@@ -121,7 +121,9 @@ class AccessControlList:
else: else:
_LOGGER.info(f"Position {position_index} is an invalid/overwrites implicit firewall rule") _LOGGER.info(f"Position {position_index} is an invalid/overwrites implicit firewall rule")
def remove_rule(self, _permission: str, _source_ip: str, _dest_ip: str, _protocol: str, _port: str) -> None: def remove_rule(
self, _permission: RulePermissionType, _source_ip: str, _dest_ip: str, _protocol: str, _port: str
) -> None:
""" """
Removes a rule. Removes a rule.

3
src/primaite/acl/acl_rule.py
View File

@@ -1,11 +1,12 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence. # 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."""
from primaite.common.enums import RulePermissionType
class ACLRule: class ACLRule:
"""Access Control List Rule class.""" """Access Control List Rule class."""
def __init__(self, _permission, _source_ip, _dest_ip, _protocol, _port): def __init__(self, _permission: RulePermissionType, _source_ip, _dest_ip, _protocol, _port):
""" """
Initialise an ACL Rule. Initialise an ACL Rule.

2
src/primaite/environment/observations.py
View File

@@ -489,7 +489,7 @@ class AccessControlList(AbstractObservationComponent):
# Map each ACL attribute from what it was to an integer to fit the observation space # Map each ACL attribute from what it was to an integer to fit the observation space
source_ip_int = None source_ip_int = None
dest_ip_int = None dest_ip_int = None
if permission == "DENY": if permission == RulePermissionType.DENY:
permission_int = 1 permission_int = 1
else: else:
permission_int = 2 permission_int = 2

5
tests/config/obs_tests/main_config_LINK_TRAFFIC_LEVELS.yaml
View File

@@ -38,6 +38,11 @@ observation_space:
# Time delay between steps (for generic agents) # Time delay between steps (for generic agents)
time_delay: 1 time_delay: 1
# Implicit ACL firewall rule at end of lists to be default action or no rule can be selected (ALLOW or DENY)
implicit_acl_rule: ALLOW
# Total number of ACL rules allowed in the environment
max_number_acl_rules: 4
# Type of session to be run (TRAINING or EVALUATION) # Type of session to be run (TRAINING or EVALUATION)
session_type: TRAIN session_type: TRAIN
# Determine whether to load an agent from file # Determine whether to load an agent from file

5
tests/config/obs_tests/main_config_NODE_LINK_TABLE.yaml
View File

@@ -36,6 +36,11 @@ observation_space:
time_delay: 1 time_delay: 1
# Filename of the scenario / laydown # Filename of the scenario / laydown
# Implicit ACL firewall rule at end of lists to be default action or no rule can be selected (ALLOW or DENY)
implicit_acl_rule: ALLOW
# Total number of ACL rules allowed in the environment
max_number_acl_rules: 4
session_type: TRAIN session_type: TRAIN
# Determine whether to load an agent from file # Determine whether to load an agent from file
load_agent: False load_agent: False

41
tests/test_acl.py
View File

@@ -3,40 +3,41 @@
from primaite.acl.access_control_list import AccessControlList from primaite.acl.access_control_list import AccessControlList
from primaite.acl.acl_rule import ACLRule from primaite.acl.acl_rule import ACLRule
from primaite.common.enums import RulePermissionType
def test_acl_address_match_1(): def test_acl_address_match_1():
"""Test that matching IP addresses produce True.""" """Test that matching IP addresses produce True."""
acl = AccessControlList("DENY", 10) acl = AccessControlList(RulePermissionType.DENY, 10)
rule = ACLRule("ALLOW", "192.168.1.1", "192.168.1.2", "TCP", "80") rule = ACLRule(RulePermissionType.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 assert acl.check_address_match(rule, "192.168.1.1", "192.168.1.2") == True
def test_acl_address_match_2(): def test_acl_address_match_2():
"""Test that mismatching IP addresses produce False.""" """Test that mismatching IP addresses produce False."""
acl = AccessControlList("DENY", 10) acl = AccessControlList(RulePermissionType.DENY, 10)
rule = ACLRule("ALLOW", "192.168.1.1", "192.168.1.2", "TCP", "80") rule = ACLRule(RulePermissionType.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 assert acl.check_address_match(rule, "192.168.1.1", "192.168.1.3") == False
def test_acl_address_match_3(): def test_acl_address_match_3():
"""Test the ANY condition for source IP addresses produce True.""" """Test the ANY condition for source IP addresses produce True."""
acl = AccessControlList("DENY", 10) acl = AccessControlList(RulePermissionType.DENY, 10)
rule = ACLRule("ALLOW", "ANY", "192.168.1.2", "TCP", "80") rule = ACLRule(RulePermissionType.ALLOW, "ANY", "192.168.1.2", "TCP", "80")
assert acl.check_address_match(rule, "192.168.1.1", "192.168.1.2") == True assert acl.check_address_match(rule, "192.168.1.1", "192.168.1.2") == True
def test_acl_address_match_4(): def test_acl_address_match_4():
"""Test the ANY condition for dest IP addresses produce True.""" """Test the ANY condition for dest IP addresses produce True."""
acl = AccessControlList("DENY", 10) acl = AccessControlList(RulePermissionType.DENY, 10)
rule = ACLRule("ALLOW", "192.168.1.1", "ANY", "TCP", "80") rule = ACLRule(RulePermissionType.ALLOW, "192.168.1.1", "ANY", "TCP", "80")
assert acl.check_address_match(rule, "192.168.1.1", "192.168.1.2") == True assert acl.check_address_match(rule, "192.168.1.1", "192.168.1.2") == True
@@ -44,10 +45,10 @@ def test_acl_address_match_4():
def test_check_acl_block_affirmative(): def test_check_acl_block_affirmative():
"""Test the block function (affirmative).""" """Test the block function (affirmative)."""
# Create the Access Control List # Create the Access Control List
acl = AccessControlList("DENY", 10) acl = AccessControlList(RulePermissionType.DENY, 10)
# Create a rule # Create a rule
acl_rule_permission = "ALLOW" acl_rule_permission = RulePermissionType.ALLOW
acl_rule_source = "192.168.1.1" acl_rule_source = "192.168.1.1"
acl_rule_destination = "192.168.1.2" acl_rule_destination = "192.168.1.2"
acl_rule_protocol = "TCP" acl_rule_protocol = "TCP"
@@ -68,10 +69,10 @@ def test_check_acl_block_affirmative():
def test_check_acl_block_negative(): def test_check_acl_block_negative():
"""Test the block function (negative).""" """Test the block function (negative)."""
# Create the Access Control List # Create the Access Control List
acl = AccessControlList("DENY", 10) acl = AccessControlList(RulePermissionType.DENY, 10)
# Create a rule # Create a rule
acl_rule_permission = "DENY" acl_rule_permission = RulePermissionType.DENY
acl_rule_source = "192.168.1.1" acl_rule_source = "192.168.1.1"
acl_rule_destination = "192.168.1.2" acl_rule_destination = "192.168.1.2"
acl_rule_protocol = "TCP" acl_rule_protocol = "TCP"
@@ -93,12 +94,12 @@ def test_check_acl_block_negative():
def test_rule_hash(): def test_rule_hash():
"""Test the rule hash.""" """Test the rule hash."""
# Create the Access Control List # Create the Access Control List
acl = AccessControlList("DENY", 10) acl = AccessControlList(RulePermissionType.DENY, 10)
rule = ACLRule("DENY", "192.168.1.1", "192.168.1.2", "TCP", "80") rule = ACLRule(RulePermissionType.DENY, "192.168.1.1", "192.168.1.2", "TCP", "80")
hash_value_local = hash(rule) 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(RulePermissionType.DENY, "192.168.1.1", "192.168.1.2", "TCP", "80")
assert hash_value_local == hash_value_remote assert hash_value_local == hash_value_remote
@@ -106,10 +107,10 @@ def test_rule_hash():
def test_delete_rule(): def test_delete_rule():
"""Adds 3 rules and deletes 1 rule and checks its deletion.""" """Adds 3 rules and deletes 1 rule and checks its deletion."""
# Create the Access Control List # Create the Access Control List
acl = AccessControlList("ALLOW", 10) acl = AccessControlList(RulePermissionType.ALLOW, 10)
# Create a first rule # Create a first rule
acl_rule_permission = "DENY" acl_rule_permission = RulePermissionType.DENY
acl_rule_source = "192.168.1.1" acl_rule_source = "192.168.1.1"
acl_rule_destination = "192.168.1.2" acl_rule_destination = "192.168.1.2"
acl_rule_protocol = "TCP" acl_rule_protocol = "TCP"
@@ -126,7 +127,7 @@ def test_delete_rule():
) )
# Create a second rule # Create a second rule
acl_rule_permission = "DENY" acl_rule_permission = RulePermissionType.DENY
acl_rule_source = "20" acl_rule_source = "20"
acl_rule_destination = "30" acl_rule_destination = "30"
acl_rule_protocol = "FTP" acl_rule_protocol = "FTP"
@@ -143,7 +144,7 @@ def test_delete_rule():
) )
# Create a third rule # Create a third rule
acl_rule_permission = "ALLOW" acl_rule_permission = RulePermissionType.ALLOW
acl_rule_source = "192.168.1.3" acl_rule_source = "192.168.1.3"
acl_rule_destination = "192.168.1.1" acl_rule_destination = "192.168.1.1"
acl_rule_protocol = "UDP" acl_rule_protocol = "UDP"
@@ -159,7 +160,7 @@ def test_delete_rule():
acl_position_in_list, acl_position_in_list,
) )
# Remove the second ACL rule added from the list # Remove the second ACL rule added from the list
acl.remove_rule("DENY", "20", "30", "FTP", "21") acl.remove_rule(RulePermissionType.DENY, "20", "30", "FTP", "21")
assert len(acl.acl) == 10 assert len(acl.acl) == 10
assert acl.acl[2] is None assert acl.acl[2] is None

21
tests/test_seeding_and_deterministic_session.py
View File

@@ -26,16 +26,16 @@ def test_seeded_learning(temp_primaite_session):
now work. If not, then you've got a bug :). now work. If not, then you've got a bug :).
""" """
expected_mean_reward_per_episode = { expected_mean_reward_per_episode = {
1: -90.703125, 1: -20.7421875,
2: -91.15234375, 2: -19.82421875,
3: -87.5, 3: -17.01171875,
4: -92.2265625, 4: -19.08203125,
5: -94.6875, 5: -21.93359375,
6: -91.19140625, 6: -20.21484375,
7: -88.984375, 7: -15.546875,
8: -88.3203125, 8: -12.08984375,
9: -112.79296875, 9: -17.59765625,
10: -100.01953125, 10: -14.6875,
} }
with temp_primaite_session as session: with temp_primaite_session as session:
@@ -44,6 +44,7 @@ def test_seeded_learning(temp_primaite_session):
), "Expected output is based upon a agent that was trained with seed 67890" ), "Expected output is based upon a agent that was trained with seed 67890"
session.learn() session.learn()
actual_mean_reward_per_episode = session.learn_av_reward_per_episode_dict() actual_mean_reward_per_episode = session.learn_av_reward_per_episode_dict()
print(actual_mean_reward_per_episode, "THISt")
assert actual_mean_reward_per_episode == expected_mean_reward_per_episode assert actual_mean_reward_per_episode == expected_mean_reward_per_episode