diff --git a/src/primaite/game/agent/actions/config.py b/src/primaite/game/agent/actions/config.py index c06ce9c8..6096a0b2 100644 --- a/src/primaite/game/agent/actions/config.py +++ b/src/primaite/game/agent/actions/config.py @@ -1,6 +1,6 @@ # © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK -from typing import Dict, Optional +from typing import Dict, List, Optional, Union from pydantic import BaseModel, ConfigDict, Field, field_validator, ValidationInfo @@ -107,3 +107,71 @@ class ConfigureC2BeaconAction(AbstractAction): ConfigureC2BeaconAction._Opts.model_validate(config) # check that options adhere to schema return ["network", "node", config.node_name, "application", "C2Beacon", "configure", config.__dict__] + + +class NodeSendRemoteCommandAction(AbstractAction): + """Action which sends a terminal command to a remote node via SSH.""" + + def __init__(self, manager: "ActionManager", **kwargs) -> None: + super().__init__(manager=manager) + + def form_request(self, node_id: int, remote_ip: str, command: RequestFormat) -> RequestFormat: + """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" + node_name = self.manager.get_node_name_by_idx(node_id) + return [ + "network", + "node", + node_name, + "service", + "Terminal", + "send_remote_command", + remote_ip, + {"command": command}, + ] + + +class TerminalC2ServerAction(AbstractAction): + """Action which causes the C2 Server to send a command to the C2 Beacon to execute the terminal command passed.""" + + class _Opts(BaseModel): + """Schema for options that can be passed to this action.""" + + commands: Union[List[RequestFormat], RequestFormat] + ip_address: Optional[str] + username: Optional[str] + password: Optional[str] + + def __init__(self, manager: "ActionManager", **kwargs) -> None: + super().__init__(manager=manager) + + def form_request(self, node_id: int, commands: List, ip_address: Optional[str], account: dict) -> 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: + return ["do_nothing"] + + command_model = { + "commands": commands, + "ip_address": ip_address, + "username": account["username"], + "password": account["password"], + } + + TerminalC2ServerAction._Opts.model_validate(command_model) + return ["network", "node", node_name, "application", "C2Server", "terminal_command", command_model] + + +class RansomwareLaunchC2ServerAction(AbstractAction): + """Action which causes the C2 Server to send a command to the C2 Beacon to launch the RansomwareScript.""" + + class ConfigSchema(AbstractAction): + """Configuration schema for RansomwareLaunchC2ServerAction.""" + node_name: str + + @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"] + # This action currently doesn't require any further configuration options. + return ["network", "node", config.node_name, "application", "C2Server", "ransomware_launch"] \ No newline at end of file diff --git a/src/primaite/game/agent/actions/host_nic.py b/src/primaite/game/agent/actions/host_nic.py index 2909772b..a4dd8d9c 100644 --- a/src/primaite/game/agent/actions/host_nic.py +++ b/src/primaite/game/agent/actions/host_nic.py @@ -32,7 +32,7 @@ class HostNICAbstractAction(AbstractAction): return ["network", "node", config.node_name, "network_interface", config.nic_num, cls.verb] -class HostNICEnableAction(HostNICAbstractAction): +class HostNICEnableAction(HostNICAbstractAction, identifier="host_nic_enable"): """Action which enables a NIC.""" class ConfigSchema(HostNICAbstractAction.ConfigSchema): @@ -41,7 +41,7 @@ class HostNICEnableAction(HostNICAbstractAction): verb: str = "enable" -class HostNICDisableAction(HostNICAbstractAction): +class HostNICDisableAction(HostNICAbstractAction, identifier="host_nic_disable"): """Action which disables a NIC.""" class ConfigSchema(HostNICAbstractAction.ConfigSchema): diff --git a/src/primaite/game/agent/actions/network.py b/src/primaite/game/agent/actions/network.py index e761fb1e..630385bd 100644 --- a/src/primaite/game/agent/actions/network.py +++ b/src/primaite/game/agent/actions/network.py @@ -1,6 +1,6 @@ # © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK -from typing import Dict, Optional +from typing import ClassVar, Dict, Optional from pydantic import BaseModel, ConfigDict @@ -8,39 +8,37 @@ from primaite.game.agent.actions.manager import AbstractAction from primaite.interface.request import RequestFormat -class NetworkPortEnableAction(AbstractAction): +class NetworkPortAbstractAction(AbstractAction): + """Base class for Network port actions""" + + class ConfigSchema(AbstractAction.ConfigSchema): + """Base configuration schema for NetworkPort actions.""" + target_nodename: str + port_id: str + + verb: ClassVar[str] + + @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.target_nodename is None or config.port_id is None: + return ["do_nothing"] + return ["network", "node", config.target_nodename, "network_interface", config.port_id, cls.verb] + + +class NetworkPortEnableAction(NetworkPortAbstractAction, identifier="network_port_enable"): """Action which enables are port on a router or a firewall.""" - def __init__(self, manager: "ActionManager", max_nics_per_node: int, **kwargs) -> None: - """Init method for NetworkPortEnableAction. + class ConfigSchema(AbstractAction.ConfigSchema): + """Configuration schema for NetworkPortEnableAction.""" - :param max_nics_per_node: Maximum number of NICs per node. - :type max_nics_per_node: int - """ - super().__init__(manager=manager) - self.shape: Dict[str, int] = {"port_id": max_nics_per_node} + verb: str = "enable" - def form_request(self, target_nodename: str, port_id: int) -> RequestFormat: - """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" - if target_nodename is None or port_id is None: - return ["do_nothing"] - return ["network", "node", target_nodename, "network_interface", port_id, "enable"] - - -class NetworkPortDisableAction(AbstractAction): +class NetworkPortDisableAction(NetworkPortAbstractAction, identifier="network_port_disable"): """Action which disables are port on a router or a firewall.""" - def __init__(self, manager: "ActionManager", max_nics_per_node: int, **kwargs) -> None: - """Init method for NetworkPortDisableAction. + class ConfigSchema(AbstractAction.ConfigSchema): + """Configuration schema for NetworkPortDisableAction""" - :param max_nics_per_node: Maximum number of NICs per node. - :type max_nics_per_node: int - """ - super().__init__(manager=manager) - self.shape: Dict[str, int] = {"port_id": max_nics_per_node} + verb: str = "disable" - def form_request(self, target_nodename: str, port_id: int) -> RequestFormat: - """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" - if target_nodename is None or port_id is None: - return ["do_nothing"] - return ["network", "node", target_nodename, "network_interface", port_id, "disable"]