Add application execution action for data manipulation bot

This commit is contained in:
Jake Walker
2023-11-24 09:25:55 +00:00
parent 5f1a5af1b4
commit b13a9d3daf
5 changed files with 25 additions and 9 deletions

View File

@@ -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

View File

@@ -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.

View File

@@ -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

View File

@@ -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

View File

@@ -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