#2869 - Remove outdated parameters from actionmanager (action map achieves the same result) [skip ci]

This commit is contained in:
Marek Wolan
2025-01-14 13:48:18 +00:00
parent fb5242bd78
commit 40d052141c
42 changed files with 36 additions and 1679 deletions

View File

@@ -32,19 +32,6 @@ agents:
2: 0.1
observation_space: null
action_space:
action_list:
- type: DONOTHING
- type: NODE_APPLICATION_EXECUTE
options:
nodes:
- node_name: client_2
applications:
- application_name: WebBrowser
- application_name: DatabaseClient
max_folders_per_node: 1
max_files_per_folder: 1
max_services_per_node: 1
max_applications_per_node: 2
action_map:
0:
action: DONOTHING
@@ -81,19 +68,6 @@ agents:
2: 0.1
observation_space: null
action_space:
action_list:
- type: DONOTHING
- type: NODE_APPLICATION_EXECUTE
options:
nodes:
- node_name: client_1
applications:
- application_name: WebBrowser
- application_name: DatabaseClient
max_folders_per_node: 1
max_files_per_folder: 1
max_services_per_node: 1
max_applications_per_node: 2
action_map:
0:
action: DONOTHING
@@ -131,20 +105,6 @@ agents:
observation_space: null
action_space:
action_list:
- type: DONOTHING
- type: NODE_APPLICATION_EXECUTE
options:
nodes:
- node_name: client_1
applications:
- application_name: DataManipulationBot
- node_name: client_2
applications:
- application_name: DataManipulationBot
max_folders_per_node: 1
max_files_per_folder: 1
max_services_per_node: 1
reward_function:
reward_components:
@@ -236,35 +196,6 @@ agents:
options: {}
action_space:
action_list:
- type: DONOTHING
- type: NODE_SERVICE_SCAN
- type: NODE_SERVICE_STOP
- type: NODE_SERVICE_START
- type: NODE_SERVICE_PAUSE
- type: NODE_SERVICE_RESUME
- type: NODE_SERVICE_RESTART
- type: NODE_SERVICE_DISABLE
- type: NODE_SERVICE_ENABLE
- type: NODE_SERVICE_FIX
- type: NODE_FILE_SCAN
- type: NODE_FILE_CHECKHASH
- type: NODE_FILE_DELETE
- type: NODE_FILE_REPAIR
- type: NODE_FILE_RESTORE
- type: NODE_FOLDER_SCAN
- type: NODE_FOLDER_CHECKHASH
- type: NODE_FOLDER_REPAIR
- type: NODE_FOLDER_RESTORE
- type: NODE_OS_SCAN
- type: NODE_SHUTDOWN
- type: NODE_STARTUP
- type: NODE_RESET
- type: ROUTER_ACL_ADDRULE
- type: ROUTER_ACL_REMOVERULE
- type: HOST_NIC_ENABLE
- type: HOST_NIC_DISABLE
action_map:
0:
action: DONOTHING
@@ -685,41 +616,6 @@ agents:
options:
nodes:
- node_name: domain_controller
- node_name: web_server
applications:
- application_name: DatabaseClient
services:
- service_name: WebServer
- node_name: database_server
folders:
- folder_name: database
files:
- file_name: database.db
services:
- service_name: DatabaseService
- node_name: backup_server
- node_name: security_suite
- node_name: client_1
- node_name: client_2
max_folders_per_node: 2
max_files_per_folder: 2
max_services_per_node: 2
max_nics_per_node: 8
max_acl_rules: 10
ip_list:
- 192.168.1.10
- 192.168.1.12
- 192.168.1.14
- 192.168.1.16
- 192.168.1.110
- 192.168.10.21
- 192.168.10.22
- 192.168.10.110
reward_function:
reward_components:

View File

@@ -28,19 +28,6 @@ agents:
2: 0.1
observation_space: null
action_space:
action_list:
- type: DONOTHING
- type: NODE_APPLICATION_EXECUTE
options:
nodes:
- node_name: client_2
applications:
- application_name: WebBrowser
- application_name: DatabaseClient
max_folders_per_node: 1
max_files_per_folder: 1
max_services_per_node: 1
max_applications_per_node: 2
action_map:
0:
action: DONOTHING
@@ -77,19 +64,6 @@ agents:
2: 0.1
observation_space: null
action_space:
action_list:
- type: DONOTHING
- type: NODE_APPLICATION_EXECUTE
options:
nodes:
- node_name: client_1
applications:
- application_name: WebBrowser
- application_name: DatabaseClient
max_folders_per_node: 1
max_files_per_folder: 1
max_services_per_node: 1
max_applications_per_node: 2
action_map:
0:
action: DONOTHING
@@ -127,23 +101,6 @@ agents:
observation_space: null
action_space:
action_list:
- type: DONOTHING
- type: NODE_APPLICATION_EXECUTE
- type: NODE_FILE_DELETE
- type: NODE_FILE_CORRUPT
- type: NODE_OS_SCAN
options:
nodes:
- node_name: client_1
applications:
- application_name: DataManipulationBot
- node_name: client_2
applications:
- application_name: DataManipulationBot
max_folders_per_node: 1
max_files_per_folder: 1
max_services_per_node: 1
reward_function:
reward_components:
@@ -230,35 +187,6 @@ agents:
options: {}
action_space:
action_list:
- type: DONOTHING
- type: NODE_SERVICE_SCAN
- type: NODE_SERVICE_STOP
- type: NODE_SERVICE_START
- type: NODE_SERVICE_PAUSE
- type: NODE_SERVICE_RESUME
- type: NODE_SERVICE_RESTART
- type: NODE_SERVICE_DISABLE
- type: NODE_SERVICE_ENABLE
- type: NODE_SERVICE_FIX
- type: NODE_FILE_SCAN
- type: NODE_FILE_CHECKHASH
- type: NODE_FILE_DELETE
- type: NODE_FILE_REPAIR
- type: NODE_FILE_RESTORE
- type: NODE_FOLDER_SCAN
- type: NODE_FOLDER_CHECKHASH
- type: NODE_FOLDER_REPAIR
- type: NODE_FOLDER_RESTORE
- type: NODE_OS_SCAN
- type: NODE_SHUTDOWN
- type: NODE_STARTUP
- type: NODE_RESET
- type: ROUTER_ACL_ADDRULE
- type: ROUTER_ACL_REMOVERULE
- type: HOST_NIC_ENABLE
- type: HOST_NIC_DISABLE
action_map:
0:
action: DONOTHING
@@ -678,42 +606,6 @@ agents:
nic_id: 0
options:
nodes:
- node_name: domain_controller
- node_name: web_server
applications:
- application_name: DatabaseClient
services:
- service_name: WebServer
- node_name: database_server
folders:
- folder_name: database
files:
- file_name: database.db
services:
- service_name: DatabaseService
- node_name: backup_server
- node_name: security_suite
- node_name: client_1
- node_name: client_2
max_folders_per_node: 2
max_files_per_folder: 2
max_services_per_node: 2
max_nics_per_node: 8
max_acl_rules: 10
ip_list:
- 192.168.1.10
- 192.168.1.12
- 192.168.1.14
- 192.168.1.16
- 192.168.1.110
- 192.168.10.21
- 192.168.10.22
- 192.168.10.110
reward_function:
reward_components:
- type: DATABASE_FILE_INTEGRITY
@@ -810,39 +702,6 @@ agents:
options: {}
action_space:
action_list:
- type: DONOTHING
- type: NODE_SERVICE_SCAN
- type: NODE_SERVICE_STOP
- type: NODE_SERVICE_START
- type: NODE_SERVICE_PAUSE
- type: NODE_SERVICE_RESUME
- type: NODE_SERVICE_RESTART
- type: NODE_SERVICE_DISABLE
- type: NODE_SERVICE_ENABLE
- type: NODE_SERVICE_FIX
- type: NODE_FILE_SCAN
- type: NODE_FILE_CHECKHASH
- type: NODE_FILE_DELETE
- type: NODE_FILE_REPAIR
- type: NODE_FILE_RESTORE
- type: NODE_FOLDER_SCAN
- type: NODE_FOLDER_CHECKHASH
- type: NODE_FOLDER_REPAIR
- type: NODE_FOLDER_RESTORE
- type: NODE_OS_SCAN
- type: NODE_SHUTDOWN
- type: NODE_STARTUP
- type: NODE_RESET
- type: ROUTER_ACL_ADDRULE
options:
target_router: router_1
- type: ROUTER_ACL_REMOVERULE
options:
target_router: router_1
- type: HOST_NIC_ENABLE
- type: HOST_NIC_DISABLE
action_map:
0:
action: DONOTHING
@@ -1263,41 +1122,6 @@ agents:
options:
nodes:
- node_name: domain_controller
- node_name: web_server
applications:
- application_name: DatabaseClient
services:
- service_name: WebServer
- node_name: database_server
folders:
- folder_name: database
files:
- file_name: database.db
services:
- service_name: DatabaseService
- node_name: backup_server
- node_name: security_suite
- node_name: client_1
- node_name: client_2
max_folders_per_node: 2
max_files_per_folder: 2
max_services_per_node: 2
max_nics_per_node: 8
max_acl_rules: 10
ip_list:
- 192.168.1.10
- 192.168.1.12
- 192.168.1.14
- 192.168.1.16
- 192.168.1.110
- 192.168.10.21
- 192.168.10.22
- 192.168.10.110
reward_function:
reward_components:
- type: DATABASE_FILE_INTEGRITY

View File

@@ -8,12 +8,6 @@ agents:
type: ProxyAgent
observation_space: null
action_space:
action_list:
- type: DONOTHING
- type: NODE_SHUTDOWN
- type: NODE_STARTUP
- type: HOST_NIC_ENABLE
- type: HOST_NIC_DISABLE
action_map:
0:
action: DONOTHING
@@ -54,18 +48,6 @@ agents:
options:
node_id: 1
nic_id: 0
options:
nodes:
- node_name: client_1
- node_name: server
max_folders_per_node: 0
max_files_per_folder: 0
max_services_per_node: 0
max_nics_per_node: 1
max_acl_rules: 0
ip_list:
- 192.168.1.2
- 192.168.1.3
reward_function:
reward_components: []

View File

@@ -8,14 +8,6 @@ agents: &greens
1: 0.8
observation_space: null
action_space:
action_list:
- type: DONOTHING
- type: NODE_APPLICATION_EXECUTE
options:
nodes:
- node_name: client
applications:
- application_name: DatabaseClient
action_map:
0:
action: DONOTHING

View File

@@ -8,14 +8,6 @@ agents: &greens
1: 0.05
observation_space: null
action_space:
action_list:
- type: DONOTHING
- type: NODE_APPLICATION_EXECUTE
options:
nodes:
- node_name: client
applications:
- application_name: DatabaseClient
action_map:
0:
action: DONOTHING

View File

@@ -6,14 +6,6 @@ reds: &reds
observation_space: null
action_space:
action_list:
- type: DONOTHING
- type: NODE_APPLICATION_EXECUTE
options:
nodes:
- node_name: client
applications:
- application_name: DataManipulationBot
reward_function:
reward_components:

View File

@@ -6,14 +6,6 @@ reds: &reds
observation_space: null
action_space:
action_list:
- type: DONOTHING
- type: NODE_APPLICATION_EXECUTE
options:
nodes:
- node_name: client
applications:
- application_name: DataManipulationBot
reward_function:
reward_components:

View File

@@ -54,12 +54,6 @@ agents:
- server:eth-1<->switch_1:eth-2
action_space:
action_list:
- type: do_nothing
- type: node_shutdown
- type: node_startup
- type: host_nic_enable
- type: host_nic_enable
action_map:
0:
action: do_nothing
@@ -100,19 +94,6 @@ agents:
options:
node_name: server
nic_id: 0
options:
nodes:
- node_name: client
- node_name: server
max_folders_per_node: 0
max_files_per_folder: 0
max_services_per_node: 0
max_nics_per_node: 1
max_acl_rules: 0
ip_list:
- 192.168.1.2
- 192.168.1.3
reward_function:
reward_components:

View File

@@ -13,7 +13,7 @@ agents:
from __future__ import annotations
from typing import Dict, List, Optional, Tuple
from typing import Dict, Optional, Tuple
from gymnasium import spaces
@@ -41,31 +41,12 @@ class DoNothingAction(AbstractAction, identifier="do_nothing"):
class ActionManager:
"""Class which manages the action space for an agent."""
def __init__(
self,
actions: List[Dict], # stores list of actions available to agent
nodes: List[Dict], # extra configuration for each node
act_map: Optional[
Dict[int, Dict]
] = None, # allows restricting set of possible actions - TODO: Refactor to be a list?
*args,
**kwargs,
) -> None:
def __init__(self, act_map: Optional[Dict[int, Dict]] = None) -> None:
"""Init method for ActionManager.
:param game: Reference to the game to which the agent belongs.
:type game: PrimaiteGame
:param actions: List of action specs which should be made available to the agent. The keys of each spec are:
'type' and 'options' for passing any options to the action class's init method
:type actions: List[dict]
:param act_map: Action map which maps integers to actions. Used for restricting the set of possible actions.
:type act_map: Optional[Dict[int, Dict]]
"""
self.actions: Dict[str, AbstractAction] = {}
for act_spec in actions:
act_type = act_spec.get("type")
self.actions[act_type] = AbstractAction._registry[act_type]
self.action_map: Dict[int, Tuple[str, Dict]] = {}
"""
Action mapping that converts an integer to a specific action and parameter choice.
@@ -73,6 +54,7 @@ class ActionManager:
For example :
{0: ("node_service_scan", {node_name:"client_1", service_name:"WebBrowser"})}
"""
# allows restricting set of possible actions - TODO: Refactor to be a list?
if act_map is None:
# raise RuntimeError("Action map must be specified in the config file.")
pass
@@ -100,39 +82,17 @@ class ActionManager:
return spaces.Discrete(len(self.action_map))
@classmethod
def from_config(cls, game: "PrimaiteGame", cfg: Dict) -> "ActionManager": # noqa: F821
def from_config(cls, cfg: Dict) -> "ActionManager":
"""
Construct an ActionManager from a config definition.
Construct an ActionManager from a config dictionary.
The action space config supports the following three sections:
1. ``action_list``
``action_list`` contains a list action components which need to be included in the action space.
Each action component has a ``type`` which maps to a subclass of AbstractAction, and additional options
which will be passed to the action class's __init__ method during initialisation.
2. ``action_map``
Since the agent uses a discrete action space which acts as a flattened version of the component-based
action space, action_map provides a mapping between an integer (chosen by the agent) and a meaningful
action and values of parameters. For example action 0 can correspond to do nothing, action 1 can
correspond to "node_service_scan" with ``node_name="server"`` and
``service_name="WebBrowser"``, action 2 can be "
3. ``options``
``options`` contains a dictionary of options which are passed to the ActionManager's __init__ method.
These options are used to calculate the shape of the action space, and to provide additional information
to the ActionManager which is required to convert the agent's action choice into a CAOS request.
The action space config supports must contain the following key:
``action_map`` - List of actions available to the agent, formatted as a dictionary where the key is the
action number between 0 - N, and the value is the CAOS-formatted action.
:param game: The Primaite Game to which the agent belongs.
:type game: PrimaiteGame
:param cfg: The action space config.
:type cfg: Dict
:return: The constructed ActionManager.
:rtype: ActionManager
"""
obj = cls(
actions=cfg["action_list"],
**cfg["options"],
protocols=game.options.protocols,
ports=game.options.ports,
act_map=cfg.get("action_map"),
)
return obj
return cls(**cfg.get("options", {}), act_map=cfg.get("action_map"))

View File

@@ -93,8 +93,8 @@ class AbstractAgent(BaseModel):
return ValueError(f"Invalid Agent Type: {config['type']}")
obj = cls(
config=cls.ConfigSchema(**config["agent_settings"]),
action_manager=ActionManager.from_config(config["game"], config["action_manager"]),
observation_manager=ObservationManager.from_config(config["observation_manager"]),
action_manager=ActionManager.from_config(config["action_space"]),
observation_manager=ObservationManager.from_config(config["observation_space"]),
reward_function=RewardFunction.from_config(config["reward_function"]),
)
return obj

View File

@@ -534,8 +534,8 @@ class PrimaiteGame:
agent_config = {
"type": agent_type,
"action_manager": action_space_cfg,
"observation_manager": observation_space_cfg,
"action_space": action_space_cfg,
"observation_space": observation_space_cfg,
"reward_function": reward_function_cfg,
"agent_settings": agent_settings,
"game": game,

View File

@@ -53,22 +53,13 @@
" type: ProxyAgent\n",
" observation_space: null\n",
" action_space:\n",
" action_list:\n",
" - type: DONOTHING\n",
" - type: NODE_APPLICATION_INSTALL\n",
" - type: NODE_APPLICATION_EXECUTE\n",
" - type: CONFIGURE_C2_BEACON\n",
" - type: C2_SERVER_RANSOMWARE_LAUNCH\n",
" - type: C2_SERVER_RANSOMWARE_CONFIGURE\n",
" - type: C2_SERVER_TERMINAL_COMMAND\n",
" - type: C2_SERVER_DATA_EXFILTRATE\n",
" options:\n",
" nodes:\n",
" - node_name: web_server\n",
" applications: \n",
" applications:\n",
" - application_name: C2Beacon\n",
" - node_name: client_1\n",
" applications: \n",
" applications:\n",
" - application_name: C2Server\n",
" max_folders_per_node: 1\n",
" max_files_per_folder: 1\n",
@@ -102,7 +93,7 @@
" action: NODE_APPLICATION_EXECUTE\n",
" options:\n",
" node_id: 0\n",
" application_id: 0 \n",
" application_id: 0\n",
" 4:\n",
" action: C2_SERVER_TERMINAL_COMMAND\n",
" options:\n",
@@ -112,7 +103,7 @@
" username: admin\n",
" password: admin\n",
" commands:\n",
" - \n",
" -\n",
" - software_manager\n",
" - application\n",
" - install\n",
@@ -134,7 +125,7 @@
" target_ip_address: 192.168.1.14\n",
" account:\n",
" username: admin\n",
" password: admin \n",
" password: admin\n",
"\n",
" 7:\n",
" action: C2_SERVER_RANSOMWARE_LAUNCH\n",
@@ -177,7 +168,7 @@
" # removing all agents & adding the custom agent.\n",
" cfg['agents'] = {}\n",
" cfg['agents'] = c2_agent_yaml\n",
" \n",
"\n",
"\n",
"env = PrimaiteGymEnv(env_config=cfg)"
]
@@ -230,10 +221,6 @@
"\n",
"```yaml\n",
" action_space:\n",
" action_list:\n",
" ...\n",
" - type: NODE_APPLICATION_INSTALL\n",
" ...\n",
" options:\n",
" nodes: # Node List\n",
" - node_name: web_server\n",
@@ -273,10 +260,6 @@
"\n",
"```yaml\n",
" action_space:\n",
" action_list:\n",
" ...\n",
" - type: CONFIGURE_C2_BEACON\n",
" ...\n",
" options:\n",
" nodes: # Node List\n",
" - node_name: web_server\n",
@@ -320,10 +303,6 @@
"\n",
"```yaml\n",
" action_space:\n",
" action_list:\n",
" ...\n",
" - type: NODE_APPLICATION_EXECUTE\n",
" ...\n",
" options:\n",
" nodes: # Node List\n",
" - node_name: web_server\n",
@@ -347,7 +326,7 @@
"metadata": {},
"outputs": [],
"source": [
"env.step(3) "
"env.step(3)"
]
},
{
@@ -390,10 +369,6 @@
"\n",
"``` yaml\n",
" action_space:\n",
" action_list:\n",
" ...\n",
" - type: C2_SERVER_TERMINAL_COMMAND\n",
" ...\n",
" options:\n",
" nodes: # Node List\n",
" ...\n",
@@ -451,10 +426,6 @@
"\n",
"``` yaml\n",
" action_space:\n",
" action_list:\n",
" ...\n",
" - type: C2_SERVER_RANSOMWARE_CONFIGURE\n",
" ...\n",
" options:\n",
" nodes: # Node List\n",
" ...\n",
@@ -507,10 +478,6 @@
"\n",
"``` yaml\n",
" action_space:\n",
" action_list:\n",
" ...\n",
" - type: C2_SERVER_DATA_EXFILTRATE\n",
" ...\n",
" options:\n",
" nodes: # Node List\n",
" ...\n",
@@ -577,10 +544,6 @@
"\n",
"``` yaml\n",
" action_space:\n",
" action_list:\n",
" ...\n",
" - type: C2_SERVER_RANSOMWARE_LAUNCH\n",
" ...\n",
" options:\n",
" nodes: # Node List\n",
" ...\n",
@@ -632,7 +595,7 @@
"metadata": {},
"outputs": [],
"source": [
"custom_blue_agent_yaml = \"\"\" \n",
"custom_blue_agent_yaml = \"\"\"\n",
" - ref: defender\n",
" team: BLUE\n",
" type: ProxyAgent\n",
@@ -715,13 +678,8 @@
" - type: \"NONE\"\n",
" label: ICS\n",
" options: {}\n",
" \n",
"\n",
" action_space:\n",
" action_list:\n",
" - type: NODE_APPLICATION_REMOVE\n",
" - type: NODE_SHUTDOWN\n",
" - type: ROUTER_ACL_ADDRULE\n",
" - type: DONOTHING\n",
" action_map:\n",
" 0:\n",
" action: DONOTHING\n",
@@ -747,7 +705,7 @@
" dest_port_id: 2\n",
" protocol_id: 1\n",
" source_wildcard_id: 0\n",
" dest_wildcard_id: 0 \n",
" dest_wildcard_id: 0\n",
"\n",
"\n",
" options:\n",
@@ -796,7 +754,7 @@
" # removing all agents & adding the custom agent.\n",
" cfg['agents'] = {}\n",
" cfg['agents'] = custom_blue\n",
" \n",
"\n",
"\n",
"blue_env = PrimaiteGymEnv(env_config=cfg)"
]
@@ -1468,7 +1426,7 @@
" # removing all agents & adding the custom agent.\n",
" cfg['agents'] = {}\n",
" cfg['agents'] = c2_agent_yaml\n",
" \n",
"\n",
"\n",
"c2_config_env = PrimaiteGymEnv(env_config=cfg)"
]
@@ -1555,7 +1513,7 @@
"source": [
"for i in range(6):\n",
" env.step(0)\n",
" \n",
"\n",
"c2_server_1.show()"
]
},
@@ -1676,7 +1634,7 @@
"metadata": {},
"outputs": [],
"source": [
"# Comparing the OBS of the default frequency to a timestep frequency of 1 \n",
"# Comparing the OBS of the default frequency to a timestep frequency of 1\n",
"for i in range(2):\n",
" keep_alive_obs, _, _, _, _ = blue_config_env.step(0)\n",
" display_obs_diffs(default_obs, keep_alive_obs, blue_config_env.game.step_counter)"
@@ -1760,7 +1718,7 @@
"metadata": {},
"outputs": [],
"source": [
"# Capturing default C2 Traffic \n",
"# Capturing default C2 Traffic\n",
"for i in range(3):\n",
" tcp_c2_obs, _, _, _, _ = blue_config_env.step(0)\n",
"\n",

View File

@@ -147,12 +147,7 @@
" nodes: {}\n",
"\n",
" action_space:\n",
"\n",
" # The agent has two action choices, either do nothing, or execute a pre-scripted attack by using \n",
" action_list:\n",
" - type: DONOTHING\n",
" - type: NODE_APPLICATION_EXECUTE\n",
"\n",
" \n",
" # The agent has access to the DataManipulationBoth on clients 1 and 2.\n",
" options:\n",
" nodes:\n",
@@ -306,19 +301,9 @@
"outputs": [],
"source": [
"change = yaml.safe_load(\"\"\"\n",
"action_space:\n",
" action_list:\n",
" - type: DONOTHING\n",
" - type: NODE_APPLICATION_EXECUTE\n",
" options:\n",
" nodes:\n",
" - node_name: client_1\n",
" applications:\n",
" - application_name: DataManipulationBot\n",
" max_folders_per_node: 1\n",
" max_files_per_folder: 1\n",
" max_services_per_node: 1\n",
"# TODO:\n",
"\"\"\")\n",
"#TODO 2869 fix\n",
"\n",
"with open(data_manipulation_config_path(), 'r') as f:\n",
" cfg = yaml.safe_load(f)\n",