From c93705867f28d2b68e3e98888a5de1cb424bc890 Mon Sep 17 00:00:00 2001 From: Jake Walker Date: Thu, 23 Nov 2023 15:53:47 +0000 Subject: [PATCH] Move configuration from agent to data manipulation bot --- .../config/_package_data/example_config.yaml | 14 +++--- src/primaite/game/agent/interface.py | 43 ------------------- src/primaite/game/session.py | 22 +++++----- .../red_services/data_manipulation_bot.py | 11 ++++- 4 files changed, 25 insertions(+), 65 deletions(-) diff --git a/src/primaite/config/_package_data/example_config.yaml b/src/primaite/config/_package_data/example_config.yaml index 700a45fd..274da7aa 100644 --- a/src/primaite/config/_package_data/example_config.yaml +++ b/src/primaite/config/_package_data/example_config.yaml @@ -59,10 +59,6 @@ game_config: team: RED type: RedDatabaseCorruptingAgent - execution_definition: - port_scan_p_of_success: 0.1 - data_manipulation_p_of_success: 0.1 - observation_space: type: UC2RedObservation options: @@ -83,11 +79,6 @@ game_config: - type: DONOTHING # "AgentExecutionDefinition": - """Construct an AgentExecutionDefinition from a config dictionary. - - :param config: A dict of options for the execution definition. - :type config: Dict - :return: The execution definition. - :rtype: AgentExecutionDefinition - """ - if config is None: - return cls() - - return cls(**config) - - class AgentStartSettings(BaseModel): """Configuration values for when an agent starts performing actions.""" @@ -81,7 +57,6 @@ class AbstractAgent(ABC): action_space: Optional[ActionManager], observation_space: Optional[ObservationSpace], reward_function: Optional[RewardFunction], - execution_definition: Optional[AgentExecutionDefinition], agent_settings: Optional[AgentSettings], ) -> None: """ @@ -100,11 +75,6 @@ class AbstractAgent(ABC): self.action_space: Optional[ActionManager] = action_space self.observation_space: Optional[ObservationSpace] = observation_space self.reward_function: Optional[RewardFunction] = reward_function - - # exection definiton converts CAOS action to Primaite simulator request, sometimes having to enrich the info - # by for example specifying target ip addresses, or converting a node ID into a uuid - self.execution_definition = execution_definition or AgentExecutionDefinition() - self.agent_settings = agent_settings or AgentSettings() def convert_state_to_obs(self, state: Dict) -> ObsType: @@ -186,19 +156,6 @@ class DataManipulationAgent(AbstractScriptedAgent): self.next_execution_timestep = self.agent_settings.start_settings.start_step - # get node ids that are part of the agent's observation space - node_ids: List[str] = [n.where[-1] for n in self.observation_space.obs.nodes] - # get all nodes from their ids - nodes: List[Node] = [n for n_id, n in self.action_space.sim.network.nodes.items() if n_id in node_ids] - - # get execution definition for data manipulation bots - for node in nodes: - bot_sw: Optional["DataManipulationBot"] = node.software_manager.software.get("DataManipulationBot") - - if bot_sw is not None: - bot_sw.execution_definition = self.execution_definition - self.data_manipulation_bots.append(bot_sw) - def get_action(self, obs: ObsType, reward: float = None) -> Tuple[str, Dict]: """Randomly sample an action from the action space. diff --git a/src/primaite/game/session.py b/src/primaite/game/session.py index 1b086c35..f675e33c 100644 --- a/src/primaite/game/session.py +++ b/src/primaite/game/session.py @@ -10,13 +10,7 @@ from pydantic import BaseModel from primaite import getLogger from primaite.game.agent.actions import ActionManager -from primaite.game.agent.interface import ( - AbstractAgent, - AgentExecutionDefinition, - AgentSettings, - DataManipulationAgent, - RandomAgent, -) +from primaite.game.agent.interface import AbstractAgent, AgentSettings, DataManipulationAgent, RandomAgent from primaite.game.agent.observations import ObservationSpace from primaite.game.agent.rewards import RewardFunction from primaite.simulator.network.hardware.base import Link, NIC, Node @@ -366,6 +360,16 @@ class PrimaiteSession: if "domain_mapping" in opt: for domain, ip in opt["domain_mapping"].items(): new_service.dns_register(domain, ip) + if service_type == "DataManipulationBot": + if "options" in service_cfg: + opt = service_cfg["options"] + new_service.configure( + server_ip_address=opt.get("server_ip"), + payload=opt.get("payload"), + port_scan_p_of_success=float(opt.get("port_scan_p_of_success", "0.1")), + data_manipulation_p_of_success=float(opt.get("data_manipulation_p_of_success", "0.1")), + ) + if "applications" in node_cfg: for application_cfg in node_cfg["applications"]: application_ref = application_cfg["ref"] @@ -444,7 +448,6 @@ class PrimaiteSession: # CREATE REWARD FUNCTION rew_function = RewardFunction.from_config(reward_function_cfg, session=sess) - execution_definition = AgentExecutionDefinition.from_config(agent_cfg.get("execution_definition")) agent_settings = AgentSettings.from_config(agent_cfg.get("agent_settings")) # CREATE AGENT @@ -455,7 +458,6 @@ class PrimaiteSession: action_space=action_space, observation_space=obs_space, reward_function=rew_function, - execution_definition=execution_definition, agent_settings=agent_settings, ) sess.agents.append(new_agent) @@ -465,7 +467,6 @@ class PrimaiteSession: action_space=action_space, observation_space=obs_space, reward_function=rew_function, - execution_definition=execution_definition, agent_settings=agent_settings, ) sess.agents.append(new_agent) @@ -476,7 +477,6 @@ class PrimaiteSession: action_space=action_space, observation_space=obs_space, reward_function=rew_function, - execution_definition=execution_definition, agent_settings=agent_settings, ) sess.agents.append(new_agent) 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 eae3f0e3..e3f5b95d 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 @@ -2,7 +2,6 @@ from enum import IntEnum from ipaddress import IPv4Address from typing import Optional -from primaite.game.agent.interface import AgentExecutionDefinition from primaite.game.science import simulate_trial from primaite.simulator.system.applications.application import ApplicationOperatingState from primaite.simulator.system.applications.database_client import DatabaseClient @@ -36,8 +35,10 @@ class DataManipulationBot(DatabaseClient): server_ip_address: Optional[IPv4Address] = None payload: Optional[str] = None server_password: Optional[str] = None + port_scan_p_of_success: float = 0.1 + data_manipulation_p_of_success: float = 0.1 + attack_stage: DataManipulationAttackStage = DataManipulationAttackStage.NOT_STARTED - execution_definition: AgentExecutionDefinition = AgentExecutionDefinition() repeat: bool = False "Whether to repeat attacking once finished." @@ -50,6 +51,8 @@ class DataManipulationBot(DatabaseClient): server_ip_address: IPv4Address, server_password: Optional[str] = None, payload: Optional[str] = None, + port_scan_p_of_success: float = 0.1, + data_manipulation_p_of_success: float = 0.1, repeat: bool = False, ): """ @@ -58,11 +61,15 @@ class DataManipulationBot(DatabaseClient): :param server_ip_address: The IP address of the Node the DatabaseService is on. :param server_password: The password on the DatabaseService. :param payload: The data manipulation query payload. + :param port_scan_p_of_success: The probability of success for the port scan stage. + :param data_manipulation_p_of_success: The probability of success for the data manipulation stage. :param repeat: Whether to repeat attacking once finished. """ self.server_ip_address = server_ip_address self.payload = payload self.server_password = server_password + self.port_scan_p_of_success = port_scan_p_of_success + self.data_manipulation_p_of_success = data_manipulation_p_of_success self.repeat = repeat self.sys_log.info( f"{self.name}: Configured the {self.name} with {server_ip_address=}, {payload=}, {server_password=}, "