Add application execution action for data manipulation bot
This commit is contained in:
@@ -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
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user