From a798d262b86eeefc90566dc3b85a55407f824ee3 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Mon, 18 Dec 2023 14:03:47 +0000 Subject: [PATCH] use names instead of uuids for requests --- .../config/_package_data/example_config.yaml | 128 ++++++----- src/primaite/game/agent/actions.py | 209 ++++++++++-------- src/primaite/game/agent/observations.py | 2 +- src/primaite/game/agent/rewards.py | 1 - src/primaite/game/game.py | 33 +-- .../simulator/file_system/file_system.py | 4 +- src/primaite/simulator/network/container.py | 2 +- .../simulator/network/hardware/base.py | 11 +- 8 files changed, 201 insertions(+), 189 deletions(-) diff --git a/src/primaite/config/_package_data/example_config.yaml b/src/primaite/config/_package_data/example_config.yaml index db0bca74..83a6de73 100644 --- a/src/primaite/config/_package_data/example_config.yaml +++ b/src/primaite/config/_package_data/example_config.yaml @@ -40,9 +40,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 @@ -76,9 +76,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 @@ -110,17 +110,15 @@ agents: - service_name: DNSServer - node_hostname: web_server services: - - service_name: web_server_database_client + - service_name: DatabaseClient - node_hostname: database_server services: - - service_name: database_service + - service_name: DatabaseService folders: - folder_name: database 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 @@ -184,10 +182,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 @@ -407,97 +405,105 @@ agents: 38: action: "NETWORK_NIC_DISABLE" options: - node_id: 1 - nic_id: 1 + node_id: 0 + nic_id: 0 39: action: "NETWORK_NIC_ENABLE" options: - node_id: 1 - nic_id: 1 + node_id: 0 + nic_id: 0 40: action: "NETWORK_NIC_DISABLE" options: - node_id: 2 - nic_id: 1 + node_id: 1 + nic_id: 0 41: action: "NETWORK_NIC_ENABLE" options: - node_id: 2 - nic_id: 1 + node_id: 1 + nic_id: 0 42: action: "NETWORK_NIC_DISABLE" options: - node_id: 3 - nic_id: 1 + node_id: 2 + nic_id: 0 43: action: "NETWORK_NIC_ENABLE" options: - node_id: 3 - nic_id: 1 + node_id: 2 + nic_id: 0 44: action: "NETWORK_NIC_DISABLE" options: - node_id: 4 - nic_id: 1 + node_id: 3 + nic_id: 0 45: action: "NETWORK_NIC_ENABLE" options: - node_id: 4 - nic_id: 1 + node_id: 3 + nic_id: 0 46: action: "NETWORK_NIC_DISABLE" options: - node_id: 5 - nic_id: 1 + node_id: 4 + nic_id: 0 47: action: "NETWORK_NIC_ENABLE" options: - node_id: 5 - nic_id: 1 + node_id: 4 + nic_id: 0 48: action: "NETWORK_NIC_DISABLE" options: - node_id: 5 - nic_id: 2 + node_id: 4 + nic_id: 1 49: action: "NETWORK_NIC_ENABLE" options: - node_id: 5 - nic_id: 2 + node_id: 4 + nic_id: 1 50: action: "NETWORK_NIC_DISABLE" options: - node_id: 6 - nic_id: 1 + node_id: 5 + nic_id: 0 51: action: "NETWORK_NIC_ENABLE" options: - node_id: 6 - nic_id: 1 + node_id: 5 + nic_id: 0 52: action: "NETWORK_NIC_DISABLE" options: - node_id: 7 - nic_id: 1 + node_id: 6 + nic_id: 0 53: action: "NETWORK_NIC_ENABLE" options: - node_id: 7 - nic_id: 1 + node_id: 6 + nic_id: 0 options: nodes: - - node_ref: router_1 - - node_ref: switch_1 - - node_ref: switch_2 - - node_ref: domain_controller - - node_ref: web_server - - node_ref: database_server - - node_ref: backup_server - - node_ref: security_suite - - node_ref: client_1 - - node_ref: client_2 + # - node_name: router_1 + # - node_name: switch_1 + # - node_name: switch_2 + - node_name: domain_controller + - node_name: web_server + applications: + - application_name: DatabaseClient + services: + - 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 @@ -518,7 +524,7 @@ agents: weight: 0.5 options: node_hostname: web_server - service_name: web_server_web_service + service_name: WebServer agent_settings: diff --git a/src/primaite/game/agent/actions.py b/src/primaite/game/agent/actions.py index 8eed3ba4..ff063dbd 100644 --- a/src/primaite/game/agent/actions.py +++ b/src/primaite/game/agent/actions.py @@ -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, "services", service_uuid, self.verb] + return ["network", "node", node_name, "services", service_name, self.verb] class NodeServiceScanAction(NodeServiceAbstractAction): @@ -172,11 +172,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): @@ -203,11 +203,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): @@ -257,12 +257,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, "files", file_uuid, self.verb] + return ["network", "node", node_name, "file_system", "folder", folder_name, "files", file_name, self.verb] class NodeFileScanAction(NodeFileAbstractAction): @@ -328,8 +328,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): @@ -370,7 +370,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, @@ -381,8 +381,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_name: hostname of the router to which the ACL rule should be added. + :type target_router_name: 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. @@ -403,7 +403,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, @@ -464,7 +464,7 @@ class NetworkACLAddRuleAction(AbstractAction): return [ "network", "node", - self.target_router_uuid, + self.target_router_name, "acl", "add_rule", permission_str, @@ -480,23 +480,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_name: Hostname of the router from which the ACL rule should be removed. + :type target_router_name: 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): @@ -523,16 +523,16 @@ 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, + node_name, "nic", - nic_uuid, + nic_num, self.verb, ] @@ -592,12 +592,11 @@ class ActionManager: 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 + 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 @@ -633,8 +632,60 @@ 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.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 n in nodes: + app_list = [a["application_name"] for a in n.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 n.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 n.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 n.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 @@ -643,15 +694,17 @@ class ActionManager: 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 nics = node_obj.nics for nic_uuid, nic_obj in nics.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, @@ -696,7 +749,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. @@ -745,34 +798,31 @@ 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]: + def get_file_name_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. :param node_idx: The index of the node. @@ -781,45 +831,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] """ - 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 + return self.service_names[node_idx][service_idx] - 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. @@ -853,23 +893,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] - nics = list(node_obj.nics.keys()) - if len(nics) <= nic_idx: - return None - return nics[nic_idx] + return nic_idx + 1 @classmethod def from_config(cls, game: "PrimaiteGame", cfg: Dict) -> "ActionManager": @@ -878,7 +913,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`` diff --git a/src/primaite/game/agent/observations.py b/src/primaite/game/agent/observations.py index ac091b77..b2e88b7c 100644 --- a/src/primaite/game/agent/observations.py +++ b/src/primaite/game/agent/observations.py @@ -704,7 +704,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.ethernet_port[nic_num] diff --git a/src/primaite/game/agent/rewards.py b/src/primaite/game/agent/rewards.py index e2c7d6fc..7a57bc2f 100644 --- a/src/primaite/game/agent/rewards.py +++ b/src/primaite/game/agent/rewards.py @@ -161,7 +161,6 @@ class WebServer404Penalty(AbstractReward): """ web_service_state = access_from_nested_dict(state, self.location_in_state) if web_service_state is NOT_PRESENT_IN_STATE: - print("error getting web service state") return 0.0 most_recent_return_code = web_service_state["last_response_status_code"] # TODO: reward needs to use the current web state. Observation should return web state at the time of last scan. diff --git a/src/primaite/game/game.py b/src/primaite/game/game.py index b6b815f1..c0885cd0 100644 --- a/src/primaite/game/game.py +++ b/src/primaite/game/game.py @@ -360,43 +360,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"] = [] - - # 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([]) - # 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 diff --git a/src/primaite/simulator/file_system/file_system.py b/src/primaite/simulator/file_system/file_system.py index c2eb0d2d..f5e734cf 100644 --- a/src/primaite/simulator/file_system/file_system.py +++ b/src/primaite/simulator/file_system/file_system.py @@ -175,7 +175,7 @@ class FileSystem(SimComponent): self.folders[folder.uuid] = folder self._folders_by_name[folder.name] = 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 @@ -282,7 +282,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]: diff --git a/src/primaite/simulator/network/container.py b/src/primaite/simulator/network/container.py index 8d8709d3..1dadd9e2 100644 --- a/src/primaite/simulator/network/container.py +++ b/src/primaite/simulator/network/container.py @@ -235,7 +235,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]: """ diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py index 95c8e570..bbcdfe37 100644 --- a/src/primaite/simulator/network/hardware/base.py +++ b/src/primaite/simulator/network/hardware/base.py @@ -1318,14 +1318,15 @@ class Node(SimComponent): """ if nic.uuid not in self.nics: self.nics[nic.uuid] = nic - self.ethernet_port[len(self.nics)] = nic + new_nic_num = len(self.nics) + self.ethernet_port[new_nic_num] = nic nic._connected_node = self - nic._port_num_on_node = len(self.nics) + nic._port_num_on_node = new_nic_num nic.parent = self self.sys_log.info(f"Connected NIC {nic}") if self.operating_state == NodeOperatingState.ON: nic.enable() - self._nic_request_manager.add_request(nic.uuid, RequestType(func=nic._request_manager)) + self._nic_request_manager.add_request(new_nic_num, RequestType(func=nic._request_manager)) else: msg = f"Cannot connect NIC {nic} as it is already connected" self.sys_log.logger.error(msg) @@ -1342,15 +1343,17 @@ class Node(SimComponent): if isinstance(nic, str): nic = self.nics.get(nic) if nic or nic.uuid in self.nics: + nic_num = -1 for port, _nic in self.ethernet_port.items(): if nic == _nic: self.ethernet_port.pop(port) + nic_num = port break self.nics.pop(nic.uuid) nic.parent = None nic.disable() self.sys_log.info(f"Disconnected NIC {nic}") - self._nic_request_manager.remove_request(nic.uuid) + self._nic_request_manager.remove_request(nic_num) else: msg = f"Cannot disconnect NIC {nic} as it is not connected" self.sys_log.logger.error(msg)