901 - changed name of enum in enums.py and added class attriubutes in access_control_list.py
This commit is contained in:
@@ -8,10 +8,12 @@ from primaite.acl.acl_rule import ACLRule
|
|||||||
class AccessControlList:
|
class AccessControlList:
|
||||||
"""Access Control List class."""
|
"""Access Control List class."""
|
||||||
|
|
||||||
def __init__(self):
|
def __init__(self, implicit_permission):
|
||||||
"""Init."""
|
"""Init."""
|
||||||
# A list of ACL Rules
|
# A list of ACL Rules
|
||||||
self.acl: List[ACLRule] = []
|
self.acl: List[ACLRule] = []
|
||||||
|
self.acl_implicit_rule = implicit_permission
|
||||||
|
self.max_acl_rules: int
|
||||||
|
|
||||||
def check_address_match(self, _rule, _source_ip_address, _dest_ip_address):
|
def check_address_match(self, _rule, _source_ip_address, _dest_ip_address):
|
||||||
"""
|
"""
|
||||||
@@ -103,6 +105,7 @@ class AccessControlList:
|
|||||||
_protocol: the protocol
|
_protocol: the protocol
|
||||||
_port: the port
|
_port: the port
|
||||||
"""
|
"""
|
||||||
|
# Add check so you cant remove implicit rule
|
||||||
rule = ACLRule(_permission, _source_ip, _dest_ip, _protocol, str(_port))
|
rule = ACLRule(_permission, _source_ip, _dest_ip, _protocol, str(_port))
|
||||||
hash_value = hash(rule)
|
hash_value = hash(rule)
|
||||||
# There will not always be something 'popable' since the agent will be trying random things
|
# There will not always be something 'popable' since the agent will be trying random things
|
||||||
|
|||||||
@@ -128,7 +128,7 @@ class LinkStatus(Enum):
|
|||||||
OVERLOAD = 4
|
OVERLOAD = 4
|
||||||
|
|
||||||
|
|
||||||
class ImplicitFirewallRule(Enum):
|
class RulePermissionType(Enum):
|
||||||
"""Implicit firewall rule."""
|
"""Implicit firewall rule."""
|
||||||
|
|
||||||
DENY = 0
|
DENY = 0
|
||||||
|
|||||||
@@ -9,7 +9,7 @@ from gym import spaces
|
|||||||
from primaite.common.enums import (
|
from primaite.common.enums import (
|
||||||
FileSystemState,
|
FileSystemState,
|
||||||
HardwareState,
|
HardwareState,
|
||||||
ImplicitFirewallRule,
|
RulePermissionType,
|
||||||
SoftwareState,
|
SoftwareState,
|
||||||
)
|
)
|
||||||
from primaite.nodes.active_node import ActiveNode
|
from primaite.nodes.active_node import ActiveNode
|
||||||
@@ -336,25 +336,16 @@ class AccessControlList(AbstractObservationComponent):
|
|||||||
|
|
||||||
_DATA_TYPE: type = np.int64
|
_DATA_TYPE: type = np.int64
|
||||||
|
|
||||||
def __init__(
|
def __init__(self, env: "Primaite"):
|
||||||
self,
|
|
||||||
env: "Primaite",
|
|
||||||
acl_implicit_rule=ImplicitFirewallRule.DENY,
|
|
||||||
max_acl_rules: int = 5,
|
|
||||||
):
|
|
||||||
super().__init__(env)
|
super().__init__(env)
|
||||||
|
|
||||||
self.acl_implicit_rule: ImplicitFirewallRule = acl_implicit_rule
|
|
||||||
self.max_acl_rules = max_acl_rules
|
|
||||||
|
|
||||||
# 1. Define the shape of your observation space component
|
# 1. Define the shape of your observation space component
|
||||||
acl_shape = [
|
acl_shape = [
|
||||||
len(ImplicitFirewallRule),
|
len(RulePermissionType),
|
||||||
len(env.nodes),
|
len(env.nodes),
|
||||||
len(env.nodes),
|
len(env.nodes),
|
||||||
len(env.services_list),
|
len(env.services_list),
|
||||||
len(env.ports_list),
|
len(env.ports_list),
|
||||||
len(env.acl),
|
|
||||||
]
|
]
|
||||||
shape = acl_shape * self.env.max_acl_rules
|
shape = acl_shape * self.env.max_acl_rules
|
||||||
|
|
||||||
@@ -394,6 +385,7 @@ class ObservationsHandler:
|
|||||||
"NODE_LINK_TABLE": NodeLinkTable,
|
"NODE_LINK_TABLE": NodeLinkTable,
|
||||||
"NODE_STATUSES": NodeStatuses,
|
"NODE_STATUSES": NodeStatuses,
|
||||||
"LINK_TRAFFIC_LEVELS": LinkTrafficLevels,
|
"LINK_TRAFFIC_LEVELS": LinkTrafficLevels,
|
||||||
|
"ACCESS_CONTROL_LIST": AccessControlList,
|
||||||
}
|
}
|
||||||
|
|
||||||
def __init__(self):
|
def __init__(self):
|
||||||
|
|||||||
@@ -21,11 +21,11 @@ from primaite.common.enums import (
|
|||||||
ActionType,
|
ActionType,
|
||||||
FileSystemState,
|
FileSystemState,
|
||||||
HardwareState,
|
HardwareState,
|
||||||
ImplicitFirewallRule,
|
|
||||||
NodePOLInitiator,
|
NodePOLInitiator,
|
||||||
NodePOLType,
|
NodePOLType,
|
||||||
NodeType,
|
NodeType,
|
||||||
Priority,
|
Priority,
|
||||||
|
RulePermissionType,
|
||||||
SoftwareState,
|
SoftwareState,
|
||||||
)
|
)
|
||||||
from primaite.common.service import Service
|
from primaite.common.service import Service
|
||||||
@@ -160,7 +160,7 @@ class Primaite(Env):
|
|||||||
|
|
||||||
# Set by main_config
|
# Set by main_config
|
||||||
# Adds a DENY ALL or ALLOW ALL to the end of the Access Control List
|
# Adds a DENY ALL or ALLOW ALL to the end of the Access Control List
|
||||||
self.acl_implicit_rule = ImplicitFirewallRule.DENY
|
self.acl_implicit_rule = RulePermissionType.DENY
|
||||||
|
|
||||||
# Sets a limit to how many ACL
|
# Sets a limit to how many ACL
|
||||||
self.max_acl_rules = 0
|
self.max_acl_rules = 0
|
||||||
@@ -1173,7 +1173,7 @@ class Primaite(Env):
|
|||||||
def create_acl_action_dict(self):
|
def create_acl_action_dict(self):
|
||||||
"""Creates a dictionary mapping each possible discrete action to more readable multidiscrete action."""
|
"""Creates a dictionary mapping each possible discrete action to more readable multidiscrete action."""
|
||||||
# reserve 0 action to be a nothing action
|
# 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
|
action_key = 1
|
||||||
# 3 possible action decisions, 0=NOTHING, 1=CREATE, 2=DELETE
|
# 3 possible action decisions, 0=NOTHING, 1=CREATE, 2=DELETE
|
||||||
@@ -1185,14 +1185,16 @@ class Primaite(Env):
|
|||||||
for dest_ip in range(self.num_nodes + 1):
|
for dest_ip in range(self.num_nodes + 1):
|
||||||
for protocol in range(self.num_services + 1):
|
for protocol in range(self.num_services + 1):
|
||||||
for port in range(self.num_ports + 1):
|
for port in range(self.num_ports + 1):
|
||||||
action = [
|
for position in range(self.max_acl_rules - 1):
|
||||||
action_decision,
|
action = [
|
||||||
action_permission,
|
action_decision,
|
||||||
source_ip,
|
action_permission,
|
||||||
dest_ip,
|
source_ip,
|
||||||
protocol,
|
dest_ip,
|
||||||
port,
|
protocol,
|
||||||
]
|
port,
|
||||||
|
position,
|
||||||
|
]
|
||||||
# Check to see if its an action we want to include as possible i.e. not a nothing action
|
# Check to see if its an action we want to include as possible i.e. not a nothing action
|
||||||
if is_valid_acl_action_extra(action):
|
if is_valid_acl_action_extra(action):
|
||||||
actions[action_key] = action
|
actions[action_key] = action
|
||||||
|
|||||||
Reference in New Issue
Block a user