From d757bd01f0f634965e5999582412801edb177f37 Mon Sep 17 00:00:00 2001 From: Charlie Crane Date: Tue, 12 Nov 2024 14:49:44 +0000 Subject: [PATCH] #2912 - Updated to include __all__ and some test fixes. New extensible actions documentation page --- .../how_to_guides/extensible_actions.rst | 66 ++++++++ src/primaite/game/agent/actions.py | 26 +++- src/primaite/game/agent/actions/__init__.py | 2 +- src/primaite/game/agent/actions/abstract.py | 20 +-- src/primaite/game/agent/actions/acl.py | 145 +++++++----------- src/primaite/game/agent/actions/config.py | 24 +++ src/primaite/game/agent/actions/file.py | 2 + src/primaite/game/agent/actions/manager.py | 6 +- src/primaite/game/agent/actions/node.py | 10 +- src/primaite/game/agent/actions/session.py | 8 +- .../notebooks/Training-an-RLLib-Agent.ipynb | 4 +- .../configs/firewall_actions_network.yaml | 38 ++--- tests/conftest.py | 2 +- .../actions/test_configure_actions.py | 10 +- .../game_layer/test_actions.py | 36 ++--- 15 files changed, 239 insertions(+), 160 deletions(-) create mode 100644 docs/source/how_to_guides/extensible_actions.rst diff --git a/docs/source/how_to_guides/extensible_actions.rst b/docs/source/how_to_guides/extensible_actions.rst new file mode 100644 index 00000000..a6c12303 --- /dev/null +++ b/docs/source/how_to_guides/extensible_actions.rst @@ -0,0 +1,66 @@ +.. only:: comment + + © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK + +.. _about: + +Extensible Actions +****************** + +Actions defined within PrimAITE have been updated to allow for easier creation of new bespoke actions, without the need to make changes to the ActionManager class within the core PrimAITE repository. + + +Developing Actions for PrimAITE +=============================== + +When developing new actions for PrimAITE, it's important to ensure new actions inherit from the AbstractAction class. This is so that the `ActionManager` has visibility +of the new action through the `AbstractAction` registry attribute. This also removes the need for actions to contain an `__init__` method. + +New actions to be used within PrimAITE require: + +#. **ConfigSchema**: + + This should be a nested class that defines the required configuration items for the new action. + + .. code-block:: python + + class ExampleAction(AbstractAction, identifier="Example_action"): + + class ConfigSchema(AbstractAction.ConfigSchema): + target_application: str + + The ConfigSchema is used when the class is called to form the action. + + +#. **Unique Identifier**: + + New actions should have a Unique identifier when declared. This is used by the `ActionManager` when forming/processing action commands from agents. See the example code block in ConfigSchema for how this should be implemented. + +#. **form_request method**: + + New actions need a `form_request()` method, to convert the action into a ``Requestformat`` that can be ingested by PrimAITE's `RequestManager`. + The below is an example of how this is done, taken from the `NodeFolderCreateAction`. + + .. code-block:: python + + @classmethod + def form_request(cls, config: ConfigSchema) -> RequestFormat: + """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" + if config.node_name is None or config.folder_name is None: + return ["do_nothing"] + return [ + "network", + "node", + config.node_name, + "file_system", + cls.model_fields["verb"].default, + "folder", + config.folder_name, + ] + +There is no longer a need for a `from_config()` method to be defined within new actions, as this is handled within the base `AbstractAction` class. + +Changes to YAML file. +===================== + +Action identifiers now follow the snake_case naming style, instead of the MACRO_CASE that has been seen previously. Please review any custom YAML files for any issues seen. This should be backwards compatible. diff --git a/src/primaite/game/agent/actions.py b/src/primaite/game/agent/actions.py index 68e42fb1..885e0238 100644 --- a/src/primaite/game/agent/actions.py +++ b/src/primaite/game/agent/actions.py @@ -175,6 +175,10 @@ class NodeApplicationInstallAction(AbstractAction): class ConfigureDatabaseClientAction(AbstractAction): """Action which sets config parameters for a database client on a node.""" + model_config = ConfigDict(extra="forbid") + server_ip_address: Optional[str] = None + server_password: Optional[str] = None + class _Opts(BaseModel): """Schema for options that can be passed to this action.""" @@ -197,9 +201,10 @@ class ConfigureDatabaseClientAction(AbstractAction): class ConfigureRansomwareScriptAction(AbstractAction): """Action which sets config parameters for a ransomware script on a node.""" - class _Opts(BaseModel): + class _Opts(BaseModel, AbstractAction.ConfigSchema): """Schema for options that can be passed to this option.""" + node_name: str model_config = ConfigDict(extra="forbid") server_ip_address: Optional[str] = None server_password: Optional[str] = None @@ -208,13 +213,20 @@ class ConfigureRansomwareScriptAction(AbstractAction): def __init__(self, manager: "ActionManager", **kwargs) -> None: super().__init__(manager=manager) - def form_request(self, node_id: int, config: Dict) -> RequestFormat: + def form_request(self, config: _Opts) -> RequestFormat: """Return the action formatted as a request that can be ingested by the simulation.""" - node_name = self.manager.get_node_name_by_idx(node_id) - if node_name is None: + if config.node_name is None: return ["do_nothing"] ConfigureRansomwareScriptAction._Opts.model_validate(config) # check that options adhere to schema - return ["network", "node", node_name, "application", "RansomwareScript", "configure", config] + return [ + "network", + "node", + config.node_name, + "application", + "RansomwareScript", + "configure", + config.model_config, + ] class ConfigureDoSBotAction(AbstractAction): @@ -285,8 +297,8 @@ class NodeFolderAbstractAction(AbstractAction): class NodeFolderScanAction(NodeFolderAbstractAction): """Action which scans a folder.""" - def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, **kwargs) -> None: - super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, **kwargs) + def __init__(self, manager: "ActionManager", node_name: str, folder_name, **kwargs) -> None: + super().__init__(manager, node_name=node_name, folder_name=folder_name, **kwargs) self.verb: str = "scan" diff --git a/src/primaite/game/agent/actions/__init__.py b/src/primaite/game/agent/actions/__init__.py index 3540b128..016a09ba 100644 --- a/src/primaite/game/agent/actions/__init__.py +++ b/src/primaite/game/agent/actions/__init__.py @@ -17,8 +17,8 @@ from primaite.game.agent.actions import ( from primaite.game.agent.actions.manager import ActionManager __all__ = ( - "acl", "abstract", + "acl", "application", "config", "file", diff --git a/src/primaite/game/agent/actions/abstract.py b/src/primaite/game/agent/actions/abstract.py index b96b14c9..2ed168d9 100644 --- a/src/primaite/game/agent/actions/abstract.py +++ b/src/primaite/game/agent/actions/abstract.py @@ -8,17 +8,18 @@ from pydantic import BaseModel, ConfigDict from primaite.interface.request import RequestFormat +# notes: +# we actually don't need to hold any state in actions, so there's no need to define any __init__ logic. +# all the init methods in the old actions are just used for holding a verb and shape, which are not really used. +# the config schema should be used to the actual parameters for formatting the action itself. +# (therefore there's no need for creating action instances, just the action class contains logic for converting +# CAOS actions to requests for simulator. Similar to the network node adder, that class also doesn't need to be +# instantiated.) + class AbstractAction(BaseModel): """Base class for actions.""" - # notes: - # we actually don't need to hold any state in actions, so there's no need to define any __init__ logic. - # all the init methods in the old actions are just used for holding a verb and shape, which are not really used. - # the config schema should be used to the actual parameters for formatting the action itself. - # (therefore there's no need for creating action instances, just the action class contains logic for converting - # CAOS actions to requests for simulator. Similar to the network node adder, that class also doesn't need to be - # instantiated.) class ConfigSchema(BaseModel, ABC): """Base configuration schema for Actions.""" @@ -41,12 +42,7 @@ class AbstractAction(BaseModel): @classmethod def from_config(cls, config: Dict) -> "AbstractAction": """Create an action component from a config dictionary.""" - # set attributes for action based off config dict - # if config["type"] not in cls._registry: - # raise ValueError(f"Invalid action reward type {config['type']}") - for attribute, value in config.items(): if not hasattr(cls.ConfigSchema, attribute): setattr(cls.ConfigSchema, attribute, value) - return cls diff --git a/src/primaite/game/agent/actions/acl.py b/src/primaite/game/agent/actions/acl.py index 8f5a79da..72a0b262 100644 --- a/src/primaite/game/agent/actions/acl.py +++ b/src/primaite/game/agent/actions/acl.py @@ -17,27 +17,35 @@ __all__ = ( class ACLAbstractAction(AbstractAction, identifier="acl_abstract_action"): """Base class for ACL actions.""" - class ConfigSchema(AbstractAction.ConfigSchema): """Configuration Schema base for ACL abstract actions.""" - - class RouterACLAddRuleAction(AbstractAction, identifier="router_acl_add_rule"): """Action which adds a rule to a router's ACL.""" + target_router: str + position: int + permission: Literal[1, 2] + src_ip: str + source_wildcard_id: int + source_port: str + dst_ip: str + dst_wildcard: int + dst_port: int + protocol_name: str + class ConfigSchema(AbstractAction.ConfigSchema): """Configuration Schema for RouterACLAddRuleAction.""" target_router: str position: int permission: Literal[1, 2] - source_ip_id: int - source_wildcard_id: int - source_port_id: int + src_ip: str + src_wildcard: int + source_port: str dst_ip: str - dst_wildcard_id: int + dst_wildcard: int dst_port: int protocol_name: str @@ -50,95 +58,54 @@ class RouterACLAddRuleAction(AbstractAction, identifier="router_acl_add_rule"): """At what position to add the rule, must be specified.""" permission: Literal[1, 2] """Whether to allow or deny traffic, must be specified. 1 = PERMIT, 2 = DENY.""" - source_ip_id: int = Field(default=1, ge=1) + src_ip: str """Rule source IP address. By default, all ip addresses.""" source_wildcard_id: int = Field(default=0, ge=0) """Rule source IP wildcard. By default, use the wildcard at index 0 from action manager.""" - source_port_id: int = Field(default=1, ge=1) + source_port: int = Field(default=1, ge=1) """Rule source port. By default, all source ports.""" dst_ip_id: int = Field(default=1, ge=1) """Rule destination IP address. By default, all ip addresses.""" - dst_wildcard_id: int = Field(default=0, ge=0) + dst_wildcard: int = Field(default=0, ge=0) """Rule destination IP wildcard. By default, use the wildcard at index 0 from action manager.""" dst_port_id: int = Field(default=1, ge=1) """Rule destination port. By default, all destination ports.""" protocol_name: str = "ALL" """Rule protocol. By default, all protocols.""" - @field_validator( - "source_ip_id", - "source_port_id", - "source_wildcard_id", - "dest_ip_id", - "dest_port_id", - "dest_wildcard_id", - "protocol_name", - mode="before", - ) - @classmethod - def not_none(cls, v: str, info: ValidationInfo) -> int: - """If None is passed, use the default value instead.""" - if v is None: - return cls.model_fields[info.field_name].default - return v + # @field_validator( + # "source_ip_id", + # "source_port_id", + # "source_wildcard_id", + # "dest_ip_id", + # "dest_port_id", + # "dest_wildcard_id", + # "protocol_name", + # mode="before", + # ) + # @classmethod + # def not_none(cls, v: str, info: ValidationInfo) -> int: + # """If None is passed, use the default value instead.""" + # if v is None: + # return cls.model_fields[info.field_name].default + # return v @classmethod def form_request(cls, config: ConfigSchema) -> List[str]: """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" # Validate incoming data. - parsed_options = RouterACLAddRuleAction.ACLRuleOptions( - target_router=config.target_router, - position=config.position, - permission=config.permission, - source_ip_id=config.source_ip_id, - source_wildcard_id=config.source_wildcard_id, - dest_ip_id=config.dst_ip, - dest_wildcard_id=config.dest_wildcard_id, - source_port_id=config.source_port_id, - dest_port_id=config.dst_port_id, - protocol=config.protocol_name, - ) - if parsed_options.permission == 1: - permission_str = "PERMIT" - elif parsed_options.permission == 2: - permission_str = "DENY" - # else: - # _LOGGER.warning(f"{self.__class__} received permission {permission}, expected 0 or 1.") - - if parsed_options.protocol_name == "ALL": - protocol = "ALL" - else: - protocol = parsed_options.protocol_name - # subtract 2 to account for UNUSED=0 and ALL=1. - - if parsed_options.source_ip_id == 1: - src_ip = "ALL" - else: - src_ip = cls.manager.get_ip_address_by_idx(parsed_options.source_ip_id - 2) - # subtract 2 to account for UNUSED=0, and ALL=1 - - src_wildcard = cls.manager.get_wildcard_by_idx(parsed_options.source_wildcard_id) - - if parsed_options.source_port_id == 1: - src_port = "ALL" - else: - src_port = cls.manager.get_port_by_idx(parsed_options.source_port_id - 2) - # subtract 2 to account for UNUSED=0, and ALL=1 - - if parsed_options.dest_ip_id == 1: - dst_ip = "ALL" - else: - dst_ip = cls.manager.get_ip_address_by_idx(parsed_options.dest_ip_id - 2) - # subtract 2 to account for UNUSED=0, and ALL=1 - dst_ip=config.dst_ip - - dst_wildcard = config.dest_wildcard_id - - if parsed_options.dest_port_id == 1: - dst_port = "ALL" - else: - dst_port = cls.manager.get_port_by_idx(parsed_options.dest_port_id - 2) - # subtract 2 to account for UNUSED=0, and ALL=1 + # parsed_options = RouterACLAddRuleAction.ACLRuleOptions( + # target_router=config.target_router, + # position=config.position, + # permission=config.permission, + # src_ip=config.src_ip, + # source_wildcard_id=config.source_wildcard_id, + # dst_ip_id=config.dst_ip, + # dst_wildcard=config.dst_wildcard, + # source_port_id=config.source_port_id, + # dest_port=config.dst_port, + # protocol=config.protocol_name, + # ) return [ "network", @@ -146,11 +113,11 @@ class RouterACLAddRuleAction(AbstractAction, identifier="router_acl_add_rule"): config.target_router, "acl", "add_rule", - config.permission_str, - protocol, - str(src_ip), + config.permission, + config.protocol_name, + config.src_ip, config.src_wildcard, - config.src_port, + config.source_port, str(config.dst_ip), config.dst_wildcard, config.dst_port, @@ -193,7 +160,6 @@ class FirewallACLAddRuleAction(ACLAbstractAction, identifier="firewall_acl_add_r num_permissions: int = 3 permission: str - def __init__( self, manager: "ActionManager", @@ -228,10 +194,8 @@ class FirewallACLAddRuleAction(ACLAbstractAction, identifier="firewall_acl_add_r "protocol_id": num_protocols, } - - @classmethod - def form_request(cls, config:ConfigSchema) -> List[str]: + def form_request(cls, config: ConfigSchema) -> List[str]: """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" if config.permission == 0: permission_str = "UNUSED" @@ -260,6 +224,7 @@ class FirewallACLAddRuleAction(ACLAbstractAction, identifier="firewall_acl_add_r else: # src_ip = self.manager.get_ip_address_by_idx(source_ip_id - 2) # subtract 2 to account for UNUSED=0, and ALL=1 + pass if config.source_port_id == 0: return ["do_nothing"] # invalid formulation @@ -286,8 +251,8 @@ class FirewallACLAddRuleAction(ACLAbstractAction, identifier="firewall_acl_add_r else: # dst_port = self.manager.get_port_by_idx(dest_port_id - 2) # subtract 2 to account for UNUSED=0, and ALL=1 - # src_wildcard = self.manager.get_wildcard_by_idx(source_wildcard_id) - # dst_wildcard = self.manager.get_wildcard_by_idx(dest_wildcard_id) + # src_wildcard = self.manager.get_wildcard_by_idx(source_wildcard_id) + # dst_wildcard = self.manager.get_wildcard_by_idx(dest_wildcard_id) pass return [ @@ -298,7 +263,7 @@ class FirewallACLAddRuleAction(ACLAbstractAction, identifier="firewall_acl_add_r config.firewall_port_direction, "acl", "add_rule", - permission_str, + config.permission, protocol, str(src_ip), config.src_wildcard, diff --git a/src/primaite/game/agent/actions/config.py b/src/primaite/game/agent/actions/config.py index beda8f27..a4247e21 100644 --- a/src/primaite/game/agent/actions/config.py +++ b/src/primaite/game/agent/actions/config.py @@ -15,6 +15,7 @@ __all__ = ( "TerminalC2ServerAction", "RansomwareLaunchC2ServerAction", "ExfiltrationC2ServerAction", + "ConfigureDatabaseClientAction", ) @@ -230,3 +231,26 @@ class ExfiltrationC2ServerAction(AbstractAction, identifier="c2_server_data_exfi } ExfiltrationC2ServerAction._Opts.model_validate(command_model) return ["network", "node", node_name, "application", "C2Server", "exfiltrate", command_model] + + +class ConfigureDatabaseClientAction(AbstractAction, identifier="configure_database_client"): + """Action which sets config parameters for a database client on a node.""" + + node_name: str + model_config: ConfigDict = ConfigDict(extra="forbid") + + class ConfigSchema(AbstractAction.ConfigSchema): + """Schema for options that can be passed to this action.""" + + node_name: str + model_config = ConfigDict(extra="forbid") + + def __init__(self, manager: "ActionManager", **kwargs) -> None: + super().__init__(manager=manager) + + def form_request(self, 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] diff --git a/src/primaite/game/agent/actions/file.py b/src/primaite/game/agent/actions/file.py index 6d1f95b9..c5ba1602 100644 --- a/src/primaite/game/agent/actions/file.py +++ b/src/primaite/game/agent/actions/file.py @@ -11,6 +11,8 @@ __all__ = ( "NodeFileRestoreAction", "NodeFileCorruptAction", "NodeFileAccessAction", + "NodeFileCheckhashAction", + "NodeFileRepairAction", ) diff --git a/src/primaite/game/agent/actions/manager.py b/src/primaite/game/agent/actions/manager.py index 26814bf2..764b5b6e 100644 --- a/src/primaite/game/agent/actions/manager.py +++ b/src/primaite/game/agent/actions/manager.py @@ -24,6 +24,8 @@ from primaite.interface.request import RequestFormat # TODO: Make sure that actions are backwards compatible where the old YAML format is used. +__all__ = "DoNothingAction" + class DoNothingAction(AbstractAction, identifier="do_nothing"): """Do Nothing Action.""" @@ -69,7 +71,7 @@ class ActionManager: Action mapping that converts an integer to a specific action and parameter choice. For example : - {0: ("NODE_SERVICE_SCAN", {node_id:0, service_id:2})} + {0: ("node_service_scan", {node_name:"client_1", service_name:"WebBrowser"})} """ if act_map is None: # raise RuntimeError("Action map must be specified in the config file.") @@ -147,7 +149,7 @@ class ActionManager: 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_id=1`` and ``service_id=1``, action 2 can be " + 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 diff --git a/src/primaite/game/agent/actions/node.py b/src/primaite/game/agent/actions/node.py index 29833b15..ab4209e5 100644 --- a/src/primaite/game/agent/actions/node.py +++ b/src/primaite/game/agent/actions/node.py @@ -5,7 +5,15 @@ from typing import ClassVar, List, Optional, Union from primaite.game.agent.actions.manager import AbstractAction from primaite.interface.request import RequestFormat -__all__ = ("NodeOSScanAction", "NodeShutdownAction", "NodeStartupAction", "NodeResetAction") +__all__ = ( + "NodeOSScanAction", + "NodeShutdownAction", + "NodeStartupAction", + "NodeResetAction", + "NodeNMAPPingScanAction", + "NodeNMAPPortScanAction", + "NodeNetworkServiceReconAction", +) class NodeAbstractAction(AbstractAction, identifier="node_abstract"): diff --git a/src/primaite/game/agent/actions/session.py b/src/primaite/game/agent/actions/session.py index 76b97cbd..79ff0705 100644 --- a/src/primaite/game/agent/actions/session.py +++ b/src/primaite/game/agent/actions/session.py @@ -4,7 +4,11 @@ from abc import abstractmethod from primaite.game.agent.actions.manager import AbstractAction from primaite.interface.request import RequestFormat -__all__ = ("NodeSessionsRemoteLoginAction", "NodeSessionsRemoteLogoutAction") +__all__ = ( + "NodeSessionsRemoteLoginAction", + "NodeSessionsRemoteLogoutAction", + "NodeAccountChangePasswordAction", +) class NodeSessionAbstractAction(AbstractAction, identifier="node_session_abstract"): @@ -73,7 +77,7 @@ class NodeSessionsRemoteLogoutAction(NodeSessionAbstractAction, identifier="node return ["network", "node", config.node_name, "service", "Terminal", "remote_logoff", config.remote_ip] -class NodeAccountsChangePasswordAction(NodeSessionAbstractAction, identifier="node_accounts_change_password"): +class NodeAccountChangePasswordAction(NodeSessionAbstractAction, identifier="node_account_change_password"): """Action which changes the password for a user.""" username: str diff --git a/src/primaite/notebooks/Training-an-RLLib-Agent.ipynb b/src/primaite/notebooks/Training-an-RLLib-Agent.ipynb index dbe8871c..0fd212f2 100644 --- a/src/primaite/notebooks/Training-an-RLLib-Agent.ipynb +++ b/src/primaite/notebooks/Training-an-RLLib-Agent.ipynb @@ -95,7 +95,7 @@ ], "metadata": { "kernelspec": { - "display_name": "venv", + "display_name": ".venv", "language": "python", "name": "python3" }, @@ -109,7 +109,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.12" + "version": "3.10.11" } }, "nbformat": 4, diff --git a/tests/assets/configs/firewall_actions_network.yaml b/tests/assets/configs/firewall_actions_network.yaml index 18e27de7..88b09a29 100644 --- a/tests/assets/configs/firewall_actions_network.yaml +++ b/tests/assets/configs/firewall_actions_network.yaml @@ -97,16 +97,16 @@ agents: action_space: action_list: - type: do_nothing - - type: FIREWALL_ACL_ADDRULE firewall_acl_add_rule - - type: FIREWALL_ACL_REMOVERULE - - type: NETWORK_PORT_DISABLE - - type: NETWORK_PORT_ENABLE + - type: firewall_acl_add_rule + - type: firewall_acl_remove_rule + - type: network_port_disable + - type: network_port_enable action_map: 0: - action: DONOTHING + action: do_nothing options: {} 1: - action: FIREWALL_ACL_ADDRULE + action: firewall_acl_add_rule options: target_firewall_nodename: firewall firewall_port_name: internal @@ -121,14 +121,14 @@ agents: source_wildcard_id: 0 dest_wildcard_id: 0 2: - action: FIREWALL_ACL_REMOVERULE + action: firewall_acl_remove_rule options: target_firewall_nodename: firewall firewall_port_name: internal firewall_port_direction: inbound position: 1 3: - action: FIREWALL_ACL_ADDRULE + action: firewall_acl_add_rule options: target_firewall_nodename: firewall firewall_port_name: internal @@ -143,14 +143,14 @@ agents: source_wildcard_id: 0 dest_wildcard_id: 0 4: - action: FIREWALL_ACL_REMOVERULE + action: firewall_acl_remove_rule options: target_firewall_nodename: firewall firewall_port_name: internal firewall_port_direction: outbound position: 1 5: - action: FIREWALL_ACL_ADDRULE + action: firewall_acl_add_rule options: target_firewall_nodename: firewall firewall_port_name: dmz @@ -165,14 +165,14 @@ agents: source_wildcard_id: 0 dest_wildcard_id: 0 6: - action: FIREWALL_ACL_REMOVERULE + action: firewall_acl_remove_rule options: target_firewall_nodename: firewall firewall_port_name: dmz firewall_port_direction: inbound position: 1 7: - action: FIREWALL_ACL_ADDRULE + action: firewall_acl_add_rule options: target_firewall_nodename: firewall firewall_port_name: dmz @@ -187,14 +187,14 @@ agents: source_wildcard_id: 0 dest_wildcard_id: 0 8: - action: FIREWALL_ACL_REMOVERULE + action: firewall_acl_remove_rule options: target_firewall_nodename: firewall firewall_port_name: dmz firewall_port_direction: outbound position: 2 9: - action: FIREWALL_ACL_ADDRULE + action: firewall_acl_add_rule options: target_firewall_nodename: firewall firewall_port_name: external @@ -209,14 +209,14 @@ agents: source_wildcard_id: 0 dest_wildcard_id: 0 10: - action: FIREWALL_ACL_REMOVERULE + action: firewall_acl_remove_rule options: target_firewall_nodename: firewall firewall_port_name: external firewall_port_direction: inbound position: 10 11: - action: FIREWALL_ACL_ADDRULE + action: firewall_acl_add_rule options: target_firewall_nodename: firewall firewall_port_name: external @@ -231,19 +231,19 @@ agents: source_wildcard_id: 0 dest_wildcard_id: 0 12: - action: FIREWALL_ACL_REMOVERULE + action: firewall_acl_remove_rule options: target_firewall_nodename: firewall firewall_port_name: external firewall_port_direction: outbound position: 1 13: - action: NETWORK_PORT_DISABLE + action: network_port_disable options: target_nodename: firewall port_id: 3 14: - action: NETWORK_PORT_ENABLE + action: network_port_enable options: target_nodename: firewall port_id: 3 diff --git a/tests/conftest.py b/tests/conftest.py index 8bfc78a4..bd1b79ee 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -463,7 +463,7 @@ def game_and_agent(): {"type": "c2_server_ransomware_configure"}, {"type": "c2_server_terminal_command"}, {"type": "c2_server_data_exfiltrate"}, - {"type": "node_accounts_change_password"}, + {"type": "node_account_change_password"}, {"type": "node_session_remote_login"}, {"type": "node_session_remote_logoff"}, {"type": "node_send_remote_command"}, 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 508bd5a4..7bf45fb4 100644 --- a/tests/integration_tests/game_layer/actions/test_configure_actions.py +++ b/tests/integration_tests/game_layer/actions/test_configure_actions.py @@ -4,7 +4,7 @@ from ipaddress import IPv4Address import pytest from pydantic import ValidationError -from primaite.game.agent.actions import ( +from primaite.game.agent.actions.config import ( ConfigureDatabaseClientAction, ConfigureDoSBotAction, ConfigureRansomwareScriptAction, @@ -35,10 +35,10 @@ class TestConfigureDatabaseAction: db_client: DatabaseClient = client_1.software_manager.software["DatabaseClient"] action = ( - "CONFIGURE_DATABASE_CLIENT", + "configure_database_client", { - "node_id": 0, - "config": { + "node_name": "client_1", + "model_config": { "server_ip_address": "192.168.1.99", "server_password": "admin123", }, @@ -53,7 +53,7 @@ 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) + 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") diff --git a/tests/integration_tests/game_layer/test_actions.py b/tests/integration_tests/game_layer/test_actions.py index 21aa31de..baa4c725 100644 --- a/tests/integration_tests/game_layer/test_actions.py +++ b/tests/integration_tests/game_layer/test_actions.py @@ -124,15 +124,15 @@ def test_router_acl_addrule_integration(game_and_agent: Tuple[PrimaiteGame, Prox "router_acl_add_rule", { "target_router": "router", - "position": 4, # 4th rule - "permission": 2, # DENY - "source_ip_id": 3, # 10.0.1.2 (client_1) - "dest_ip_id": 6, # 10.0.2.3 (server_2) - "dest_port_id": 1, # ALL - "source_port_id": 1, # ALL - "protocol_name": "ALL", # ALL - "source_wildcard_id": 0, - "dest_wildcard_id": 0, + "position": 4, + "permission": "DENY", + "src_ip": "10.0.1.2", + "src_wildcard": 0, + "source_port": "ALL", + "dst_ip": "10.0.2.3", + "dst_wildcard": 0, + "dst_port": "ALL", + "protocol_name": "ALL", }, ) agent.store_action(action) @@ -148,18 +148,18 @@ def test_router_acl_addrule_integration(game_and_agent: Tuple[PrimaiteGame, Prox # 4: Add a rule to block server_1 from reaching server_2 on router (this should not affect comms as they are on same subnet) action = ( - "ROUTER_ACL_ADDRULE", + "router_acl_add_rule", { "target_router": "router", "position": 5, # 5th rule - "permission": 2, # DENY - "source_ip_id": 5, # 10.0.2.2 (server_1) - "dest_ip_id": 6, # 10.0.2.3 (server_2) - "dest_port_id": 1, # ALL - "source_port_id": 1, # ALL - "protocol_id": 1, # ALL - "source_wildcard_id": 0, - "dest_wildcard_id": 0, + "permission": "DENY", # DENY + "src_ip": "10.0.2.2", # 10.0.2.2 (server_1) + "src_wildcard": 0, + "source_port": "ALL", # ALL + "dst_ip": "10.0.2.3", # 10.0.2.3 (server_2) + "dst_wildcard": 0, + "dst_port": "ALL", # ALL + "protocol_name": "ALL", # ALL }, ) agent.store_action(action)