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

901 - changed ACL instantiation and changed acl t private _acl (list not dict) attribute, added laydown_ACL.yaml for testing, fixed encoding of acl rules to integers for obs space, added ACL position to node action space and added generic test where agents adds two ACL rules.

Browse Source
This commit is contained in:
SunilSamra
2023-06-20 11:47:20 +01:00
parent 52d759bcd9
commit 9c17b54073
9 changed files with 305 additions and 107 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -11,21 +11,43 @@ _LOGGER: Final[logging.Logger] = logging.getLogger(__name__)
class AccessControlList:
"""Access Control List class."""
def __init__(self, implicit_permission, max_acl_rules):
def __init__(self, apply_implicit_rule, implicit_permission, max_acl_rules):
"""Init."""
# Bool option in main_config to decide to use implicit rule or not
self.apply_implicit_rule: bool = apply_implicit_rule
# Implicit ALLOW or DENY firewall spec
# Last rule in the ACL list
self.acl_implicit_rule = implicit_permission
# Create implicit rule based on input
if self.acl_implicit_rule == "DENY":
implicit_rule = ACLRule("DENY", "ANY", "ANY", "ANY", "ANY")
else:
implicit_rule = ACLRule("ALLOW", "ANY", "ANY", "ANY", "ANY")
self.acl_implicit_permission = implicit_permission
# Maximum number of ACL Rules in ACL
self.max_acl_rules: int = max_acl_rules
# A list of ACL Rules
self.acl: List[ACLRule] = [implicit_rule]
self._acl: List[ACLRule] = []
# Implicit rule
@property
def acl_implicit_rule(self):
"""ACL implicit rule class attribute with added logic to change it depending on option in main_config."""
# Create implicit rule based on input
if self.apply_implicit_rule:
if self.acl_implicit_permission == "DENY":
return ACLRule("DENY", "ANY", "ANY", "ANY", "ANY")
elif self.acl_implicit_permission == "ALLOW":
return ACLRule("ALLOW", "ANY", "ANY", "ANY", "ANY")
else:
return None
else:
return None
@property
def acl(self):
"""Public access method for private _acl.
Adds implicit rule to end of acl list and
Pads out rest of list (if empty) with -1.
"""
if self.acl_implicit_rule is not None:
acl_list = self._acl + [self.acl_implicit_rule]
return acl_list + [-1] * (self.max_acl_rules - len(acl_list))
def check_address_match(self, _rule, _source_ip_address, _dest_ip_address):
"""
@@ -85,7 +107,9 @@ class AccessControlList:
# If there has been no rule to allow the IER through, it will return a blocked signal by default
return True
def add_rule(self, _permission, _source_ip, _dest_ip, _protocol, _port, _position):
def add_rule(
self, _permission, _source_ip, _dest_ip, _protocol, _port, _position=None
):
"""
Adds a new rule.
@@ -99,18 +123,22 @@ class AccessControlList:
"""
position_index = int(_position)
new_rule = ACLRule(_permission, _source_ip, _dest_ip, _protocol, str(_port))
if len(self.acl) < self.max_acl_rules:
if len(self.acl) > position_index > -1:
try:
self.acl.insert(position_index, new_rule)
except Exception:
print(len(self._acl))
if len(self._acl) + 1 < self.max_acl_rules:
if _position is not None:
if self.max_acl_rules - 1 > position_index > -1:
try:
self._acl.insert(position_index, new_rule)
except Exception:
_LOGGER.info(
f"New Rule could NOT be added to list at position {position_index}."
)
else:
_LOGGER.info(
f"New Rule could NOT be added to list at position {position_index}."
f"Position {position_index} is an invalid index for list/overwrites implicit firewall rule"
)
else:
_LOGGER.info(
f"Position {position_index} is an invalid index for list and/or overwrites implicit firewall rule"
)
self.acl.append(new_rule)
else:
_LOGGER.info(
f"The ACL list is FULL."

2
src/primaite/config/training_config.py
View File

@@ -1,5 +1,4 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
import logging
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, Final, Optional, Union
@@ -10,7 +9,6 @@ from primaite import USERS_CONFIG_DIR, getLogger
from primaite.common.enums import ActionType
_LOGGER = getLogger(__name__)
logging.basicConfig(level=logging.DEBUG, format="%(message)s")
_EXAMPLE_TRAINING: Final[Path] = USERS_CONFIG_DIR / "example_config" / "training"

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

@@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Dict, Final, List, Tuple, Union
import numpy as np
from gym import spaces
from primaite.acl.acl_rule import ACLRule
from primaite.common.enums import (
FileSystemState,
HardwareState,
@@ -22,7 +23,6 @@ from primaite.nodes.service_node import ServiceNode
if TYPE_CHECKING:
from primaite.environment.primaite_env import Primaite
_LOGGER = logging.getLogger(__name__)
@@ -346,16 +346,19 @@ class AccessControlList(AbstractObservationComponent):
# 1. Define the shape of your observation space component
acl_shape = [
len(RulePermissionType),
len(env.nodes),
len(env.nodes),
len(env.nodes) + 1,
len(env.nodes) + 1,
len(env.services_list),
len(env.ports_list),
env.max_number_acl_rules,
]
len(acl_shape)
# shape = acl_shape
shape = acl_shape * self.env.max_number_acl_rules
# 2. Create Observation space
self.space = spaces.MultiDiscrete(shape)
print("obs space:", self.space)
# 3. Initialise observation with zeroes
self.current_observation = np.zeros(len(shape), dtype=self._DATA_TYPE)
@@ -365,67 +368,85 @@ class AccessControlList(AbstractObservationComponent):
The structure of the observation space is described in :class:`.AccessControlList`
"""
obs = []
for acl_rule in self.env.acl.acl:
permission = acl_rule.permission
source_ip = acl_rule.source_ip
dest_ip = acl_rule.dest_ip
protocol = acl_rule.protocol
port = acl_rule.port
position = self.env.acl.acl.index(acl_rule)
if permission == "DENY":
permission_int = 0
else:
permission_int = 1
if source_ip == "ANY":
source_ip_int = 0
else:
source_ip_int = self.obtain_node_id_using_ip(source_ip)
if dest_ip == "ANY":
dest_ip_int = 0
else:
dest_ip_int = self.obtain_node_id_using_ip(dest_ip)
if protocol == "ANY":
protocol_int = 0
else:
try:
protocol_int = Protocol[protocol].value
except AttributeError:
_LOGGER.info(f"Service {protocol} could not be found")
protocol_int = -1
if port == "ANY":
port_int = 0
else:
if port in self.env.ports_list:
port_int = self.env.ports_list.index(port)
for index in range(len(self.env.acl.acl)):
acl_rule = self.env.acl.acl[index]
if isinstance(acl_rule, ACLRule):
permission = acl_rule.permission
source_ip = acl_rule.source_ip
dest_ip = acl_rule.dest_ip
protocol = acl_rule.protocol
port = acl_rule.port
position = index
source_ip_int = -1
dest_ip_int = -1
if permission == "DENY":
permission_int = 0
else:
_LOGGER.info(f"Port {port} could not be found.")
permission_int = 1
if source_ip == "ANY":
source_ip_int = 0
else:
nodes = list(self.env.nodes.values())
for node in nodes:
# print(node.ip_address, source_ip, node.ip_address == source_ip)
if (
isinstance(node, ServiceNode)
or isinstance(node, ActiveNode)
) and node.ip_address == source_ip:
source_ip_int = node.node_id
break
if dest_ip == "ANY":
dest_ip_int = 0
else:
nodes = list(self.env.nodes.values())
for node in nodes:
if (
isinstance(node, ServiceNode)
or isinstance(node, ActiveNode)
) and node.ip_address == dest_ip:
dest_ip_int = node.node_id
if protocol == "ANY":
protocol_int = 0
else:
try:
protocol_int = Protocol[protocol].value
except AttributeError:
_LOGGER.info(f"Service {protocol} could not be found")
protocol_int = -1
if port == "ANY":
port_int = 0
else:
if port in self.env.ports_list:
port_int = self.env.ports_list.index(port)
else:
_LOGGER.info(f"Port {port} could not be found.")
print(permission_int, source_ip, dest_ip, protocol_int, port_int, position)
obs.extend(
[
permission_int,
source_ip_int,
dest_ip_int,
protocol_int,
port_int,
position,
]
)
# Either do the multiply on the obs space
# Change the obs to
if source_ip_int != -1 and dest_ip_int != -1:
items_to_add = [
permission_int,
source_ip_int,
dest_ip_int,
protocol_int,
port_int,
position,
]
position = position * 6
for item in items_to_add:
obs.insert(position, int(item))
position += 1
else:
items_to_add = [-1, -1, -1, -1, -1, index]
position = index * 6
for item in items_to_add:
obs.insert(position, int(item))
position += 1
self.current_observation[:] = obs
def obtain_node_id_using_ip(self, ip_address):
"""Uses IP address of Nodes to find the ID.
Resolves IP address -> x (node id e.g. 1 or 2 or 3 or 4) for observation space
"""
print(type(self.env.nodes))
for key, node in self.env.nodes.items():
if isinstance(node, ActiveNode) or isinstance(node, ServiceNode):
if node.ip_address == ip_address:
return key
_LOGGER.info(f"Node ID was not found from IP Address {ip_address}")
return -1
self.current_observation = obs
print("current observation space:", self.current_observation)
class ObservationsHandler:

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

@@ -45,7 +45,7 @@ from primaite.pol.red_agent_pol import apply_red_agent_iers, apply_red_agent_nod
from primaite.transactions.transaction import Transaction
_LOGGER = logging.getLogger(__name__)
_LOGGER.setLevel(logging.INFO)
# _LOGGER.setLevel(logging.INFO)
class Primaite(Env):
@@ -119,6 +119,7 @@ class Primaite(Env):
# Create the Access Control List
self.acl = AccessControlList(
self.training_config.apply_implicit_rule,
self.training_config.implicit_acl_rule,
self.training_config.max_number_acl_rules,
)
@@ -546,6 +547,7 @@ class Primaite(Env):
action_destination_ip = readable_action[3]
action_protocol = readable_action[4]
action_port = readable_action[5]
acl_rule_position = readable_action[6]
if action_decision == 0:
# It's decided to do nothing
@@ -595,6 +597,7 @@ class Primaite(Env):
acl_rule_destination,
acl_rule_protocol,
acl_rule_port,
acl_rule_position,
)
elif action_decision == 2:
# Remove the rule
@@ -1172,13 +1175,9 @@ class Primaite(Env):
# [0, num ports] - Port (0 = any, then 1 -> x resolving to port)
# [0, max acl rules - 1] - Position (0 = first index, then 1 -> x index resolving to acl rule in acl list)
# reserve 0 action to be a nothing action
actions = {0: [0, 0, 0, 0, 0, 0]}
actions = {0: [0, 0, 0, 0, 0, 0, 0]}
action_key = 1
print(
"what is this primaite_env.py 1177",
self.training_config.max_number_acl_rules - 1,
)
# 3 possible action decisions, 0=NOTHING, 1=CREATE, 2=DELETE
for action_decision in range(3):
# 2 possible action permissions 0 = DENY, 1 = CREATE
@@ -1188,9 +1187,7 @@ class Primaite(Env):
for dest_ip in range(self.num_nodes + 1):
for protocol in range(self.num_services + 1):
for port in range(self.num_ports + 1):
for position in range(
self.training_config.max_number_acl_rules - 1
):
for position in range(self.max_number_acl_rules - 1):
action = [
action_decision,
action_permission,