Move configuration from agent to data manipulation bot
This commit is contained in:
@@ -59,10 +59,6 @@ game_config:
|
|||||||
team: RED
|
team: RED
|
||||||
type: RedDatabaseCorruptingAgent
|
type: RedDatabaseCorruptingAgent
|
||||||
|
|
||||||
execution_definition:
|
|
||||||
port_scan_p_of_success: 0.1
|
|
||||||
data_manipulation_p_of_success: 0.1
|
|
||||||
|
|
||||||
observation_space:
|
observation_space:
|
||||||
type: UC2RedObservation
|
type: UC2RedObservation
|
||||||
options:
|
options:
|
||||||
@@ -83,11 +79,6 @@ game_config:
|
|||||||
- type: DONOTHING
|
- type: DONOTHING
|
||||||
#<not yet implemented
|
#<not yet implemented
|
||||||
# - type: NODE_APPLICATION_EXECUTE
|
# - type: NODE_APPLICATION_EXECUTE
|
||||||
# options:
|
|
||||||
# execution_definition:
|
|
||||||
# server_ip: 192.168.1.14
|
|
||||||
# payload: "DROP TABLE IF EXISTS user;"
|
|
||||||
# success_rate: 80%
|
|
||||||
- type: NODE_FILE_DELETE
|
- type: NODE_FILE_DELETE
|
||||||
- type: NODE_FILE_CORRUPT
|
- type: NODE_FILE_CORRUPT
|
||||||
# - type: NODE_FOLDER_DELETE
|
# - type: NODE_FOLDER_DELETE
|
||||||
@@ -662,6 +653,11 @@ simulation:
|
|||||||
services:
|
services:
|
||||||
- ref: data_manipulation_bot
|
- ref: data_manipulation_bot
|
||||||
type: DataManipulationBot
|
type: DataManipulationBot
|
||||||
|
options:
|
||||||
|
port_scan_p_of_success: 0.1
|
||||||
|
data_manipulation_p_of_success: 0.1
|
||||||
|
payload: "DROP TABLE IF EXISTS user;"
|
||||||
|
server_ip: 192.168.1.14
|
||||||
- ref: client_1_dns_client
|
- ref: client_1_dns_client
|
||||||
type: DNSClient
|
type: DNSClient
|
||||||
|
|
||||||
|
|||||||
@@ -9,7 +9,6 @@ from pydantic import BaseModel
|
|||||||
from primaite.game.agent.actions import ActionManager
|
from primaite.game.agent.actions import ActionManager
|
||||||
from primaite.game.agent.observations import ObservationSpace
|
from primaite.game.agent.observations import ObservationSpace
|
||||||
from primaite.game.agent.rewards import RewardFunction
|
from primaite.game.agent.rewards import RewardFunction
|
||||||
from primaite.simulator.network.hardware.base import Node
|
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from primaite.simulator.system.services.red_services.data_manipulation_bot import DataManipulationBot
|
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]
|
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):
|
class AgentStartSettings(BaseModel):
|
||||||
"""Configuration values for when an agent starts performing actions."""
|
"""Configuration values for when an agent starts performing actions."""
|
||||||
|
|
||||||
@@ -81,7 +57,6 @@ class AbstractAgent(ABC):
|
|||||||
action_space: Optional[ActionManager],
|
action_space: Optional[ActionManager],
|
||||||
observation_space: Optional[ObservationSpace],
|
observation_space: Optional[ObservationSpace],
|
||||||
reward_function: Optional[RewardFunction],
|
reward_function: Optional[RewardFunction],
|
||||||
execution_definition: Optional[AgentExecutionDefinition],
|
|
||||||
agent_settings: Optional[AgentSettings],
|
agent_settings: Optional[AgentSettings],
|
||||||
) -> None:
|
) -> None:
|
||||||
"""
|
"""
|
||||||
@@ -100,11 +75,6 @@ class AbstractAgent(ABC):
|
|||||||
self.action_space: Optional[ActionManager] = action_space
|
self.action_space: Optional[ActionManager] = action_space
|
||||||
self.observation_space: Optional[ObservationSpace] = observation_space
|
self.observation_space: Optional[ObservationSpace] = observation_space
|
||||||
self.reward_function: Optional[RewardFunction] = reward_function
|
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()
|
self.agent_settings = agent_settings or AgentSettings()
|
||||||
|
|
||||||
def convert_state_to_obs(self, state: Dict) -> ObsType:
|
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
|
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]:
|
def get_action(self, obs: ObsType, reward: float = None) -> Tuple[str, Dict]:
|
||||||
"""Randomly sample an action from the action space.
|
"""Randomly sample an action from the action space.
|
||||||
|
|
||||||
|
|||||||
@@ -10,13 +10,7 @@ from pydantic import BaseModel
|
|||||||
|
|
||||||
from primaite import getLogger
|
from primaite import getLogger
|
||||||
from primaite.game.agent.actions import ActionManager
|
from primaite.game.agent.actions import ActionManager
|
||||||
from primaite.game.agent.interface import (
|
from primaite.game.agent.interface import AbstractAgent, AgentSettings, DataManipulationAgent, RandomAgent
|
||||||
AbstractAgent,
|
|
||||||
AgentExecutionDefinition,
|
|
||||||
AgentSettings,
|
|
||||||
DataManipulationAgent,
|
|
||||||
RandomAgent,
|
|
||||||
)
|
|
||||||
from primaite.game.agent.observations import ObservationSpace
|
from primaite.game.agent.observations import ObservationSpace
|
||||||
from primaite.game.agent.rewards import RewardFunction
|
from primaite.game.agent.rewards import RewardFunction
|
||||||
from primaite.simulator.network.hardware.base import Link, NIC, Node
|
from primaite.simulator.network.hardware.base import Link, NIC, Node
|
||||||
@@ -366,6 +360,16 @@ class PrimaiteSession:
|
|||||||
if "domain_mapping" in opt:
|
if "domain_mapping" in opt:
|
||||||
for domain, ip in opt["domain_mapping"].items():
|
for domain, ip in opt["domain_mapping"].items():
|
||||||
new_service.dns_register(domain, ip)
|
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:
|
if "applications" in node_cfg:
|
||||||
for application_cfg in node_cfg["applications"]:
|
for application_cfg in node_cfg["applications"]:
|
||||||
application_ref = application_cfg["ref"]
|
application_ref = application_cfg["ref"]
|
||||||
@@ -444,7 +448,6 @@ class PrimaiteSession:
|
|||||||
# CREATE REWARD FUNCTION
|
# CREATE REWARD FUNCTION
|
||||||
rew_function = RewardFunction.from_config(reward_function_cfg, session=sess)
|
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"))
|
agent_settings = AgentSettings.from_config(agent_cfg.get("agent_settings"))
|
||||||
|
|
||||||
# CREATE AGENT
|
# CREATE AGENT
|
||||||
@@ -455,7 +458,6 @@ class PrimaiteSession:
|
|||||||
action_space=action_space,
|
action_space=action_space,
|
||||||
observation_space=obs_space,
|
observation_space=obs_space,
|
||||||
reward_function=rew_function,
|
reward_function=rew_function,
|
||||||
execution_definition=execution_definition,
|
|
||||||
agent_settings=agent_settings,
|
agent_settings=agent_settings,
|
||||||
)
|
)
|
||||||
sess.agents.append(new_agent)
|
sess.agents.append(new_agent)
|
||||||
@@ -465,7 +467,6 @@ class PrimaiteSession:
|
|||||||
action_space=action_space,
|
action_space=action_space,
|
||||||
observation_space=obs_space,
|
observation_space=obs_space,
|
||||||
reward_function=rew_function,
|
reward_function=rew_function,
|
||||||
execution_definition=execution_definition,
|
|
||||||
agent_settings=agent_settings,
|
agent_settings=agent_settings,
|
||||||
)
|
)
|
||||||
sess.agents.append(new_agent)
|
sess.agents.append(new_agent)
|
||||||
@@ -476,7 +477,6 @@ class PrimaiteSession:
|
|||||||
action_space=action_space,
|
action_space=action_space,
|
||||||
observation_space=obs_space,
|
observation_space=obs_space,
|
||||||
reward_function=rew_function,
|
reward_function=rew_function,
|
||||||
execution_definition=execution_definition,
|
|
||||||
agent_settings=agent_settings,
|
agent_settings=agent_settings,
|
||||||
)
|
)
|
||||||
sess.agents.append(new_agent)
|
sess.agents.append(new_agent)
|
||||||
|
|||||||
@@ -2,7 +2,6 @@ from enum import IntEnum
|
|||||||
from ipaddress import IPv4Address
|
from ipaddress import IPv4Address
|
||||||
from typing import Optional
|
from typing import Optional
|
||||||
|
|
||||||
from primaite.game.agent.interface import AgentExecutionDefinition
|
|
||||||
from primaite.game.science import simulate_trial
|
from primaite.game.science import simulate_trial
|
||||||
from primaite.simulator.system.applications.application import ApplicationOperatingState
|
from primaite.simulator.system.applications.application import ApplicationOperatingState
|
||||||
from primaite.simulator.system.applications.database_client import DatabaseClient
|
from primaite.simulator.system.applications.database_client import DatabaseClient
|
||||||
@@ -36,8 +35,10 @@ class DataManipulationBot(DatabaseClient):
|
|||||||
server_ip_address: Optional[IPv4Address] = None
|
server_ip_address: Optional[IPv4Address] = None
|
||||||
payload: Optional[str] = None
|
payload: Optional[str] = None
|
||||||
server_password: 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
|
attack_stage: DataManipulationAttackStage = DataManipulationAttackStage.NOT_STARTED
|
||||||
execution_definition: AgentExecutionDefinition = AgentExecutionDefinition()
|
|
||||||
repeat: bool = False
|
repeat: bool = False
|
||||||
"Whether to repeat attacking once finished."
|
"Whether to repeat attacking once finished."
|
||||||
|
|
||||||
@@ -50,6 +51,8 @@ class DataManipulationBot(DatabaseClient):
|
|||||||
server_ip_address: IPv4Address,
|
server_ip_address: IPv4Address,
|
||||||
server_password: Optional[str] = None,
|
server_password: Optional[str] = None,
|
||||||
payload: 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,
|
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_ip_address: The IP address of the Node the DatabaseService is on.
|
||||||
:param server_password: The password on the DatabaseService.
|
:param server_password: The password on the DatabaseService.
|
||||||
:param payload: The data manipulation query payload.
|
: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.
|
:param repeat: Whether to repeat attacking once finished.
|
||||||
"""
|
"""
|
||||||
self.server_ip_address = server_ip_address
|
self.server_ip_address = server_ip_address
|
||||||
self.payload = payload
|
self.payload = payload
|
||||||
self.server_password = server_password
|
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.repeat = repeat
|
||||||
self.sys_log.info(
|
self.sys_log.info(
|
||||||
f"{self.name}: Configured the {self.name} with {server_ip_address=}, {payload=}, {server_password=}, "
|
f"{self.name}: Configured the {self.name} with {server_ip_address=}, {payload=}, {server_password=}, "
|
||||||
|
|||||||
Reference in New Issue
Block a user