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

#2912 - Updated to include __all__ and some test fixes. New extensible actions documentation page

Browse Source
This commit is contained in:
Charlie Crane
2024-11-12 14:49:44 +00:00
parent 844a3a60fa
commit d757bd01f0
15 changed files with 239 additions and 160 deletions
Download Patch File Download Diff File Expand all files Collapse all files

66
docs/source/how_to_guides/extensible_actions.rst Normal file
View File

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

26
src/primaite/game/agent/actions.py
View File

@@ -175,6 +175,10 @@ class NodeApplicationInstallAction(AbstractAction):
class ConfigureDatabaseClientAction(AbstractAction): class ConfigureDatabaseClientAction(AbstractAction):
"""Action which sets config parameters for a database client on a node.""" """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): class _Opts(BaseModel):
"""Schema for options that can be passed to this action.""" """Schema for options that can be passed to this action."""
@@ -197,9 +201,10 @@ class ConfigureDatabaseClientAction(AbstractAction):
class ConfigureRansomwareScriptAction(AbstractAction): class ConfigureRansomwareScriptAction(AbstractAction):
"""Action which sets config parameters for a ransomware script on a node.""" """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.""" """Schema for options that can be passed to this option."""
node_name: str
model_config = ConfigDict(extra="forbid") model_config = ConfigDict(extra="forbid")
server_ip_address: Optional[str] = None server_ip_address: Optional[str] = None
server_password: Optional[str] = None server_password: Optional[str] = None
@@ -208,13 +213,20 @@ class ConfigureRansomwareScriptAction(AbstractAction):
def __init__(self, manager: "ActionManager", **kwargs) -> None: def __init__(self, manager: "ActionManager", **kwargs) -> None:
super().__init__(manager=manager) 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.""" """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 config.node_name is None:
if node_name is None:
return ["do_nothing"] return ["do_nothing"]
ConfigureRansomwareScriptAction._Opts.model_validate(config) # check that options adhere to schema 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): class ConfigureDoSBotAction(AbstractAction):
@@ -285,8 +297,8 @@ class NodeFolderAbstractAction(AbstractAction):
class NodeFolderScanAction(NodeFolderAbstractAction): class NodeFolderScanAction(NodeFolderAbstractAction):
"""Action which scans a folder.""" """Action which scans a folder."""
def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, **kwargs) -> None: def __init__(self, manager: "ActionManager", node_name: str, folder_name, **kwargs) -> None:
super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, **kwargs) super().__init__(manager, node_name=node_name, folder_name=folder_name, **kwargs)
self.verb: str = "scan" self.verb: str = "scan"

2
src/primaite/game/agent/actions/__init__.py
View File

@@ -17,8 +17,8 @@ from primaite.game.agent.actions import (
from primaite.game.agent.actions.manager import ActionManager from primaite.game.agent.actions.manager import ActionManager
__all__ = ( __all__ = (
"acl",
"abstract", "abstract",
"acl",
"application", "application",
"config", "config",
"file", "file",

20
src/primaite/game/agent/actions/abstract.py
View File

@@ -8,17 +8,18 @@ from pydantic import BaseModel, ConfigDict
from primaite.interface.request import RequestFormat 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): class AbstractAction(BaseModel):
"""Base class for actions.""" """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): class ConfigSchema(BaseModel, ABC):
"""Base configuration schema for Actions.""" """Base configuration schema for Actions."""
@@ -41,12 +42,7 @@ class AbstractAction(BaseModel):
@classmethod @classmethod
def from_config(cls, config: Dict) -> "AbstractAction": def from_config(cls, config: Dict) -> "AbstractAction":
"""Create an action component from a config dictionary.""" """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(): for attribute, value in config.items():
if not hasattr(cls.ConfigSchema, attribute): if not hasattr(cls.ConfigSchema, attribute):
setattr(cls.ConfigSchema, attribute, value) setattr(cls.ConfigSchema, attribute, value)
return cls return cls

145
src/primaite/game/agent/actions/acl.py
View File

@@ -17,27 +17,35 @@ __all__ = (
class ACLAbstractAction(AbstractAction, identifier="acl_abstract_action"): class ACLAbstractAction(AbstractAction, identifier="acl_abstract_action"):
"""Base class for ACL actions.""" """Base class for ACL actions."""
class ConfigSchema(AbstractAction.ConfigSchema): class ConfigSchema(AbstractAction.ConfigSchema):
"""Configuration Schema base for ACL abstract actions.""" """Configuration Schema base for ACL abstract actions."""
class RouterACLAddRuleAction(AbstractAction, identifier="router_acl_add_rule"): class RouterACLAddRuleAction(AbstractAction, identifier="router_acl_add_rule"):
"""Action which adds a rule to a router's ACL.""" """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): class ConfigSchema(AbstractAction.ConfigSchema):
"""Configuration Schema for RouterACLAddRuleAction.""" """Configuration Schema for RouterACLAddRuleAction."""
target_router: str target_router: str
position: int position: int
permission: Literal[1, 2] permission: Literal[1, 2]
source_ip_id: int src_ip: str
source_wildcard_id: int src_wildcard: int
source_port_id: int source_port: str
dst_ip: str dst_ip: str
dst_wildcard_id: int dst_wildcard: int
dst_port: int dst_port: int
protocol_name: str 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.""" """At what position to add the rule, must be specified."""
permission: Literal[1, 2] permission: Literal[1, 2]
"""Whether to allow or deny traffic, must be specified. 1 = PERMIT, 2 = DENY.""" """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.""" """Rule source IP address. By default, all ip addresses."""
source_wildcard_id: int = Field(default=0, ge=0) source_wildcard_id: int = Field(default=0, ge=0)
"""Rule source IP wildcard. By default, use the wildcard at index 0 from action manager.""" """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.""" """Rule source port. By default, all source ports."""
dst_ip_id: int = Field(default=1, ge=1) dst_ip_id: int = Field(default=1, ge=1)
"""Rule destination IP address. By default, all ip addresses.""" """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.""" """Rule destination IP wildcard. By default, use the wildcard at index 0 from action manager."""
dst_port_id: int = Field(default=1, ge=1) dst_port_id: int = Field(default=1, ge=1)
"""Rule destination port. By default, all destination ports.""" """Rule destination port. By default, all destination ports."""
protocol_name: str = "ALL" protocol_name: str = "ALL"
"""Rule protocol. By default, all protocols.""" """Rule protocol. By default, all protocols."""
@field_validator( # @field_validator(
"source_ip_id", # "source_ip_id",
"source_port_id", # "source_port_id",
"source_wildcard_id", # "source_wildcard_id",
"dest_ip_id", # "dest_ip_id",
"dest_port_id", # "dest_port_id",
"dest_wildcard_id", # "dest_wildcard_id",
"protocol_name", # "protocol_name",
mode="before", # mode="before",
) # )
@classmethod # @classmethod
def not_none(cls, v: str, info: ValidationInfo) -> int: # def not_none(cls, v: str, info: ValidationInfo) -> int:
"""If None is passed, use the default value instead.""" # """If None is passed, use the default value instead."""
if v is None: # if v is None:
return cls.model_fields[info.field_name].default # return cls.model_fields[info.field_name].default
return v # return v
@classmethod @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.""" """Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
# Validate incoming data. # Validate incoming data.
parsed_options = RouterACLAddRuleAction.ACLRuleOptions( # parsed_options = RouterACLAddRuleAction.ACLRuleOptions(
target_router=config.target_router, # target_router=config.target_router,
position=config.position, # position=config.position,
permission=config.permission, # permission=config.permission,
source_ip_id=config.source_ip_id, # src_ip=config.src_ip,
source_wildcard_id=config.source_wildcard_id, # source_wildcard_id=config.source_wildcard_id,
dest_ip_id=config.dst_ip, # dst_ip_id=config.dst_ip,
dest_wildcard_id=config.dest_wildcard_id, # dst_wildcard=config.dst_wildcard,
source_port_id=config.source_port_id, # source_port_id=config.source_port_id,
dest_port_id=config.dst_port_id, # dest_port=config.dst_port,
protocol=config.protocol_name, # 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
return [ return [
"network", "network",
@@ -146,11 +113,11 @@ class RouterACLAddRuleAction(AbstractAction, identifier="router_acl_add_rule"):
config.target_router, config.target_router,
"acl", "acl",
"add_rule", "add_rule",
config.permission_str, config.permission,
protocol, config.protocol_name,
str(src_ip), config.src_ip,
config.src_wildcard, config.src_wildcard,
config.src_port, config.source_port,
str(config.dst_ip), str(config.dst_ip),
config.dst_wildcard, config.dst_wildcard,
config.dst_port, config.dst_port,
@@ -193,7 +160,6 @@ class FirewallACLAddRuleAction(ACLAbstractAction, identifier="firewall_acl_add_r
num_permissions: int = 3 num_permissions: int = 3
permission: str permission: str
def __init__( def __init__(
self, self,
manager: "ActionManager", manager: "ActionManager",
@@ -228,10 +194,8 @@ class FirewallACLAddRuleAction(ACLAbstractAction, identifier="firewall_acl_add_r
"protocol_id": num_protocols, "protocol_id": num_protocols,
} }
@classmethod @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.""" """Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
if config.permission == 0: if config.permission == 0:
permission_str = "UNUSED" permission_str = "UNUSED"
@@ -260,6 +224,7 @@ class FirewallACLAddRuleAction(ACLAbstractAction, identifier="firewall_acl_add_r
else: else:
# src_ip = self.manager.get_ip_address_by_idx(source_ip_id - 2) # src_ip = self.manager.get_ip_address_by_idx(source_ip_id - 2)
# subtract 2 to account for UNUSED=0, and ALL=1 # subtract 2 to account for UNUSED=0, and ALL=1
pass
if config.source_port_id == 0: if config.source_port_id == 0:
return ["do_nothing"] # invalid formulation return ["do_nothing"] # invalid formulation
@@ -286,8 +251,8 @@ class FirewallACLAddRuleAction(ACLAbstractAction, identifier="firewall_acl_add_r
else: else:
# dst_port = self.manager.get_port_by_idx(dest_port_id - 2) # dst_port = self.manager.get_port_by_idx(dest_port_id - 2)
# subtract 2 to account for UNUSED=0, and ALL=1 # subtract 2 to account for UNUSED=0, and ALL=1
# src_wildcard = self.manager.get_wildcard_by_idx(source_wildcard_id) # src_wildcard = self.manager.get_wildcard_by_idx(source_wildcard_id)
# dst_wildcard = self.manager.get_wildcard_by_idx(dest_wildcard_id) # dst_wildcard = self.manager.get_wildcard_by_idx(dest_wildcard_id)
pass pass
return [ return [
@@ -298,7 +263,7 @@ class FirewallACLAddRuleAction(ACLAbstractAction, identifier="firewall_acl_add_r
config.firewall_port_direction, config.firewall_port_direction,
"acl", "acl",
"add_rule", "add_rule",
permission_str, config.permission,
protocol, protocol,
str(src_ip), str(src_ip),
config.src_wildcard, config.src_wildcard,

24
src/primaite/game/agent/actions/config.py
View File

@@ -15,6 +15,7 @@ __all__ = (
"TerminalC2ServerAction", "TerminalC2ServerAction",
"RansomwareLaunchC2ServerAction", "RansomwareLaunchC2ServerAction",
"ExfiltrationC2ServerAction", "ExfiltrationC2ServerAction",
"ConfigureDatabaseClientAction",
) )
@@ -230,3 +231,26 @@ class ExfiltrationC2ServerAction(AbstractAction, identifier="c2_server_data_exfi
} }
ExfiltrationC2ServerAction._Opts.model_validate(command_model) ExfiltrationC2ServerAction._Opts.model_validate(command_model)
return ["network", "node", node_name, "application", "C2Server", "exfiltrate", 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]

2
src/primaite/game/agent/actions/file.py
View File

@@ -11,6 +11,8 @@ __all__ = (
"NodeFileRestoreAction", "NodeFileRestoreAction",
"NodeFileCorruptAction", "NodeFileCorruptAction",
"NodeFileAccessAction", "NodeFileAccessAction",
"NodeFileCheckhashAction",
"NodeFileRepairAction",
) )

6
src/primaite/game/agent/actions/manager.py
View File

@@ -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. # TODO: Make sure that actions are backwards compatible where the old YAML format is used.
__all__ = "DoNothingAction"
class DoNothingAction(AbstractAction, identifier="do_nothing"): class DoNothingAction(AbstractAction, identifier="do_nothing"):
"""Do Nothing Action.""" """Do Nothing Action."""
@@ -69,7 +71,7 @@ class ActionManager:
Action mapping that converts an integer to a specific action and parameter choice. Action mapping that converts an integer to a specific action and parameter choice.
For example : 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: if act_map is None:
# raise RuntimeError("Action map must be specified in the config file.") # 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 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 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 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`` 3. ``options``
``options`` contains a dictionary of options which are passed to the ActionManager's __init__ method. ``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 These options are used to calculate the shape of the action space, and to provide additional information

10
src/primaite/game/agent/actions/node.py
View File

@@ -5,7 +5,15 @@ from typing import ClassVar, List, Optional, Union
from primaite.game.agent.actions.manager import AbstractAction from primaite.game.agent.actions.manager import AbstractAction
from primaite.interface.request import RequestFormat 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"): class NodeAbstractAction(AbstractAction, identifier="node_abstract"):

8
src/primaite/game/agent/actions/session.py
View File

@@ -4,7 +4,11 @@ from abc import abstractmethod
from primaite.game.agent.actions.manager import AbstractAction from primaite.game.agent.actions.manager import AbstractAction
from primaite.interface.request import RequestFormat from primaite.interface.request import RequestFormat
__all__ = ("NodeSessionsRemoteLoginAction", "NodeSessionsRemoteLogoutAction") __all__ = (
"NodeSessionsRemoteLoginAction",
"NodeSessionsRemoteLogoutAction",
"NodeAccountChangePasswordAction",
)
class NodeSessionAbstractAction(AbstractAction, identifier="node_session_abstract"): 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] 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.""" """Action which changes the password for a user."""
username: str username: str

4
src/primaite/notebooks/Training-an-RLLib-Agent.ipynb
View File

@@ -95,7 +95,7 @@
], ],
"metadata": { "metadata": {
"kernelspec": { "kernelspec": {
"display_name": "venv", "display_name": ".venv",
"language": "python", "language": "python",
"name": "python3" "name": "python3"
}, },
@@ -109,7 +109,7 @@
"name": "python", "name": "python",
"nbconvert_exporter": "python", "nbconvert_exporter": "python",
"pygments_lexer": "ipython3", "pygments_lexer": "ipython3",
"version": "3.10.12" "version": "3.10.11"
} }
}, },
"nbformat": 4, "nbformat": 4,

38
tests/assets/configs/firewall_actions_network.yaml
View File

@@ -97,16 +97,16 @@ agents:
action_space: action_space:
action_list: action_list:
- type: do_nothing - type: do_nothing
- type: FIREWALL_ACL_ADDRULE firewall_acl_add_rule - type: firewall_acl_add_rule
- type: FIREWALL_ACL_REMOVERULE - type: firewall_acl_remove_rule
- type: NETWORK_PORT_DISABLE - type: network_port_disable
- type: NETWORK_PORT_ENABLE - type: network_port_enable
action_map: action_map:
0: 0:
action: DONOTHING action: do_nothing
options: {} options: {}
1: 1:
action: FIREWALL_ACL_ADDRULE action: firewall_acl_add_rule
options: options:
target_firewall_nodename: firewall target_firewall_nodename: firewall
firewall_port_name: internal firewall_port_name: internal
@@ -121,14 +121,14 @@ agents:
source_wildcard_id: 0 source_wildcard_id: 0
dest_wildcard_id: 0 dest_wildcard_id: 0
2: 2:
action: FIREWALL_ACL_REMOVERULE action: firewall_acl_remove_rule
options: options:
target_firewall_nodename: firewall target_firewall_nodename: firewall
firewall_port_name: internal firewall_port_name: internal
firewall_port_direction: inbound firewall_port_direction: inbound
position: 1 position: 1
3: 3:
action: FIREWALL_ACL_ADDRULE action: firewall_acl_add_rule
options: options:
target_firewall_nodename: firewall target_firewall_nodename: firewall
firewall_port_name: internal firewall_port_name: internal
@@ -143,14 +143,14 @@ agents:
source_wildcard_id: 0 source_wildcard_id: 0
dest_wildcard_id: 0 dest_wildcard_id: 0
4: 4:
action: FIREWALL_ACL_REMOVERULE action: firewall_acl_remove_rule
options: options:
target_firewall_nodename: firewall target_firewall_nodename: firewall
firewall_port_name: internal firewall_port_name: internal
firewall_port_direction: outbound firewall_port_direction: outbound
position: 1 position: 1
5: 5:
action: FIREWALL_ACL_ADDRULE action: firewall_acl_add_rule
options: options:
target_firewall_nodename: firewall target_firewall_nodename: firewall
firewall_port_name: dmz firewall_port_name: dmz
@@ -165,14 +165,14 @@ agents:
source_wildcard_id: 0 source_wildcard_id: 0
dest_wildcard_id: 0 dest_wildcard_id: 0
6: 6:
action: FIREWALL_ACL_REMOVERULE action: firewall_acl_remove_rule
options: options:
target_firewall_nodename: firewall target_firewall_nodename: firewall
firewall_port_name: dmz firewall_port_name: dmz
firewall_port_direction: inbound firewall_port_direction: inbound
position: 1 position: 1
7: 7:
action: FIREWALL_ACL_ADDRULE action: firewall_acl_add_rule
options: options:
target_firewall_nodename: firewall target_firewall_nodename: firewall
firewall_port_name: dmz firewall_port_name: dmz
@@ -187,14 +187,14 @@ agents:
source_wildcard_id: 0 source_wildcard_id: 0
dest_wildcard_id: 0 dest_wildcard_id: 0
8: 8:
action: FIREWALL_ACL_REMOVERULE action: firewall_acl_remove_rule
options: options:
target_firewall_nodename: firewall target_firewall_nodename: firewall
firewall_port_name: dmz firewall_port_name: dmz
firewall_port_direction: outbound firewall_port_direction: outbound
position: 2 position: 2
9: 9:
action: FIREWALL_ACL_ADDRULE action: firewall_acl_add_rule
options: options:
target_firewall_nodename: firewall target_firewall_nodename: firewall
firewall_port_name: external firewall_port_name: external
@@ -209,14 +209,14 @@ agents:
source_wildcard_id: 0 source_wildcard_id: 0
dest_wildcard_id: 0 dest_wildcard_id: 0
10: 10:
action: FIREWALL_ACL_REMOVERULE action: firewall_acl_remove_rule
options: options:
target_firewall_nodename: firewall target_firewall_nodename: firewall
firewall_port_name: external firewall_port_name: external
firewall_port_direction: inbound firewall_port_direction: inbound
position: 10 position: 10
11: 11:
action: FIREWALL_ACL_ADDRULE action: firewall_acl_add_rule
options: options:
target_firewall_nodename: firewall target_firewall_nodename: firewall
firewall_port_name: external firewall_port_name: external
@@ -231,19 +231,19 @@ agents:
source_wildcard_id: 0 source_wildcard_id: 0
dest_wildcard_id: 0 dest_wildcard_id: 0
12: 12:
action: FIREWALL_ACL_REMOVERULE action: firewall_acl_remove_rule
options: options:
target_firewall_nodename: firewall target_firewall_nodename: firewall
firewall_port_name: external firewall_port_name: external
firewall_port_direction: outbound firewall_port_direction: outbound
position: 1 position: 1
13: 13:
action: NETWORK_PORT_DISABLE action: network_port_disable
options: options:
target_nodename: firewall target_nodename: firewall
port_id: 3 port_id: 3
14: 14:
action: NETWORK_PORT_ENABLE action: network_port_enable
options: options:
target_nodename: firewall target_nodename: firewall
port_id: 3 port_id: 3

2
tests/conftest.py
View File

@@ -463,7 +463,7 @@ def game_and_agent():
{"type": "c2_server_ransomware_configure"}, {"type": "c2_server_ransomware_configure"},
{"type": "c2_server_terminal_command"}, {"type": "c2_server_terminal_command"},
{"type": "c2_server_data_exfiltrate"}, {"type": "c2_server_data_exfiltrate"},
{"type": "node_accounts_change_password"}, {"type": "node_account_change_password"},
{"type": "node_session_remote_login"}, {"type": "node_session_remote_login"},
{"type": "node_session_remote_logoff"}, {"type": "node_session_remote_logoff"},
{"type": "node_send_remote_command"}, {"type": "node_send_remote_command"},

10
tests/integration_tests/game_layer/actions/test_configure_actions.py
View File

@@ -4,7 +4,7 @@ from ipaddress import IPv4Address
import pytest import pytest
from pydantic import ValidationError from pydantic import ValidationError
from primaite.game.agent.actions import ( from primaite.game.agent.actions.config import (
ConfigureDatabaseClientAction, ConfigureDatabaseClientAction,
ConfigureDoSBotAction, ConfigureDoSBotAction,
ConfigureRansomwareScriptAction, ConfigureRansomwareScriptAction,
@@ -35,10 +35,10 @@ class TestConfigureDatabaseAction:
db_client: DatabaseClient = client_1.software_manager.software["DatabaseClient"] db_client: DatabaseClient = client_1.software_manager.software["DatabaseClient"]
action = ( action = (
"CONFIGURE_DATABASE_CLIENT", "configure_database_client",
{ {
"node_id": 0, "node_name": "client_1",
"config": { "model_config": {
"server_ip_address": "192.168.1.99", "server_ip_address": "192.168.1.99",
"server_password": "admin123", "server_password": "admin123",
}, },
@@ -53,7 +53,7 @@ class TestConfigureDatabaseAction:
def test_configure_ip(self, game_and_agent): def test_configure_ip(self, game_and_agent):
game, agent = game_and_agent game, agent = game_and_agent
agent: ControlledAgent 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 # make sure there is a database client on this node
client_1 = game.simulation.network.get_node_by_hostname("client_1") client_1 = game.simulation.network.get_node_by_hostname("client_1")

36
tests/integration_tests/game_layer/test_actions.py
View File

@@ -124,15 +124,15 @@ def test_router_acl_addrule_integration(game_and_agent: Tuple[PrimaiteGame, Prox
"router_acl_add_rule", "router_acl_add_rule",
{ {
"target_router": "router", "target_router": "router",
"position": 4, # 4th rule "position": 4,
"permission": 2, # DENY "permission": "DENY",
"source_ip_id": 3, # 10.0.1.2 (client_1) "src_ip": "10.0.1.2",
"dest_ip_id": 6, # 10.0.2.3 (server_2) "src_wildcard": 0,
"dest_port_id": 1, # ALL "source_port": "ALL",
"source_port_id": 1, # ALL "dst_ip": "10.0.2.3",
"protocol_name": "ALL", # ALL "dst_wildcard": 0,
"source_wildcard_id": 0, "dst_port": "ALL",
"dest_wildcard_id": 0, "protocol_name": "ALL",
}, },
) )
agent.store_action(action) 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) # 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 = ( action = (
"ROUTER_ACL_ADDRULE", "router_acl_add_rule",
{ {
"target_router": "router", "target_router": "router",
"position": 5, # 5th rule "position": 5, # 5th rule
"permission": 2, # DENY "permission": "DENY", # DENY
"source_ip_id": 5, # 10.0.2.2 (server_1) "src_ip": "10.0.2.2", # 10.0.2.2 (server_1)
"dest_ip_id": 6, # 10.0.2.3 (server_2) "src_wildcard": 0,
"dest_port_id": 1, # ALL "source_port": "ALL", # ALL
"source_port_id": 1, # ALL "dst_ip": "10.0.2.3", # 10.0.2.3 (server_2)
"protocol_id": 1, # ALL "dst_wildcard": 0,
"source_wildcard_id": 0, "dst_port": "ALL", # ALL
"dest_wildcard_id": 0, "protocol_name": "ALL", # ALL
}, },
) )
agent.store_action(action) agent.store_action(action)