From 495d847a71ddf29164952f927df6139a00a85003 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Fri, 15 Dec 2023 13:04:18 +0000 Subject: [PATCH 01/18] Use service and app name for node software requests --- .../simulator/network/hardware/base.py | 28 +++++++++++-------- 1 file changed, 16 insertions(+), 12 deletions(-) diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py index ad3d73aa..95c8e570 100644 --- a/src/primaite/simulator/network/hardware/base.py +++ b/src/primaite/simulator/network/hardware/base.py @@ -1441,14 +1441,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.sys_log.info(f"Installed service {service.name}") - _LOGGER.info(f"Added service {service.uuid} to node {self.uuid}") - self._service_request_manager.add_request(service.uuid, RequestType(func=service._request_manager)) + _LOGGER.info(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: """ @@ -1458,14 +1458,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: """ @@ -1475,13 +1475,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.sys_log.info(f"Installed application {application.name}") - _LOGGER.info(f"Added application {application.uuid} to node {self.uuid}") - self._application_request_manager.add_request(application.uuid, RequestType(func=application._request_manager)) + _LOGGER.info(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: """ @@ -1491,13 +1493,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.""" From a798d262b86eeefc90566dc3b85a55407f824ee3 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Mon, 18 Dec 2023 14:03:47 +0000 Subject: [PATCH 02/18] 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) From a1dcfa291bcd700640563b721d51aafaebdec8b7 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Thu, 21 Dec 2023 09:25:54 +0000 Subject: [PATCH 03/18] Update test configs with new action spec --- .../assets/configs/bad_primaite_session.yaml | 48 +++++----- .../configs/eval_only_primaite_session.yaml | 48 +++++----- tests/assets/configs/multi_agent_session.yaml | 90 +++++++++---------- .../assets/configs/test_primaite_session.yaml | 46 +++++----- .../configs/train_only_primaite_session.yaml | 46 +++++----- .../_file_system/test_file_actions.py | 2 +- 6 files changed, 136 insertions(+), 144 deletions(-) diff --git a/tests/assets/configs/bad_primaite_session.yaml b/tests/assets/configs/bad_primaite_session.yaml index 478cbfae..4c1d7ce7 100644 --- a/tests/assets/configs/bad_primaite_session.yaml +++ b/tests/assets/configs/bad_primaite_session.yaml @@ -29,7 +29,7 @@ agents: - type: DONOTHING options: nodes: - - node_ref: client_2 + - node_hostname: client_2 max_folders_per_node: 1 max_files_per_folder: 1 max_services_per_node: 1 @@ -64,9 +64,9 @@ agents: - type: NODE_OS_SCAN options: nodes: - - node_ref: client_1 + - node_hostname: client_1 applications: - - application_ref: data_manipulation_bot + - application_name: data_manipulation_bot max_folders_per_node: 1 max_files_per_folder: 1 max_services_per_node: 1 @@ -107,8 +107,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 @@ -128,21 +126,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 @@ -172,10 +170,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 @@ -476,16 +474,16 @@ agents: 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_hostname: router_1 + - node_hostname: switch_1 + - node_hostname: switch_2 + - node_hostname: domain_controller + - node_hostname: web_server + - node_hostname: database_server + - node_hostname: backup_server + - node_hostname: security_suite + - node_hostname: client_1 + - node_hostname: client_2 max_folders_per_node: 2 max_files_per_folder: 2 max_services_per_node: 2 diff --git a/tests/assets/configs/eval_only_primaite_session.yaml b/tests/assets/configs/eval_only_primaite_session.yaml index ec6bfb63..29b7937b 100644 --- a/tests/assets/configs/eval_only_primaite_session.yaml +++ b/tests/assets/configs/eval_only_primaite_session.yaml @@ -34,7 +34,7 @@ agents: options: nodes: - - node_ref: client_2 + - node_name: client_2 max_folders_per_node: 1 max_files_per_folder: 1 max_services_per_node: 1 @@ -69,9 +69,9 @@ agents: - type: NODE_OS_SCAN options: nodes: - - node_ref: client_1 + - node_hostname: client_1 applications: - - application_ref: data_manipulation_bot + - application_name: data_manipulation_bot max_folders_per_node: 1 max_files_per_folder: 1 max_services_per_node: 1 @@ -111,8 +111,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 @@ -132,21 +130,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 @@ -176,10 +174,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 @@ -480,16 +478,16 @@ agents: 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_hostname: router_1 + - node_hostname: switch_1 + - node_hostname: switch_2 + - node_hostname: domain_controller + - node_hostname: web_server + - node_hostname: database_server + - node_hostname: backup_server + - node_hostname: security_suite + - node_hostname: client_1 + - node_hostname: client_2 max_folders_per_node: 2 max_files_per_folder: 2 max_services_per_node: 2 diff --git a/tests/assets/configs/multi_agent_session.yaml b/tests/assets/configs/multi_agent_session.yaml index 3671b809..54727790 100644 --- a/tests/assets/configs/multi_agent_session.yaml +++ b/tests/assets/configs/multi_agent_session.yaml @@ -40,7 +40,7 @@ agents: options: nodes: - - node_ref: client_2 + - node_hostname: client_2 max_folders_per_node: 1 max_files_per_folder: 1 max_services_per_node: 1 @@ -75,9 +75,9 @@ agents: - type: NODE_OS_SCAN options: nodes: - - node_ref: client_1 + - node_hostname: client_1 applications: - - application_ref: data_manipulation_bot + - application_name: data_manipulation_bot max_folders_per_node: 1 max_files_per_folder: 1 max_services_per_node: 1 @@ -118,8 +118,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 @@ -139,21 +137,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 @@ -183,10 +181,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 @@ -487,16 +485,16 @@ agents: 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_hostname: router_1 + - node_hostname: switch_1 + - node_hostname: switch_2 + - node_hostname: domain_controller + - node_hostname: web_server + - node_hostname: database_server + - node_hostname: backup_server + - node_hostname: security_suite + - node_hostname: client_1 + - node_hostname: client_2 max_folders_per_node: 2 max_files_per_folder: 2 max_services_per_node: 2 @@ -549,8 +547,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 @@ -570,21 +566,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 @@ -614,10 +610,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 @@ -918,16 +914,16 @@ agents: 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_hostname: router_1 + - node_hostname: switch_1 + - node_hostname: switch_2 + - node_hostname: domain_controller + - node_hostname: web_server + - node_hostname: database_server + - node_hostname: backup_server + - node_hostname: security_suite + - node_hostname: client_1 + - node_hostname: client_2 max_folders_per_node: 2 max_files_per_folder: 2 max_services_per_node: 2 diff --git a/tests/assets/configs/test_primaite_session.yaml b/tests/assets/configs/test_primaite_session.yaml index cc198a64..f677b4e0 100644 --- a/tests/assets/configs/test_primaite_session.yaml +++ b/tests/assets/configs/test_primaite_session.yaml @@ -38,7 +38,7 @@ agents: options: nodes: - - node_ref: client_2 + - node_hostname: client_2 max_folders_per_node: 1 max_files_per_folder: 1 max_services_per_node: 1 @@ -73,9 +73,9 @@ agents: - type: NODE_OS_SCAN options: nodes: - - node_ref: client_1 + - node_hostname: client_1 applications: - - application_ref: data_manipulation_bot + - application_hostname: data_manipulation_bot max_folders_per_node: 1 max_files_per_folder: 1 max_services_per_node: 1 @@ -137,21 +137,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 @@ -181,10 +181,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 @@ -485,16 +485,16 @@ agents: 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_hostname: router_1 + - node_hostname: switch_1 + - node_hostname: switch_2 + - node_hostname: domain_controller + - node_hostname: web_server + - node_hostname: database_server + - node_hostname: backup_server + - node_hostname: security_suite + - node_hostname: client_1 + - node_hostname: client_2 max_folders_per_node: 2 max_files_per_folder: 2 max_services_per_node: 2 diff --git a/tests/assets/configs/train_only_primaite_session.yaml b/tests/assets/configs/train_only_primaite_session.yaml index ebef7f6a..b788e33f 100644 --- a/tests/assets/configs/train_only_primaite_session.yaml +++ b/tests/assets/configs/train_only_primaite_session.yaml @@ -41,7 +41,7 @@ agents: options: nodes: - - node_ref: client_2 + - node_hostname: client_2 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_hostname: client_1 applications: - - application_ref: data_manipulation_bot + - application_name: data_manipulation_bot max_folders_per_node: 1 max_files_per_folder: 1 max_services_per_node: 1 @@ -138,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 @@ -182,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 @@ -486,16 +486,16 @@ agents: 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_hostname: router_1 + - node_hostname: switch_1 + - node_hostname: switch_2 + - node_hostname: domain_controller + - node_hostname: web_server + - node_hostname: database_server + - node_hostname: backup_server + - node_hostname: security_suite + - node_hostname: client_1 + - node_hostname: client_2 max_folders_per_node: 2 max_files_per_folder: 2 max_services_per_node: 2 diff --git a/tests/unit_tests/_primaite/_simulator/_file_system/test_file_actions.py b/tests/unit_tests/_primaite/_simulator/_file_system/test_file_actions.py index aa8faa90..8590153a 100644 --- a/tests/unit_tests/_primaite/_simulator/_file_system/test_file_actions.py +++ b/tests/unit_tests/_primaite/_simulator/_file_system/test_file_actions.py @@ -25,7 +25,7 @@ def test_file_scan_request(populated_file_system): assert file.health_status == FileSystemItemHealthStatus.CORRUPT assert file.visible_health_status == FileSystemItemHealthStatus.GOOD - fs.apply_request(request=["file", file.uuid, "scan"]) + fs.apply_request(request=["file", file.name, "scan"]) assert file.health_status == FileSystemItemHealthStatus.CORRUPT assert file.visible_health_status == FileSystemItemHealthStatus.CORRUPT From 48f1d13fd86c39f2b78b6c41dc00b7709a502c3e Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Wed, 3 Jan 2024 16:23:44 +0000 Subject: [PATCH 04/18] Minor refactor and add comment --- src/primaite/game/agent/actions.py | 21 ++++++++++----------- 1 file changed, 10 insertions(+), 11 deletions(-) diff --git a/src/primaite/game/agent/actions.py b/src/primaite/game/agent/actions.py index ff063dbd..9b6b63cc 100644 --- a/src/primaite/game/agent/actions.py +++ b/src/primaite/game/agent/actions.py @@ -527,14 +527,7 @@ class NetworkNICAbstractAction(AbstractAction): 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_name, - "nic", - nic_num, - self.verb, - ] + return ["network", "node", node_name, "nic", nic_num, self.verb] class NetworkNICEnableAction(NetworkNICAbstractAction): @@ -610,8 +603,8 @@ class ActionManager: :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 nodes: Extra configuration for each node. + :type nodes: 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. @@ -690,6 +683,13 @@ class ActionManager: 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: @@ -936,7 +936,6 @@ class ActionManager: obj = cls( game=game, actions=cfg["action_list"], - # node_uuids=cfg["options"]["node_uuids"], **cfg["options"], protocols=game.options.protocols, ports=game.options.ports, From 25c8ec2ec9f085a5b90870b77c90b44167ba9c1e Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Wed, 3 Jan 2024 18:19:10 +0000 Subject: [PATCH 05/18] Add skeleton for action integration and unit tests --- src/primaite/game/agent/actions.py | 11 +- .../game_layer/test_actions.py | 185 ++++++++++++++++++ tests/unit_tests/_primaite/_game/__init__.py | 0 .../_primaite/_game/_agent/__init__.py | 0 .../_primaite/_game/_agent/test_actions.py | 90 +++++++++ 5 files changed, 281 insertions(+), 5 deletions(-) create mode 100644 tests/integration_tests/game_layer/test_actions.py create mode 100644 tests/unit_tests/_primaite/_game/__init__.py create mode 100644 tests/unit_tests/_primaite/_game/_agent/__init__.py create mode 100644 tests/unit_tests/_primaite/_game/_agent/test_actions.py diff --git a/src/primaite/game/agent/actions.py b/src/primaite/game/agent/actions.py index 9b6b63cc..d65cd8d0 100644 --- a/src/primaite/game/agent/actions.py +++ b/src/primaite/game/agent/actions.py @@ -549,7 +549,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, @@ -584,7 +584,7 @@ class ActionManager: def __init__( self, game: "PrimaiteGame", # reference to game for information lookup - actions: List[str], # stores list of actions available to agent + 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 @@ -601,8 +601,9 @@ 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 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: Dict :param max_folders_per_node: Maximum number of folders per node. Used for calculating action shape. @@ -728,7 +729,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]] = {} """ diff --git a/tests/integration_tests/game_layer/test_actions.py b/tests/integration_tests/game_layer/test_actions.py new file mode 100644 index 00000000..37a680c8 --- /dev/null +++ b/tests/integration_tests/game_layer/test_actions.py @@ -0,0 +1,185 @@ +# Plan for creating integration tests for the actions: +# I need to test that the requests coming out of the actions have the intended effect on the simulation. +# I can do this by creating a simulation, and then running the action on the simulation, and then checking +# the state of the simulation. + +# Steps for creating the integration tests: +# 1. Create a fixture which creates a simulation. +# 2. Create a fixture which creates a game, including a simple agent with some actions. +# 3. Get the agent to perform an action of my choosing. +# 4. Check that the simulation has changed in the way that I expect. +# 5. Repeat for all actions. + +import pytest + +from primaite.game.agent.actions import ActionManager +from primaite.game.agent.interface import ProxyAgent +from primaite.game.agent.observations import ObservationManager +from primaite.game.agent.rewards import RewardFunction +from primaite.game.game import PrimaiteGame +from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState +from primaite.simulator.network.hardware.nodes.computer import Computer +from primaite.simulator.network.hardware.nodes.router import ACLAction, Router +from primaite.simulator.network.hardware.nodes.server import Server +from primaite.simulator.network.hardware.nodes.switch import Switch +from primaite.simulator.network.transmission.network_layer import IPProtocol +from primaite.simulator.network.transmission.transport_layer import Port +from primaite.simulator.sim_container import Simulation +from primaite.simulator.system.applications.web_browser import WebBrowser +from primaite.simulator.system.services.dns.dns_client import DNSClient +from primaite.simulator.system.services.dns.dns_server import DNSServer +from primaite.simulator.system.services.web_server.web_server import WebServer + + +def install_stuff_to_sim(sim: Simulation): + """Create a simulation with a three computers, two switches, and a router.""" + + # 0: Pull out the network + network = sim.network + + # 1: Set up network hardware + # 1.1: Configure the router + router = Router(hostname="router", num_ports=3) + router.power_on() + router.configure_port(port=1, ip_address="10.0.1.1", subnet_mask="255.255.255.0") + router.configure_port(port=2, ip_address="10.0.2.1", subnet_mask="255.255.255.0") + + # 1.2: Create and connect switches + switch_1 = Switch(hostname="switch_1", num_ports=6) + switch_1.power_on() + network.connect(endpoint_a=router.ethernet_ports[1], endpoint_b=switch_1.switch_ports[6]) + router.enable_port(1) + switch_2 = Switch(hostname="switch_2", num_ports=6) + switch_2.power_on() + network.connect(endpoint_a=router.ethernet_ports[2], endpoint_b=switch_2.switch_ports[6]) + router.enable_port(2) + + # 1.3: Create and connect computer + client_1 = Computer( + hostname="client_1", + ip_address="10.0.1.2", + subnet_mask="255.255.255.0", + default_gateway="10.0.1.1", + operating_state=NodeOperatingState.ON, + ) + client_1.power_on() + network.connect( + endpoint_a=client_1.ethernet_port[1], + endpoint_b=switch_1.switch_ports[1], + ) + + # 1.4: Create and connect servers + server_1 = Server( + hostname="server_1", + ip_address="10.0.2.2", + subnet_mask="255.255.255.0", + default_gateway="10.0.2.1", + operating_state=NodeOperatingState.ON, + ) + server_1.power_on() + network.connect(endpoint_a=server_1.ethernet_port[1], endpoint_b=switch_2.switch_ports[1]) + + server_2 = Server( + hostname="server_2", + ip_address="10.0.2.3", + subnet_mask="255.255.255.0", + default_gateway="10.0.2.1", + operating_state=NodeOperatingState.ON, + ) + server_2.power_on() + network.connect(endpoint_a=server_2.ethernet_port[1], endpoint_b=switch_2.switch_ports[2]) + + # 2: Configure base ACL + router.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.ARP, dst_port=Port.ARP, position=22) + router.acl.add_rule(action=ACLAction.PERMIT, protocol=IPProtocol.ICMP, position=23) + router.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.DNS, dst_port=Port.DNS, position=1) + router.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.HTTP, dst_port=Port.HTTP, position=3) + + # 3: Install server software + server_1.software_manager.install(DNSServer) + dns_service: DNSServer = server_1.software_manager.software.get("DNSServer") # noqa + dns_service.dns_register("example.com", server_2.ip_address) + server_2.software_manager.install(WebServer) + + # 4: Check that client came pre-installed with web browser and dns client + assert isinstance(client_1.software_manager.software.get("WebBrowser"), WebBrowser) + assert isinstance(client_1.software_manager.software.get("DNSClient"), DNSClient) + + # 5: Return the simulation + return sim + + +@pytest.fixture +def game(): + """Create a game with a simple agent that can be controlled by the tests.""" + game = PrimaiteGame() + sim = game.simulation + install_stuff_to_sim(sim) + + actions = [ + {"type": "DONOTHING"}, + {"type": "NODE_SERVICE_SCAN"}, + {"type": "NODE_SERVICE_STOP"}, + {"type": "NODE_SERVICE_START"}, + {"type": "NODE_SERVICE_PAUSE"}, + {"type": "NODE_SERVICE_RESUME"}, + {"type": "NODE_SERVICE_RESTART"}, + {"type": "NODE_SERVICE_DISABLE"}, + {"type": "NODE_SERVICE_ENABLE"}, + {"type": "NODE_APPLICATION_EXECUTE"}, + {"type": "NODE_FILE_SCAN"}, + {"type": "NODE_FILE_CHECKHASH"}, + {"type": "NODE_FILE_DELETE"}, + {"type": "NODE_FILE_REPAIR"}, + {"type": "NODE_FILE_RESTORE"}, + {"type": "NODE_FILE_CORRUPT"}, + {"type": "NODE_FOLDER_SCAN"}, + {"type": "NODE_FOLDER_CHECKHASH"}, + {"type": "NODE_FOLDER_REPAIR"}, + {"type": "NODE_FOLDER_RESTORE"}, + {"type": "NODE_OS_SCAN"}, + {"type": "NODE_SHUTDOWN"}, + {"type": "NODE_STARTUP"}, + {"type": "NODE_RESET"}, + {"type": "NETWORK_ACL_ADDRULE", "options": {"target_router_hostname": "router"}}, + {"type": "NETWORK_ACL_REMOVERULE", "options": {"target_router_hostname": "router"}}, + {"type": "NETWORK_NIC_ENABLE"}, + {"type": "NETWORK_NIC_DISABLE"}, + ] + + action_space = ActionManager( + game=game, + actions=actions, # ALL POSSIBLE ACTIONS + nodes=[ + {"node_name": "client_1", "applications": [{"application_name": "WebBrowser"}]}, + {"node_name": "server_1", "services": [{"service_name": "DNSServer"}]}, + {"node_name": "server_2", "services": [{"service_name": "WebServer"}]}, + ], + max_folders_per_node=2, + max_files_per_folder=2, + max_services_per_node=2, + max_applications_per_node=2, + max_nics_per_node=2, + max_acl_rules=10, + protocols=["TCP", "UDP", "ICMP"], + ports=["HTTP", "DNS", "ARP"], + ip_address_list=["10.0.1.1", "10.0.1.2", "10.0.2.1", "10.0.2.2", "10.0.2.3"], + act_map={}, + ) + observation_space = None + reward_function = None + + test_agent = ProxyAgent( + agent_name="test_agent", + action_space=action_space, + observation_space=observation_space, + reward_function=reward_function, + ) + + game.agents.append(test_agent) + + return game, test_agent + + +def test_test(game): + assert True diff --git a/tests/unit_tests/_primaite/_game/__init__.py b/tests/unit_tests/_primaite/_game/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit_tests/_primaite/_game/_agent/__init__.py b/tests/unit_tests/_primaite/_game/_agent/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit_tests/_primaite/_game/_agent/test_actions.py b/tests/unit_tests/_primaite/_game/_agent/test_actions.py new file mode 100644 index 00000000..9b641fe2 --- /dev/null +++ b/tests/unit_tests/_primaite/_game/_agent/test_actions.py @@ -0,0 +1,90 @@ +from unittest.mock import Mock + +import pytest + +from primaite.game.agent.actions import ( + ActionManager, + DoNothingAction, + NodeServiceDisableAction, + NodeServiceEnableAction, + NodeServicePauseAction, + NodeServiceRestartAction, + NodeServiceResumeAction, + NodeServiceScanAction, + NodeServiceStartAction, + NodeServiceStopAction, +) + + +def test_do_nothing_action_form_request(): + """Test that the DoNothingAction can form a request and that it is correct.""" + manager = Mock() + + action = DoNothingAction(manager=manager) + + request = action.form_request() + + assert request == ["do_nothing"] + + +@pytest.mark.parametrize( + "action_class, action_verb", + [ + (NodeServiceScanAction, "scan"), + (NodeServiceStopAction, "stop"), + (NodeServiceStartAction, "start"), + (NodeServicePauseAction, "pause"), + (NodeServiceResumeAction, "resume"), + (NodeServiceRestartAction, "restart"), + (NodeServiceDisableAction, "disable"), + (NodeServiceEnableAction, "enable"), + ], +) # flake8: noqa +@pytest.mark.parametrize( + "node_name, service_name, expect_to_do_nothing", + [ + ("pc_1", "chrome", False), + (None, "chrome", True), + ("pc_1", None, True), + (None, None, True), + ], +) # flake8: noqa +def test_service_action_form_request(node_name, service_name, expect_to_do_nothing, action_class, action_verb): + """Test that the ServiceScanAction can form a request and that it is correct.""" + manager: ActionManager = Mock() + manager.get_node_name_by_idx.return_value = node_name + manager.get_service_name_by_idx.return_value = service_name + + action = action_class(manager=manager, num_nodes=1, num_services=1) + + request = action.form_request(node_id=0, service_id=0) + + if expect_to_do_nothing: + assert request == ["do_nothing"] + else: + assert request == ["network", "node", node_name, "services", service_name, action_verb] + + +@pytest.mark.parametrize( + "node_name, service_name, expect_to_do_nothing", + [ + ("pc_1", "chrome", False), + (None, "chrome", True), + ("pc_1", None, True), + (None, None, True), + ], +) # flake8: noqa +def test_service_scan_form_request(node_name, service_name, expect_to_do_nothing): + """Test that the ServiceScanAction can form a request and that it is correct.""" + manager: ActionManager = Mock() + manager.get_node_name_by_idx.return_value = node_name + manager.get_service_name_by_idx.return_value = service_name + + action = NodeServiceScanAction(manager=manager, num_nodes=1, num_services=1) + + request = action.form_request(node_id=0, service_id=0) + + if expect_to_do_nothing: + assert request == ["do_nothing"] + else: + assert request == ["network", "node", node_name, "services", service_name, "scan"] From 528e3b22a9c9549aebc1ae232ed96e8a99773474 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Thu, 4 Jan 2024 12:47:35 +0000 Subject: [PATCH 06/18] Add integration tests --- .../system/applications/web_browser.py | 1 + .../game_layer/test_actions.py | 124 ++++++++++++++++-- 2 files changed, 115 insertions(+), 10 deletions(-) diff --git a/src/primaite/simulator/system/applications/web_browser.py b/src/primaite/simulator/system/applications/web_browser.py index 7533f6f3..a5738d76 100644 --- a/src/primaite/simulator/system/applications/web_browser.py +++ b/src/primaite/simulator/system/applications/web_browser.py @@ -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.""" diff --git a/tests/integration_tests/game_layer/test_actions.py b/tests/integration_tests/game_layer/test_actions.py index 37a680c8..85660796 100644 --- a/tests/integration_tests/game_layer/test_actions.py +++ b/tests/integration_tests/game_layer/test_actions.py @@ -10,11 +10,13 @@ # 4. Check that the simulation has changed in the way that I expect. # 5. Repeat for all actions. +from typing import Dict, Tuple + import pytest from primaite.game.agent.actions import ActionManager -from primaite.game.agent.interface import ProxyAgent -from primaite.game.agent.observations import ObservationManager +from primaite.game.agent.interface import AbstractAgent, ProxyAgent +from primaite.game.agent.observations import ICSObservation, ObservationManager from primaite.game.agent.rewards import RewardFunction from primaite.game.game import PrimaiteGame from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState @@ -29,6 +31,34 @@ from primaite.simulator.system.applications.web_browser import WebBrowser from primaite.simulator.system.services.dns.dns_client import DNSClient from primaite.simulator.system.services.dns.dns_server import DNSServer from primaite.simulator.system.services.web_server.web_server import WebServer +from primaite.simulator.system.software import SoftwareHealthState + + +class ControlledAgent(AbstractAgent): + """Agent that can be controlled by the tests.""" + + def __init__( + self, + agent_name: str, + action_space: ActionManager, + observation_space: ObservationManager, + reward_function: RewardFunction, + ) -> None: + super().__init__( + agent_name=agent_name, + action_space=action_space, + observation_space=observation_space, + reward_function=reward_function, + ) + self.most_recent_action: Tuple[str, Dict] + + def get_action(self, obs: None, reward: float = 0.0) -> Tuple[str, Dict]: + """Return the agent's most recent action, formatted in CAOS format.""" + return self.most_recent_action + + def store_action(self, action: Tuple[str, Dict]): + """Store the most recent action.""" + self.most_recent_action = action def install_stuff_to_sim(sim: Simulation): @@ -105,12 +135,47 @@ def install_stuff_to_sim(sim: Simulation): assert isinstance(client_1.software_manager.software.get("WebBrowser"), WebBrowser) assert isinstance(client_1.software_manager.software.get("DNSClient"), DNSClient) - # 5: Return the simulation + # 5: Assert that the simulation starts off in the state that we expect + assert len(sim.network.nodes) == 6 + assert len(sim.network.links) == 5 + # 5.1: Assert the router is correctly configured + r = sim.network.routers[0] + for i, acl_rule in enumerate(r.acl.acl): + if i == 1: + assert acl_rule.src_port == acl_rule.dst_port == Port.DNS + elif i == 3: + assert acl_rule.src_port == acl_rule.dst_port == Port.HTTP + elif i == 22: + assert acl_rule.src_port == acl_rule.dst_port == Port.ARP + elif i == 23: + assert acl_rule.protocol == IPProtocol.ICMP + elif i == 24: + ... + else: + assert acl_rule is None + + # 5.2: Assert the client is correctly configured + c: Computer = [node for node in sim.network.nodes.values() if node.hostname == "client_1"][0] + assert c.software_manager.software.get("WebBrowser") is not None + assert c.software_manager.software.get("DNSClient") is not None + assert str(c.ethernet_port[1].ip_address) == "10.0.1.2" + + # 5.3: Assert that server_1 is correctly configured + s1: Server = [node for node in sim.network.nodes.values() if node.hostname == "server_1"][0] + assert str(s1.ethernet_port[1].ip_address) == "10.0.2.2" + assert s1.software_manager.software.get("DNSServer") is not None + + # 5.4: Assert that server_2 is correctly configured + s2: Server = [node for node in sim.network.nodes.values() if node.hostname == "server_2"][0] + assert str(s2.ethernet_port[1].ip_address) == "10.0.2.3" + assert s2.software_manager.software.get("WebServer") is not None + + # 6: Return the simulation return sim @pytest.fixture -def game(): +def game_and_agent(): """Create a game with a simple agent that can be controlled by the tests.""" game = PrimaiteGame() sim = game.simulation @@ -166,10 +231,10 @@ def game(): ip_address_list=["10.0.1.1", "10.0.1.2", "10.0.2.1", "10.0.2.2", "10.0.2.3"], act_map={}, ) - observation_space = None - reward_function = None + observation_space = ObservationManager(ICSObservation()) + reward_function = RewardFunction() - test_agent = ProxyAgent( + test_agent = ControlledAgent( agent_name="test_agent", action_space=action_space, observation_space=observation_space, @@ -178,8 +243,47 @@ def game(): game.agents.append(test_agent) - return game, test_agent + return (game, test_agent) -def test_test(game): - assert True +# def test_test(game_and_agent:Tuple[PrimaiteGame, ProxyAgent]): +# game, agent = game_and_agent + + +def test_do_nothing_integration(game_and_agent: Tuple[PrimaiteGame, ProxyAgent]): + """Test that the DoNothingAction can form a request and that it is accepted by the simulation.""" + game, agent = game_and_agent + + action = ("DONOTHING", {}) + agent.store_action(action) + game.step() + + +@pytest.mark.skip(reason="Waiting to merge ticket 2160") +def test_node_service_scan_integration(game_and_agent: Tuple[PrimaiteGame, ProxyAgent]): + """ + Test that the NodeServiceScanAction can form a request and that it is accepted by the simulation. + + The health status of applications is not always updated in the state dict, rather the agent needs to perform a scan. + Therefore, we set the web browser to be corrupted, check the state is still good, then perform a scan, and check + that the state changes to the true value. + """ + game, agent = game_and_agent + + browser = game.simulation.network.get_node_by_hostname("client_1").software_manager.software.get("WebBrowser") + browser.health_state_actual = SoftwareHealthState.COMPROMISED + + state_before = game.get_sim_state() + assert ( + game.get_sim_state()["network"]["nodes"]["client_1"]["applications"]["WebBrowser"]["health_state"] + == SoftwareHealthState.GOOD + ) + action = ("NODE_SERVICE_SCAN", {"node_id": 0, "service_id": 0}) + agent.store_action(action) + game.step() + state_after = game.get_sim_state() + pass + assert ( + game.get_sim_state()["network"]["nodes"]["client_1"]["services"]["WebBrowser"]["health_state"] + == SoftwareHealthState.COMPROMISED + ) From 27f70204ed009c9fa18ea6de935c082792caeea7 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Mon, 8 Jan 2024 13:28:34 +0000 Subject: [PATCH 07/18] Fix minor issues in actions --- src/primaite/game/agent/actions.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/primaite/game/agent/actions.py b/src/primaite/game/agent/actions.py index 47cd2394..d99d3818 100644 --- a/src/primaite/game/agent/actions.py +++ b/src/primaite/game/agent/actions.py @@ -89,7 +89,7 @@ class NodeServiceAbstractAction(AbstractAction): 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_name, "services", service_name, self.verb] + return ["network", "node", node_name, "service", service_name, self.verb] class NodeServiceScanAction(NodeServiceAbstractAction): @@ -460,13 +460,13 @@ class NetworkACLAddRuleAction(AbstractAction): dst_ip = "ALL" return ["do_nothing"] # NOT SUPPORTED, JUST DO NOTHING IF WE COME ACROSS THIS else: - dst_ip = self.manager.get_ip_address_by_idx(dest_ip_id) + dst_ip = self.manager.get_ip_address_by_idx(dest_ip_id - 2) # subtract 2 to account for UNUSED=0, and ALL=1 if dest_port_id == 1: dst_port = "ALL" else: - dst_port = self.manager.get_port_by_idx(dest_port_id) + dst_port = self.manager.get_port_by_idx(dest_port_id - 2) # subtract 2 to account for UNUSED=0, and ALL=1 return [ From 294c8b982f48556c053dceecc19337ca100a4e42 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Mon, 8 Jan 2024 13:28:55 +0000 Subject: [PATCH 08/18] Add convenience method for router acl --- src/primaite/simulator/network/hardware/nodes/router.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/src/primaite/simulator/network/hardware/nodes/router.py b/src/primaite/simulator/network/hardware/nodes/router.py index 0234934d..0e6bc946 100644 --- a/src/primaite/simulator/network/hardware/nodes/router.py +++ b/src/primaite/simulator/network/hardware/nodes/router.py @@ -319,6 +319,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): """ From 534f84ccd1382cadfdbc3c7017a3e2b0c6597e02 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Mon, 8 Jan 2024 13:29:17 +0000 Subject: [PATCH 09/18] Add action tests. --- .../game_layer/test_actions.py | 133 +++++++++++++++--- .../_primaite/_game/_agent/test_actions.py | 4 +- 2 files changed, 118 insertions(+), 19 deletions(-) diff --git a/tests/integration_tests/game_layer/test_actions.py b/tests/integration_tests/game_layer/test_actions.py index 85660796..b756553e 100644 --- a/tests/integration_tests/game_layer/test_actions.py +++ b/tests/integration_tests/game_layer/test_actions.py @@ -69,17 +69,17 @@ def install_stuff_to_sim(sim: Simulation): # 1: Set up network hardware # 1.1: Configure the router - router = Router(hostname="router", num_ports=3) + router = Router(hostname="router", num_ports=3, operating_state=NodeOperatingState.ON) router.power_on() router.configure_port(port=1, ip_address="10.0.1.1", subnet_mask="255.255.255.0") router.configure_port(port=2, ip_address="10.0.2.1", subnet_mask="255.255.255.0") # 1.2: Create and connect switches - switch_1 = Switch(hostname="switch_1", num_ports=6) + switch_1 = Switch(hostname="switch_1", num_ports=6, operating_state=NodeOperatingState.ON) switch_1.power_on() network.connect(endpoint_a=router.ethernet_ports[1], endpoint_b=switch_1.switch_ports[6]) router.enable_port(1) - switch_2 = Switch(hostname="switch_2", num_ports=6) + switch_2 = Switch(hostname="switch_2", num_ports=6, operating_state=NodeOperatingState.ON) switch_2.power_on() network.connect(endpoint_a=router.ethernet_ports[2], endpoint_b=switch_2.switch_ports[6]) router.enable_port(2) @@ -191,6 +191,7 @@ def game_and_agent(): {"type": "NODE_SERVICE_RESTART"}, {"type": "NODE_SERVICE_DISABLE"}, {"type": "NODE_SERVICE_ENABLE"}, + {"type": "NODE_SERVICE_PATCH"}, {"type": "NODE_APPLICATION_EXECUTE"}, {"type": "NODE_FILE_SCAN"}, {"type": "NODE_FILE_CHECKHASH"}, @@ -259,31 +260,129 @@ def test_do_nothing_integration(game_and_agent: Tuple[PrimaiteGame, ProxyAgent]) game.step() -@pytest.mark.skip(reason="Waiting to merge ticket 2160") +# @pytest.mark.skip(reason="Waiting to merge ticket 2166") def test_node_service_scan_integration(game_and_agent: Tuple[PrimaiteGame, ProxyAgent]): """ Test that the NodeServiceScanAction can form a request and that it is accepted by the simulation. The health status of applications is not always updated in the state dict, rather the agent needs to perform a scan. - Therefore, we set the web browser to be corrupted, check the state is still good, then perform a scan, and check + Therefore, we set a service to be compromised, check the state is still good, then perform a scan, and check that the state changes to the true value. """ game, agent = game_and_agent - browser = game.simulation.network.get_node_by_hostname("client_1").software_manager.software.get("WebBrowser") - browser.health_state_actual = SoftwareHealthState.COMPROMISED + # 1: Check that the service starts off in a good state, and that visible state is hidden until first scan + svc = game.simulation.network.get_node_by_hostname("client_1").software_manager.software.get("DNSClient") + assert svc.health_state_actual == SoftwareHealthState.GOOD + assert svc.health_state_visible == SoftwareHealthState.UNUSED - state_before = game.get_sim_state() - assert ( - game.get_sim_state()["network"]["nodes"]["client_1"]["applications"]["WebBrowser"]["health_state"] - == SoftwareHealthState.GOOD - ) + # 2: Scan and check that the visible state is now correct action = ("NODE_SERVICE_SCAN", {"node_id": 0, "service_id": 0}) agent.store_action(action) game.step() - state_after = game.get_sim_state() - pass - assert ( - game.get_sim_state()["network"]["nodes"]["client_1"]["services"]["WebBrowser"]["health_state"] - == SoftwareHealthState.COMPROMISED + assert svc.health_state_actual == SoftwareHealthState.GOOD + assert svc.health_state_visible == SoftwareHealthState.GOOD + + # 3: Corrupt the service and check that the visible state is still good + svc.health_state_actual = SoftwareHealthState.COMPROMISED + assert svc.health_state_visible == SoftwareHealthState.GOOD + + # 4: Scan and check that the visible state is now correct + action = ("NODE_SERVICE_SCAN", {"node_id": 0, "service_id": 0}) + agent.store_action(action) + game.step() + assert svc.health_state_actual == SoftwareHealthState.COMPROMISED + assert svc.health_state_visible == SoftwareHealthState.COMPROMISED + + +def test_node_service_patch_integration(game_and_agent: Tuple[PrimaiteGame, ProxyAgent]): + """ + Test that the NodeServicePatchAction can form a request and that it is accepted by the simulation. + + When you initiate a patch action, the software health state turns to PATCHING, then after a few steps, it goes + to GOOD. + """ + game, agent = game_and_agent + + # 1: Corrupt the service + svc = game.simulation.network.get_node_by_hostname("server_1").software_manager.software.get("DNSServer") + svc.health_state_actual = SoftwareHealthState.COMPROMISED + + # 2: Apply a patch action + action = ("NODE_SERVICE_PATCH", {"node_id": 1, "service_id": 0}) + agent.store_action(action) + game.step() + + # 3: Check that the service is now in the patching state + assert svc.health_state_actual == SoftwareHealthState.PATCHING + + # 4: perform a few do-nothing steps and check that the service is now in the good state + action = ("DONOTHING", {}) + agent.store_action(action) + game.step() + assert svc.health_state_actual == SoftwareHealthState.GOOD + + +def test_network_acl_addrule_integration(game_and_agent: Tuple[PrimaiteGame, ProxyAgent]): + """ + Test that the NetworkACLAddRuleAction can form a request and that it is accepted by the simulation. + + The ACL starts off with 4 rules, and we add a rule, and check that the ACL now has 5 rules. + """ + game, agent = game_and_agent + + # 1: Check that traffic is normal and acl starts off with 4 rules. + client_1 = game.simulation.network.get_node_by_hostname("client_1") + server_1 = game.simulation.network.get_node_by_hostname("server_1") + server_2 = game.simulation.network.get_node_by_hostname("server_2") + router = game.simulation.network.get_node_by_hostname("router") + assert router.acl.num_rules == 4 + assert client_1.ping("10.0.2.3") # client_1 can ping server_2 + assert server_2.ping("10.0.1.2") # server_2 can ping client_1 + + # 2: Add a rule to block client 1 from reaching server 2 on router + action = ( + "NETWORK_ACL_ADDRULE", + { + "position": 4, # 4th rule + "permission": 2, # DENY + "source_ip_id": 3, # 10.0.1.2 (client_1) + "dest_ip_id": 6, # 10.0.2.3 (server_2) + "dest_port_id": 1, # ALL + "source_port_id": 1, # ALL + "protocol_id": 1, # ALL + }, ) + agent.store_action(action) + game.step() + + # 3: Check that the ACL now has 5 rules, and that client 1 cannot ping server 2 + assert router.acl.num_rules == 5 + assert not client_1.ping("10.0.2.3") # Cannot ping server_2 + assert client_1.ping("10.0.2.2") # Can ping server_1 + assert not server_2.ping( + "10.0.1.2" + ) # Server 2 can't ping client_1 (although rule is one-way, the ping response is blocked) + + # 4: Add a rule to block server_1 from reaching server_2 on router (this should not affect comms as they are on same subnet) + action = ( + "NETWORK_ACL_ADDRULE", + { + "position": 5, # 5th rule + "permission": 2, # DENY + "source_ip_id": 5, # 10.0.2.2 (server_1) + "dest_ip_id": 6, # 10.0.2.3 (server_2) + "dest_port_id": 1, # ALL + "source_port_id": 1, # ALL + "protocol_id": 1, # ALL + }, + ) + agent.store_action(action) + game.step() + + # 5: Check that the ACL now has 6 rules, but that server_1 can still ping server_2 + assert router.acl.num_rules == 6 + assert server_1.ping("10.0.2.3") # Can ping server_2 + + +# def test_network_acl_removerule_integration() diff --git a/tests/unit_tests/_primaite/_game/_agent/test_actions.py b/tests/unit_tests/_primaite/_game/_agent/test_actions.py index 9b641fe2..b41e22c9 100644 --- a/tests/unit_tests/_primaite/_game/_agent/test_actions.py +++ b/tests/unit_tests/_primaite/_game/_agent/test_actions.py @@ -62,7 +62,7 @@ def test_service_action_form_request(node_name, service_name, expect_to_do_nothi if expect_to_do_nothing: assert request == ["do_nothing"] else: - assert request == ["network", "node", node_name, "services", service_name, action_verb] + assert request == ["network", "node", node_name, "service", service_name, action_verb] @pytest.mark.parametrize( @@ -87,4 +87,4 @@ def test_service_scan_form_request(node_name, service_name, expect_to_do_nothing if expect_to_do_nothing: assert request == ["do_nothing"] else: - assert request == ["network", "node", node_name, "services", service_name, "scan"] + assert request == ["network", "node", node_name, "service", service_name, "scan"] From 7c0ff8e3f06e8d7f56cd78a685a6d292b2f5d757 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Mon, 8 Jan 2024 16:24:09 +0000 Subject: [PATCH 10/18] Add acl remove rule integration test. --- .../game_layer/test_actions.py | 100 +++++++++++++++++- 1 file changed, 97 insertions(+), 3 deletions(-) diff --git a/tests/integration_tests/game_layer/test_actions.py b/tests/integration_tests/game_layer/test_actions.py index b756553e..2bc3b095 100644 --- a/tests/integration_tests/game_layer/test_actions.py +++ b/tests/integration_tests/game_layer/test_actions.py @@ -128,9 +128,13 @@ def install_stuff_to_sim(sim: Simulation): # 3: Install server software server_1.software_manager.install(DNSServer) dns_service: DNSServer = server_1.software_manager.software.get("DNSServer") # noqa - dns_service.dns_register("example.com", server_2.ip_address) + dns_service.dns_register("www.example.com", server_2.ip_address) server_2.software_manager.install(WebServer) + # 3.1: Ensure that the dns clients are configured correctly + client_1.software_manager.software.get("DNSClient").dns_server = server_1.ethernet_port[1].ip_address + server_2.software_manager.software.get("DNSClient").dns_server = server_1.ethernet_port[1].ip_address + # 4: Check that client came pre-installed with web browser and dns client assert isinstance(client_1.software_manager.software.get("WebBrowser"), WebBrowser) assert isinstance(client_1.software_manager.software.get("DNSClient"), DNSClient) @@ -260,7 +264,7 @@ def test_do_nothing_integration(game_and_agent: Tuple[PrimaiteGame, ProxyAgent]) game.step() -# @pytest.mark.skip(reason="Waiting to merge ticket 2166") +@pytest.mark.skip(reason="Waiting to merge ticket 2166") def test_node_service_scan_integration(game_and_agent: Tuple[PrimaiteGame, ProxyAgent]): """ Test that the NodeServiceScanAction can form a request and that it is accepted by the simulation. @@ -385,4 +389,94 @@ def test_network_acl_addrule_integration(game_and_agent: Tuple[PrimaiteGame, Pro assert server_1.ping("10.0.2.3") # Can ping server_2 -# def test_network_acl_removerule_integration() +def test_network_acl_removerule_integration(game_and_agent: Tuple[PrimaiteGame, ProxyAgent]): + """Test that the NetworkACLRemoveRuleAction can form a request and that it is accepted by the simulation.""" + game, agent = game_and_agent + + # 1: Check that http traffic is going across the network nicely. + client_1 = game.simulation.network.get_node_by_hostname("client_1") + server_1 = game.simulation.network.get_node_by_hostname("server_1") + router = game.simulation.network.get_node_by_hostname("router") + + browser: WebBrowser = client_1.software_manager.software.get("WebBrowser") + browser.run() + browser.target_url = "http://www.example.com" + assert browser.get_webpage() # check that the browser can access example.com before we block it + + # 2: Remove rule that allows HTTP traffic across the network + action = ( + "NETWORK_ACL_REMOVERULE", + { + "position": 3, # 4th rule + }, + ) + agent.store_action(action) + game.step() + + # 3: Check that the ACL now has 3 rules, and that client 1 cannot access example.com + assert router.acl.num_rules == 3 + assert not browser.get_webpage() + client_1.software_manager.software.get("DNSClient").dns_cache.clear() + assert client_1.ping("10.0.2.2") # pinging still works because ICMP is allowed + assert client_1.ping("10.0.2.3") + + +def test_network_nic_disable_integration(game_and_agent: Tuple[PrimaiteGame, ProxyAgent]): + """Test that the NetworkNICDisableAction can form a request and that it is accepted by the simulation.""" + game, agent = game_and_agent + + # 1: Check that client_1 can access the network + client_1 = game.simulation.network.get_node_by_hostname("client_1") + server_1 = game.simulation.network.get_node_by_hostname("server_1") + server_2 = game.simulation.network.get_node_by_hostname("server_2") + + browser: WebBrowser = client_1.software_manager.software.get("WebBrowser") + browser.run() + browser.target_url = "http://www.example.com" + assert browser.get_webpage() # check that the browser can access example.com before we block it + + # 2: Disable the NIC on client_1 + action = ( + "NETWORK_NIC_DISABLE", + { + "node_id": 0, # client_1 + "nic_id": 0, # the only nic (eth-1) + }, + ) + agent.store_action(action) + game.step() + + # 3: Check that the NIC is disabled, and that client 1 cannot access example.com + assert client_1.ethernet_port[1].enabled == False + assert not browser.get_webpage() + assert not client_1.ping("10.0.2.2") + assert not client_1.ping("10.0.2.3") + + # 4: check that servers can still communicate + assert server_1.ping("10.0.2.3") + + +def test_network_nic_enable_integration(game_and_agent: Tuple[PrimaiteGame, ProxyAgent]): + """Test that the NetworkNICEnableAction can form a request and that it is accepted by the simulation.""" + + game, agent = game_and_agent + + # 1: Disable client_1 nic + client_1 = game.simulation.network.get_node_by_hostname("client_1") + client_1.ethernet_port[1].disable() + assert not client_1.ping("10.0.2.2") + + # 2: Use action to enable nic + action = ( + "NETWORK_NIC_ENABLE", + { + "node_id": 0, # client_1 + "nic_id": 0, # the only nic (eth-1) + }, + ) + agent.store_action(action) + game.step() + + # 3: Check that the NIC is enabled, and that client 1 can ping again + assert client_1.ethernet_port[1].enabled == True + assert client_1.ping("10.0.2.3") From 9f993dda5746aaea0e19466eeb1045f0ce53cba2 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Wed, 31 Jan 2024 10:48:40 +0000 Subject: [PATCH 11/18] Fix test config discrepancies --- .../config/_package_data/example_config.yaml | 16 +- src/primaite/game/agent/actions.py | 16 +- .../assets/configs/bad_primaite_session.yaml | 221 +++++---- .../configs/eval_only_primaite_session.yaml | 217 +++++---- tests/assets/configs/multi_agent_session.yaml | 433 +++++++++--------- .../assets/configs/test_primaite_session.yaml | 220 +++++---- .../configs/train_only_primaite_session.yaml | 218 +++++---- 7 files changed, 660 insertions(+), 681 deletions(-) diff --git a/src/primaite/config/_package_data/example_config.yaml b/src/primaite/config/_package_data/example_config.yaml index 68aa9106..7a286931 100644 --- a/src/primaite/config/_package_data/example_config.yaml +++ b/src/primaite/config/_package_data/example_config.yaml @@ -508,21 +508,21 @@ agents: 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 diff --git a/src/primaite/game/agent/actions.py b/src/primaite/game/agent/actions.py index 40c40077..3b1fb926 100644 --- a/src/primaite/game/agent/actions.py +++ b/src/primaite/game/agent/actions.py @@ -298,13 +298,13 @@ 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): @@ -849,7 +849,7 @@ class ActionManager: return self.folder_names[node_idx][folder_idx] 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. + """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 @@ -962,9 +962,9 @@ 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.ethernet_port[nic_num].ip_address ip_address_list.append(ip_address) diff --git a/tests/assets/configs/bad_primaite_session.yaml b/tests/assets/configs/bad_primaite_session.yaml index 3e9be3bb..f5d7850c 100644 --- a/tests/assets/configs/bad_primaite_session.yaml +++ b/tests/assets/configs/bad_primaite_session.yaml @@ -29,7 +29,7 @@ agents: - type: DONOTHING options: nodes: - - node_hostname: client_2 + - node_name: client_2 max_folders_per_node: 1 max_files_per_folder: 1 max_services_per_node: 1 @@ -64,7 +64,7 @@ agents: - type: NODE_OS_SCAN options: nodes: - - node_hostname: client_1 + - node_name: client_1 applications: - application_name: data_manipulation_bot max_folders_per_node: 1 @@ -185,168 +185,167 @@ agents: 1: action: NODE_SERVICE_SCAN options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 # stop webapp service 2: action: NODE_SERVICE_STOP options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 # start webapp service 3: action: "NODE_SERVICE_START" options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 4: action: "NODE_SERVICE_PAUSE" options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 5: action: "NODE_SERVICE_RESUME" options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 6: action: "NODE_SERVICE_RESTART" options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 7: action: "NODE_SERVICE_DISABLE" options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 8: action: "NODE_SERVICE_ENABLE" options: - node_id: 2 - service_id: 1 - 9: + node_id: 1 + service_id: 0 + 9: # check database.db file action: "NODE_FILE_SCAN" options: - node_id: 3 + node_id: 2 folder_id: 1 - file_id: 1 + file_id: 0 10: action: "NODE_FILE_CHECKHASH" options: - node_id: 3 + node_id: 2 folder_id: 1 - file_id: 1 + file_id: 0 11: action: "NODE_FILE_DELETE" options: - node_id: 3 + node_id: 2 folder_id: 1 - file_id: 1 + file_id: 0 12: action: "NODE_FILE_REPAIR" options: - node_id: 3 + node_id: 2 folder_id: 1 - file_id: 1 + file_id: 0 13: - action: "NODE_FILE_RESTORE" + action: "NODE_SERVICE_PATCH" options: - node_id: 3 - folder_id: 1 - file_id: 1 + node_id: 2 + service_id: 0 14: action: "NODE_FOLDER_SCAN" options: - node_id: 3 + node_id: 2 folder_id: 1 15: action: "NODE_FOLDER_CHECKHASH" options: - node_id: 3 + node_id: 2 folder_id: 1 16: action: "NODE_FOLDER_REPAIR" options: - node_id: 3 + node_id: 2 folder_id: 1 17: action: "NODE_FOLDER_RESTORE" options: - node_id: 3 + node_id: 2 folder_id: 1 18: action: "NODE_OS_SCAN" options: - node_id: 3 - 19: + node_id: 2 + 19: # shutdown client 1 action: "NODE_SHUTDOWN" options: - node_id: 6 + node_id: 5 20: action: "NODE_STARTUP" options: - node_id: 6 + node_id: 5 21: action: "NODE_RESET" options: - node_id: 6 - 22: + node_id: 5 + 22: # "ACL: ADDRULE - Block outgoing traffic from client 1" (not supported in Primaite) action: "NETWORK_ACL_ADDRULE" options: position: 1 permission: 2 - source_ip_id: 7 - dest_ip_id: 1 + source_ip_id: 7 # client 1 + dest_ip_id: 1 # ALL source_port_id: 1 dest_port_id: 1 protocol_id: 1 - 23: + 23: # "ACL: ADDRULE - Block outgoing traffic from client 2" (not supported in Primaite) action: "NETWORK_ACL_ADDRULE" options: - position: 1 + position: 2 permission: 2 - source_ip_id: 8 - dest_ip_id: 1 + source_ip_id: 8 # client 2 + dest_ip_id: 1 # ALL source_port_id: 1 dest_port_id: 1 protocol_id: 1 - 24: + 24: # block tcp traffic from client 1 to web app action: "NETWORK_ACL_ADDRULE" options: - position: 1 + position: 3 permission: 2 - source_ip_id: 7 - dest_ip_id: 3 + source_ip_id: 7 # client 1 + dest_ip_id: 3 # web server source_port_id: 1 dest_port_id: 1 protocol_id: 3 - 25: + 25: # block tcp traffic from client 2 to web app action: "NETWORK_ACL_ADDRULE" options: - position: 1 + position: 4 permission: 2 - source_ip_id: 8 - dest_ip_id: 3 + source_ip_id: 8 # client 2 + dest_ip_id: 3 # web server source_port_id: 1 dest_port_id: 1 protocol_id: 3 26: action: "NETWORK_ACL_ADDRULE" options: - position: 1 + position: 5 permission: 2 - source_ip_id: 7 - dest_ip_id: 4 + source_ip_id: 7 # client 1 + dest_ip_id: 4 # database source_port_id: 1 dest_port_id: 1 protocol_id: 3 27: action: "NETWORK_ACL_ADDRULE" options: - position: 1 + position: 6 permission: 2 - source_ip_id: 8 - dest_ip_id: 4 + source_ip_id: 8 # client 2 + dest_ip_id: 4 # database source_port_id: 1 dest_port_id: 1 protocol_id: 3 @@ -393,118 +392,116 @@ 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_hostname: router_1 - - node_hostname: switch_1 - - node_hostname: switch_2 - - node_hostname: domain_controller - - node_hostname: web_server - - node_hostname: database_server - - node_hostname: backup_server - - node_hostname: security_suite - - node_hostname: client_1 - - node_hostname: client_2 + - node_name: domain_controller + - node_name: web_server + - node_name: database_server + - 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 reward_function: @@ -622,7 +619,7 @@ simulation: dns_server: 192.168.1.10 services: - ref: backup_service - type: DatabaseBackup + type: FTPServer - ref: security_suite type: server diff --git a/tests/assets/configs/eval_only_primaite_session.yaml b/tests/assets/configs/eval_only_primaite_session.yaml index 0c3872b0..b46c2f8d 100644 --- a/tests/assets/configs/eval_only_primaite_session.yaml +++ b/tests/assets/configs/eval_only_primaite_session.yaml @@ -69,7 +69,7 @@ agents: - type: NODE_OS_SCAN options: nodes: - - node_hostname: client_1 + - node_name: client_1 applications: - application_name: data_manipulation_bot max_folders_per_node: 1 @@ -189,168 +189,167 @@ agents: 1: action: NODE_SERVICE_SCAN options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 # stop webapp service 2: action: NODE_SERVICE_STOP options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 # start webapp service 3: action: "NODE_SERVICE_START" options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 4: action: "NODE_SERVICE_PAUSE" options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 5: action: "NODE_SERVICE_RESUME" options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 6: action: "NODE_SERVICE_RESTART" options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 7: action: "NODE_SERVICE_DISABLE" options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 8: action: "NODE_SERVICE_ENABLE" options: - node_id: 2 - service_id: 1 - 9: + node_id: 1 + service_id: 0 + 9: # check database.db file action: "NODE_FILE_SCAN" options: - node_id: 3 + node_id: 2 folder_id: 1 - file_id: 1 + file_id: 0 10: action: "NODE_FILE_CHECKHASH" options: - node_id: 3 + node_id: 2 folder_id: 1 - file_id: 1 + file_id: 0 11: action: "NODE_FILE_DELETE" options: - node_id: 3 + node_id: 2 folder_id: 1 - file_id: 1 + file_id: 0 12: action: "NODE_FILE_REPAIR" options: - node_id: 3 + node_id: 2 folder_id: 1 - file_id: 1 + file_id: 0 13: - action: "NODE_FILE_RESTORE" + action: "NODE_SERVICE_PATCH" options: - node_id: 3 - folder_id: 1 - file_id: 1 + node_id: 2 + service_id: 0 14: action: "NODE_FOLDER_SCAN" options: - node_id: 3 + node_id: 2 folder_id: 1 15: action: "NODE_FOLDER_CHECKHASH" options: - node_id: 3 + node_id: 2 folder_id: 1 16: action: "NODE_FOLDER_REPAIR" options: - node_id: 3 + node_id: 2 folder_id: 1 17: action: "NODE_FOLDER_RESTORE" options: - node_id: 3 + node_id: 2 folder_id: 1 18: action: "NODE_OS_SCAN" options: - node_id: 3 - 19: + node_id: 2 + 19: # shutdown client 1 action: "NODE_SHUTDOWN" options: - node_id: 6 + node_id: 5 20: action: "NODE_STARTUP" options: - node_id: 6 + node_id: 5 21: action: "NODE_RESET" options: - node_id: 6 - 22: + node_id: 5 + 22: # "ACL: ADDRULE - Block outgoing traffic from client 1" (not supported in Primaite) action: "NETWORK_ACL_ADDRULE" options: position: 1 permission: 2 - source_ip_id: 7 - dest_ip_id: 1 + source_ip_id: 7 # client 1 + dest_ip_id: 1 # ALL source_port_id: 1 dest_port_id: 1 protocol_id: 1 - 23: + 23: # "ACL: ADDRULE - Block outgoing traffic from client 2" (not supported in Primaite) action: "NETWORK_ACL_ADDRULE" options: - position: 1 + position: 2 permission: 2 - source_ip_id: 8 - dest_ip_id: 1 + source_ip_id: 8 # client 2 + dest_ip_id: 1 # ALL source_port_id: 1 dest_port_id: 1 protocol_id: 1 - 24: + 24: # block tcp traffic from client 1 to web app action: "NETWORK_ACL_ADDRULE" options: - position: 1 + position: 3 permission: 2 - source_ip_id: 7 - dest_ip_id: 3 + source_ip_id: 7 # client 1 + dest_ip_id: 3 # web server source_port_id: 1 dest_port_id: 1 protocol_id: 3 - 25: + 25: # block tcp traffic from client 2 to web app action: "NETWORK_ACL_ADDRULE" options: - position: 1 + position: 4 permission: 2 - source_ip_id: 8 - dest_ip_id: 3 + source_ip_id: 8 # client 2 + dest_ip_id: 3 # web server source_port_id: 1 dest_port_id: 1 protocol_id: 3 26: action: "NETWORK_ACL_ADDRULE" options: - position: 1 + position: 5 permission: 2 - source_ip_id: 7 - dest_ip_id: 4 + source_ip_id: 7 # client 1 + dest_ip_id: 4 # database source_port_id: 1 dest_port_id: 1 protocol_id: 3 27: action: "NETWORK_ACL_ADDRULE" options: - position: 1 + position: 6 permission: 2 - source_ip_id: 8 - dest_ip_id: 4 + source_ip_id: 8 # client 2 + dest_ip_id: 4 # database source_port_id: 1 dest_port_id: 1 protocol_id: 3 @@ -397,118 +396,116 @@ 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_hostname: router_1 - - node_hostname: switch_1 - - node_hostname: switch_2 - - node_hostname: domain_controller - - node_hostname: web_server - - node_hostname: database_server - - node_hostname: backup_server - - node_hostname: security_suite - - node_hostname: client_1 - - node_hostname: client_2 + - node_name: domain_controller + - node_name: web_server + - node_name: database_server + - 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 reward_function: diff --git a/tests/assets/configs/multi_agent_session.yaml b/tests/assets/configs/multi_agent_session.yaml index 87bcc14f..23bd46c2 100644 --- a/tests/assets/configs/multi_agent_session.yaml +++ b/tests/assets/configs/multi_agent_session.yaml @@ -40,7 +40,7 @@ agents: options: nodes: - - node_hostname: client_2 + - node_name: client_2 max_folders_per_node: 1 max_files_per_folder: 1 max_services_per_node: 1 @@ -75,7 +75,7 @@ agents: - type: NODE_OS_SCAN options: nodes: - - node_hostname: client_1 + - node_name: client_1 applications: - application_name: data_manipulation_bot max_folders_per_node: 1 @@ -196,168 +196,167 @@ agents: 1: action: NODE_SERVICE_SCAN options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 # stop webapp service 2: action: NODE_SERVICE_STOP options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 # start webapp service 3: action: "NODE_SERVICE_START" options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 4: action: "NODE_SERVICE_PAUSE" options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 5: action: "NODE_SERVICE_RESUME" options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 6: action: "NODE_SERVICE_RESTART" options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 7: action: "NODE_SERVICE_DISABLE" options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 8: action: "NODE_SERVICE_ENABLE" options: - node_id: 2 - service_id: 1 - 9: + node_id: 1 + service_id: 0 + 9: # check database.db file action: "NODE_FILE_SCAN" options: - node_id: 3 + node_id: 2 folder_id: 1 - file_id: 1 + file_id: 0 10: action: "NODE_FILE_CHECKHASH" options: - node_id: 3 + node_id: 2 folder_id: 1 - file_id: 1 + file_id: 0 11: action: "NODE_FILE_DELETE" options: - node_id: 3 + node_id: 2 folder_id: 1 - file_id: 1 + file_id: 0 12: action: "NODE_FILE_REPAIR" options: - node_id: 3 + node_id: 2 folder_id: 1 - file_id: 1 + file_id: 0 13: - action: "NODE_FILE_RESTORE" + action: "NODE_SERVICE_PATCH" options: - node_id: 3 - folder_id: 1 - file_id: 1 + node_id: 2 + service_id: 0 14: action: "NODE_FOLDER_SCAN" options: - node_id: 3 + node_id: 2 folder_id: 1 15: action: "NODE_FOLDER_CHECKHASH" options: - node_id: 3 + node_id: 2 folder_id: 1 16: action: "NODE_FOLDER_REPAIR" options: - node_id: 3 + node_id: 2 folder_id: 1 17: action: "NODE_FOLDER_RESTORE" options: - node_id: 3 + node_id: 2 folder_id: 1 18: action: "NODE_OS_SCAN" options: - node_id: 3 - 19: + node_id: 2 + 19: # shutdown client 1 action: "NODE_SHUTDOWN" options: - node_id: 6 + node_id: 5 20: action: "NODE_STARTUP" options: - node_id: 6 + node_id: 5 21: action: "NODE_RESET" options: - node_id: 6 - 22: + node_id: 5 + 22: # "ACL: ADDRULE - Block outgoing traffic from client 1" (not supported in Primaite) action: "NETWORK_ACL_ADDRULE" options: position: 1 permission: 2 - source_ip_id: 7 - dest_ip_id: 1 + source_ip_id: 7 # client 1 + dest_ip_id: 1 # ALL source_port_id: 1 dest_port_id: 1 protocol_id: 1 - 23: + 23: # "ACL: ADDRULE - Block outgoing traffic from client 2" (not supported in Primaite) action: "NETWORK_ACL_ADDRULE" options: - position: 1 + position: 2 permission: 2 - source_ip_id: 8 - dest_ip_id: 1 + source_ip_id: 8 # client 2 + dest_ip_id: 1 # ALL source_port_id: 1 dest_port_id: 1 protocol_id: 1 - 24: + 24: # block tcp traffic from client 1 to web app action: "NETWORK_ACL_ADDRULE" options: - position: 1 + position: 3 permission: 2 - source_ip_id: 7 - dest_ip_id: 3 + source_ip_id: 7 # client 1 + dest_ip_id: 3 # web server source_port_id: 1 dest_port_id: 1 protocol_id: 3 - 25: + 25: # block tcp traffic from client 2 to web app action: "NETWORK_ACL_ADDRULE" options: - position: 1 + position: 4 permission: 2 - source_ip_id: 8 - dest_ip_id: 3 + source_ip_id: 8 # client 2 + dest_ip_id: 3 # web server source_port_id: 1 dest_port_id: 1 protocol_id: 3 26: action: "NETWORK_ACL_ADDRULE" options: - position: 1 + position: 5 permission: 2 - source_ip_id: 7 - dest_ip_id: 4 + source_ip_id: 7 # client 1 + dest_ip_id: 4 # database source_port_id: 1 dest_port_id: 1 protocol_id: 3 27: action: "NETWORK_ACL_ADDRULE" options: - position: 1 + position: 6 permission: 2 - source_ip_id: 8 - dest_ip_id: 4 + source_ip_id: 8 # client 2 + dest_ip_id: 4 # database source_port_id: 1 dest_port_id: 1 protocol_id: 3 @@ -404,118 +403,116 @@ 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_hostname: router_1 - - node_hostname: switch_1 - - node_hostname: switch_2 - - node_hostname: domain_controller - - node_hostname: web_server - - node_hostname: database_server - - node_hostname: backup_server - - node_hostname: security_suite - - node_hostname: client_1 - - node_hostname: client_2 + - node_name: domain_controller + - node_name: web_server + - node_name: database_server + - 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 reward_function: @@ -642,168 +639,167 @@ agents: 1: action: NODE_SERVICE_SCAN options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 # stop webapp service 2: action: NODE_SERVICE_STOP options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 # start webapp service 3: action: "NODE_SERVICE_START" options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 4: action: "NODE_SERVICE_PAUSE" options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 5: action: "NODE_SERVICE_RESUME" options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 6: action: "NODE_SERVICE_RESTART" options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 7: action: "NODE_SERVICE_DISABLE" options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 8: action: "NODE_SERVICE_ENABLE" options: - node_id: 2 - service_id: 1 - 9: + node_id: 1 + service_id: 0 + 9: # check database.db file action: "NODE_FILE_SCAN" options: - node_id: 3 + node_id: 2 folder_id: 1 - file_id: 1 + file_id: 0 10: action: "NODE_FILE_CHECKHASH" options: - node_id: 3 + node_id: 2 folder_id: 1 - file_id: 1 + file_id: 0 11: action: "NODE_FILE_DELETE" options: - node_id: 3 + node_id: 2 folder_id: 1 - file_id: 1 + file_id: 0 12: action: "NODE_FILE_REPAIR" options: - node_id: 3 + node_id: 2 folder_id: 1 - file_id: 1 + file_id: 0 13: - action: "NODE_FILE_RESTORE" + action: "NODE_SERVICE_PATCH" options: - node_id: 3 - folder_id: 1 - file_id: 1 + node_id: 2 + service_id: 0 14: action: "NODE_FOLDER_SCAN" options: - node_id: 3 + node_id: 2 folder_id: 1 15: action: "NODE_FOLDER_CHECKHASH" options: - node_id: 3 + node_id: 2 folder_id: 1 16: action: "NODE_FOLDER_REPAIR" options: - node_id: 3 + node_id: 2 folder_id: 1 17: action: "NODE_FOLDER_RESTORE" options: - node_id: 3 + node_id: 2 folder_id: 1 18: action: "NODE_OS_SCAN" options: - node_id: 3 - 19: + node_id: 2 + 19: # shutdown client 1 action: "NODE_SHUTDOWN" options: - node_id: 6 + node_id: 5 20: action: "NODE_STARTUP" options: - node_id: 6 + node_id: 5 21: action: "NODE_RESET" options: - node_id: 6 - 22: + node_id: 5 + 22: # "ACL: ADDRULE - Block outgoing traffic from client 1" (not supported in Primaite) action: "NETWORK_ACL_ADDRULE" options: position: 1 permission: 2 - source_ip_id: 7 - dest_ip_id: 1 + source_ip_id: 7 # client 1 + dest_ip_id: 1 # ALL source_port_id: 1 dest_port_id: 1 protocol_id: 1 - 23: + 23: # "ACL: ADDRULE - Block outgoing traffic from client 2" (not supported in Primaite) action: "NETWORK_ACL_ADDRULE" options: - position: 1 + position: 2 permission: 2 - source_ip_id: 8 - dest_ip_id: 1 + source_ip_id: 8 # client 2 + dest_ip_id: 1 # ALL source_port_id: 1 dest_port_id: 1 protocol_id: 1 - 24: + 24: # block tcp traffic from client 1 to web app action: "NETWORK_ACL_ADDRULE" options: - position: 1 + position: 3 permission: 2 - source_ip_id: 7 - dest_ip_id: 3 + source_ip_id: 7 # client 1 + dest_ip_id: 3 # web server source_port_id: 1 dest_port_id: 1 protocol_id: 3 - 25: + 25: # block tcp traffic from client 2 to web app action: "NETWORK_ACL_ADDRULE" options: - position: 1 + position: 4 permission: 2 - source_ip_id: 8 - dest_ip_id: 3 + source_ip_id: 8 # client 2 + dest_ip_id: 3 # web server source_port_id: 1 dest_port_id: 1 protocol_id: 3 26: action: "NETWORK_ACL_ADDRULE" options: - position: 1 + position: 5 permission: 2 - source_ip_id: 7 - dest_ip_id: 4 + source_ip_id: 7 # client 1 + dest_ip_id: 4 # database source_port_id: 1 dest_port_id: 1 protocol_id: 3 27: action: "NETWORK_ACL_ADDRULE" options: - position: 1 + position: 6 permission: 2 - source_ip_id: 8 - dest_ip_id: 4 + source_ip_id: 8 # client 2 + dest_ip_id: 4 # database source_port_id: 1 dest_port_id: 1 protocol_id: 3 @@ -850,118 +846,115 @@ 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_hostname: router_1 - - node_hostname: switch_1 - - node_hostname: switch_2 - - node_hostname: domain_controller - - node_hostname: web_server - - node_hostname: database_server - - node_hostname: backup_server - - node_hostname: security_suite - - node_hostname: client_1 - - node_hostname: client_2 + - node_name: domain_controller + - node_name: web_server + - node_name: database_server + - 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 reward_function: diff --git a/tests/assets/configs/test_primaite_session.yaml b/tests/assets/configs/test_primaite_session.yaml index 84b1c15f..1aa6ad71 100644 --- a/tests/assets/configs/test_primaite_session.yaml +++ b/tests/assets/configs/test_primaite_session.yaml @@ -38,7 +38,7 @@ agents: options: nodes: - - node_hostname: client_2 + - node_name: client_2 max_folders_per_node: 1 max_files_per_folder: 1 max_services_per_node: 1 @@ -73,9 +73,9 @@ agents: - type: NODE_OS_SCAN options: nodes: - - node_hostname: client_1 + - node_name: client_1 applications: - - application_hostname: data_manipulation_bot + - application_name: data_manipulation_bot max_folders_per_node: 1 max_files_per_folder: 1 max_services_per_node: 1 @@ -196,168 +196,167 @@ agents: 1: action: NODE_SERVICE_SCAN options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 # stop webapp service 2: action: NODE_SERVICE_STOP options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 # start webapp service 3: action: "NODE_SERVICE_START" options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 4: action: "NODE_SERVICE_PAUSE" options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 5: action: "NODE_SERVICE_RESUME" options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 6: action: "NODE_SERVICE_RESTART" options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 7: action: "NODE_SERVICE_DISABLE" options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 8: action: "NODE_SERVICE_ENABLE" options: - node_id: 2 - service_id: 1 - 9: + node_id: 1 + service_id: 0 + 9: # check database.db file action: "NODE_FILE_SCAN" options: - node_id: 3 + node_id: 2 folder_id: 1 - file_id: 1 + file_id: 0 10: action: "NODE_FILE_CHECKHASH" options: - node_id: 3 + node_id: 2 folder_id: 1 - file_id: 1 + file_id: 0 11: action: "NODE_FILE_DELETE" options: - node_id: 3 + node_id: 2 folder_id: 1 - file_id: 1 + file_id: 0 12: action: "NODE_FILE_REPAIR" options: - node_id: 3 + node_id: 2 folder_id: 1 - file_id: 1 + file_id: 0 13: - action: "NODE_FILE_RESTORE" + action: "NODE_SERVICE_PATCH" options: - node_id: 3 - folder_id: 1 - file_id: 1 + node_id: 2 + service_id: 0 14: action: "NODE_FOLDER_SCAN" options: - node_id: 3 + node_id: 2 folder_id: 1 15: action: "NODE_FOLDER_CHECKHASH" options: - node_id: 3 + node_id: 2 folder_id: 1 16: action: "NODE_FOLDER_REPAIR" options: - node_id: 3 + node_id: 2 folder_id: 1 17: action: "NODE_FOLDER_RESTORE" options: - node_id: 3 + node_id: 2 folder_id: 1 18: action: "NODE_OS_SCAN" options: - node_id: 3 - 19: + node_id: 2 + 19: # shutdown client 1 action: "NODE_SHUTDOWN" options: - node_id: 6 + node_id: 5 20: action: "NODE_STARTUP" options: - node_id: 6 + node_id: 5 21: action: "NODE_RESET" options: - node_id: 6 - 22: + node_id: 5 + 22: # "ACL: ADDRULE - Block outgoing traffic from client 1" (not supported in Primaite) action: "NETWORK_ACL_ADDRULE" options: position: 1 permission: 2 - source_ip_id: 7 - dest_ip_id: 1 + source_ip_id: 7 # client 1 + dest_ip_id: 1 # ALL source_port_id: 1 dest_port_id: 1 protocol_id: 1 - 23: + 23: # "ACL: ADDRULE - Block outgoing traffic from client 2" (not supported in Primaite) action: "NETWORK_ACL_ADDRULE" options: - position: 1 + position: 2 permission: 2 - source_ip_id: 8 - dest_ip_id: 1 + source_ip_id: 8 # client 2 + dest_ip_id: 1 # ALL source_port_id: 1 dest_port_id: 1 protocol_id: 1 - 24: + 24: # block tcp traffic from client 1 to web app action: "NETWORK_ACL_ADDRULE" options: - position: 1 + position: 3 permission: 2 - source_ip_id: 7 - dest_ip_id: 3 + source_ip_id: 7 # client 1 + dest_ip_id: 3 # web server source_port_id: 1 dest_port_id: 1 protocol_id: 3 - 25: + 25: # block tcp traffic from client 2 to web app action: "NETWORK_ACL_ADDRULE" options: - position: 1 + position: 4 permission: 2 - source_ip_id: 8 - dest_ip_id: 3 + source_ip_id: 8 # client 2 + dest_ip_id: 3 # web server source_port_id: 1 dest_port_id: 1 protocol_id: 3 26: action: "NETWORK_ACL_ADDRULE" options: - position: 1 + position: 5 permission: 2 - source_ip_id: 7 - dest_ip_id: 4 + source_ip_id: 7 # client 1 + dest_ip_id: 4 # database source_port_id: 1 dest_port_id: 1 protocol_id: 3 27: action: "NETWORK_ACL_ADDRULE" options: - position: 1 + position: 6 permission: 2 - source_ip_id: 8 - dest_ip_id: 4 + source_ip_id: 8 # client 2 + dest_ip_id: 4 # database source_port_id: 1 dest_port_id: 1 protocol_id: 3 @@ -404,118 +403,115 @@ 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_hostname: router_1 - - node_hostname: switch_1 - - node_hostname: switch_2 - - node_hostname: domain_controller - - node_hostname: web_server - - node_hostname: database_server - - node_hostname: backup_server - - node_hostname: security_suite - - node_hostname: client_1 - - node_hostname: client_2 + - node_name: domain_controller + - node_name: web_server + - node_name: database_server + - 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 reward_function: diff --git a/tests/assets/configs/train_only_primaite_session.yaml b/tests/assets/configs/train_only_primaite_session.yaml index 62826cd4..115b6c85 100644 --- a/tests/assets/configs/train_only_primaite_session.yaml +++ b/tests/assets/configs/train_only_primaite_session.yaml @@ -41,7 +41,7 @@ agents: options: nodes: - - node_hostname: client_2 + - node_name: client_2 max_folders_per_node: 1 max_files_per_folder: 1 max_services_per_node: 1 @@ -76,7 +76,7 @@ agents: - type: NODE_OS_SCAN options: nodes: - - node_hostname: client_1 + - node_name: client_1 applications: - application_name: data_manipulation_bot max_folders_per_node: 1 @@ -197,168 +197,167 @@ agents: 1: action: NODE_SERVICE_SCAN options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 # stop webapp service 2: action: NODE_SERVICE_STOP options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 # start webapp service 3: action: "NODE_SERVICE_START" options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 4: action: "NODE_SERVICE_PAUSE" options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 5: action: "NODE_SERVICE_RESUME" options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 6: action: "NODE_SERVICE_RESTART" options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 7: action: "NODE_SERVICE_DISABLE" options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 8: action: "NODE_SERVICE_ENABLE" options: - node_id: 2 - service_id: 1 - 9: + node_id: 1 + service_id: 0 + 9: # check database.db file action: "NODE_FILE_SCAN" options: - node_id: 3 + node_id: 2 folder_id: 1 - file_id: 1 + file_id: 0 10: action: "NODE_FILE_CHECKHASH" options: - node_id: 3 + node_id: 2 folder_id: 1 - file_id: 1 + file_id: 0 11: action: "NODE_FILE_DELETE" options: - node_id: 3 + node_id: 2 folder_id: 1 - file_id: 1 + file_id: 0 12: action: "NODE_FILE_REPAIR" options: - node_id: 3 + node_id: 2 folder_id: 1 - file_id: 1 + file_id: 0 13: - action: "NODE_FILE_RESTORE" + action: "NODE_SERVICE_PATCH" options: - node_id: 3 - folder_id: 1 - file_id: 1 + node_id: 2 + service_id: 0 14: action: "NODE_FOLDER_SCAN" options: - node_id: 3 + node_id: 2 folder_id: 1 15: action: "NODE_FOLDER_CHECKHASH" options: - node_id: 3 + node_id: 2 folder_id: 1 16: action: "NODE_FOLDER_REPAIR" options: - node_id: 3 + node_id: 2 folder_id: 1 17: action: "NODE_FOLDER_RESTORE" options: - node_id: 3 + node_id: 2 folder_id: 1 18: action: "NODE_OS_SCAN" options: - node_id: 3 - 19: + node_id: 2 + 19: # shutdown client 1 action: "NODE_SHUTDOWN" options: - node_id: 6 + node_id: 5 20: action: "NODE_STARTUP" options: - node_id: 6 + node_id: 5 21: action: "NODE_RESET" options: - node_id: 6 - 22: + node_id: 5 + 22: # "ACL: ADDRULE - Block outgoing traffic from client 1" (not supported in Primaite) action: "NETWORK_ACL_ADDRULE" options: position: 1 permission: 2 - source_ip_id: 7 - dest_ip_id: 1 + source_ip_id: 7 # client 1 + dest_ip_id: 1 # ALL source_port_id: 1 dest_port_id: 1 protocol_id: 1 - 23: + 23: # "ACL: ADDRULE - Block outgoing traffic from client 2" (not supported in Primaite) action: "NETWORK_ACL_ADDRULE" options: - position: 1 + position: 2 permission: 2 - source_ip_id: 8 - dest_ip_id: 1 + source_ip_id: 8 # client 2 + dest_ip_id: 1 # ALL source_port_id: 1 dest_port_id: 1 protocol_id: 1 - 24: + 24: # block tcp traffic from client 1 to web app action: "NETWORK_ACL_ADDRULE" options: - position: 1 + position: 3 permission: 2 - source_ip_id: 7 - dest_ip_id: 3 + source_ip_id: 7 # client 1 + dest_ip_id: 3 # web server source_port_id: 1 dest_port_id: 1 protocol_id: 3 - 25: + 25: # block tcp traffic from client 2 to web app action: "NETWORK_ACL_ADDRULE" options: - position: 1 + position: 4 permission: 2 - source_ip_id: 8 - dest_ip_id: 3 + source_ip_id: 8 # client 2 + dest_ip_id: 3 # web server source_port_id: 1 dest_port_id: 1 protocol_id: 3 26: action: "NETWORK_ACL_ADDRULE" options: - position: 1 + position: 5 permission: 2 - source_ip_id: 7 - dest_ip_id: 4 + source_ip_id: 7 # client 1 + dest_ip_id: 4 # database source_port_id: 1 dest_port_id: 1 protocol_id: 3 27: action: "NETWORK_ACL_ADDRULE" options: - position: 1 + position: 6 permission: 2 - source_ip_id: 8 - dest_ip_id: 4 + source_ip_id: 8 # client 2 + dest_ip_id: 4 # database source_port_id: 1 dest_port_id: 1 protocol_id: 3 @@ -405,118 +404,115 @@ 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_hostname: router_1 - - node_hostname: switch_1 - - node_hostname: switch_2 - - node_hostname: domain_controller - - node_hostname: web_server - - node_hostname: database_server - - node_hostname: backup_server - - node_hostname: security_suite - - node_hostname: client_1 - - node_hostname: client_2 + - node_name: domain_controller + - node_name: web_server + - node_name: database_server + - 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 reward_function: From 6aa6383fb73bdecae063f20f14f4d7204baeb5e8 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Wed, 31 Jan 2024 11:44:56 +0000 Subject: [PATCH 12/18] Fix broken test configs --- tests/assets/configs/bad_primaite_session.yaml | 3 ++- tests/assets/configs/eval_only_primaite_session.yaml | 3 ++- tests/assets/configs/multi_agent_session.yaml | 4 +++- tests/assets/configs/test_primaite_session.yaml | 3 ++- tests/assets/configs/train_only_primaite_session.yaml | 3 ++- 5 files changed, 11 insertions(+), 5 deletions(-) diff --git a/tests/assets/configs/bad_primaite_session.yaml b/tests/assets/configs/bad_primaite_session.yaml index f5d7850c..552351b2 100644 --- a/tests/assets/configs/bad_primaite_session.yaml +++ b/tests/assets/configs/bad_primaite_session.yaml @@ -66,7 +66,7 @@ agents: nodes: - node_name: client_1 applications: - - application_name: data_manipulation_bot + - application_name: DataManipulationBot max_folders_per_node: 1 max_files_per_folder: 1 max_services_per_node: 1 @@ -155,6 +155,7 @@ agents: - type: NODE_SERVICE_RESTART - type: NODE_SERVICE_DISABLE - type: NODE_SERVICE_ENABLE + - type: NODE_SERVICE_PATCH - type: NODE_FILE_SCAN - type: NODE_FILE_CHECKHASH - type: NODE_FILE_DELETE diff --git a/tests/assets/configs/eval_only_primaite_session.yaml b/tests/assets/configs/eval_only_primaite_session.yaml index b46c2f8d..d49562c8 100644 --- a/tests/assets/configs/eval_only_primaite_session.yaml +++ b/tests/assets/configs/eval_only_primaite_session.yaml @@ -71,7 +71,7 @@ agents: nodes: - node_name: client_1 applications: - - application_name: data_manipulation_bot + - application_name: DataManipulationBot max_folders_per_node: 1 max_files_per_folder: 1 max_services_per_node: 1 @@ -159,6 +159,7 @@ agents: - type: NODE_SERVICE_RESTART - type: NODE_SERVICE_DISABLE - type: NODE_SERVICE_ENABLE + - type: NODE_SERVICE_PATCH - type: NODE_FILE_SCAN - type: NODE_FILE_CHECKHASH - type: NODE_FILE_DELETE diff --git a/tests/assets/configs/multi_agent_session.yaml b/tests/assets/configs/multi_agent_session.yaml index 23bd46c2..29f0ae7f 100644 --- a/tests/assets/configs/multi_agent_session.yaml +++ b/tests/assets/configs/multi_agent_session.yaml @@ -77,7 +77,7 @@ agents: nodes: - node_name: client_1 applications: - - application_name: data_manipulation_bot + - application_name: DataManipulationBot max_folders_per_node: 1 max_files_per_folder: 1 max_services_per_node: 1 @@ -166,6 +166,7 @@ agents: - type: NODE_SERVICE_RESTART - type: NODE_SERVICE_DISABLE - type: NODE_SERVICE_ENABLE + - type: NODE_SERVICE_PATCH - type: NODE_FILE_SCAN - type: NODE_FILE_CHECKHASH - type: NODE_FILE_DELETE @@ -609,6 +610,7 @@ agents: - type: NODE_SERVICE_RESTART - type: NODE_SERVICE_DISABLE - type: NODE_SERVICE_ENABLE + - type: NODE_SERVICE_PATCH - type: NODE_FILE_SCAN - type: NODE_FILE_CHECKHASH - type: NODE_FILE_DELETE diff --git a/tests/assets/configs/test_primaite_session.yaml b/tests/assets/configs/test_primaite_session.yaml index 1aa6ad71..0c70840d 100644 --- a/tests/assets/configs/test_primaite_session.yaml +++ b/tests/assets/configs/test_primaite_session.yaml @@ -75,7 +75,7 @@ agents: nodes: - node_name: client_1 applications: - - application_name: data_manipulation_bot + - application_name: DataManipulationBot max_folders_per_node: 1 max_files_per_folder: 1 max_services_per_node: 1 @@ -166,6 +166,7 @@ agents: - type: NODE_SERVICE_RESTART - type: NODE_SERVICE_DISABLE - type: NODE_SERVICE_ENABLE + - type: NODE_SERVICE_PATCH - type: NODE_FILE_SCAN - type: NODE_FILE_CHECKHASH - type: NODE_FILE_DELETE diff --git a/tests/assets/configs/train_only_primaite_session.yaml b/tests/assets/configs/train_only_primaite_session.yaml index 115b6c85..0466a5ac 100644 --- a/tests/assets/configs/train_only_primaite_session.yaml +++ b/tests/assets/configs/train_only_primaite_session.yaml @@ -78,7 +78,7 @@ agents: nodes: - node_name: client_1 applications: - - application_name: data_manipulation_bot + - application_name: DataManipulationBot max_folders_per_node: 1 max_files_per_folder: 1 max_services_per_node: 1 @@ -167,6 +167,7 @@ agents: - type: NODE_SERVICE_RESTART - type: NODE_SERVICE_DISABLE - type: NODE_SERVICE_ENABLE + - type: NODE_SERVICE_PATCH - type: NODE_FILE_SCAN - type: NODE_FILE_CHECKHASH - type: NODE_FILE_DELETE From d6a83fd1fb33140b3f4b3952069d1c4c3889c606 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Wed, 31 Jan 2024 11:55:38 +0000 Subject: [PATCH 13/18] Update action tests to use name, not uuid --- src/primaite/simulator/network/container.py | 6 +++--- .../test_action_integration.py | 4 +--- .../_simulator/_file_system/test_file_actions.py | 16 ++++++++-------- .../_file_system/test_folder_actions.py | 16 ++++++++-------- 4 files changed, 20 insertions(+), 22 deletions(-) diff --git a/src/primaite/simulator/network/container.py b/src/primaite/simulator/network/container.py index 1dadd9e2..8989a60f 100644 --- a/src/primaite/simulator/network/container.py +++ b/src/primaite/simulator/network/container.py @@ -260,7 +260,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(): @@ -268,8 +268,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[NIC, SwitchPort], endpoint_b: Union[NIC, SwitchPort], **kwargs diff --git a/tests/integration_tests/component_creation/test_action_integration.py b/tests/integration_tests/component_creation/test_action_integration.py index a2be923b..b506544d 100644 --- a/tests/integration_tests/component_creation/test_action_integration.py +++ b/tests/integration_tests/component_creation/test_action_integration.py @@ -48,8 +48,6 @@ def test_passing_actions_down(monkeypatch) -> None: assert not action_invoked # call the patched method - sim.apply_request( - ["network", "node", pc1.uuid, "file_system", "folder", pc1.file_system.get_folder("downloads").uuid, "repair"] - ) + sim.apply_request(["network", "node", pc1.hostname, "file_system", "folder", "downloads", "repair"]) assert action_invoked diff --git a/tests/unit_tests/_primaite/_simulator/_file_system/test_file_actions.py b/tests/unit_tests/_primaite/_simulator/_file_system/test_file_actions.py index 8590153a..f43652d8 100644 --- a/tests/unit_tests/_primaite/_simulator/_file_system/test_file_actions.py +++ b/tests/unit_tests/_primaite/_simulator/_file_system/test_file_actions.py @@ -35,12 +35,12 @@ def test_file_checkhash_request(populated_file_system): """Test that an agent can request a file hash check.""" fs, folder, file = populated_file_system - fs.apply_request(request=["file", file.uuid, "checkhash"]) + fs.apply_request(request=["file", file.name, "checkhash"]) assert file.health_status == FileSystemItemHealthStatus.GOOD file.sim_size = 0 - fs.apply_request(request=["file", file.uuid, "checkhash"]) + fs.apply_request(request=["file", file.name, "checkhash"]) assert file.health_status == FileSystemItemHealthStatus.CORRUPT @@ -52,7 +52,7 @@ def test_file_repair_request(populated_file_system): file.corrupt() assert file.health_status == FileSystemItemHealthStatus.CORRUPT - fs.apply_request(request=["file", file.uuid, "repair"]) + fs.apply_request(request=["file", file.name, "repair"]) assert file.health_status == FileSystemItemHealthStatus.GOOD @@ -69,7 +69,7 @@ def test_file_restore_request(populated_file_system): assert fs.get_file(folder_name=folder.name, file_name=file.name) is not None assert fs.get_file(folder_name=folder.name, file_name=file.name).deleted is False - fs.apply_request(request=["file", file.uuid, "corrupt"]) + fs.apply_request(request=["file", file.name, "corrupt"]) assert fs.get_file(folder_name=folder.name, file_name=file.name).health_status == FileSystemItemHealthStatus.CORRUPT fs.apply_request(request=["restore", "file", folder.uuid, file.uuid]) @@ -79,7 +79,7 @@ def test_file_restore_request(populated_file_system): def test_file_corrupt_request(populated_file_system): """Test that an agent can request a file corruption.""" fs, folder, file = populated_file_system - fs.apply_request(request=["file", file.uuid, "corrupt"]) + fs.apply_request(request=["file", file.name, "corrupt"]) assert file.health_status == FileSystemItemHealthStatus.CORRUPT @@ -88,7 +88,7 @@ def test_deleted_file_cannot_be_interacted_with(populated_file_system): fs, folder, file = populated_file_system assert fs.get_file(folder_name=folder.name, file_name=file.name) is not None - fs.apply_request(request=["file", file.uuid, "corrupt"]) + fs.apply_request(request=["file", file.name, "corrupt"]) assert fs.get_file(folder_name=folder.name, file_name=file.name).health_status == FileSystemItemHealthStatus.CORRUPT assert ( fs.get_file(folder_name=folder.name, file_name=file.name).visible_health_status @@ -98,8 +98,8 @@ def test_deleted_file_cannot_be_interacted_with(populated_file_system): fs.apply_request(request=["delete", "file", folder.uuid, file.uuid]) assert fs.get_file(folder_name=folder.name, file_name=file.name) is None - fs.apply_request(request=["file", file.uuid, "repair"]) - fs.apply_request(request=["file", file.uuid, "scan"]) + fs.apply_request(request=["file", file.name, "repair"]) + fs.apply_request(request=["file", file.name, "scan"]) file = folder.deleted_files.get(file.uuid) diff --git a/tests/unit_tests/_primaite/_simulator/_file_system/test_folder_actions.py b/tests/unit_tests/_primaite/_simulator/_file_system/test_folder_actions.py index efa74e1f..398af0cc 100644 --- a/tests/unit_tests/_primaite/_simulator/_file_system/test_folder_actions.py +++ b/tests/unit_tests/_primaite/_simulator/_file_system/test_folder_actions.py @@ -31,7 +31,7 @@ def test_folder_scan_request(populated_file_system): assert file1.visible_health_status == FileSystemItemHealthStatus.GOOD assert file2.visible_health_status == FileSystemItemHealthStatus.GOOD - fs.apply_request(request=["folder", folder.uuid, "scan"]) + fs.apply_request(request=["folder", folder.name, "scan"]) folder.apply_timestep(timestep=0) @@ -53,12 +53,12 @@ def test_folder_checkhash_request(populated_file_system): """Test that an agent can request a folder hash check.""" fs, folder, file = populated_file_system - fs.apply_request(request=["folder", folder.uuid, "checkhash"]) + fs.apply_request(request=["folder", folder.name, "checkhash"]) assert folder.health_status == FileSystemItemHealthStatus.GOOD file.sim_size = 0 - fs.apply_request(request=["folder", folder.uuid, "checkhash"]) + fs.apply_request(request=["folder", folder.name, "checkhash"]) assert folder.health_status == FileSystemItemHealthStatus.CORRUPT @@ -70,7 +70,7 @@ def test_folder_repair_request(populated_file_system): assert file.health_status == FileSystemItemHealthStatus.CORRUPT assert folder.health_status == FileSystemItemHealthStatus.CORRUPT - fs.apply_request(request=["folder", folder.uuid, "repair"]) + fs.apply_request(request=["folder", folder.name, "repair"]) assert file.health_status == FileSystemItemHealthStatus.GOOD assert folder.health_status == FileSystemItemHealthStatus.GOOD @@ -116,7 +116,7 @@ def test_folder_restore_request(populated_file_system): assert fs.get_file_by_id(folder_uuid=folder.uuid, file_uuid=file.uuid, include_deleted=True).deleted is False # corrupt folder - fs.apply_request(request=["folder", folder.uuid, "corrupt"]) + fs.apply_request(request=["folder", folder.name, "corrupt"]) assert fs.get_folder(folder_name=folder.name).health_status == FileSystemItemHealthStatus.CORRUPT assert fs.get_file(folder_name=folder.name, file_name=file.name).health_status == FileSystemItemHealthStatus.CORRUPT @@ -143,7 +143,7 @@ def test_folder_restore_request(populated_file_system): def test_folder_corrupt_request(populated_file_system): """Test that an agent can request a folder corruption.""" fs, folder, file = populated_file_system - fs.apply_request(request=["folder", folder.uuid, "corrupt"]) + fs.apply_request(request=["folder", folder.name, "corrupt"]) assert file.health_status == FileSystemItemHealthStatus.CORRUPT assert folder.health_status == FileSystemItemHealthStatus.CORRUPT @@ -153,13 +153,13 @@ def test_deleted_folder_and_its_files_cannot_be_interacted_with(populated_file_s fs, folder, file = populated_file_system assert fs.get_file(folder_name=folder.name, file_name=file.name) is not None - fs.apply_request(request=["file", file.uuid, "corrupt"]) + fs.apply_request(request=["file", file.name, "corrupt"]) assert fs.get_file(folder_name=folder.name, file_name=file.name).health_status == FileSystemItemHealthStatus.CORRUPT fs.apply_request(request=["delete", "folder", folder.uuid]) assert fs.get_file(folder_name=folder.name, file_name=file.name) is None - fs.apply_request(request=["file", file.uuid, "repair"]) + fs.apply_request(request=["file", file.name, "repair"]) deleted_folder = fs.deleted_folders.get(folder.uuid) deleted_file = deleted_folder.deleted_files.get(file.uuid) From 41a7f838873745467d0e86480406d92327e61c33 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Wed, 31 Jan 2024 13:29:46 +0000 Subject: [PATCH 14/18] Add file scan test --- src/primaite/game/agent/actions.py | 2 +- src/primaite/simulator/file_system/folder.py | 2 +- .../game_layer/test_actions.py | 50 ++++++++++++++++++- 3 files changed, 50 insertions(+), 4 deletions(-) diff --git a/src/primaite/game/agent/actions.py b/src/primaite/game/agent/actions.py index 3b1fb926..dbf1b6a2 100644 --- a/src/primaite/game/agent/actions.py +++ b/src/primaite/game/agent/actions.py @@ -270,7 +270,7 @@ class NodeFileAbstractAction(AbstractAction): 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_name, "file_system", "folder", folder_name, "files", file_name, self.verb] + return ["network", "node", node_name, "file_system", "folder", folder_name, "file", file_name, self.verb] class NodeFileScanAction(NodeFileAbstractAction): diff --git a/src/primaite/simulator/file_system/folder.py b/src/primaite/simulator/file_system/folder.py index dae32cd5..a93b2927 100644 --- a/src/primaite/simulator/file_system/folder.py +++ b/src/primaite/simulator/file_system/folder.py @@ -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]): diff --git a/tests/integration_tests/game_layer/test_actions.py b/tests/integration_tests/game_layer/test_actions.py index 2bc3b095..7527151c 100644 --- a/tests/integration_tests/game_layer/test_actions.py +++ b/tests/integration_tests/game_layer/test_actions.py @@ -19,6 +19,7 @@ from primaite.game.agent.interface import AbstractAgent, ProxyAgent from primaite.game.agent.observations import ICSObservation, ObservationManager from primaite.game.agent.rewards import RewardFunction from primaite.game.game import PrimaiteGame +from primaite.simulator.file_system.file_system_item_abc import FileSystemItemHealthStatus from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState from primaite.simulator.network.hardware.nodes.computer import Computer from primaite.simulator.network.hardware.nodes.router import ACLAction, Router @@ -62,7 +63,7 @@ class ControlledAgent(AbstractAgent): def install_stuff_to_sim(sim: Simulation): - """Create a simulation with a three computers, two switches, and a router.""" + """Create a simulation with a computer, two servers, two switches, and a router.""" # 0: Pull out the network network = sim.network @@ -139,6 +140,9 @@ def install_stuff_to_sim(sim: Simulation): assert isinstance(client_1.software_manager.software.get("WebBrowser"), WebBrowser) assert isinstance(client_1.software_manager.software.get("DNSClient"), DNSClient) + # 4.1: Create a file on the computer + client_1.file_system.create_file("cat.png", 300, folder_name="downloads") + # 5: Assert that the simulation starts off in the state that we expect assert len(sim.network.nodes) == 6 assert len(sim.network.links) == 5 @@ -221,7 +225,11 @@ def game_and_agent(): game=game, actions=actions, # ALL POSSIBLE ACTIONS nodes=[ - {"node_name": "client_1", "applications": [{"application_name": "WebBrowser"}]}, + { + "node_name": "client_1", + "applications": [{"application_name": "WebBrowser"}], + "folders": [{"folder_name": "downloads", "files": [{"file_name": "cat.png"}]}], + }, {"node_name": "server_1", "services": [{"service_name": "DNSServer"}]}, {"node_name": "server_2", "services": [{"service_name": "WebServer"}]}, ], @@ -480,3 +488,41 @@ def test_network_nic_enable_integration(game_and_agent: Tuple[PrimaiteGame, Prox # 3: Check that the NIC is enabled, and that client 1 can ping again assert client_1.ethernet_port[1].enabled == True assert client_1.ping("10.0.2.3") + + +def test_node_file_scan_integration(game_and_agent: Tuple[PrimaiteGame, ProxyAgent]): + """Test that a when a file is scanned, it's visible health status gets set to the actual health status.""" + + game, agent = game_and_agent + + # 1: assert file is healthy + client_1 = game.simulation.network.get_node_by_hostname("client_1") + file = client_1.file_system.get_file("downloads", "cat.png") + assert file.health_status == FileSystemItemHealthStatus.GOOD + assert file.visible_health_status == FileSystemItemHealthStatus.GOOD + + # 2: perform a scan and make sure nothing has changed + action = ( + "NODE_FILE_SCAN", + { + "node_id": 0, # client_1, + "folder_id": 0, # downloads, + "file_id": 0, # cat.png + }, + ) + agent.store_action(action) + game.step() + + assert file.health_status == FileSystemItemHealthStatus.GOOD + assert file.visible_health_status == FileSystemItemHealthStatus.GOOD + + # 3: Set the file to corrupted, and check that only actual updates, not visible. + file.health_status = FileSystemItemHealthStatus.CORRUPT + assert file.health_status == FileSystemItemHealthStatus.CORRUPT + assert file.visible_health_status == FileSystemItemHealthStatus.GOOD + + # 4: Perform a scan and check that it updates + agent.store_action(action) + game.step() + assert file.health_status == FileSystemItemHealthStatus.CORRUPT + assert file.visible_health_status == FileSystemItemHealthStatus.CORRUPT From 83db5b1eb59712cb952462f8e913ec6d56fe0700 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Wed, 31 Jan 2024 13:50:20 +0000 Subject: [PATCH 15/18] Fix node file delete action --- src/primaite/game/agent/actions.py | 3 +- .../simulator/file_system/file_system.py | 4 +-- .../game_layer/test_actions.py | 28 +++++++++++++++++++ 3 files changed, 31 insertions(+), 4 deletions(-) diff --git a/src/primaite/game/agent/actions.py b/src/primaite/game/agent/actions.py index dbf1b6a2..6123da50 100644 --- a/src/primaite/game/agent/actions.py +++ b/src/primaite/game/agent/actions.py @@ -303,8 +303,7 @@ class NodeFileDeleteAction(NodeFileAbstractAction): 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_name, "file_system", "delete", "file", folder_name, file_name] + return ["network", "node", node_name, "file_system", "delete", "file", folder_name, file_name] class NodeFileRepairAction(NodeFileAbstractAction): diff --git a/src/primaite/simulator/file_system/file_system.py b/src/primaite/simulator/file_system/file_system.py index 48ea587d..2ab3b005 100644 --- a/src/primaite/simulator/file_system/file_system.py +++ b/src/primaite/simulator/file_system/file_system.py @@ -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", diff --git a/tests/integration_tests/game_layer/test_actions.py b/tests/integration_tests/game_layer/test_actions.py index 7527151c..e771dbd2 100644 --- a/tests/integration_tests/game_layer/test_actions.py +++ b/tests/integration_tests/game_layer/test_actions.py @@ -526,3 +526,31 @@ def test_node_file_scan_integration(game_and_agent: Tuple[PrimaiteGame, ProxyAge game.step() assert file.health_status == FileSystemItemHealthStatus.CORRUPT assert file.visible_health_status == FileSystemItemHealthStatus.CORRUPT + + +def test_node_file_delete_integration(game_and_agent: Tuple[PrimaiteGame, ProxyAgent]): + """Test that a file can be deleted by the agent.""" + game, agent = game_and_agent + + # 1: assert the file is there + client_1 = game.simulation.network.get_node_by_hostname("client_1") + file = client_1.file_system.get_file("downloads", "cat.png") + assert file is not None + assert not file.deleted + + # 2: delete the file + action = ( + "NODE_FILE_DELETE", + { + "node_id": 0, # client_1 + "folder_id": 0, # downloads + "file_id": 0, # cat.png + }, + ) + agent.store_action(action) + game.step() + + # 3. Check that the file is not there any more + assert not client_1.file_system.get_file("downloads", "cat.png") + # 3.1 (but with the reference to the original file, we can check that deleted flag is True ) + assert file.deleted From 261423bc78c7ef1cbe18aa83dfadc55ff4f88680 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Wed, 31 Jan 2024 13:54:15 +0000 Subject: [PATCH 16/18] Update changelog --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5d8706ad..fdb530e2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ## [Unreleased] +- Refactored actions and observations to be configurable via object name, instead of UUID. - Fixed a bug where ACL rules were not resetting on episode reset. - Fixed a bug where blue agent's ACL actions were being applied against the wrong IP addresses - Fixed a bug where deleted files and folders did not reset correctly on episode reset. From c163fb37ea4389fcace1a79b2f1723ea2c624c78 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Thu, 1 Feb 2024 11:19:42 +0000 Subject: [PATCH 17/18] Update request system docs. --- docs/source/request_system.rst | 35 +++++++++++++++++++--------------- 1 file changed, 20 insertions(+), 15 deletions(-) diff --git a/docs/source/request_system.rst b/docs/source/request_system.rst index 1b06e2d9..392bc792 100644 --- a/docs/source/request_system.rst +++ b/docs/source/request_system.rst @@ -7,26 +7,29 @@ Request System ``SimComponent`` objects in the simulation are decoupled from the agent training logic. However, they still need a managed means of accepting requests to perform actions. For this, they use ``RequestManager`` and ``RequestType``. -Just like other aspects of SimComponent, the request types are not managed centrally for the whole simulation, but instead they are dynamically created and updated based on the nodes, links, and other components that currently exist. This was achieved in the following way: +Just like other aspects of SimComponent, the request types are not managed centrally for the whole simulation, but instead they are dynamically created and updated based on the nodes, links, and other components that currently exist in the simulation. This is achieved in the following way: - API - A ``RequestType`` contains two elements: + When requesting an action within the simulation, these two arguments must be provided: - 1. ``request`` - selects which action you want to take on this ``SimComponent``. This is formatted as a list of strings such as `['network', 'node', '', 'service', '', 'restart']`. + 1. ``request`` - selects which action you want to take on this ``SimComponent``. This is formatted as a list of strings such as `['network', 'node', '', 'service', '', 'restart']`. 2. ``context`` - optional extra information that can be used to decide how to process the request. This is formatted as a dictionary. For example, if the request requires authentication, the context can include information about the user that initiated the request to decide if their permissions are sufficient. -- request +- ``request`` detail The request is a list of strings which help specify who should handle the request. The strings in the request list help RequestManagers traverse the 'ownership tree' of SimComponent. The example given above would be handled in the following way: - 1. ``Simulation`` receives `['network', 'node', '', 'service', '', 'restart']`. + 1. ``Simulation`` receives `['network', 'node', '', 'service', '', 'restart']`. The first element of the request is ``network``, therefore it passes the request down to its network. - 2. ``Network`` receives `['node', '', 'service', '', 'restart']`. - The first element of the request is ``node``, therefore the network looks at the node uuid and passes the request down to the node with that uuid. - 3. ``Node`` receives `['service', '', 'restart']`. - The first element of the request is ``service``, therefore the node looks at the service uuid and passes the rest of the request to the service with that uuid. + 2. ``Network`` receives `['node', '', 'service', '', 'restart']`. + The first element of the request is ``node``, therefore the network looks at the node name and passes the request down to the node with that name. + 3. ``Node`` receives `['service', '', 'restart']`. + The first element of the request is ``service``, therefore the node looks at the service name and passes the rest of the request to the service with that name. 4. ``Service`` receives ``['restart']``. Since ``restart`` is a defined request type in the service's own RequestManager, the service performs a restart. +- ``context`` detail + The context is not used by any of the currently implemented components or requests. + Technical Detail ---------------- @@ -75,16 +78,18 @@ An example of how this works is in the :py:class:`primaite.simulator.network.har request_manager.add_request("turn_on", RequestType(func=lambda request, context: self.turn_on())) # if the Node receives a request where the first word is 'service', it will use a dummy manager - # called self._service_request_manager to pass on the reqeust to the relevant service. This dummy - # manager is simply here to map the service UUID that that service's own action manager. This is - # done because the next string after "service" is always the uuid of that service, so we need an + # called self._service_request_manager to pass on the request to the relevant service. This dummy + # manager is simply here to map the service name that that service's own action manager. This is + # done because the next string after "service" is always the name of that service, so we need an # RequestManager to pop that string before sending it onto the relevant service's RequestManager. self._service_request_manager = RequestManager() request_manager.add_request("service", RequestType(func=self._service_request_manager)) ... def install_service(self, service): - self.services[service.uuid] = service + self.services[service.name] = service ... - # Here, the service UUID is registered to allow passing actions between the node and the service. - self._service_request_manager.add_request(service.uuid, RequestType(func=service._request_manager)) + # Here, the service name is registered to allow passing actions between the node and the service. + self._service_request_manager.add_request(service.name, RequestType(func=service._request_manager)) + +This process is repeated until the request word corresponds to a callable function rather than another ``RequestManager`` . From 4c6ae135cdf93a2ccc83f94ab578654112001661 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Fri, 2 Feb 2024 14:48:49 +0000 Subject: [PATCH 18/18] Fix typos --- src/primaite/game/agent/actions.py | 20 +++++++++---------- src/primaite/game/agent/rewards.py | 4 ++-- .../simulator/network/hardware/base.py | 3 ++- 3 files changed, 14 insertions(+), 13 deletions(-) diff --git a/src/primaite/game/agent/actions.py b/src/primaite/game/agent/actions.py index 6123da50..d5b78a16 100644 --- a/src/primaite/game/agent/actions.py +++ b/src/primaite/game/agent/actions.py @@ -398,8 +398,8 @@ class NetworkACLAddRuleAction(AbstractAction): :param manager: Reference to the ActionManager which created this action. :type manager: ActionManager - :param target_router_name: hostname of the router to which the ACL rule should be added. - :type target_router_name: 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. @@ -508,8 +508,8 @@ class NetworkACLRemoveRuleAction(AbstractAction): :param manager: Reference to the ActionManager which created this action. :type manager: ActionManager - :param target_router_name: Hostname of the router from which the ACL rule should be removed. - :type target_router_name: 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 """ @@ -629,7 +629,7 @@ class ActionManager: '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: Dict + :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. @@ -679,24 +679,24 @@ class ActionManager: """ # 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", [])] + 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 n.get("services", [])] + 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 n.get("folders", [])] + 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 n.get("folders", [{"files": []}]): + 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) diff --git a/src/primaite/game/agent/rewards.py b/src/primaite/game/agent/rewards.py index 0b292bcb..8944a184 100644 --- a/src/primaite/game/agent/rewards.py +++ b/src/primaite/game/agent/rewards.py @@ -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 diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py index 8e4c1d76..5334021a 100644 --- a/src/primaite/simulator/network/hardware/base.py +++ b/src/primaite/simulator/network/hardware/base.py @@ -1388,7 +1388,8 @@ class Node(SimComponent): nic.parent = None nic.disable() self.sys_log.info(f"Disconnected NIC {nic}") - self._nic_request_manager.remove_request(nic_num) + if nic_num != -1: + 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)