From b13a9d3daf34f38992b19f7854cbbf0eeb3e2723 Mon Sep 17 00:00:00 2001 From: Jake Walker Date: Fri, 24 Nov 2023 09:25:55 +0000 Subject: [PATCH] Add application execution action for data manipulation bot --- .../config/_package_data/example_config.yaml | 9 ++++++--- src/primaite/game/agent/actions.py | 7 +++---- src/primaite/game/session.py | 14 ++++++++++++++ .../system/applications/database_client.py | 2 +- .../services/red_services/data_manipulation_bot.py | 2 +- 5 files changed, 25 insertions(+), 9 deletions(-) diff --git a/src/primaite/config/_package_data/example_config.yaml b/src/primaite/config/_package_data/example_config.yaml index aff54d62..8ea1c83c 100644 --- a/src/primaite/config/_package_data/example_config.yaml +++ b/src/primaite/config/_package_data/example_config.yaml @@ -67,8 +67,8 @@ game_config: observations: - logon_status - operating_status - services: - - service_ref: data_manipulation_bot + applications: + - application_ref: data_manipulation_bot observations: operating_status health_status @@ -89,6 +89,8 @@ game_config: options: nodes: - node_ref: client_1 + applications: + - application_ref: data_manipulation_bot max_folders_per_node: 1 max_files_per_folder: 1 max_services_per_node: 1 @@ -650,7 +652,7 @@ simulation: subnet_mask: 255.255.255.0 default_gateway: 192.168.10.1 dns_server: 192.168.1.10 - services: + applications: - ref: data_manipulation_bot type: DataManipulationBot options: @@ -658,6 +660,7 @@ simulation: data_manipulation_p_of_success: 0.1 payload: "DROP TABLE IF EXISTS user;" server_ip: 192.168.1.14 + services: - ref: client_1_dns_client type: DNSClient diff --git a/src/primaite/game/agent/actions.py b/src/primaite/game/agent/actions.py index 0c78dac7..64d89722 100644 --- a/src/primaite/game/agent/actions.py +++ b/src/primaite/game/agent/actions.py @@ -594,6 +594,7 @@ class ActionManager: session: "PrimaiteSession", # reference to session for looking up stuff 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 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 @@ -635,6 +636,7 @@ class ActionManager: self.session: "PrimaiteSession" = session self.sim: Simulation = self.session.simulation self.node_uuids: List[str] = node_uuids + self.application_uuids: List[List[str]] = application_uuids self.protocols: List[str] = protocols self.ports: List[str] = ports @@ -819,10 +821,7 @@ class ActionManager: :return: The UUID 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.sim.network.nodes[node_uuid] - application_uuids = list(node.applications.keys()) - return application_uuids[application_idx] if len(application_uuids) > application_idx else None + return self.application_uuids[node_idx][application_idx] def get_internet_protocol_by_idx(self, protocol_idx: int) -> str: """Get the internet protocol corresponding to the given index. diff --git a/src/primaite/game/session.py b/src/primaite/game/session.py index f675e33c..cc4036ef 100644 --- a/src/primaite/game/session.py +++ b/src/primaite/game/session.py @@ -426,11 +426,25 @@ class PrimaiteSession: # 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 = sess.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: remove inconsistency with the above nodes + application_uuid = sess.ref_map_applications[application_option["application_ref"]].uuid + 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 diff --git a/src/primaite/simulator/system/applications/database_client.py b/src/primaite/simulator/system/applications/database_client.py index e15249e3..9d85221e 100644 --- a/src/primaite/simulator/system/applications/database_client.py +++ b/src/primaite/simulator/system/applications/database_client.py @@ -141,7 +141,7 @@ class DatabaseClient(Application): :param sql: The SQL query. :return: True if the query was successful, otherwise False. """ - if self.connected and self.operating_state.RUNNING: + if self.connected and self.operating_state == ApplicationOperatingState.RUNNING: query_id = str(uuid4()) # Initialise the tracker of this ID to False diff --git a/src/primaite/simulator/system/services/red_services/data_manipulation_bot.py b/src/primaite/simulator/system/services/red_services/data_manipulation_bot.py index f4b31cb1..0ec64950 100644 --- a/src/primaite/simulator/system/services/red_services/data_manipulation_bot.py +++ b/src/primaite/simulator/system/services/red_services/data_manipulation_bot.py @@ -50,7 +50,7 @@ class DataManipulationBot(DatabaseClient): def _init_request_manager(self) -> RequestManager: rm = super()._init_request_manager() - rm.add_request(name="execute", request_type=RequestType(func=self.execute)) + rm.add_request(name="execute", request_type=RequestType(func=lambda request, context: self.execute())) return rm