oxbrbtx/PrimAITE
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Move configuration from agent to data manipulation bot

Browse Source
This commit is contained in:
Jake Walker
2023-11-23 15:53:47 +00:00
parent 061e508187
commit c93705867f
4 changed files with 25 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files

43
src/primaite/game/agent/interface.py
View File

@@ -9,7 +9,6 @@ from pydantic import BaseModel
from primaite.game.agent.actions import ActionManager
from primaite.game.agent.observations import ObservationSpace
from primaite.game.agent.rewards import RewardFunction
from primaite.simulator.network.hardware.base import Node
if TYPE_CHECKING:
from primaite.simulator.system.services.red_services.data_manipulation_bot import DataManipulationBot
@@ -17,29 +16,6 @@ if TYPE_CHECKING:
ObsType: TypeAlias = Union[Dict, np.ndarray]
class AgentExecutionDefinition(BaseModel):
"""Additional configuration for agents."""
port_scan_p_of_success: float = 0.1
"The probability of a port scan succeeding."
data_manipulation_p_of_success: float = 0.1
"The probability of data manipulation succeeding."
@classmethod
def from_config(cls, config: Optional[Dict]) -> "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.

22
src/primaite/game/session.py
View File

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