From 858406c4a355392c23cbec69c45fd22061bccc57 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Fri, 17 Jan 2025 14:38:59 +0000 Subject: [PATCH] update more tests to new actions schema --- docs/index.rst | 2 +- docs/source/about.rst | 4 +- docs/source/action_masking.rst | 8 +-- .../_package_data/data_manipulation.yaml | 20 +++--- .../_package_data/data_manipulation_marl.yaml | 40 ++++++------ .../scenario_with_placeholders/scenario.yaml | 8 +-- src/primaite/game/agent/actions/abstract.py | 2 +- src/primaite/game/agent/actions/acl.py | 10 +-- .../game/agent/actions/application.py | 2 +- src/primaite/game/agent/actions/host_nic.py | 4 +- src/primaite/game/agent/actions/manager.py | 2 +- src/primaite/game/agent/actions/node.py | 23 +++---- src/primaite/game/agent/actions/session.py | 2 +- src/primaite/game/agent/actions/software.py | 64 ++++++++++--------- src/primaite/game/agent/interface.py | 4 +- src/primaite/game/agent/rewards.py | 4 +- .../scripted_agents/probabilistic_agent.py | 3 +- ...ommand-and-Control-E2E-Demonstration.ipynb | 2 +- ...ege-Escalation-and Data-Loss-Example.ipynb | 6 +- .../system/services/terminal/terminal.py | 2 +- tests/assets/configs/action_penalty.yaml | 20 +++--- tests/assets/configs/basic_firewall.yaml | 4 -- .../configs/basic_switched_network.yaml | 3 - tests/assets/configs/data_manipulation.yaml | 20 +++--- tests/assets/configs/dmz_network.yaml | 2 +- tests/assets/configs/extended_config.yaml | 32 +++++----- .../configs/firewall_actions_network.yaml | 35 +++++----- .../configs/install_and_configure_apps.yaml | 28 ++++---- tests/assets/configs/multi_agent_session.yaml | 40 ++++++------ .../nmap_ping_scan_red_agent_config.yaml | 2 +- tests/assets/configs/shared_rewards.yaml | 32 +++++----- .../configs/test_application_install.yaml | 27 ++++---- .../test_application_request_permission.py | 8 +-- .../actions/test_c2_suite_actions.py | 38 +++++------ .../actions/test_configure_actions.py | 50 +++++---------- .../actions/test_file_request_permission.py | 14 ++-- .../actions/test_folder_request_permission.py | 14 ++-- .../actions/test_nic_request_permission.py | 16 ++--- .../actions/test_node_request_permission.py | 14 ++-- .../test_service_request_permission.py | 18 +++--- .../actions/test_terminal_actions.py | 24 +++---- .../game_layer/test_action_mask.py | 4 +- .../game_layer/test_actions.py | 31 +++------ .../game_layer/test_rewards.py | 8 +-- .../_primaite/_game/_agent/test_actions.py | 40 ++++-------- .../_primaite/_game/_agent/test_agent.py | 2 + .../_game/_agent/test_observations.py | 2 +- .../_game/_agent/test_sticky_rewards.py | 24 +++---- 48 files changed, 351 insertions(+), 413 deletions(-) diff --git a/docs/index.rst b/docs/index.rst index 42cc1d6d..aa7d16e0 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -70,7 +70,7 @@ PrimAITE incorporates the following features: - Architected with a separate Simulation layer and Game layer. This separation of concerns defines a clear path towards transfer learning with environments of differing fidelity; - Ability to reconfigure an RL reward function based on (a) the ability to counter the modelled adversarial cyber-attack, and (b) the ability to ensure success for green agents; -- Access Control List (ACL) functions for network devices (routers and firewalls), following standard ACL rule format (e.g., DENY / ALLOW, source / destination IP addresses, protocol and port); +- Access Control List (ACL) functions for network devices (routers and firewalls), following standard ACL rule format (e.g., DENY / PERMIT, source / destination IP addresses, protocol and port); - Application of traffic to the links of the system laydown adheres to the ACL rulesets and routing tables contained within each network device; - Provides RL environments adherent to the Farama Foundation Gymnasium (Previously OpenAI Gym) API, allowing integration with any compliant RL Agent frameworks; - Provides RL environments adherent to Ray RLlib environment specifications for single-agent and multi-agent scenarios; diff --git a/docs/source/about.rst b/docs/source/about.rst index da87102a..839bbb0b 100644 --- a/docs/source/about.rst +++ b/docs/source/about.rst @@ -184,7 +184,7 @@ Head over to the :ref:`getting-started` page to install and setup PrimAITE! - 192.168.1.5 - ANY - ANY - All ACL rules are considered when applying an IER. Logic follows the order of rules, so a DENY or ALLOW for the same parameters will override an earlier entry. + All ACL rules are considered when applying an IER. Logic follows the order of rules, so a DENY or PERMIT for the same parameters will override an earlier entry. Observation Spaces ****************** The observation space provides the blue agent with information about the current status of nodes and links. @@ -331,7 +331,7 @@ Head over to the :ref:`getting-started` page to install and setup PrimAITE! * Dictionary item {... ,1: [x1, x2, x3, x4, x5, x6] ...} The placeholders inside the list under the key '1' mean the following: * [0, 2] - Action (0 = do nothing, 1 = create rule, 2 = delete rule) - * [0, 1] - Permission (0 = DENY, 1 = ALLOW) + * [0, 1] - Permission (0 = DENY, 1 = PERMIT) * [0, num nodes] - Source IP (0 = any, then 1 -> x resolving to IP addresses) * [0, num nodes] - Dest IP (0 = any, then 1 -> x resolving to IP addresses) * [0, num services] - Protocol (0 = any, then 1 -> x resolving to protocol) diff --git a/docs/source/action_masking.rst b/docs/source/action_masking.rst index c6e4ca59..359ad452 100644 --- a/docs/source/action_masking.rst +++ b/docs/source/action_masking.rst @@ -134,15 +134,15 @@ The following logic is applied: +------------------------------------------+---------------------------------------------------------------------+ | **C2_SERVER_RANSOMWARE_CONFIGURE** | Node is on. | +------------------------------------------+---------------------------------------------------------------------+ -| **C2_SERVER_TERMINAL_COMMAND** | Node is on. | +| **c2_server_terminal_command** | Node is on. | +------------------------------------------+---------------------------------------------------------------------+ | **C2_SERVER_DATA_EXFILTRATE** | Node is on. | +------------------------------------------+---------------------------------------------------------------------+ -| **NODE_ACCOUNTS_CHANGE_PASSWORD** | Node is on. | +| **node_account_change_password** | Node is on. | +------------------------------------------+---------------------------------------------------------------------+ -| **SSH_TO_REMOTE** | Node is on. | +| **node_session_remote_login** | Node is on. | +------------------------------------------+---------------------------------------------------------------------+ -| **SESSIONS_REMOTE_LOGOFF** | Node is on. | +| **node_session_remote_logoff** | Node is on. | +------------------------------------------+---------------------------------------------------------------------+ | **NODE_SEND_REMOTE_COMMAND** | Node is on. | +------------------------------------------+---------------------------------------------------------------------+ diff --git a/src/primaite/config/_package_data/data_manipulation.yaml b/src/primaite/config/_package_data/data_manipulation.yaml index 2f6e24b3..fa10a463 100644 --- a/src/primaite/config/_package_data/data_manipulation.yaml +++ b/src/primaite/config/_package_data/data_manipulation.yaml @@ -478,52 +478,52 @@ agents: 52: # old action num: 28 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 0 53: # old action num: 29 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 1 54: # old action num: 30 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 2 55: # old action num: 31 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 3 56: # old action num: 32 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 4 57: # old action num: 33 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 5 58: # old action num: 34 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 6 59: # old action num: 35 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 7 60: # old action num: 36 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 8 61: # old action num: 37 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 9 62: # old action num: 38 action: "host_nic_disable" diff --git a/src/primaite/config/_package_data/data_manipulation_marl.yaml b/src/primaite/config/_package_data/data_manipulation_marl.yaml index 53ff0634..b0131c8c 100644 --- a/src/primaite/config/_package_data/data_manipulation_marl.yaml +++ b/src/primaite/config/_package_data/data_manipulation_marl.yaml @@ -470,52 +470,52 @@ agents: 52: # old action num: 28 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 0 53: # old action num: 29 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 1 54: # old action num: 30 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 2 55: # old action num: 31 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 3 56: # old action num: 32 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 4 57: # old action num: 33 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 5 58: # old action num: 34 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 6 59: # old action num: 35 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 7 60: # old action num: 36 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 8 61: # old action num: 37 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 9 62: # old action num: 38 action: "host_nic_disable" @@ -985,52 +985,52 @@ agents: 52: # old action num: 28 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 0 53: # old action num: 29 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 1 54: # old action num: 30 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 2 55: # old action num: 31 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 3 56: # old action num: 32 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 4 57: # old action num: 33 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 5 58: # old action num: 34 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 6 59: # old action num: 35 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 7 60: # old action num: 36 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 8 61: # old action num: 37 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 9 62: # old action num: 38 action: "host_nic_disable" diff --git a/src/primaite/config/_package_data/scenario_with_placeholders/scenario.yaml b/src/primaite/config/_package_data/scenario_with_placeholders/scenario.yaml index 0223beb6..c692c725 100644 --- a/src/primaite/config/_package_data/scenario_with_placeholders/scenario.yaml +++ b/src/primaite/config/_package_data/scenario_with_placeholders/scenario.yaml @@ -78,22 +78,22 @@ agents: action: host_nic_disable options: node_name: client_1 - nic_id: 0 + nic_num: 1 6: action: host_nic_disable options: node_name: server - nic_id: 0 + nic_num: 1 7: action: host_nic_enable options: node_name: client_1 - nic_id: 0 + nic_num: 1 8: action: host_nic_enable options: node_name: server - nic_id: 0 + nic_num: 1 reward_function: reward_components: diff --git a/src/primaite/game/agent/actions/abstract.py b/src/primaite/game/agent/actions/abstract.py index c570119b..1cda4360 100644 --- a/src/primaite/game/agent/actions/abstract.py +++ b/src/primaite/game/agent/actions/abstract.py @@ -18,7 +18,7 @@ class AbstractAction(BaseModel, ABC): """Base configuration schema for Actions.""" model_config = ConfigDict(extra="forbid") - type: str + type: str = "" _registry: ClassVar[Dict[str, Type[AbstractAction]]] = {} diff --git a/src/primaite/game/agent/actions/acl.py b/src/primaite/game/agent/actions/acl.py index ee5ed292..a097b906 100644 --- a/src/primaite/game/agent/actions/acl.py +++ b/src/primaite/game/agent/actions/acl.py @@ -28,7 +28,7 @@ class ACLAddRuleAbstractAction(AbstractAction, ABC): src_ip: IPV4Address protocol_name: Union[IPProtocol, Literal["ALL"]] - permission: Literal["ALLOW", "DENY"] + permission: Literal["PERMIT", "DENY"] position: int dst_ip: Union[IPV4Address, Literal["ALL"]] src_port: Union[Port, Literal["ALL"]] @@ -70,10 +70,10 @@ class RouterACLAddRuleAction(ACLAddRuleAbstractAction, identifier="router_acl_ad config.permission, config.protocol_name, str(config.src_ip), - config.src_wildcard, + str(config.src_wildcard), config.src_port, str(config.dst_ip), - config.dst_wildcard, + str(config.dst_wildcard), config.dst_port, config.position, ] @@ -121,10 +121,10 @@ class FirewallACLAddRuleAction(ACLAddRuleAbstractAction, identifier="firewall_ac config.permission, config.protocol_name, str(config.src_ip), - config.src_wildcard, + str(config.src_wildcard), config.src_port, str(config.dst_ip), - config.dst_wildcard, + str(config.dst_wildcard), config.dst_port, config.position, ] diff --git a/src/primaite/game/agent/actions/application.py b/src/primaite/game/agent/actions/application.py index 223effc4..f6ce0624 100644 --- a/src/primaite/game/agent/actions/application.py +++ b/src/primaite/game/agent/actions/application.py @@ -19,7 +19,7 @@ class NodeApplicationAbstractAction(AbstractAction, ABC): """ Base class for application actions. - Any action which applies to an application and uses node_id and application_id as its only two parameters can + Any action which applies to an application and uses node_name and application_name as its only two parameters can inherit from this base class. """ diff --git a/src/primaite/game/agent/actions/host_nic.py b/src/primaite/game/agent/actions/host_nic.py index b9206b9c..0ca816f3 100644 --- a/src/primaite/game/agent/actions/host_nic.py +++ b/src/primaite/game/agent/actions/host_nic.py @@ -12,8 +12,8 @@ class HostNICAbstractAction(AbstractAction, ABC): """ Abstract base class for NIC actions. - Any action which applies to a NIC and uses node_id and nic_id as its only two parameters can inherit from this base - class. + Any action which applies to a NIC and uses node_name and nic_num as its only two parameters can inherit from this + base class. """ config: "HostNICAbstractAction.ConfigSchema" diff --git a/src/primaite/game/agent/actions/manager.py b/src/primaite/game/agent/actions/manager.py index fefa22b8..a6e235c5 100644 --- a/src/primaite/game/agent/actions/manager.py +++ b/src/primaite/game/agent/actions/manager.py @@ -84,7 +84,7 @@ class ActionManager(BaseModel): def form_request(self, action_identifier: str, action_options: Dict) -> RequestFormat: """Take action in CAOS format and use the execution definition to change it into PrimAITE request format.""" act_class = AbstractAction._registry[action_identifier] - config = act_class.ConfigSchema(type=action_identifier, **action_options) + config = act_class.ConfigSchema(**action_options) return act_class.form_request(config=config) @property diff --git a/src/primaite/game/agent/actions/node.py b/src/primaite/game/agent/actions/node.py index 4a7f725e..95bf5c34 100644 --- a/src/primaite/game/agent/actions/node.py +++ b/src/primaite/game/agent/actions/node.py @@ -4,6 +4,8 @@ from typing import ClassVar, List, Optional, Union from primaite.game.agent.actions.manager import AbstractAction from primaite.interface.request import RequestFormat +from primaite.utils.validation.ip_protocol import IPProtocol +from primaite.utils.validation.port import Port __all__ = ( "NodeOSScanAction", @@ -92,7 +94,7 @@ class NodeNMAPAbstractAction(AbstractAction, identifier="node_nmap_abstract_acti target_ip_address: Union[str, List[str]] show: bool = False - node_name: str + source_node: str @classmethod @abstractmethod @@ -107,18 +109,13 @@ class NodeNMAPPingScanAction(NodeNMAPAbstractAction, identifier="node_nmap_ping_ config: "NodeNMAPPingScanAction.ConfigSchema" - class ConfigSchema(NodeNMAPAbstractAction.ConfigSchema): - """Configuration schema for NodeNMAPPingScanAction.""" - - pass - @classmethod - def form_request(cls, config: ConfigSchema) -> List[str]: # noqa + def form_request(cls, config: "NodeNMAPPingScanAction.ConfigSchema") -> List[str]: # noqa """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" return [ "network", "node", - config.node_name, + config.source_node, "application", "NMAP", "ping_scan", @@ -135,8 +132,8 @@ class NodeNMAPPortScanAction(NodeNMAPAbstractAction, identifier="node_nmap_port_ """Configuration Schema for NodeNMAPPortScanAction.""" source_node: str - target_protocol: Optional[Union[str, List[str]]] = (None,) - target_port: Optional[Union[str, List[str]]] = (None,) + target_protocol: Optional[Union[IPProtocol, List[IPProtocol]]] = None + target_port: Optional[Union[Port, List[Port]]] = None show: Optional[bool] = (False,) @classmethod @@ -166,11 +163,11 @@ class NodeNetworkServiceReconAction(NodeNMAPAbstractAction, identifier="node_net config: "NodeNetworkServiceReconAction.ConfigSchema" - class ConfigSchema(AbstractAction.ConfigSchema): + class ConfigSchema(NodeNMAPAbstractAction.ConfigSchema): """Configuration schema for NodeNetworkServiceReconAction.""" - target_protocol: Optional[Union[str, List[str]]] = (None,) - target_port: Optional[Union[str, List[str]]] = (None,) + target_protocol: Optional[Union[IPProtocol, List[IPProtocol]]] = None + target_port: Optional[Union[Port, List[Port]]] = None show: Optional[bool] = (False,) @classmethod diff --git a/src/primaite/game/agent/actions/session.py b/src/primaite/game/agent/actions/session.py index 1191987b..9720d371 100644 --- a/src/primaite/game/agent/actions/session.py +++ b/src/primaite/game/agent/actions/session.py @@ -55,7 +55,7 @@ class NodeSessionsRemoteLoginAction(NodeSessionAbstractAction, identifier="node_ config.node_name, "service", "Terminal", - "ssh_to_remote", + "node_session_remote_login", config.username, config.password, config.remote_ip, diff --git a/src/primaite/game/agent/actions/software.py b/src/primaite/game/agent/actions/software.py index 760e8dfa..23fbd70d 100644 --- a/src/primaite/game/agent/actions/software.py +++ b/src/primaite/game/agent/actions/software.py @@ -4,7 +4,7 @@ from typing import List, Optional, Union from pydantic import ConfigDict, Field, field_validator, ValidationInfo -from primaite.game.agent.actions.manager import AbstractAction, ActionManager +from primaite.game.agent.actions.manager import AbstractAction from primaite.interface.request import RequestFormat __all__ = ( @@ -28,36 +28,31 @@ class ConfigureRansomwareScriptAction(AbstractAction, identifier="c2_server_rans """Configuration schema for ConfigureRansomwareScriptAction.""" node_name: str - server_ip_address: Optional[str] - server_password: Optional[str] - payload: Optional[str] + server_ip_address: Optional[str] = None + server_password: Optional[str] = None + payload: Optional[str] = None @classmethod def form_request(cls, config: ConfigSchema) -> RequestFormat: """Return the action formatted as a request that can be ingested by the simulation.""" if config.node_name is None: return ["do_nothing"] - return [ - "network", - "node", - config.node_name, - "application", - "RansomwareScript", - "configure", - config.model_config, - ] + data = dict( + server_ip_address=config.server_ip_address, + server_password=config.server_password, + payload=config.payload, + ) + return ["network", "node", config.node_name, "application", "RansomwareScript", "configure", data] class ConfigureDoSBotAction(AbstractAction, identifier="configure_dos_bot"): """Action which sets config parameters for a DoS bot on a node.""" - config: "ConfigureDoSBotAction.ConfigSchema" - class ConfigSchema(AbstractAction.ConfigSchema): """Schema for options that can be passed to this action.""" - node_name: str model_config = ConfigDict(extra="forbid") + node_name: str target_ip_address: Optional[str] = None target_port: Optional[str] = None payload: Optional[str] = None @@ -66,22 +61,24 @@ class ConfigureDoSBotAction(AbstractAction, identifier="configure_dos_bot"): dos_intensity: Optional[float] = None max_sessions: Optional[int] = None - def __init__(self, manager: "ActionManager", **kwargs) -> None: - super().__init__(manager=manager) - - def form_request(self, config: ConfigSchema) -> RequestFormat: + @classmethod + def form_request(config: ConfigSchema) -> RequestFormat: """Return the action formatted as a request that can be ingested by the simulation.""" - if config.node_name is None: - return ["do_nothing"] - self.ConfigSchema.model_validate(config) # check that options adhere to schema - return ["network", "node", config.node_name, "application", "DoSBot", "configure", config] + data = dict( + target_ip_address=config.target_ip_address, + target_port=config.target_port, + payload=config.payload, + repeat=config.repeat, + port_scan_p_of_success=config.port_scan_p_of_success, + dos_intensity=config.dos_intensity, + max_sessions=config.max_sessions, + ) + return ["network", "node", config.node_name, "application", "DoSBot", "configure", data] class ConfigureC2BeaconAction(AbstractAction, identifier="configure_c2_beacon"): """Action which configures a C2 Beacon based on the parameters given.""" - config: "ConfigureC2BeaconAction.ConfigSchema" - class ConfigSchema(AbstractAction.ConfigSchema): """Configuration schema for ConfigureC2BeaconAction.""" @@ -91,6 +88,7 @@ class ConfigureC2BeaconAction(AbstractAction, identifier="configure_c2_beacon"): masquerade_protocol: str = Field(default="TCP") masquerade_port: str = Field(default="HTTP") + # TODO: this validator should not be needed anymore, test what happens if removed. @field_validator( "c2_server_ip_address", "keep_alive_frequency", @@ -108,7 +106,13 @@ class ConfigureC2BeaconAction(AbstractAction, identifier="configure_c2_beacon"): @classmethod def form_request(self, config: ConfigSchema) -> RequestFormat: """Return the action formatted as a request that can be ingested by the simulation.""" - return ["network", "node", config.node_name, "application", "C2Beacon", "configure", config] + data = dict( + c2_server_ip_address=config.c2_server_ip_address, + keep_alive_frequency=config.keep_alive_frequency, + masquerade_protocol=config.masquerade_protocol, + masquerade_port=config.masquerade_port, + ) + return ["network", "node", config.node_name, "application", "C2Beacon", "configure", data] class NodeSendRemoteCommandAction(AbstractAction, identifier="node_send_remote_command"): @@ -228,11 +232,13 @@ class ConfigureDatabaseClientAction(AbstractAction, identifier="configure_databa """Schema for options that can be passed to this action.""" node_name: str - model_config = ConfigDict(extra="forbid") + server_ip_address: Optional[str] = None + server_password: Optional[str] = None @classmethod def form_request(cls, config: ConfigSchema) -> RequestFormat: """Return the action formatted as a request that can be ingested by the simulation.""" if config.node_name is None: return ["do_nothing"] - return ["network", "node", config.node_name, "application", "DatabaseClient", "configure", config.model_config] + data = {"server_ip_address": config.server_ip_address, "server_password": config.server_password} + return ["network", "node", config.node_name, "application", "DatabaseClient", "configure", data] diff --git a/src/primaite/game/agent/interface.py b/src/primaite/game/agent/interface.py index f5714644..b58cdf29 100644 --- a/src/primaite/game/agent/interface.py +++ b/src/primaite/game/agent/interface.py @@ -58,9 +58,9 @@ class AbstractAgent(BaseModel, ABC): model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True) type: str - ref: str + ref: str = "" """name of the agent.""" - team: Optional[Literal["BLUE", "GREEN", "RED"]] + team: Optional[Literal["BLUE", "GREEN", "RED"]] = None agent_settings: AbstractAgent.AgentSettingsSchema = Field(default=lambda: AbstractAgent.AgentSettingsSchema()) action_space: ActionManager.ConfigSchema = Field(default_factory=lambda: ActionManager.ConfigSchema()) observation_space: ObservationManager.ConfigSchema = Field( diff --git a/src/primaite/game/agent/rewards.py b/src/primaite/game/agent/rewards.py index 2881f967..80be14ef 100644 --- a/src/primaite/game/agent/rewards.py +++ b/src/primaite/game/agent/rewards.py @@ -51,7 +51,7 @@ class AbstractReward(BaseModel): class ConfigSchema(BaseModel, ABC): """Config schema for AbstractReward.""" - type: str + type: str = "" _registry: ClassVar[Dict[str, Type["AbstractReward"]]] = {} @@ -404,7 +404,7 @@ class ActionPenalty(AbstractReward, identifier="ACTION_PENALTY"): :rtype: float """ if last_action_response.action == "do_nothing": - return self.do_nothing_penalty + return self.config.do_nothing_penalty else: return self.config.action_penalty diff --git a/src/primaite/game/agent/scripted_agents/probabilistic_agent.py b/src/primaite/game/agent/scripted_agents/probabilistic_agent.py index 20924a95..959eaadc 100644 --- a/src/primaite/game/agent/scripted_agents/probabilistic_agent.py +++ b/src/primaite/game/agent/scripted_agents/probabilistic_agent.py @@ -16,7 +16,6 @@ __all__ = "ProbabilisticAgent" class ProbabilisticAgent(AbstractScriptedAgent, identifier="ProbabilisticAgent"): """Scripted agent which randomly samples its action space with prescribed probabilities for each action.""" - config: "ProbabilisticAgent.ConfigSchema" = Field(default_factory=lambda: ProbabilisticAgent.ConfigSchema()) rng: Generator = np.random.default_rng(np.random.randint(0, 65535)) class AgentSettingsSchema(AbstractScriptedAgent.AgentSettingsSchema): @@ -52,6 +51,8 @@ class ProbabilisticAgent(AbstractScriptedAgent, identifier="ProbabilisticAgent") default_factory=lambda: ProbabilisticAgent.AgentSettingsSchema() ) + config: "ProbabilisticAgent.ConfigSchema" = Field(default_factory=lambda: ProbabilisticAgent.ConfigSchema()) + @property def probabilities(self) -> Dict[str, int]: """Convenience method to view the probabilities of the Agent.""" diff --git a/src/primaite/notebooks/Command-and-Control-E2E-Demonstration.ipynb b/src/primaite/notebooks/Command-and-Control-E2E-Demonstration.ipynb index 278fb3dc..66a684de 100644 --- a/src/primaite/notebooks/Command-and-Control-E2E-Demonstration.ipynb +++ b/src/primaite/notebooks/Command-and-Control-E2E-Demonstration.ipynb @@ -95,7 +95,7 @@ " node_id: 0\n", " application_id: 0\n", " 4:\n", - " action: C2_SERVER_TERMINAL_COMMAND\n", + " action: c2_server_terminal_command\n", " options:\n", " node_id: 1\n", " ip_address:\n", diff --git a/src/primaite/notebooks/Privilege-Escalation-and Data-Loss-Example.ipynb b/src/primaite/notebooks/Privilege-Escalation-and Data-Loss-Example.ipynb index c751edfd..fcda4dbd 100644 --- a/src/primaite/notebooks/Privilege-Escalation-and Data-Loss-Example.ipynb +++ b/src/primaite/notebooks/Privilege-Escalation-and Data-Loss-Example.ipynb @@ -201,7 +201,7 @@ "source": [ "caos_action = [\n", " \"network\", \"node\", \"some_tech_jnr_dev_pc\", \n", - " \"service\", \"Terminal\", \"ssh_to_remote\", \"admin\", \"admin\", str(some_tech_storage_srv.network_interface[1].ip_address)\n", + " \"service\", \"Terminal\", \"node_session_remote_login\", \"admin\", \"admin\", str(some_tech_storage_srv.network_interface[1].ip_address)\n", "]\n", "game.simulation.apply_request(caos_action)" ] @@ -259,7 +259,7 @@ "source": [ "caos_action = [\n", " \"network\", \"node\", \"some_tech_jnr_dev_pc\", \n", - " \"service\", \"Terminal\", \"ssh_to_remote\", \"admin\", \"admin\", str(some_tech_rt.network_interface[4].ip_address)\n", + " \"service\", \"Terminal\", \"node_session_remote_login\", \"admin\", \"admin\", str(some_tech_rt.network_interface[4].ip_address)\n", "]\n", "game.simulation.apply_request(caos_action)" ] @@ -396,7 +396,7 @@ "source": [ "caos_action = [\n", " \"network\", \"node\", \"some_tech_jnr_dev_pc\", \n", - " \"service\", \"Terminal\", \"ssh_to_remote\", \"admin\", \"admin\", str(some_tech_storage_srv.network_interface[1].ip_address)\n", + " \"service\", \"Terminal\", \"node_session_remote_login\", \"admin\", \"admin\", str(some_tech_storage_srv.network_interface[1].ip_address)\n", "]\n", "game.simulation.apply_request(caos_action)" ] diff --git a/src/primaite/simulator/system/services/terminal/terminal.py b/src/primaite/simulator/system/services/terminal/terminal.py index e26e77f6..1c249ebb 100644 --- a/src/primaite/simulator/system/services/terminal/terminal.py +++ b/src/primaite/simulator/system/services/terminal/terminal.py @@ -179,7 +179,7 @@ class Terminal(Service): return RequestResponse(status="failure", data={}) rm.add_request( - "ssh_to_remote", + "node_session_remote_login", request_type=RequestType(func=_remote_login), ) diff --git a/tests/assets/configs/action_penalty.yaml b/tests/assets/configs/action_penalty.yaml index 2828b5aa..9ab13036 100644 --- a/tests/assets/configs/action_penalty.yaml +++ b/tests/assets/configs/action_penalty.yaml @@ -386,52 +386,52 @@ agents: 52: # old action num: 28 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 0 53: # old action num: 29 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 1 54: # old action num: 30 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 2 55: # old action num: 31 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 3 56: # old action num: 32 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 4 57: # old action num: 33 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 5 58: # old action num: 34 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 6 59: # old action num: 35 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 7 60: # old action num: 36 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 8 61: # old action num: 37 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 9 62: # old action num: 38 action: "host_nic_disable" diff --git a/tests/assets/configs/basic_firewall.yaml b/tests/assets/configs/basic_firewall.yaml index 154956d3..3a62c75c 100644 --- a/tests/assets/configs/basic_firewall.yaml +++ b/tests/assets/configs/basic_firewall.yaml @@ -39,10 +39,6 @@ agents: node_name: client_2 application_name: WebBrowser - reward_function: - reward_components: - - type: DUMMY - agent_settings: action_probabilities: 0: 0.4 diff --git a/tests/assets/configs/basic_switched_network.yaml b/tests/assets/configs/basic_switched_network.yaml index e74a6a4e..10a92d7a 100644 --- a/tests/assets/configs/basic_switched_network.yaml +++ b/tests/assets/configs/basic_switched_network.yaml @@ -51,9 +51,6 @@ agents: 0: 0.4 1: 0.6 - - - - ref: defender team: BLUE type: ProxyAgent diff --git a/tests/assets/configs/data_manipulation.yaml b/tests/assets/configs/data_manipulation.yaml index 5e12f1c6..328fe413 100644 --- a/tests/assets/configs/data_manipulation.yaml +++ b/tests/assets/configs/data_manipulation.yaml @@ -477,52 +477,52 @@ agents: 52: # old action num: 28 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 0 53: # old action num: 29 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 1 54: # old action num: 30 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 2 55: # old action num: 31 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 3 56: # old action num: 32 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 4 57: # old action num: 33 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 5 58: # old action num: 34 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 6 59: # old action num: 35 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 7 60: # old action num: 36 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 8 61: # old action num: 37 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 9 62: # old action num: 38 action: "host_nic_disable" diff --git a/tests/assets/configs/dmz_network.yaml b/tests/assets/configs/dmz_network.yaml index 6cdae6a5..b4d018c8 100644 --- a/tests/assets/configs/dmz_network.yaml +++ b/tests/assets/configs/dmz_network.yaml @@ -62,7 +62,7 @@ agents: action: node_application_execute options: node_name: client_1 - application_id: WebBrowser + application_name: WebBrowser agent_settings: action_probabilities: diff --git a/tests/assets/configs/extended_config.yaml b/tests/assets/configs/extended_config.yaml index 97d9299a..0ec0c91f 100644 --- a/tests/assets/configs/extended_config.yaml +++ b/tests/assets/configs/extended_config.yaml @@ -529,82 +529,82 @@ agents: action: "host_nic_disable" options: node_name: domain_controller - nic_id: 0 + nic_num: 1 63: # old action num: 39 action: "host_nic_enable" options: node_name: domain_controller - nic_id: 0 + nic_num: 1 64: # old action num: 40 action: "host_nic_disable" options: node_name: web_server - nic_id: 0 + nic_num: 1 65: # old action num: 41 action: "host_nic_enable" options: node_name: web_server - nic_id: 0 + nic_num: 1 66: # old action num: 42 action: "host_nic_disable" options: node_name: database_server - nic_id: 0 + nic_num: 1 67: # old action num: 43 action: "host_nic_enable" options: node_name: database_server - nic_id: 0 + nic_num: 1 68: # old action num: 44 action: "host_nic_disable" options: node_name: backup_server - nic_id: 0 + nic_num: 1 69: # old action num: 45 action: "host_nic_enable" options: node_name: backup_server - nic_id: 0 + nic_num: 1 70: # old action num: 46 action: "host_nic_disable" options: node_name: security_suite - nic_id: 0 + nic_num: 1 71: # old action num: 47 action: "host_nic_enable" options: node_name: security_suite - nic_id: 0 + nic_num: 1 72: # old action num: 48 action: "host_nic_disable" options: node_name: security_suite - nic_id: 1 + nic_num: 2 73: # old action num: 49 action: "host_nic_enable" options: node_name: security_suite - nic_id: 1 + nic_num: 2 74: # old action num: 50 action: "host_nic_disable" options: node_name: client_1 - nic_id: 0 + nic_num: 1 75: # old action num: 51 action: "host_nic_enable" options: node_name: client_1 - nic_id: 0 + nic_num: 1 76: # old action num: 52 action: "host_nic_disable" options: node_name: client_2 - nic_id: 0 + nic_num: 1 77: # old action num: 53 action: "host_nic_enable" options: node_name: client_2 - nic_id: 0 + nic_num: 1 diff --git a/tests/assets/configs/firewall_actions_network.yaml b/tests/assets/configs/firewall_actions_network.yaml index 41b856fc..ff8e784d 100644 --- a/tests/assets/configs/firewall_actions_network.yaml +++ b/tests/assets/configs/firewall_actions_network.yaml @@ -109,12 +109,12 @@ agents: position: 1 permission: PERMIT src_ip: 192.168.0.10 - dst_ip: 0.0.0.0 + dst_ip: ALL src_port: 80 dst_port: HTTP protocol_name: TCP - src_wildcard: 0 - dst_wildcard: 0 + src_wildcard: NONE + dst_wildcard: NONE 2: action: firewall_acl_remove_rule options: @@ -135,8 +135,8 @@ agents: src_port: ARP dst_port: DNS protocol_name: ICMP - source_wildcard_id: 0 - dest_wildcard_id: 0 + src_wildcard: NONE + dst_wildcard: NONE 4: action: firewall_acl_remove_rule options: @@ -157,8 +157,8 @@ agents: src_port: HTTP dst_port: HTTP protocol_name: UDP - source_wildcard_id: 0 - dest_wildcard_id: 0 + src_wildcard: NONE + dst_wildcard: NONE 6: action: firewall_acl_remove_rule options: @@ -179,8 +179,8 @@ agents: src_port: HTTP dst_port: HTTP protocol_name: TCP - source_wildcard_id: 0 - dest_wildcard_id: 0 + src_wildcard: NONE + dst_wildcard: NONE 8: action: firewall_acl_remove_rule options: @@ -201,8 +201,8 @@ agents: src_port: POSTGRES_SERVER dst_port: POSTGRES_SERVER protocol_name: ICMP - source_wildcard_id: 0 - dest_wildcard_id: 0 + src_wildcard: NONE + dst_wildcard: NONE 10: action: firewall_acl_remove_rule options: @@ -223,8 +223,8 @@ agents: src_port: NONE dst_port: NONE protocol_name: none - source_wildcard_id: 0 - dest_wildcard_id: 0 + src_wildcard: NONE + dst_wildcard: NONE 12: action: firewall_acl_remove_rule options: @@ -237,17 +237,14 @@ agents: options: type: network_port_disable target_nodename: firewall - port_id: 3 + port_num: 3 14: action: network_port_enable options: type: network_port_enable target_nodename: firewall - port_id: 3 - agent_settings: - start_step: 5 - frequency: 4 - variance: 3 + port_num: 3 + diff --git a/tests/assets/configs/install_and_configure_apps.yaml b/tests/assets/configs/install_and_configure_apps.yaml index e4b8805e..ecc81668 100644 --- a/tests/assets/configs/install_and_configure_apps.yaml +++ b/tests/assets/configs/install_and_configure_apps.yaml @@ -43,34 +43,30 @@ agents: action: configure_database_client options: node_name: client_1 - config: - server_ip_address: 10.0.0.5 + server_ip_address: 10.0.0.5 5: action: configure_database_client options: node_name: client_1 - config: - server_password: correct_password + server_password: correct_password 6: action: c2_server_ransomware_configure options: node_name: client_2 - config: - server_ip_address: 10.0.0.5 - server_password: correct_password - payload: ENCRYPT + server_ip_address: 10.0.0.5 + server_password: correct_password + payload: ENCRYPT 7: action: configure_dos_bot options: node_name: client_3 - config: - target_ip_address: 10.0.0.5 - target_port: POSTGRES_SERVER - payload: DELETE - repeat: true - port_scan_p_of_success: 1.0 - dos_intensity: 1.0 - max_sessions: 1000 + target_ip_address: 10.0.0.5 + target_port: POSTGRES_SERVER + payload: DELETE + repeat: true + port_scan_p_of_success: 1.0 + dos_intensity: 1.0 + max_sessions: 1000 8: action: node_application_install options: diff --git a/tests/assets/configs/multi_agent_session.yaml b/tests/assets/configs/multi_agent_session.yaml index bc1f1b69..3b746273 100644 --- a/tests/assets/configs/multi_agent_session.yaml +++ b/tests/assets/configs/multi_agent_session.yaml @@ -469,52 +469,52 @@ agents: 52: # old action num: 28 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 0 53: # old action num: 29 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 1 54: # old action num: 30 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 2 55: # old action num: 31 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 3 56: # old action num: 32 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 4 57: # old action num: 33 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 5 58: # old action num: 34 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 6 59: # old action num: 35 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 7 60: # old action num: 36 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 8 61: # old action num: 37 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 9 62: # old action num: 38 action: "host_nic_disable" @@ -983,52 +983,52 @@ agents: 52: # old action num: 28 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 0 53: # old action num: 29 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 1 54: # old action num: 30 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 2 55: # old action num: 31 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 3 56: # old action num: 32 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 4 57: # old action num: 33 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 5 58: # old action num: 34 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 6 59: # old action num: 35 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 7 60: # old action num: 36 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 8 61: # old action num: 37 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 9 62: # old action num: 38 action: "host_nic_disable" diff --git a/tests/assets/configs/nmap_ping_scan_red_agent_config.yaml b/tests/assets/configs/nmap_ping_scan_red_agent_config.yaml index cd485ced..f6d549e8 100644 --- a/tests/assets/configs/nmap_ping_scan_red_agent_config.yaml +++ b/tests/assets/configs/nmap_ping_scan_red_agent_config.yaml @@ -28,7 +28,7 @@ agents: 0: action: node_nmap_ping_scan options: - node_name: client_1 + source_node: client_1 target_ip_address: 192.168.1.0/24 show: False diff --git a/tests/assets/configs/shared_rewards.yaml b/tests/assets/configs/shared_rewards.yaml index 96dada07..7ad5371d 100644 --- a/tests/assets/configs/shared_rewards.yaml +++ b/tests/assets/configs/shared_rewards.yaml @@ -519,82 +519,82 @@ agents: action: "host_nic_disable" options: node_name: domain_controller - nic_id: 0 + nic_num: 1 63: # old action num: 39 action: "host_nic_enable" options: node_name: domain_controller - nic_id: 0 + nic_num: 1 64: # old action num: 40 action: "host_nic_disable" options: node_name: web_server - nic_id: 0 + nic_num: 1 65: # old action num: 41 action: "host_nic_enable" options: node_name: web_server - nic_id: 0 + nic_num: 1 66: # old action num: 42 action: "host_nic_disable" options: node_name: database_server - nic_id: 0 + nic_num: 1 67: # old action num: 43 action: "host_nic_enable" options: node_name: database_server - nic_id: 0 + nic_num: 1 68: # old action num: 44 action: "host_nic_disable" options: node_name: backup_server - nic_id: 0 + nic_num: 1 69: # old action num: 45 action: "host_nic_enable" options: node_name: backup_server - nic_id: 0 + nic_num: 1 70: # old action num: 46 action: "host_nic_disable" options: node_name: security_suite - nic_id: 0 + nic_num: 1 71: # old action num: 47 action: "host_nic_enable" options: node_name: security_suite - nic_id: 0 + nic_num: 1 72: # old action num: 48 action: "host_nic_disable" options: node_name: security_suite - nic_id: 1 + nic_num: 2 73: # old action num: 49 action: "host_nic_enable" options: node_name: security_suite - nic_id: 1 + nic_num: 2 74: # old action num: 50 action: "host_nic_disable" options: node_name: client_1 - nic_id: 0 + nic_num: 1 75: # old action num: 51 action: "host_nic_enable" options: node_name: client_1 - nic_id: 0 + nic_num: 1 76: # old action num: 52 action: "host_nic_disable" options: node_name: client_2 - nic_id: 0 + nic_num: 1 77: # old action num: 53 action: "host_nic_enable" options: node_name: client_2 - nic_id: 0 + nic_num: 1 reward_function: reward_components: diff --git a/tests/assets/configs/test_application_install.yaml b/tests/assets/configs/test_application_install.yaml index 55c4afd3..cafcc72b 100644 --- a/tests/assets/configs/test_application_install.yaml +++ b/tests/assets/configs/test_application_install.yaml @@ -472,52 +472,52 @@ agents: 52: # old action num: 28 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 0 53: # old action num: 29 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 1 54: # old action num: 30 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 2 55: # old action num: 31 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 3 56: # old action num: 32 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 4 57: # old action num: 33 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 5 58: # old action num: 34 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 6 59: # old action num: 35 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 7 60: # old action num: 36 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 8 61: # old action num: 37 action: "router_acl_remove_rule" options: - target_router_hostname: router_1 + target_router: router_1 position: 9 62: # old action num: 38 action: "host_nic_disable" @@ -618,14 +618,13 @@ agents: action: node_application_execute options: node_name: domain_controller - application_id: 0 + application_name: DoSBot 82: action: configure_dos_bot options: node_name: domain_controller - config: - target_ip_address: 192.168.1.14 - target_port: POSTGRES_SERVER + target_ip_address: 192.168.1.14 + target_port: POSTGRES_SERVER reward_function: reward_components: diff --git a/tests/integration_tests/game_layer/actions/test_application_request_permission.py b/tests/integration_tests/game_layer/actions/test_application_request_permission.py index c0c039f6..c47b617b 100644 --- a/tests/integration_tests/game_layer/actions/test_application_request_permission.py +++ b/tests/integration_tests/game_layer/actions/test_application_request_permission.py @@ -33,22 +33,22 @@ def test_application_cannot_perform_actions_unless_running(game_and_agent_fixtur browser.close() assert browser.operating_state == ApplicationOperatingState.CLOSED - action = ("node_application_scan", {"node_id": 0, "application_id": 0}) + action = ("node_application_scan", {"node_name": "client_1", "application_name": "WebBrowser"}) agent.store_action(action) game.step() assert browser.operating_state == ApplicationOperatingState.CLOSED - action = ("node_application_close", {"node_id": 0, "application_id": 0}) + action = ("node_application_close", {"node_name": "client_1", "application_name": "WebBrowser"}) agent.store_action(action) game.step() assert browser.operating_state == ApplicationOperatingState.CLOSED - action = ("node_application_fix", {"node_id": 0, "application_id": 0}) + action = ("node_application_fix", {"node_name": "client_1", "application_name": "WebBrowser"}) agent.store_action(action) game.step() assert browser.operating_state == ApplicationOperatingState.CLOSED - action = ("node_application_execute", {"node_id": 0, "application_id": 0}) + action = ("node_application_execute", {"node_name": "client_1", "application_name": "WebBrowser"}) agent.store_action(action) game.step() assert browser.operating_state == ApplicationOperatingState.CLOSED diff --git a/tests/integration_tests/game_layer/actions/test_c2_suite_actions.py b/tests/integration_tests/game_layer/actions/test_c2_suite_actions.py index 2984429a..c52c5761 100644 --- a/tests/integration_tests/game_layer/actions/test_c2_suite_actions.py +++ b/tests/integration_tests/game_layer/actions/test_c2_suite_actions.py @@ -47,7 +47,7 @@ def test_c2_beacon_default(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAgen action = ( "node_application_install", - {"node_id": 1, "application_name": "C2Beacon"}, + {"node_name": "server_1", "application_name": "C2Beacon"}, ) agent.store_action(action) game.step() @@ -56,13 +56,11 @@ def test_c2_beacon_default(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAgen action = ( "configure_c2_beacon", { - "node_id": 1, - "config": { - "c2_server_ip_address": "10.0.1.2", - "keep_alive_frequency": 5, - "masquerade_protocol": "TCP", - "masquerade_port": "HTTP", - }, + "node_name": "server_1", + "c2_server_ip_address": "10.0.1.2", + "keep_alive_frequency": 5, + "masquerade_protocol": "TCP", + "masquerade_port": "HTTP", }, ) agent.store_action(action) @@ -71,7 +69,7 @@ def test_c2_beacon_default(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAgen action = ( "node_application_execute", - {"node_id": 1, "application_id": 0}, + {"node_name": "server_1", "application_name": "C2Beacon"}, ) agent.store_action(action) game.step() @@ -103,14 +101,12 @@ def test_c2_server_ransomware(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyA # C2 Action 1: Installing the RansomwareScript & Database client via Terminal action = ( - "C2_SERVER_TERMINAL_COMMAND", + "c2_server_terminal_command", { - "node_id": 0, + "node_name": "client_1", "ip_address": None, - "account": { - "username": "admin", - "password": "admin", - }, + "username": "admin", + "password": "admin", "commands": [ ["software_manager", "application", "install", "RansomwareScript"], ["software_manager", "application", "install", "DatabaseClient"], @@ -124,7 +120,7 @@ def test_c2_server_ransomware(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyA action = ( "c2_server_ransomware_configure", { - "node_id": 0, + "node_name": "client_1", "config": {"server_ip_address": "10.0.2.3", "payload": "ENCRYPT"}, }, ) @@ -143,7 +139,7 @@ def test_c2_server_ransomware(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyA action = ( "c2_server_ransomware_launch", { - "node_id": 0, + "node_name": "client_1", }, ) agent.store_action(action) @@ -183,15 +179,13 @@ def test_c2_server_data_exfiltration(game_and_agent_fixture: Tuple[PrimaiteGame, action = ( "c2_server_data_exfiltrate", { - "node_id": 0, + "node_name": "client_1", "target_file_name": "database.db", "target_folder_name": "database", "exfiltration_folder_name": "spoils", "target_ip_address": "10.0.2.3", - "account": { - "username": "admin", - "password": "admin", - }, + "username": "admin", + "password": "admin", }, ) agent.store_action(action) diff --git a/tests/integration_tests/game_layer/actions/test_configure_actions.py b/tests/integration_tests/game_layer/actions/test_configure_actions.py index 0e1a4873..5c9f09e4 100644 --- a/tests/integration_tests/game_layer/actions/test_configure_actions.py +++ b/tests/integration_tests/game_layer/actions/test_configure_actions.py @@ -27,7 +27,6 @@ class TestConfigureDatabaseAction: def test_configure_ip_password(self, game_and_agent): game, agent = game_and_agent agent: ControlledAgent - agent.action_manager.actions["configure_database_client"] = ConfigureDatabaseClientAction(agent.action_manager) # make sure there is a database client on this node client_1 = game.simulation.network.get_node_by_hostname("client_1") @@ -38,10 +37,8 @@ class TestConfigureDatabaseAction: "configure_database_client", { "node_name": "client_1", - "model_config": { - "server_ip_address": "192.168.1.99", - "server_password": "admin123", - }, + "server_ip_address": "192.168.1.99", + "server_password": "admin123", }, ) agent.store_action(action) @@ -53,7 +50,6 @@ class TestConfigureDatabaseAction: def test_configure_ip(self, game_and_agent): game, agent = game_and_agent agent: ControlledAgent - agent.action_manager.actions["configure_database_client"] = ConfigureDatabaseClientAction(agent.action_manager) # make sure there is a database client on this node client_1 = game.simulation.network.get_node_by_hostname("client_1") @@ -63,10 +59,8 @@ class TestConfigureDatabaseAction: action = ( "configure_database_client", { - "node_id": 0, - "config": { - "server_ip_address": "192.168.1.99", - }, + "node_name": "client_1", + "server_ip_address": "192.168.1.99", }, ) agent.store_action(action) @@ -78,7 +72,6 @@ class TestConfigureDatabaseAction: def test_configure_password(self, game_and_agent): game, agent = game_and_agent agent: ControlledAgent - agent.action_manager.actions["configure_database_client"] = ConfigureDatabaseClientAction(agent.action_manager) # make sure there is a database client on this node client_1 = game.simulation.network.get_node_by_hostname("client_1") @@ -89,10 +82,8 @@ class TestConfigureDatabaseAction: action = ( "configure_database_client", { - "node_id": 0, - "config": { - "server_password": "admin123", - }, + "node_name": "client_1", + "server_password": "admin123", }, ) agent.store_action(action) @@ -120,9 +111,6 @@ class TestConfigureRansomwareScriptAction: def test_configure_ip_password(self, game_and_agent, config): game, agent = game_and_agent agent: ControlledAgent - agent.action_manager.actions["c2_server_ransomware_configure"] = ConfigureRansomwareScriptAction( - agent.action_manager - ) # make sure there is a database client on this node client_1 = game.simulation.network.get_node_by_hostname("client_1") @@ -135,7 +123,7 @@ class TestConfigureRansomwareScriptAction: action = ( "c2_server_ransomware_configure", - {"node_id": 0, "config": config}, + {"node_name": "client_1", **config}, ) agent.store_action(action) game.step() @@ -151,9 +139,6 @@ class TestConfigureRansomwareScriptAction: def test_invalid_config(self, game_and_agent): game, agent = game_and_agent agent: ControlledAgent - agent.action_manager.actions["c2_server_ransomware_configure"] = ConfigureRansomwareScriptAction( - agent.action_manager - ) # make sure there is a database client on this node client_1 = game.simulation.network.get_node_by_hostname("client_1") @@ -162,7 +147,7 @@ class TestConfigureRansomwareScriptAction: action = ( "c2_server_ransomware_configure", { - "node_id": 0, + "node_name": "client_1", "config": {"server_password": "admin123", "bad_option": 70}, }, ) @@ -175,7 +160,6 @@ class TestConfigureDoSBot: def test_configure_dos_bot(self, game_and_agent): game, agent = game_and_agent agent: ControlledAgent - agent.action_manager.actions["configure_dos_bot"] = ConfigureDoSBotAction(agent.action_manager) client_1 = game.simulation.network.get_node_by_hostname("client_1") client_1.software_manager.install(DoSBot) @@ -184,16 +168,14 @@ class TestConfigureDoSBot: action = ( "configure_dos_bot", { - "node_id": 0, - "config": { - "target_ip_address": "192.168.1.99", - "target_port": "POSTGRES_SERVER", - "payload": "HACC", - "repeat": False, - "port_scan_p_of_success": 0.875, - "dos_intensity": 0.75, - "max_sessions": 50, - }, + "node_name": "client_1", + "target_ip_address": "192.168.1.99", + "target_port": "POSTGRES_SERVER", + "payload": "HACC", + "repeat": False, + "port_scan_p_of_success": 0.875, + "dos_intensity": 0.75, + "max_sessions": 50, }, ) agent.store_action(action) diff --git a/tests/integration_tests/game_layer/actions/test_file_request_permission.py b/tests/integration_tests/game_layer/actions/test_file_request_permission.py index 2ed76063..0976abdc 100644 --- a/tests/integration_tests/game_layer/actions/test_file_request_permission.py +++ b/tests/integration_tests/game_layer/actions/test_file_request_permission.py @@ -34,7 +34,7 @@ def test_create_file(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAgent]): action = ( "node_file_create", - {"node_id": 0, "folder_name": random_folder, "file_name": random_file}, + {"node_name": "client_1", "folder_name": random_folder, "file_name": random_file}, ) agent.store_action(action) game.step() @@ -52,7 +52,7 @@ def test_file_delete_action(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAge action = ( "node_file_delete", - {"node_id": 0, "folder_id": 0, "file_id": 0}, + {"node_name": "client_1", "folder_name": "downloads", "file_name": "cat.png"}, ) agent.store_action(action) game.step() @@ -73,7 +73,7 @@ def test_file_scan_action(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAgent action = ( "node_file_scan", - {"node_id": 0, "folder_id": 0, "file_id": 0}, + {"node_name": "client_1", "folder_name": "downloads", "file_name": "cat.png"}, ) agent.store_action(action) game.step() @@ -94,7 +94,7 @@ def test_file_repair_action(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAge action = ( "node_file_repair", - {"node_id": 0, "folder_id": 0, "file_id": 0}, + {"node_name": "client_1", "folder_name": "downloads", "file_name": "cat.png"}, ) agent.store_action(action) game.step() @@ -114,7 +114,7 @@ def test_file_restore_action(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAg action = ( "node_file_restore", - {"node_id": 0, "folder_id": 0, "file_id": 0}, + {"node_name": "client_1", "folder_name": "downloads", "file_name": "cat.png"}, ) agent.store_action(action) game.step() @@ -133,7 +133,7 @@ def test_file_corrupt_action(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAg action = ( "node_file_corrupt", - {"node_id": 0, "folder_id": 0, "file_id": 0}, + {"node_name": "client_1", "folder_name": "downloads", "file_name": "cat.png"}, ) agent.store_action(action) game.step() @@ -151,7 +151,7 @@ def test_file_access_action(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAge action = ( "node_file_access", - {"node_id": 0, "folder_name": file.folder_name, "file_name": file.name}, + {"node_name": "client_1", "folder_name": file.folder_name, "file_name": file.name}, ) agent.store_action(action) game.step() diff --git a/tests/integration_tests/game_layer/actions/test_folder_request_permission.py b/tests/integration_tests/game_layer/actions/test_folder_request_permission.py index 1c3cca7b..9cd4bfcf 100644 --- a/tests/integration_tests/game_layer/actions/test_folder_request_permission.py +++ b/tests/integration_tests/game_layer/actions/test_folder_request_permission.py @@ -34,7 +34,7 @@ def test_create_folder(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAgent]): action = ( "node_folder_create", { - "node_id": 0, + "node_name": "client_1", "folder_name": random_folder, }, ) @@ -62,8 +62,8 @@ def test_folder_scan_action(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAge action = ( "node_folder_scan", { - "node_id": 0, # client_1, - "folder_id": 0, # downloads + "node_name": "client_1", # client_1, + "folder_name": "downloads", # downloads }, ) agent.store_action(action) @@ -89,8 +89,8 @@ def test_folder_repair_action(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyA action = ( "node_folder_repair", { - "node_id": 0, # client_1, - "folder_id": 0, # downloads + "node_name": "client_1", # client_1, + "folder_name": "downloads", # downloads }, ) agent.store_action(action) @@ -113,8 +113,8 @@ def test_folder_restore_action(game_and_agent_fixture: Tuple[PrimaiteGame, Proxy action = ( "node_folder_restore", { - "node_id": 0, # client_1, - "folder_id": 0, # downloads + "node_name": "client_1", # client_1, + "folder_name": "downloads", # downloads }, ) agent.store_action(action) diff --git a/tests/integration_tests/game_layer/actions/test_nic_request_permission.py b/tests/integration_tests/game_layer/actions/test_nic_request_permission.py index ac92205b..11e39c7e 100644 --- a/tests/integration_tests/game_layer/actions/test_nic_request_permission.py +++ b/tests/integration_tests/game_layer/actions/test_nic_request_permission.py @@ -31,8 +31,8 @@ def test_nic_cannot_be_turned_off_if_not_on(game_and_agent_fixture: Tuple[Primai action = ( "host_nic_disable", { - "node_id": 0, # client_1 - "nic_id": 0, # the only nic (eth-1) + "node_name": "client_1", # client_1 + "nic_num": 1, # the only nic (eth-1) }, ) agent.store_action(action) @@ -52,8 +52,8 @@ def test_nic_cannot_be_turned_on_if_already_on(game_and_agent_fixture: Tuple[Pri action = ( "host_nic_enable", { - "node_id": 0, # client_1 - "nic_id": 0, # the only nic (eth-1) + "node_name": "client_1", # client_1 + "nic_num": 1, # the only nic (eth-1) }, ) agent.store_action(action) @@ -73,8 +73,8 @@ def test_that_a_nic_can_be_enabled_and_disabled(game_and_agent_fixture: Tuple[Pr action = ( "host_nic_disable", { - "node_id": 0, # client_1 - "nic_id": 0, # the only nic (eth-1) + "node_name": "client_1", # client_1 + "nic_num": 1, # the only nic (eth-1) }, ) agent.store_action(action) @@ -85,8 +85,8 @@ def test_that_a_nic_can_be_enabled_and_disabled(game_and_agent_fixture: Tuple[Pr action = ( "host_nic_enable", { - "node_id": 0, # client_1 - "nic_id": 0, # the only nic (eth-1) + "node_name": "client_1", # client_1 + "nic_num": 1, # the only nic (eth-1) }, ) agent.store_action(action) diff --git a/tests/integration_tests/game_layer/actions/test_node_request_permission.py b/tests/integration_tests/game_layer/actions/test_node_request_permission.py index 997a9282..8a438673 100644 --- a/tests/integration_tests/game_layer/actions/test_node_request_permission.py +++ b/tests/integration_tests/game_layer/actions/test_node_request_permission.py @@ -29,28 +29,28 @@ def test_node_startup_shutdown(game_and_agent_fixture: Tuple[PrimaiteGame, Proxy assert client_1.operating_state == NodeOperatingState.ON # turn it off - action = ("node_shutdown", {"node_id": 0}) + action = ("node_shutdown", {"node_name": "client_1"}) agent.store_action(action) game.step() assert client_1.operating_state == NodeOperatingState.SHUTTING_DOWN for i in range(client_1.shut_down_duration + 1): - action = ("do_nothing", {"node_id": 0}) + action = ("do_nothing", {}) agent.store_action(action) game.step() assert client_1.operating_state == NodeOperatingState.OFF # turn it on - action = ("node_startup", {"node_id": 0}) + action = ("node_startup", {"node_name": "client_1"}) agent.store_action(action) game.step() assert client_1.operating_state == NodeOperatingState.BOOTING for i in range(client_1.start_up_duration + 1): - action = ("do_nothing", {"node_id": 0}) + action = ("do_nothing", {}) agent.store_action(action) game.step() @@ -65,7 +65,7 @@ def test_node_cannot_be_started_up_if_node_is_already_on(game_and_agent_fixture: assert client_1.operating_state == NodeOperatingState.ON # turn it on - action = ("node_startup", {"node_id": 0}) + action = ("node_startup", {"node_name": "client_1"}) agent.store_action(action) game.step() @@ -80,14 +80,14 @@ def test_node_cannot_be_shut_down_if_node_is_already_off(game_and_agent_fixture: client_1.power_off() for i in range(client_1.shut_down_duration + 1): - action = ("do_nothing", {"node_id": 0}) + action = ("do_nothing", {}) agent.store_action(action) game.step() assert client_1.operating_state == NodeOperatingState.OFF # turn it ff - action = ("node_shutdown", {"node_id": 0}) + action = ("node_shutdown", {"node_name": "client_1"}) agent.store_action(action) game.step() diff --git a/tests/integration_tests/game_layer/actions/test_service_request_permission.py b/tests/integration_tests/game_layer/actions/test_service_request_permission.py index dad67d10..80e68131 100644 --- a/tests/integration_tests/game_layer/actions/test_service_request_permission.py +++ b/tests/integration_tests/game_layer/actions/test_service_request_permission.py @@ -31,7 +31,7 @@ def test_service_start(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAgent]): dns_server.pause() assert dns_server.operating_state == ServiceOperatingState.PAUSED - action = ("node_service_start", {"node_id": 1, "service_id": 0}) + action = ("node_service_start", {"node_name": "server_1", "service_name": "DNSServer"}) agent.store_action(action) game.step() assert dns_server.operating_state == ServiceOperatingState.PAUSED @@ -40,7 +40,7 @@ def test_service_start(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAgent]): assert dns_server.operating_state == ServiceOperatingState.STOPPED - action = ("node_service_start", {"node_id": 1, "service_id": 0}) + action = ("node_service_start", {"node_name": "server_1", "service_name": "DNSServer"}) agent.store_action(action) game.step() @@ -54,7 +54,7 @@ def test_service_resume(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAgent]) server_1: Server = game.simulation.network.get_node_by_hostname("server_1") dns_server = server_1.software_manager.software.get("DNSServer") - action = ("node_service_resume", {"node_id": 1, "service_id": 0}) + action = ("node_service_resume", {"node_name": "server_1", "service_name": "DNSServer"}) agent.store_action(action) game.step() assert dns_server.operating_state == ServiceOperatingState.RUNNING @@ -63,7 +63,7 @@ def test_service_resume(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAgent]) assert dns_server.operating_state == ServiceOperatingState.PAUSED - action = ("node_service_resume", {"node_id": 1, "service_id": 0}) + action = ("node_service_resume", {"node_name": "server_1", "service_name": "DNSServer"}) agent.store_action(action) game.step() @@ -80,27 +80,27 @@ def test_service_cannot_perform_actions_unless_running(game_and_agent_fixture: T dns_server.stop() assert dns_server.operating_state == ServiceOperatingState.STOPPED - action = ("node_service_scan", {"node_id": 1, "service_id": 0}) + action = ("node_service_scan", {"node_name": "server_1", "service_name": "DNSServer"}) agent.store_action(action) game.step() assert dns_server.operating_state == ServiceOperatingState.STOPPED - action = ("node_service_pause", {"node_id": 1, "service_id": 0}) + action = ("node_service_pause", {"node_name": "server_1", "service_name": "DNSServer"}) agent.store_action(action) game.step() assert dns_server.operating_state == ServiceOperatingState.STOPPED - action = ("node_service_resume", {"node_id": 1, "service_id": 0}) + action = ("node_service_resume", {"node_name": "server_1", "service_name": "DNSServer"}) agent.store_action(action) game.step() assert dns_server.operating_state == ServiceOperatingState.STOPPED - action = ("node_service_restart", {"node_id": 1, "service_id": 0}) + action = ("node_service_restart", {"node_name": "server_1", "service_name": "DNSServer"}) agent.store_action(action) game.step() assert dns_server.operating_state == ServiceOperatingState.STOPPED - action = ("node_service_fix", {"node_id": 1, "service_id": 0}) + action = ("node_service_fix", {"node_name": "server_1", "service_name": "DNSServer"}) agent.store_action(action) game.step() assert dns_server.operating_state == ServiceOperatingState.STOPPED diff --git a/tests/integration_tests/game_layer/actions/test_terminal_actions.py b/tests/integration_tests/game_layer/actions/test_terminal_actions.py index beaec5da..f15f7156 100644 --- a/tests/integration_tests/game_layer/actions/test_terminal_actions.py +++ b/tests/integration_tests/game_layer/actions/test_terminal_actions.py @@ -36,9 +36,9 @@ def test_remote_login(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAgent]): server_1_usm.add_user("user123", "password", is_admin=True) action = ( - "SSH_TO_REMOTE", + "node_session_remote_login", { - "node_id": 0, + "node_name": "client_1", "username": "user123", "password": "password", "remote_ip": str(server_1.network_interface[1].ip_address), @@ -68,9 +68,9 @@ def test_remote_login_wrong_password(game_and_agent_fixture: Tuple[PrimaiteGame, server_1_usm.add_user("user123", "password", is_admin=True) action = ( - "SSH_TO_REMOTE", + "node_session_remote_login", { - "node_id": 0, + "node_name": "client_1", "username": "user123", "password": "wrong_password", "remote_ip": str(server_1.network_interface[1].ip_address), @@ -100,12 +100,13 @@ def test_remote_login_change_password(game_and_agent_fixture: Tuple[PrimaiteGame server_1_um.add_user("user123", "password", is_admin=True) action = ( - "node_accounts_change_password", + "node_account_change_password", { - "node_id": 1, # server_1 + "node_name": "server_1", # server_1 "username": "user123", "current_password": "password", "new_password": "different_password", + "remote_ip": str(server_1.network_interface[1].ip_address), }, ) agent.store_action(action) @@ -126,9 +127,9 @@ def test_change_password_logs_out_user(game_and_agent_fixture: Tuple[PrimaiteGam # Log in remotely action = ( - "SSH_TO_REMOTE", + "node_session_remote_login", { - "node_id": 0, + "node_name": "client_1", "username": "user123", "password": "password", "remote_ip": str(server_1.network_interface[1].ip_address), @@ -139,12 +140,13 @@ def test_change_password_logs_out_user(game_and_agent_fixture: Tuple[PrimaiteGam # Change password action = ( - "node_accounts_change_password", + "node_account_change_password", { - "node_id": 1, # server_1 + "node_name": "server_1", # server_1 "username": "user123", "current_password": "password", "new_password": "different_password", + "remote_ip": str(server_1.network_interface[1].ip_address), }, ) agent.store_action(action) @@ -154,7 +156,7 @@ def test_change_password_logs_out_user(game_and_agent_fixture: Tuple[PrimaiteGam action = ( "node_send_remote_command", { - "node_id": 0, + "node_name": "client_1", "remote_ip": str(server_1.network_interface[1].ip_address), "command": ["file_system", "create", "file", "folder123", "doggo.pdf", False], }, diff --git a/tests/integration_tests/game_layer/test_action_mask.py b/tests/integration_tests/game_layer/test_action_mask.py index 485ad138..75965f16 100644 --- a/tests/integration_tests/game_layer/test_action_mask.py +++ b/tests/integration_tests/game_layer/test_action_mask.py @@ -25,7 +25,7 @@ def test_mask_contents_correct(): if act_type == "node_nic_enable": node_name = act_params["node_name"] node_obj = net.get_node_by_hostname(node_name) - nic_obj = node_obj.network_interface[act_params["nic_id"] + 1] + nic_obj = node_obj.network_interface[act_params["nic_num"]] assert nic_obj.enabled assert not mask[action_num] nic_obj.disable() @@ -36,7 +36,7 @@ def test_mask_contents_correct(): if act_type == "node_nic_disable": node_name = act_params["node_name"] node_obj = net.get_node_by_hostname(node_name) - nic_obj = node_obj.network_interface[act_params["nic_id"] + 1] + nic_obj = node_obj.network_interface[act_params["nic_num"]] assert nic_obj.enabled assert mask[action_num] nic_obj.disable() diff --git a/tests/integration_tests/game_layer/test_actions.py b/tests/integration_tests/game_layer/test_actions.py index ff86dbf0..cf8a33ce 100644 --- a/tests/integration_tests/game_layer/test_actions.py +++ b/tests/integration_tests/game_layer/test_actions.py @@ -56,7 +56,7 @@ def test_node_service_scan_integration(game_and_agent: Tuple[PrimaiteGame, Proxy assert svc.health_state_visible == SoftwareHealthState.UNUSED # 2: Scan and check that the visible state is now correct - action = ("node_service_scan", {"type": "node_service_scan", "node_name": "server_1", "service_name": "DNSServer"}) + action = ("node_service_scan", {"node_name": "server_1", "service_name": "DNSServer"}) agent.store_action(action) game.step() assert svc.health_state_actual == SoftwareHealthState.GOOD @@ -67,7 +67,7 @@ def test_node_service_scan_integration(game_and_agent: Tuple[PrimaiteGame, Proxy assert svc.health_state_visible == SoftwareHealthState.GOOD # 4: Scan and check that the visible state is now correct - action = ("node_service_scan", {"type": "node_service_scan", "node_name": "server_1", "service_name": "DNSServer"}) + action = ("node_service_scan", {"node_name": "server_1", "service_name": "DNSServer"}) agent.store_action(action) game.step() assert svc.health_state_actual == SoftwareHealthState.COMPROMISED @@ -88,7 +88,7 @@ def test_node_service_fix_integration(game_and_agent: Tuple[PrimaiteGame, ProxyA svc.health_state_actual = SoftwareHealthState.COMPROMISED # 2: Apply a patch action - action = ("node_service_fix", {"type": "node_service_fix", "node_name": "server_1", "service_name": "DNSServer"}) + action = ("node_service_fix", {"node_name": "server_1", "service_name": "DNSServer"}) agent.store_action(action) game.step() @@ -123,7 +123,6 @@ def test_router_acl_addrule_integration(game_and_agent: Tuple[PrimaiteGame, Prox action = ( "router_acl_add_rule", { - "type": "router_acl_add_rule", "target_router": "router", "position": 4, "permission": "DENY", @@ -151,7 +150,6 @@ def test_router_acl_addrule_integration(game_and_agent: Tuple[PrimaiteGame, Prox action = ( "router_acl_add_rule", { - "type": "router_acl_add_rule", "target_router": "router", "position": 5, # 5th rule "permission": "DENY", # DENY @@ -192,7 +190,6 @@ def test_router_acl_removerule_integration(game_and_agent: Tuple[PrimaiteGame, P action = ( "router_acl_remove_rule", { - "type": "router_acl_remove_rule", "target_router": "router", "position": 3, # 4th rule }, @@ -226,7 +223,6 @@ def test_host_nic_disable_integration(game_and_agent: Tuple[PrimaiteGame, ProxyA action = ( "host_nic_disable", { - "type": "host_nic_disable", "node_name": "client_1", # client_1 "nic_num": 1, # the only nic (eth-1) }, @@ -258,7 +254,6 @@ def test_host_nic_enable_integration(game_and_agent: Tuple[PrimaiteGame, ProxyAg action = ( "host_nic_enable", { - "type": "host_nic_enable", "node_name": "client_1", # client_1 "nic_num": 1, # the only nic (eth-1) }, @@ -286,7 +281,6 @@ def test_node_file_scan_integration(game_and_agent: Tuple[PrimaiteGame, ProxyAge action = ( "node_file_scan", { - "type": "node_file_scan", "node_name": "client_1", # client_1, "folder_name": "downloads", # downloads, "file_name": "cat.png", # cat.png @@ -324,7 +318,6 @@ def test_node_file_delete_integration(game_and_agent: Tuple[PrimaiteGame, ProxyA action = ( "node_file_delete", { - "type": "node_file_delete", "node_name": "client_1", # client_1 "folder_name": "downloads", # downloads "file_name": "cat.png", # cat.png @@ -348,7 +341,6 @@ def test_node_file_create(game_and_agent: Tuple[PrimaiteGame, ProxyAgent]): action = ( "node_file_create", { - "type": "node_file_create", "node_name": "client_1", "folder_name": "test", "file_name": "file.txt", @@ -370,7 +362,6 @@ def test_node_file_access(game_and_agent: Tuple[PrimaiteGame, ProxyAgent]): action = ( "node_file_create", { - "type": "node_file_create", "node_name": "client_1", "folder_name": "test", "file_name": "file.txt", @@ -384,7 +375,6 @@ def test_node_file_access(game_and_agent: Tuple[PrimaiteGame, ProxyAgent]): action = ( "node_file_access", { - "type": "node_file_access", "node_name": "client_1", "folder_name": "test", "file_name": "file.txt", @@ -405,7 +395,6 @@ def test_node_folder_create(game_and_agent: Tuple[PrimaiteGame, ProxyAgent]): action = ( "node_folder_create", { - "type": "node_folder_create", "node_name": "client_1", "folder_name": "test", }, @@ -434,7 +423,6 @@ def test_network_router_port_disable_integration(game_and_agent: Tuple[PrimaiteG action = ( "network_port_disable", { - "type": "network_port_disable", "target_nodename": "router", # router "port_id": 1, # port 1 }, @@ -467,7 +455,6 @@ def test_network_router_port_enable_integration(game_and_agent: Tuple[PrimaiteGa action = ( "network_port_enable", { - "type": "network_port_enable", "target_nodename": "router", # router "port_id": 1, # port 1 }, @@ -498,7 +485,7 @@ def test_node_application_scan_integration(game_and_agent: Tuple[PrimaiteGame, P # 2: Scan and check that the visible state is now correct action = ( "node_application_scan", - {"type": "node_application_scan", "node_name": "client_1", "application_name": "WebBrowser"}, + {"node_name": "client_1", "application_name": "WebBrowser"}, ) agent.store_action(action) game.step() @@ -512,7 +499,7 @@ def test_node_application_scan_integration(game_and_agent: Tuple[PrimaiteGame, P # 4: Scan and check that the visible state is now correct action = ( "node_application_scan", - {"type": "node_application_scan", "node_name": "client_1", "application_name": "WebBrowser"}, + {"node_name": "client_1", "application_name": "WebBrowser"}, ) agent.store_action(action) game.step() @@ -536,7 +523,7 @@ def test_node_application_fix_integration(game_and_agent: Tuple[PrimaiteGame, Pr # 2: Apply a fix action action = ( "node_application_fix", - {"type": "node_application_fix", "node_name": "client_1", "application_name": "WebBrowser"}, + {"node_name": "client_1", "application_name": "WebBrowser"}, ) agent.store_action(action) game.step() @@ -565,7 +552,7 @@ def test_node_application_close_integration(game_and_agent: Tuple[PrimaiteGame, # 2: Apply a close action action = ( "node_application_close", - {"type": "node_application_close", "node_name": "client_1", "application_name": "WebBrowser"}, + {"node_name": "client_1", "application_name": "WebBrowser"}, ) agent.store_action(action) game.step() @@ -587,7 +574,7 @@ def test_node_application_install_and_uninstall_integration(game_and_agent: Tupl action = ( "node_application_install", - {"type": "node_application_install", "node_name": "client_1", "application_name": "DoSBot"}, + {"node_name": "client_1", "application_name": "DoSBot"}, ) agent.store_action(action) game.step() @@ -596,7 +583,7 @@ def test_node_application_install_and_uninstall_integration(game_and_agent: Tupl action = ( "node_application_remove", - {"type": "node_application_remove", "node_name": "client_1", "application_name": "DoSBot"}, + {"node_name": "client_1", "application_name": "DoSBot"}, ) agent.store_action(action) game.step() diff --git a/tests/integration_tests/game_layer/test_rewards.py b/tests/integration_tests/game_layer/test_rewards.py index 3d360313..91a022d5 100644 --- a/tests/integration_tests/game_layer/test_rewards.py +++ b/tests/integration_tests/game_layer/test_rewards.py @@ -38,7 +38,7 @@ def test_WebpageUnavailablePenalty(game_and_agent: tuple[PrimaiteGame, Controlle assert agent.reward_function.current_reward == 0.0 # Check that successfully fetching the webpage yields a reward of 0.7 - agent.store_action(("node_application_execute", {"node_id": 0, "application_id": 0})) + agent.store_action(("node_application_execute", {"node_name": "client_1", "application_name": "WebBrowser"})) game.step() assert agent.reward_function.current_reward == 0.7 @@ -50,7 +50,7 @@ def test_WebpageUnavailablePenalty(game_and_agent: tuple[PrimaiteGame, Controlle src_port=PORT_LOOKUP["HTTP"], dst_port=PORT_LOOKUP["HTTP"], ) - agent.store_action(("node_application_execute", {"node_id": 0, "application_id": 0})) + agent.store_action(("node_application_execute", {"node_name": "client_1", "application_name": "WebBrowser"})) game.step() assert agent.reward_function.current_reward == -0.7 @@ -160,7 +160,7 @@ def test_action_penalty(): last_action_response=AgentHistoryItem( timestep=0, action="node_application_execute", - parameters={"node_id": 0, "application_id": 1}, + parameters={"node_name": "client", "application_name": "WebBrowser"}, request=["execute"], response=RequestResponse.from_bool(True), ), @@ -197,7 +197,7 @@ def test_action_penalty_e2e(game_and_agent: tuple[PrimaiteGame, ControlledAgent] game.step() assert agent.reward_function.current_reward == 0.125 - action = ("node_file_scan", {"node_id": 0, "folder_id": 0, "file_id": 0}) + action = ("node_file_scan", {"node_name": "client", "folder_name": "downloads", "file_name": "document.pdf"}) agent.store_action(action) game.step() assert agent.reward_function.current_reward == -0.75 diff --git a/tests/unit_tests/_primaite/_game/_agent/test_actions.py b/tests/unit_tests/_primaite/_game/_agent/test_actions.py index dd8d5678..5750befd 100644 --- a/tests/unit_tests/_primaite/_game/_agent/test_actions.py +++ b/tests/unit_tests/_primaite/_game/_agent/test_actions.py @@ -2,6 +2,7 @@ from unittest.mock import Mock import pytest +from pydantic import ValidationError from primaite.game.agent.actions import ActionManager from primaite.game.agent.actions.manager import DoNothingAction @@ -37,7 +38,7 @@ def test_do_nothing_action_form_request(): ], ) # flake8: noqa @pytest.mark.parametrize( - "node_name, service_name, expect_to_do_nothing", + "node_name, service_name, expect_failure", [ ("pc_1", "chrome", False), (None, "chrome", True), @@ -45,34 +46,15 @@ def test_do_nothing_action_form_request(): (None, None, True), ], ) # flake8: noqa -def test_service_action_form_request(node_name, service_name, expect_to_do_nothing, action_class, action_verb): +def test_service_action_form_request(node_name, service_name, expect_failure, action_class, action_verb): """Test that the ServiceScanAction can form a request and that it is correct.""" - request = action_class.form_request( - config=action_class.ConfigSchema(node_name=node_name, service_name=service_name) - ) - - if expect_to_do_nothing: - assert request == ["do_nothing"] + if expect_failure: + with pytest.raises(ValidationError): + request = action_class.form_request( + config=action_class.ConfigSchema(node_name=node_name, service_name=service_name) + ) else: + request = action_class.form_request( + config=action_class.ConfigSchema(node_name=node_name, service_name=service_name) + ) assert request == ["network", "node", node_name, "service", service_name, action_verb] - - -@pytest.mark.parametrize( - "node_name, service_name, expect_to_do_nothing", - [ - ("pc_1", "chrome", False), - (None, "chrome", True), - ("pc_1", None, True), - (None, None, True), - ], -) # flake8: noqa -def test_service_scan_form_request(node_name, service_name, expect_to_do_nothing): - """Test that the ServiceScanAction can form a request and that it is correct.""" - request = NodeServiceScanAction.form_request( - NodeServiceScanAction.ConfigSchema(node_id=node_name, service_id=service_name) - ) - - if expect_to_do_nothing: - assert request == ["do_nothing"] - else: - assert request == ["network", "node", node_name, "service", service_name, "scan"] diff --git a/tests/unit_tests/_primaite/_game/_agent/test_agent.py b/tests/unit_tests/_primaite/_game/_agent/test_agent.py index 5f3b4fc0..7956a44f 100644 --- a/tests/unit_tests/_primaite/_game/_agent/test_agent.py +++ b/tests/unit_tests/_primaite/_game/_agent/test_agent.py @@ -39,6 +39,8 @@ def test_creating_agent_from_dict(): } agent = RandomAgent( config={ + "ref": "random_agent", + "team": "BLUE", "action_space": action_config, "observation_space": observation_config, "reward_function": reward_config, diff --git a/tests/unit_tests/_primaite/_game/_agent/test_observations.py b/tests/unit_tests/_primaite/_game/_agent/test_observations.py index 1888e9c1..5170bcf3 100644 --- a/tests/unit_tests/_primaite/_game/_agent/test_observations.py +++ b/tests/unit_tests/_primaite/_game/_agent/test_observations.py @@ -98,7 +98,7 @@ class TestFileSystemRequiresScan: """ cfg = yaml.safe_load(obs_cfg_yaml) - manager = ObservationManager(cfg) + manager = ObservationManager(config=cfg) hosts: List[HostObservation] = manager.obs.components["NODES"].hosts for i, host in enumerate(hosts): diff --git a/tests/unit_tests/_primaite/_game/_agent/test_sticky_rewards.py b/tests/unit_tests/_primaite/_game/_agent/test_sticky_rewards.py index 289d3941..5af71319 100644 --- a/tests/unit_tests/_primaite/_game/_agent/test_sticky_rewards.py +++ b/tests/unit_tests/_primaite/_game/_agent/test_sticky_rewards.py @@ -92,7 +92,7 @@ class TestWebpageUnavailabilitySticky: # agent did a successful fetch action = "node_application_execute" - params = {"node_id": 0, "application_id": 0} + params = {"node_name": "computer", "application_name": "WebBrowser"} request = ["network", "node", "computer", "application", "WebBrowser", "execute"] response = RequestResponse(status="success", data={}) browser_history.append({"outcome": 200}) @@ -115,7 +115,7 @@ class TestWebpageUnavailabilitySticky: # agent fails to fetch, get a -1.0 reward action = "node_application_execute" - params = {"node_id": 0, "application_id": 0} + params = {"node_name": "computer", "application_name": "WebBrowser"} request = ["network", "node", "computer", "application", "WebBrowser", "execute"] response = RequestResponse(status="failure", data={}) browser_history.append({"outcome": 404}) @@ -127,7 +127,7 @@ class TestWebpageUnavailabilitySticky: # agent fails again to fetch, get a -1.0 reward again action = "node_application_execute" - params = {"node_id": 0, "application_id": 0} + params = {"node_name": "computer", "application_name": "WebBrowser"} request = ["network", "node", "computer", "application", "WebBrowser", "execute"] response = RequestResponse(status="failure", data={}) browser_history.append({"outcome": 404}) @@ -153,7 +153,7 @@ class TestWebpageUnavailabilitySticky: # agent did a successful fetch action = "node_application_execute" - params = {"node_id": 0, "application_id": 0} + params = {"node_name": "computer", "application_name": "WebBrowser"} request = ["network", "node", "computer", "application", "WebBrowser", "execute"] response = RequestResponse(status="success", data={}) browser_history.append({"outcome": 200}) @@ -175,7 +175,7 @@ class TestWebpageUnavailabilitySticky: # agent fails to fetch, get a -1.0 reward action = "node_application_execute" - params = {"node_id": 0, "application_id": 0} + params = {"node_name": "computer", "application_name": "WebBrowser"} request = ["network", "node", "computer", "application", "WebBrowser", "execute"] response = RequestResponse(status="failure", data={}) browser_history.append({"outcome": 404}) @@ -187,7 +187,7 @@ class TestWebpageUnavailabilitySticky: # agent fails again to fetch, get a -1.0 reward again action = "node_application_execute" - params = {"node_id": 0, "application_id": 0} + params = {"node_name": "computer", "application_name": "WebBrowser"} request = ["network", "node", "computer", "application", "WebBrowser", "execute"] response = RequestResponse(status="failure", data={}) browser_history.append({"outcome": 404}) @@ -217,7 +217,7 @@ class TestGreenAdminDatabaseUnreachableSticky: # agent did a successful fetch action = "node_application_execute" - params = {"node_id": 0, "application_id": 0} + params = {"node_name": "computer", "application_name": "DatabaseClient"} request = ["network", "node", "computer", "application", "DatabaseClient", "execute"] response = RequestResponse(status="success", data={}) state = {"network": {"nodes": {"computer": {"applications": {"DatabaseClient": {}}}}}} @@ -238,7 +238,7 @@ class TestGreenAdminDatabaseUnreachableSticky: # agent fails to fetch, get a -1.0 reward action = "node_application_execute" - params = {"node_id": 0, "application_id": 0} + params = {"node_name": "computer", "application_name": "DatabaseClient"} request = ["network", "node", "computer", "application", "DatabaseClient", "execute"] response = RequestResponse(status="failure", data={}) state = {"network": {"nodes": {"computer": {"applications": {"DatabaseClient": {}}}}}} @@ -249,7 +249,7 @@ class TestGreenAdminDatabaseUnreachableSticky: # agent fails again to fetch, get a -1.0 reward again action = "node_application_execute" - params = {"node_id": 0, "application_id": 0} + params = {"node_name": "computer", "application_name": "DatabaseClient"} request = ["network", "node", "computer", "application", "DatabaseClient", "execute"] response = RequestResponse(status="failure", data={}) state = {"network": {"nodes": {"computer": {"applications": {"DatabaseClient": {}}}}}} @@ -276,7 +276,7 @@ class TestGreenAdminDatabaseUnreachableSticky: # agent did a successful fetch action = "node_application_execute" - params = {"node_id": 0, "application_id": 0} + params = {"node_name": "computer", "application_name": "DatabaseClient"} request = ["network", "node", "computer", "application", "DatabaseClient", "execute"] response = RequestResponse(status="success", data={}) state = {"network": {"nodes": {"computer": {"applications": {"DatabaseClient": {}}}}}} @@ -297,7 +297,7 @@ class TestGreenAdminDatabaseUnreachableSticky: # agent fails to fetch, get a -1.0 reward action = "node_application_execute" - params = {"node_id": 0, "application_id": 0} + params = {"node_name": "computer", "application_name": "DatabaseClient"} request = ["network", "node", "computer", "application", "DatabaseClient", "execute"] response = RequestResponse(status="failure", data={}) state = {"network": {"nodes": {"computer": {"applications": {"DatabaseClient": {}}}}}} @@ -308,7 +308,7 @@ class TestGreenAdminDatabaseUnreachableSticky: # agent fails again to fetch, get a -1.0 reward again action = "node_application_execute" - params = {"node_id": 0, "application_id": 0} + params = {"node_name": "computer", "application_name": "DatabaseClient"} request = ["network", "node", "computer", "application", "DatabaseClient", "execute"] response = RequestResponse(status="failure", data={}) state = {"network": {"nodes": {"computer": {"applications": {"DatabaseClient": {}}}}}}