#2912 - Additional actions added to config.py, refactor of HostNIC
This commit is contained in:
@@ -1,6 +1,6 @@
|
|||||||
# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
|
# © 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
|
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
|
ConfigureC2BeaconAction._Opts.model_validate(config) # check that options adhere to schema
|
||||||
|
|
||||||
return ["network", "node", config.node_name, "application", "C2Beacon", "configure", config.__dict__]
|
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"]
|
||||||
@@ -32,7 +32,7 @@ class HostNICAbstractAction(AbstractAction):
|
|||||||
return ["network", "node", config.node_name, "network_interface", config.nic_num, cls.verb]
|
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."""
|
"""Action which enables a NIC."""
|
||||||
|
|
||||||
class ConfigSchema(HostNICAbstractAction.ConfigSchema):
|
class ConfigSchema(HostNICAbstractAction.ConfigSchema):
|
||||||
@@ -41,7 +41,7 @@ class HostNICEnableAction(HostNICAbstractAction):
|
|||||||
verb: str = "enable"
|
verb: str = "enable"
|
||||||
|
|
||||||
|
|
||||||
class HostNICDisableAction(HostNICAbstractAction):
|
class HostNICDisableAction(HostNICAbstractAction, identifier="host_nic_disable"):
|
||||||
"""Action which disables a NIC."""
|
"""Action which disables a NIC."""
|
||||||
|
|
||||||
class ConfigSchema(HostNICAbstractAction.ConfigSchema):
|
class ConfigSchema(HostNICAbstractAction.ConfigSchema):
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
|
# © 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
|
from pydantic import BaseModel, ConfigDict
|
||||||
|
|
||||||
@@ -8,39 +8,37 @@ from primaite.game.agent.actions.manager import AbstractAction
|
|||||||
from primaite.interface.request import RequestFormat
|
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."""
|
"""Action which enables are port on a router or a firewall."""
|
||||||
|
|
||||||
def __init__(self, manager: "ActionManager", max_nics_per_node: int, **kwargs) -> None:
|
class ConfigSchema(AbstractAction.ConfigSchema):
|
||||||
"""Init method for NetworkPortEnableAction.
|
"""Configuration schema for NetworkPortEnableAction."""
|
||||||
|
|
||||||
:param max_nics_per_node: Maximum number of NICs per node.
|
verb: str = "enable"
|
||||||
:type max_nics_per_node: int
|
|
||||||
"""
|
|
||||||
super().__init__(manager=manager)
|
|
||||||
self.shape: Dict[str, int] = {"port_id": max_nics_per_node}
|
|
||||||
|
|
||||||
def form_request(self, target_nodename: str, port_id: int) -> RequestFormat:
|
class NetworkPortDisableAction(NetworkPortAbstractAction, identifier="network_port_disable"):
|
||||||
"""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):
|
|
||||||
"""Action which disables are port on a router or a firewall."""
|
"""Action which disables are port on a router or a firewall."""
|
||||||
|
|
||||||
def __init__(self, manager: "ActionManager", max_nics_per_node: int, **kwargs) -> None:
|
class ConfigSchema(AbstractAction.ConfigSchema):
|
||||||
"""Init method for NetworkPortDisableAction.
|
"""Configuration schema for NetworkPortDisableAction"""
|
||||||
|
|
||||||
:param max_nics_per_node: Maximum number of NICs per node.
|
verb: str = "disable"
|
||||||
:type max_nics_per_node: int
|
|
||||||
"""
|
|
||||||
super().__init__(manager=manager)
|
|
||||||
self.shape: Dict[str, int] = {"port_id": max_nics_per_node}
|
|
||||||
|
|
||||||
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"]
|
|
||||||
|
|||||||
Reference in New Issue
Block a user