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

#2248 - synced wth dev

Browse Source
This commit is contained in:
Chris McCarthy
2024-02-08 15:27:02 +00:00
parent 411f0a320f cfc254753f
commit 114fb8c077
35 changed files with 1688 additions and 1035 deletions
Download Patch File Download Diff File Expand all files Collapse all files

105
src/primaite/config/_package_data/example_config.yaml
View File

@@ -42,9 +42,9 @@ agents:
- type: NODE_APPLICATION_EXECUTE
options:
nodes:
- node_ref: client_2
- node_name: client_2
applications:
- application_ref: client_2_web_browser
- application_name: WebBrowser
max_folders_per_node: 1
max_files_per_folder: 1
max_services_per_node: 1
@@ -78,9 +78,9 @@ agents:
- type: NODE_OS_SCAN
options:
nodes:
- node_ref: client_1
- node_name: client_1
applications:
- application_ref: data_manipulation_bot
- application_name: DataManipulationBot
max_folders_per_node: 1
max_files_per_folder: 1
max_services_per_node: 1
@@ -119,8 +119,6 @@ agents:
files:
- file_name: database.db
- node_hostname: backup_server
# services:
# - service_ref: backup_service
- node_hostname: security_suite
- node_hostname: client_1
- node_hostname: client_2
@@ -140,21 +138,21 @@ agents:
max_acl_rules: 10
router_hostname: router_1
ip_address_order:
- node_ref: domain_controller
- node_hostname: domain_controller
nic_num: 1
- node_ref: web_server
- node_hostname: web_server
nic_num: 1
- node_ref: database_server
- node_hostname: database_server
nic_num: 1
- node_ref: backup_server
- node_hostname: backup_server
nic_num: 1
- node_ref: security_suite
- node_hostname: security_suite
nic_num: 1
- node_ref: client_1
- node_hostname: client_1
nic_num: 1
- node_ref: client_2
- node_hostname: client_2
nic_num: 1
- node_ref: security_suite
- node_hostname: security_suite
nic_num: 2
ics: null
@@ -185,10 +183,10 @@ agents:
- type: NODE_RESET
- type: NETWORK_ACL_ADDRULE
options:
target_router_ref: router_1
target_router_hostname: router_1
- type: NETWORK_ACL_REMOVERULE
options:
target_router_ref: router_1
target_router_hostname: router_1
- type: NETWORK_NIC_ENABLE
- type: NETWORK_NIC_DISABLE
@@ -408,118 +406,123 @@ agents:
action: "NETWORK_NIC_DISABLE"
options:
node_id: 0
nic_id: 1
nic_id: 0
39:
action: "NETWORK_NIC_ENABLE"
options:
node_id: 0
nic_id: 1
nic_id: 0
40:
action: "NETWORK_NIC_DISABLE"
options:
node_id: 1
nic_id: 1
nic_id: 0
41:
action: "NETWORK_NIC_ENABLE"
options:
node_id: 1
nic_id: 1
nic_id: 0
42:
action: "NETWORK_NIC_DISABLE"
options:
node_id: 2
nic_id: 1
nic_id: 0
43:
action: "NETWORK_NIC_ENABLE"
options:
node_id: 2
nic_id: 1
nic_id: 0
44:
action: "NETWORK_NIC_DISABLE"
options:
node_id: 3
nic_id: 1
nic_id: 0
45:
action: "NETWORK_NIC_ENABLE"
options:
node_id: 3
nic_id: 1
nic_id: 0
46:
action: "NETWORK_NIC_DISABLE"
options:
node_id: 4
nic_id: 1
nic_id: 0
47:
action: "NETWORK_NIC_ENABLE"
options:
node_id: 4
nic_id: 1
nic_id: 0
48:
action: "NETWORK_NIC_DISABLE"
options:
node_id: 4
nic_id: 2
nic_id: 1
49:
action: "NETWORK_NIC_ENABLE"
options:
node_id: 4
nic_id: 2
nic_id: 1
50:
action: "NETWORK_NIC_DISABLE"
options:
node_id: 5
nic_id: 1
nic_id: 0
51:
action: "NETWORK_NIC_ENABLE"
options:
node_id: 5
nic_id: 1
nic_id: 0
52:
action: "NETWORK_NIC_DISABLE"
options:
node_id: 6
nic_id: 1
nic_id: 0
53:
action: "NETWORK_NIC_ENABLE"
options:
node_id: 6
nic_id: 1
nic_id: 0
options:
nodes:
- node_ref: domain_controller
- node_ref: web_server
- node_name: domain_controller
- node_name: web_server
applications:
- application_name: DatabaseClient
services:
- service_ref: web_server_web_service
- node_ref: database_server
services:
- service_ref: database_service
- node_ref: backup_server
- node_ref: security_suite
- node_ref: client_1
- node_ref: client_2
- service_name: WebServer
- node_name: database_server
folders:
- folder_name: database
files:
- file_name: database.db
- node_name: backup_server
- node_name: security_suite
- node_name: client_1
- node_name: client_2
max_folders_per_node: 2
max_files_per_folder: 2
max_services_per_node: 2
max_nics_per_node: 8
max_acl_rules: 10
ip_address_order:
- node_ref: domain_controller
- node_name: domain_controller
nic_num: 1
- node_ref: web_server
- node_name: web_server
nic_num: 1
- node_ref: database_server
- node_name: database_server
nic_num: 1
- node_ref: backup_server
- node_name: backup_server
nic_num: 1
- node_ref: security_suite
- node_name: security_suite
nic_num: 1
- node_ref: client_1
- node_name: client_1
nic_num: 1
- node_ref: client_2
- node_name: client_2
nic_num: 1
- node_ref: security_suite
- node_name: security_suite
nic_num: 2
@@ -537,7 +540,7 @@ agents:
weight: 0.5
options:
node_hostname: web_server
service_name: web_server_web_service
service_name: WebServer
agent_settings:

261
src/primaite/game/agent/actions.py
View File

@@ -33,7 +33,7 @@ class AbstractAction(ABC):
All action init functions should accept **kwargs as a way of ignoring extra arguments.
Since many parameters are defined for the action space as a whole (such as max files per folder, max services
per node), we need to pass those options to every action that gets created. To pervent verbosity, these
per node), we need to pass those options to every action that gets created. To prevent verbosity, these
parameters are just broadcasted to all actions and the actions can pay attention to the ones that apply.
"""
self.name: str = ""
@@ -85,11 +85,11 @@ class NodeServiceAbstractAction(AbstractAction):
def form_request(self, node_id: int, service_id: int) -> List[str]:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
node_uuid = self.manager.get_node_uuid_by_idx(node_id)
service_uuid = self.manager.get_service_uuid_by_idx(node_id, service_id)
if node_uuid is None or service_uuid is None:
node_name = self.manager.get_node_name_by_idx(node_id)
service_name = self.manager.get_service_name_by_idx(node_id, service_id)
if node_name is None or service_name is None:
return ["do_nothing"]
return ["network", "node", node_uuid, "service", service_uuid, self.verb]
return ["network", "node", node_name, "service", service_name, self.verb]
class NodeServiceScanAction(NodeServiceAbstractAction):
@@ -180,11 +180,11 @@ class NodeApplicationAbstractAction(AbstractAction):
def form_request(self, node_id: int, application_id: int) -> List[str]:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
node_uuid = self.manager.get_node_uuid_by_idx(node_id)
application_uuid = self.manager.get_application_uuid_by_idx(node_id, application_id)
if node_uuid is None or application_uuid is None:
node_name = self.manager.get_node_name_by_idx(node_id)
application_name = self.manager.get_application_name_by_idx(node_id, application_id)
if node_name is None or application_name is None:
return ["do_nothing"]
return ["network", "node", node_uuid, "application", application_uuid, self.verb]
return ["network", "node", node_name, "application", application_name, self.verb]
class NodeApplicationExecuteAction(NodeApplicationAbstractAction):
@@ -211,11 +211,11 @@ class NodeFolderAbstractAction(AbstractAction):
def form_request(self, node_id: int, folder_id: int) -> List[str]:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
node_uuid = self.manager.get_node_uuid_by_idx(node_id)
folder_uuid = self.manager.get_folder_uuid_by_idx(node_idx=node_id, folder_idx=folder_id)
if node_uuid is None or folder_uuid is None:
node_name = self.manager.get_node_name_by_idx(node_id)
folder_name = self.manager.get_folder_name_by_idx(node_idx=node_id, folder_idx=folder_id)
if node_name is None or folder_name is None:
return ["do_nothing"]
return ["network", "node", node_uuid, "file_system", "folder", folder_uuid, self.verb]
return ["network", "node", node_name, "file_system", "folder", folder_name, self.verb]
class NodeFolderScanAction(NodeFolderAbstractAction):
@@ -265,12 +265,12 @@ class NodeFileAbstractAction(AbstractAction):
def form_request(self, node_id: int, folder_id: int, file_id: int) -> List[str]:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
node_uuid = self.manager.get_node_uuid_by_idx(node_id)
folder_uuid = self.manager.get_folder_uuid_by_idx(node_idx=node_id, folder_idx=folder_id)
file_uuid = self.manager.get_file_uuid_by_idx(node_idx=node_id, folder_idx=folder_id, file_idx=file_id)
if node_uuid is None or folder_uuid is None or file_uuid is None:
node_name = self.manager.get_node_name_by_idx(node_id)
folder_name = self.manager.get_folder_name_by_idx(node_idx=node_id, folder_idx=folder_id)
file_name = self.manager.get_file_name_by_idx(node_idx=node_id, folder_idx=folder_id, file_idx=file_id)
if node_name is None or folder_name is None or file_name is None:
return ["do_nothing"]
return ["network", "node", node_uuid, "file_system", "folder", folder_uuid, "file", file_uuid, self.verb]
return ["network", "node", node_name, "file_system", "folder", folder_name, "file", file_name, self.verb]
class NodeFileScanAction(NodeFileAbstractAction):
@@ -298,13 +298,12 @@ class NodeFileDeleteAction(NodeFileAbstractAction):
def form_request(self, node_id: int, folder_id: int, file_id: int) -> List[str]:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
node_uuid = self.manager.get_node_uuid_by_idx(node_id)
folder_uuid = self.manager.get_folder_uuid_by_idx(node_idx=node_id, folder_idx=folder_id)
file_uuid = self.manager.get_file_uuid_by_idx(node_idx=node_id, folder_idx=folder_id, file_idx=file_id)
if node_uuid is None or folder_uuid is None or file_uuid is None:
node_name = self.manager.get_node_name_by_idx(node_id)
folder_name = self.manager.get_folder_name_by_idx(node_idx=node_id, folder_idx=folder_id)
file_name = self.manager.get_file_name_by_idx(node_idx=node_id, folder_idx=folder_id, file_idx=file_id)
if node_name is None or folder_name is None or file_name is None:
return ["do_nothing"]
return ["do_nothing"]
# return ["network", "node", node_uuid, "file_system", "delete", "file", folder_uuid, file_uuid]
return ["network", "node", node_name, "file_system", "delete", "file", folder_name, file_name]
class NodeFileRepairAction(NodeFileAbstractAction):
@@ -346,8 +345,8 @@ class NodeAbstractAction(AbstractAction):
def form_request(self, node_id: int) -> List[str]:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
node_uuid = self.manager.get_node_uuid_by_idx(node_id)
return ["network", "node", node_uuid, self.verb]
node_name = self.manager.get_node_name_by_idx(node_id)
return ["network", "node", node_name, self.verb]
class NodeOSScanAction(NodeAbstractAction):
@@ -388,7 +387,7 @@ class NetworkACLAddRuleAction(AbstractAction):
def __init__(
self,
manager: "ActionManager",
target_router_uuid: str,
target_router_hostname: str,
max_acl_rules: int,
num_ips: int,
num_ports: int,
@@ -399,8 +398,8 @@ class NetworkACLAddRuleAction(AbstractAction):
:param manager: Reference to the ActionManager which created this action.
:type manager: ActionManager
:param target_router_uuid: UUID of the router to which the ACL rule should be added.
:type target_router_uuid: str
:param target_router_hostname: hostname of the router to which the ACL rule should be added.
:type target_router_hostname: str
:param max_acl_rules: Maximum number of ACL rules that can be added to the router.
:type max_acl_rules: int
:param num_ips: Number of IP addresses in the simulation.
@@ -421,7 +420,7 @@ class NetworkACLAddRuleAction(AbstractAction):
"dest_port_id": num_ports,
"protocol_id": num_protocols,
}
self.target_router_uuid: str = target_router_uuid
self.target_router_name: str = target_router_hostname
def form_request(
self,
@@ -488,7 +487,7 @@ class NetworkACLAddRuleAction(AbstractAction):
return [
"network",
"node",
self.target_router_uuid,
self.target_router_name,
"acl",
"add_rule",
permission_str,
@@ -504,23 +503,23 @@ class NetworkACLAddRuleAction(AbstractAction):
class NetworkACLRemoveRuleAction(AbstractAction):
"""Action which removes a rule from a router's ACL."""
def __init__(self, manager: "ActionManager", target_router_uuid: str, max_acl_rules: int, **kwargs) -> None:
def __init__(self, manager: "ActionManager", target_router_hostname: str, max_acl_rules: int, **kwargs) -> None:
"""Init method for NetworkACLRemoveRuleAction.
:param manager: Reference to the ActionManager which created this action.
:type manager: ActionManager
:param target_router_uuid: UUID of the router from which the ACL rule should be removed.
:type target_router_uuid: str
:param target_router_hostname: Hostname of the router from which the ACL rule should be removed.
:type target_router_hostname: str
:param max_acl_rules: Maximum number of ACL rules that can be added to the router.
:type max_acl_rules: int
"""
super().__init__(manager=manager)
self.shape: Dict[str, int] = {"position": max_acl_rules}
self.target_router_uuid: str = target_router_uuid
self.target_router_name: str = target_router_hostname
def form_request(self, position: int) -> List[str]:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
return ["network", "node", self.target_router_uuid, "acl", "remove_rule", position]
return ["network", "node", self.target_router_name, "acl", "remove_rule", position]
class NetworkNICAbstractAction(AbstractAction):
@@ -547,18 +546,11 @@ class NetworkNICAbstractAction(AbstractAction):
def form_request(self, node_id: int, nic_id: int) -> List[str]:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
node_uuid = self.manager.get_node_uuid_by_idx(node_idx=node_id)
nic_uuid = self.manager.get_nic_uuid_by_idx(node_idx=node_id, nic_idx=nic_id)
if node_uuid is None or nic_uuid is None:
node_name = self.manager.get_node_name_by_idx(node_idx=node_id)
nic_num = self.manager.get_nic_num_by_idx(node_idx=node_id, nic_idx=nic_id)
if node_name is None or nic_num is None:
return ["do_nothing"]
return [
"network",
"node",
node_uuid,
"network_interface",
nic_uuid,
self.verb,
]
return ["network", "node", node_name, "network_interface", nic_num, self.verb]
class NetworkNICEnableAction(NetworkNICAbstractAction):
@@ -580,7 +572,7 @@ class NetworkNICDisableAction(NetworkNICAbstractAction):
class ActionManager:
"""Class which manages the action space for an agent."""
__act_class_identifiers: Dict[str, type] = {
_act_class_identifiers: Dict[str, type] = {
"DONOTHING": DoNothingAction,
"NODE_SERVICE_SCAN": NodeServiceScanAction,
"NODE_SERVICE_STOP": NodeServiceStopAction,
@@ -616,14 +608,12 @@ class ActionManager:
def __init__(
self,
game: "PrimaiteGame", # reference to game for information lookup
actions: List[str], # stores list of actions available to agent
node_uuids: List[str], # allows mapping index to node
application_uuids: List[List[str]], # allows mapping index to application
service_uuids: List[List[str]], # allows mapping index to service
actions: List[Dict], # stores list of actions available to agent
nodes: List[Dict], # extra configuration for each node
max_folders_per_node: int = 2, # allows calculating shape
max_files_per_folder: int = 2, # allows calculating shape
max_services_per_node: int = 2, # allows calculating shape
max_applications_per_node: int = 10, # allows calculating shape
max_applications_per_node: int = 2, # allows calculating shape
max_nics_per_node: int = 8, # allows calculating shape
max_acl_rules: int = 10, # allows calculating shape
protocols: List[str] = ["TCP", "UDP", "ICMP"], # allow mapping index to protocol
@@ -635,10 +625,11 @@ class ActionManager:
:param game: Reference to the game to which the agent belongs.
:type game: PrimaiteGame
:param actions: List of action types which should be made available to the agent.
:type actions: List[str]
:param node_uuids: List of node UUIDs that this agent can act on.
:type node_uuids: List[str]
:param actions: List of action specs which should be made available to the agent. The keys of each spec are:
'type' and 'options' for passing any options to the action class's init method
:type actions: List[dict]
:param nodes: Extra configuration for each node.
:type nodes: List[Dict]
:param max_folders_per_node: Maximum number of folders per node. Used for calculating action shape.
:type max_folders_per_node: int
:param max_files_per_folder: Maximum number of files per folder. Used for calculating action shape.
@@ -659,26 +650,86 @@ class ActionManager:
:type act_map: Optional[Dict[int, Dict]]
"""
self.game: "PrimaiteGame" = game
self.node_uuids: List[str] = node_uuids
self.application_uuids: List[List[str]] = application_uuids
self.service_uuids: List[List[str]] = service_uuids
self.node_names: List[str] = [n["node_name"] for n in nodes]
"""List of node names in this action space. The list order is the mapping between node index and node name."""
self.application_names: List[List[str]] = []
"""
List of applications per node. The list order gives the two-index mapping between (node_id, app_id) to app name.
The first index corresponds to node id, the second index is the app id on that particular node.
For instance, self.application_names[0][2] is the name of the third application on the first node.
"""
self.service_names: List[List[str]] = []
"""
List of services per node. The list order gives the two-index mapping between (node_id, svc_id) to svc name.
The first index corresponds to node id, the second index is the service id on that particular node.
For instance, self.service_names[0][2] is the name of the third service on the first node.
"""
self.folder_names: List[List[str]] = []
"""
List of folders per node. The list order gives the two-index mapping between (node_id, folder_id) to folder
name. The first index corresponds to node id, the second index is the folder id on that particular node.
For instance, self.folder_names[0][2] is the name of the third folder on the first node.
"""
self.file_names: List[List[List[str]]] = []
"""
List of files per folder per node. The list order gives the three-index mapping between
(node_id, folder_id, file_id) to file name. The first index corresponds to node id, the second index is the
folder id on that particular node, and the third index is the file id in that particular folder.
For instance, self.file_names[0][2][1] is the name of the second file in the third folder on the first node.
"""
# Populate lists of apps, services, files, folders, etc on nodes.
for node in nodes:
app_list = [a["application_name"] for a in node.get("applications", [])]
while len(app_list) < max_applications_per_node:
app_list.append(None)
self.application_names.append(app_list)
svc_list = [s["service_name"] for s in node.get("services", [])]
while len(svc_list) < max_services_per_node:
svc_list.append(None)
self.service_names.append(svc_list)
folder_list = [f["folder_name"] for f in node.get("folders", [])]
while len(folder_list) < max_folders_per_node:
folder_list.append(None)
self.folder_names.append(folder_list)
file_sublist = []
for folder in node.get("folders", [{"files": []}]):
file_list = [f["file_name"] for f in folder.get("files", [])]
while len(file_list) < max_files_per_folder:
file_list.append(None)
file_sublist.append(file_list)
while len(file_sublist) < max_folders_per_node:
file_sublist.append([None] * max_files_per_folder)
self.file_names.append(file_sublist)
self.protocols: List[str] = protocols
self.ports: List[str] = ports
self.ip_address_list: List[str]
# If the user has provided a list of IP addresses, use that. Otherwise, generate a list of IP addresses from
# the nodes in the simulation.
# TODO: refactor. Options:
# 1: This should be pulled out into it's own function for clarity
# 2: The simulation itself should be able to provide a list of IP addresses with its API, rather than having to
# go through the nodes here.
if ip_address_list is not None:
self.ip_address_list = ip_address_list
else:
self.ip_address_list = []
for node_uuid in self.node_uuids:
node_obj = self.game.simulation.network.nodes[node_uuid]
for node_name in self.node_names:
node_obj = self.game.simulation.network.get_node_by_hostname(node_name)
if node_obj is None:
continue
network_interfaces = node_obj.network_interfaces
for nic_uuid, nic_obj in network_interfaces.items():
self.ip_address_list.append(nic_obj.ip_address)
# action_args are settings which are applied to the action space as a whole.
global_action_args = {
"num_nodes": len(node_uuids),
"num_nodes": len(self.node_names),
"num_folders": max_folders_per_node,
"num_files": max_files_per_folder,
"num_services": max_services_per_node,
@@ -702,7 +753,7 @@ class ActionManager:
# and `options` is an optional dict of options to pass to the init method of the action class
act_type = act_spec.get("type")
act_options = act_spec.get("options", {})
self.actions[act_type] = self.__act_class_identifiers[act_type](self, **global_action_args, **act_options)
self.actions[act_type] = self._act_class_identifiers[act_type](self, **global_action_args, **act_options)
self.action_map: Dict[int, Tuple[str, Dict]] = {}
"""
@@ -723,7 +774,7 @@ class ActionManager:
) -> Dict[int, Tuple[str, Dict]]:
"""Generate a list of all the possible actions that could be taken.
This enumerates all actions all combinations of parametes you could choose for those actions. The output
This enumerates all actions all combinations of parameters you could choose for those actions. The output
of this function is intended to populate the self.action_map parameter in the situation where the user provides
a list of action types, but doesn't specify any subset of actions that should be made available to the agent.
@@ -772,35 +823,32 @@ class ActionManager:
"""Return the gymnasium action space for this agent."""
return spaces.Discrete(len(self.action_map))
def get_node_uuid_by_idx(self, node_idx: int) -> str:
def get_node_name_by_idx(self, node_idx: int) -> str:
"""
Get the node UUID corresponding to the given index.
Get the node name corresponding to the given index.
:param node_idx: The index of the node to retrieve.
:type node_idx: int
:return: The node UUID.
:return: The node hostname.
:rtype: str
"""
return self.node_uuids[node_idx]
return self.node_names[node_idx]
def get_folder_uuid_by_idx(self, node_idx: int, folder_idx: int) -> Optional[str]:
def get_folder_name_by_idx(self, node_idx: int, folder_idx: int) -> Optional[str]:
"""
Get the folder UUID corresponding to the given node and folder indices.
Get the folder name corresponding to the given node and folder indices.
:param node_idx: The index of the node.
:type node_idx: int
:param folder_idx: The index of the folder on the node.
:type folder_idx: int
:return: The UUID of the folder. Or None if the node has fewer folders than the given index.
:return: The name of the folder. Or None if the node has fewer folders than the given index.
:rtype: Optional[str]
"""
node_uuid = self.get_node_uuid_by_idx(node_idx)
node = self.game.simulation.network.nodes[node_uuid]
folder_uuids = list(node.file_system.folders.keys())
return folder_uuids[folder_idx] if len(folder_uuids) > folder_idx else None
return self.folder_names[node_idx][folder_idx]
def get_file_uuid_by_idx(self, node_idx: int, folder_idx: int, file_idx: int) -> Optional[str]:
"""Get the file UUID corresponding to the given node, folder, and file indices.
def get_file_name_by_idx(self, node_idx: int, folder_idx: int, file_idx: int) -> Optional[str]:
"""Get the file name corresponding to the given node, folder, and file indices.
:param node_idx: The index of the node.
:type node_idx: int
@@ -808,50 +856,35 @@ class ActionManager:
:type folder_idx: int
:param file_idx: The index of the file in the folder.
:type file_idx: int
:return: The UUID of the file. Or None if the node has fewer folders than the given index, or the folder has
:return: The name of the file. Or None if the node has fewer folders than the given index, or the folder has
fewer files than the given index.
:rtype: Optional[str]
"""
node_uuid = self.get_node_uuid_by_idx(node_idx)
node = self.game.simulation.network.nodes[node_uuid]
folder_uuids = list(node.file_system.folders.keys())
if len(folder_uuids) <= folder_idx:
return None
folder = node.file_system.folders[folder_uuids[folder_idx]]
file_uuids = list(folder.files.keys())
return file_uuids[file_idx] if len(file_uuids) > file_idx else None
return self.file_names[node_idx][folder_idx][file_idx]
def get_service_uuid_by_idx(self, node_idx: int, service_idx: int) -> Optional[str]:
"""Get the service UUID corresponding to the given node and service indices.
def get_service_name_by_idx(self, node_idx: int, service_idx: int) -> Optional[str]:
"""Get the service name corresponding to the given node and service indices.
:param node_idx: The index of the node.
:type node_idx: int
:param service_idx: The index of the service on the node.
:type service_idx: int
:return: The UUID of the service. Or None if the node has fewer services than the given index.
:return: The name of the service. Or None if the node has fewer services than the given index.
:rtype: Optional[str]
"""
# if a mapping was specified, use that mapping, otherwise just use the list of all installed services
if self.service_uuids:
if self.service_uuids[node_idx]:
return self.service_uuids[node_idx][service_idx]
return self.service_names[node_idx][service_idx]
node_uuid = self.get_node_uuid_by_idx(node_idx)
node = self.game.simulation.network.nodes[node_uuid]
service_uuids = list(node.services.keys())
return service_uuids[service_idx] if len(service_uuids) > service_idx else None
def get_application_uuid_by_idx(self, node_idx: int, application_idx: int) -> Optional[str]:
"""Get the application UUID corresponding to the given node and service indices.
def get_application_name_by_idx(self, node_idx: int, application_idx: int) -> Optional[str]:
"""Get the application name corresponding to the given node and service indices.
:param node_idx: The index of the node.
:type node_idx: int
:param application_idx: The index of the service on the node.
:type application_idx: int
:return: The UUID of the service. Or None if the node has fewer services than the given index.
:return: The name of the service. Or None if the node has fewer services than the given index.
:rtype: Optional[str]
"""
return self.application_uuids[node_idx][application_idx]
return self.application_names[node_idx][application_idx]
def get_internet_protocol_by_idx(self, protocol_idx: int) -> str:
"""Get the internet protocol corresponding to the given index.
@@ -885,23 +918,18 @@ class ActionManager:
"""
return self.ports[port_idx]
def get_nic_uuid_by_idx(self, node_idx: int, nic_idx: int) -> str:
def get_nic_num_by_idx(self, node_idx: int, nic_idx: int) -> int:
"""
Get the NIC UUID corresponding to the given node and NIC indices.
Get the NIC number corresponding to the given node and NIC indices.
:param node_idx: The index of the node.
:type node_idx: int
:param nic_idx: The index of the NIC on the node.
:type nic_idx: int
:return: The NIC UUID.
:rtype: str
:return: The NIC number.
:rtype: int
"""
node_uuid = self.get_node_uuid_by_idx(node_idx)
node_obj = self.game.simulation.network.nodes[node_uuid]
network_interfaces = list(node_obj.network_interfaces.keys())
if len(network_interfaces) <= nic_idx:
return None
return network_interfaces[nic_idx]
return nic_idx + 1
@classmethod
def from_config(cls, game: "PrimaiteGame", cfg: Dict) -> "ActionManager":
@@ -910,7 +938,7 @@ class ActionManager:
The action space config supports the following three sections:
1. ``action_list``
``action_list`` contians a list action components which need to be included in the action space.
``action_list`` contains a list action components which need to be included in the action space.
Each action component has a ``type`` which maps to a subclass of AbstractAction, and additional options
which will be passed to the action class's __init__ method during initialisation.
2. ``action_map``
@@ -933,16 +961,15 @@ class ActionManager:
ip_address_order = cfg["options"].pop("ip_address_order", {})
ip_address_list = []
for entry in ip_address_order:
node_ref = entry["node_ref"]
node_name = entry["node_name"]
nic_num = entry["nic_num"]
node_obj = game.simulation.network.get_node_by_hostname(node_ref)
node_obj = game.simulation.network.get_node_by_hostname(node_name)
ip_address = node_obj.network_interface[nic_num].ip_address
ip_address_list.append(ip_address)
obj = cls(
game=game,
actions=cfg["action_list"],
# node_uuids=cfg["options"]["node_uuids"],
**cfg["options"],
protocols=game.options.protocols,
ports=game.options.ports,

2
src/primaite/game/agent/observations.py
View File

@@ -729,7 +729,7 @@ class AclObservation(AbstractObservation):
max_acl_rules = config["options"]["max_acl_rules"]
node_ip_to_idx = {}
for ip_idx, ip_map_config in enumerate(config["ip_address_order"]):
node_ref = ip_map_config["node_ref"]
node_ref = ip_map_config["node_hostname"]
nic_num = ip_map_config["nic_num"]
node_obj = game.simulation.network.nodes[game.ref_map_nodes[node_ref]]
nic_obj = node_obj.network_interface[nic_num]

4
src/primaite/game/agent/rewards.py
View File

@@ -218,7 +218,7 @@ class RewardFunction:
self.current_reward: float = 0.0
self.total_reward: float = 0.0
def regsiter_component(self, component: AbstractReward, weight: float = 1.0) -> None:
def register_component(self, component: AbstractReward, weight: float = 1.0) -> None:
"""Add a reward component to the reward function.
:param component: Instance of a reward component.
@@ -260,5 +260,5 @@ class RewardFunction:
weight = rew_component_cfg.get("weight", 1.0)
rew_class = cls.__rew_class_identifiers[rew_type]
rew_instance = rew_class.from_config(config=rew_component_cfg.get("options", {}), game=game)
new.regsiter_component(component=rew_instance, weight=weight)
new.register_component(component=rew_instance, weight=weight)
return new

45
src/primaite/game/game.py
View File

@@ -343,55 +343,12 @@ class PrimaiteGame:
obs_space = ObservationManager.from_config(observation_space_cfg, game)
# CREATE ACTION SPACE
action_space_cfg["options"]["node_uuids"] = []
action_space_cfg["options"]["application_uuids"] = []
action_space_cfg["options"]["service_uuids"] = []
# if a list of nodes is defined, convert them from node references to node UUIDs
for action_node_option in action_space_cfg.get("options", {}).pop("nodes", {}):
if "node_ref" in action_node_option:
node_uuid = game.ref_map_nodes[action_node_option["node_ref"]]
action_space_cfg["options"]["node_uuids"].append(node_uuid)
if "applications" in action_node_option:
node_application_uuids = []
for application_option in action_node_option["applications"]:
# TODO: fix inconsistency with node uuids and application uuids. The node object get added to
# node_uuid, whereas here the application gets added by uuid.
application_uuid = game.ref_map_applications[application_option["application_ref"]]
node_application_uuids.append(application_uuid)
action_space_cfg["options"]["application_uuids"].append(node_application_uuids)
else:
action_space_cfg["options"]["application_uuids"].append([])
if "services" in action_node_option:
node_service_uuids = []
for service_option in action_node_option["services"]:
service_uuid = game.ref_map_services[service_option["service_ref"]]
node_service_uuids.append(service_uuid)
action_space_cfg["options"]["service_uuids"].append(node_service_uuids)
else:
action_space_cfg["options"]["service_uuids"].append([])
# Each action space can potentially have a different list of nodes that it can apply to. Therefore,
# we will pass node_uuids as a part of the action space config.
# However, it's not possible to specify the node uuids directly in the config, as they are generated
# dynamically, so we have to translate node references to uuids before passing this config on.
if "action_list" in action_space_cfg:
for action_config in action_space_cfg["action_list"]:
if "options" in action_config:
if "target_router_ref" in action_config["options"]:
_target = action_config["options"]["target_router_ref"]
action_config["options"]["target_router_uuid"] = game.ref_map_nodes[_target]
action_space = ActionManager.from_config(game, action_space_cfg)
# CREATE REWARD FUNCTION
rew_function = RewardFunction.from_config(reward_function_cfg, game=game)
# OTHER AGENT SETTINGS
agent_settings = AgentSettings.from_config(agent_cfg.get("agent_settings"))
# CREATE AGENT

8
src/primaite/simulator/file_system/file_system.py
View File

@@ -78,12 +78,12 @@ class FileSystem(SimComponent):
self._delete_manager.add_request(
name="file",
request_type=RequestType(
func=lambda request, context: self.delete_file_by_id(folder_uuid=request[0], file_uuid=request[1])
func=lambda request, context: self.delete_file(folder_name=request[0], file_name=request[1])
),
)
self._delete_manager.add_request(
name="folder",
request_type=RequestType(func=lambda request, context: self.delete_folder_by_id(folder_uuid=request[0])),
request_type=RequestType(func=lambda request, context: self.delete_folder(folder_name=request[0])),
)
rm.add_request(
name="delete",
@@ -171,7 +171,7 @@ class FileSystem(SimComponent):
self.folders[folder.uuid] = folder
self._folder_request_manager.add_request(
name=folder.uuid, request_type=RequestType(func=folder._request_manager)
name=folder.name, request_type=RequestType(func=folder._request_manager)
)
return folder
@@ -280,7 +280,7 @@ class FileSystem(SimComponent):
sys_log=self.sys_log,
)
folder.add_file(file)
self._file_request_manager.add_request(name=file.uuid, request_type=RequestType(func=file._request_manager))
self._file_request_manager.add_request(name=file.name, request_type=RequestType(func=file._request_manager))
return file
def get_file(self, folder_name: str, file_name: str) -> Optional[File]:

2
src/primaite/simulator/file_system/folder.py
View File

@@ -257,7 +257,7 @@ class Folder(FileSystemItemABC):
# add to list
self.files[file.uuid] = file
self._file_request_manager.add_request(file.uuid, RequestType(func=file._request_manager))
self._file_request_manager.add_request(file.name, RequestType(func=file._request_manager))
file.folder = self
def remove_file(self, file: Optional[File]):

8
src/primaite/simulator/network/container.py
View File

@@ -236,7 +236,7 @@ class Network(SimComponent):
node.parent = self
self._nx_graph.add_node(node.hostname)
_LOGGER.debug(f"Added node {node.uuid} to Network {self.uuid}")
self._node_request_manager.add_request(name=node.uuid, request_type=RequestType(func=node._request_manager))
self._node_request_manager.add_request(name=node.hostname, request_type=RequestType(func=node._request_manager))
def get_node_by_hostname(self, hostname: str) -> Optional[Node]:
"""
@@ -261,7 +261,7 @@ class Network(SimComponent):
:type node: Node
"""
if node not in self:
_LOGGER.warning(f"Can't remove node {node.uuid}. It's not in the network.")
_LOGGER.warning(f"Can't remove node {node.hostname}. It's not in the network.")
return
self.nodes.pop(node.uuid)
for i, _node in self._node_id_map.items():
@@ -269,8 +269,8 @@ class Network(SimComponent):
self._node_id_map.pop(i)
break
node.parent = None
_LOGGER.info(f"Removed node {node.uuid} from network {self.uuid}")
self._node_request_manager.remove_request(name=node.uuid)
self._node_request_manager.remove_request(name=node.hostname)
_LOGGER.info(f"Removed node {node.hostname} from network {self.uuid}")
def connect(
self, endpoint_a: Union[WiredNetworkInterface], endpoint_b: Union[WiredNetworkInterface], **kwargs

50
src/primaite/simulator/network/hardware/base.py
View File

@@ -1056,18 +1056,17 @@ class Node(SimComponent):
:param network_interface: The NIC to connect.
:raise NetworkError: If the NIC is already connected.
"""
if network_interface.uuid not in self.network_interfaces:
if network_interface.uuid not in self.network_interface:
self.network_interfaces[network_interface.uuid] = network_interface
self.network_interface[len(self.network_interfaces)] = network_interface
new_nic_num = len(self.network_interfaces)
self.network_interface[new_nic_num] = network_interface
network_interface._connected_node = self
network_interface.port_num = len(self.network_interfaces)
network_interface._port_num_on_node = new_nic_num
network_interface.parent = self
self.sys_log.info(f"Connected Network Interface {network_interface}")
if self.operating_state == NodeOperatingState.ON:
network_interface.enable()
self._nic_request_manager.add_request(
network_interface.uuid, RequestType(func=network_interface._request_manager)
)
self._nic_request_manager.add_request(new_nic_num, RequestType(func=network_interface._request_manager))
else:
msg = f"Cannot connect NIC {network_interface} as it is already connected"
self.sys_log.logger.error(msg)
@@ -1083,17 +1082,20 @@ class Node(SimComponent):
if isinstance(network_interface, str):
network_interface = self.network_interfaces.get(network_interface)
if network_interface or network_interface.uuid in self.network_interfaces:
for port, _nic in self.network_interface.items():
if network_interface == _nic:
network_interface_num = -1
for port, _network_interface in self.network_interface.items():
if network_interface == _network_interface:
self.network_interface.pop(port)
network_interface_num = port
break
self.network_interfaces.pop(network_interface.uuid)
network_interface.parent = None
network_interface.disable()
self.sys_log.info(f"Disconnected Network Interface {network_interface}")
self._nic_request_manager.remove_request(network_interface.uuid)
if network_interface_num != -1:
self._nic_request_manager.remove_request(network_interface_num)
else:
msg = f"Cannot disconnect NIC {network_interface} as it is not connected"
msg = f"Cannot disconnect Network Interface {network_interface} as it is not connected"
self.sys_log.logger.error(msg)
raise NetworkError(msg)
@@ -1142,12 +1144,14 @@ class Node(SimComponent):
:type service: Service
"""
if service in self:
_LOGGER.warning(f"Can't add service {service.uuid} to node {self.uuid}. It's already installed.")
_LOGGER.warning(f"Can't add service {service.name} to node {self.hostname}. It's already installed.")
return
self.services[service.uuid] = service
service.parent = self
service.install() # Perform any additional setup, such as creating files for this service on the node.
self._service_request_manager.add_request(service.uuid, RequestType(func=service._request_manager))
self.sys_log.info(f"Installed service {service.name}")
_LOGGER.debug(f"Added service {service.name} to node {self.hostname}")
self._service_request_manager.add_request(service.name, RequestType(func=service._request_manager))
def uninstall_service(self, service: Service) -> None:
"""
@@ -1157,14 +1161,14 @@ class Node(SimComponent):
:type service: Service
"""
if service not in self:
_LOGGER.warning(f"Can't remove service {service.uuid} from node {self.uuid}. It's not installed.")
_LOGGER.warning(f"Can't remove service {service.name} from node {self.hostname}. It's not installed.")
return
service.uninstall() # Perform additional teardown, such as removing files or restarting the machine.
self.services.pop(service.uuid)
service.parent = None
self.sys_log.info(f"Uninstalled service {service.name}")
_LOGGER.info(f"Removed service {service.uuid} from node {self.uuid}")
self._service_request_manager.remove_request(service.uuid)
_LOGGER.info(f"Removed service {service.name} from node {self.hostname}")
self._service_request_manager.remove_request(service.name)
def install_application(self, application: Application) -> None:
"""
@@ -1174,11 +1178,15 @@ class Node(SimComponent):
:type application: Application
"""
if application in self:
_LOGGER.warning(f"Can't add application {application.uuid} to node {self.uuid}. It's already installed.")
_LOGGER.warning(
f"Can't add application {application.name} to node {self.hostname}. It's already installed."
)
return
self.applications[application.uuid] = application
application.parent = self
self._application_request_manager.add_request(application.uuid, RequestType(func=application._request_manager))
self.sys_log.info(f"Installed application {application.name}")
_LOGGER.debug(f"Added application {application.name} to node {self.hostname}")
self._application_request_manager.add_request(application.name, RequestType(func=application._request_manager))
def uninstall_application(self, application: Application) -> None:
"""
@@ -1188,13 +1196,15 @@ class Node(SimComponent):
:type application: Application
"""
if application not in self:
_LOGGER.warning(f"Can't remove application {application.uuid} from node {self.uuid}. It's not installed.")
_LOGGER.warning(
f"Can't remove application {application.name} from node {self.hostname}. It's not installed."
)
return
self.applications.pop(application.uuid)
application.parent = None
self.sys_log.info(f"Uninstalled application {application.name}")
_LOGGER.info(f"Removed application {application.uuid} from node {self.uuid}")
self._application_request_manager.remove_request(application.uuid)
_LOGGER.info(f"Removed application {application.name} from node {self.hostname}")
self._application_request_manager.remove_request(application.name)
def _shut_down_actions(self):
"""Actions to perform when the node is shut down."""

18
src/primaite/simulator/network/hardware/nodes/network/router.py
View File

@@ -201,6 +201,15 @@ class AccessControlList(SimComponent):
"""
return self._acl
@property
def num_rules(self) -> int:
"""
Get the number of rules in the ACL.
:return: The number of rules in the ACL.
"""
return len([rule for rule in self._acl if rule is not None])
def add_rule(
self,
action: ACLAction,
@@ -357,6 +366,15 @@ class AccessControlList(SimComponent):
)
print(table)
@property
def num_rules(self) -> int:
"""
Get the number of rules in the ACL.
:return: The number of rules in the ACL.
"""
return len([rule for rule in self._acl if rule is not None])
class RouteEntry(SimComponent):
"""

1
src/primaite/simulator/system/applications/web_browser.py
View File

@@ -72,6 +72,7 @@ class WebBrowser(Application):
"""
state = super().describe_state()
state["last_response_status_code"] = self.latest_response.status_code if self.latest_response else None
return state
def reset_component_for_episode(self, episode: int):
"""Reset the original state of the SimComponent."""