Files
PrimAITE/src/primaite/game/agent/actions.py

1247 lines
50 KiB
Python
Raw Normal View History

# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
2023-10-19 01:56:40 +01:00
"""
This module contains the ActionManager class which belongs to the Agent class.
An agent's action space is made up of a collection of actions. Each action is an instance of a subclass of
AbstractAction. The ActionManager is responsible for:
1. Creating the action space from a list of action types.
2. Converting an integer action choice into a specific action and parameter choice.
3. Converting an action and parameter choice into a request which can be ingested by the PrimAITE simulation. This
ensures that requests conform to the simulator's request format.
"""
2024-10-11 15:00:26 +01:00
from abc import abstractmethod
from typing import Dict, List, Literal, Optional, Union
2024-07-01 14:41:41 +01:00
from pydantic import BaseModel, ConfigDict, Field, field_validator, ValidationInfo
2023-10-09 17:29:50 +01:00
from primaite import getLogger
from primaite.game.agent.actions.manager import AbstractAction, ActionManager
2024-10-11 15:00:26 +01:00
from primaite.game.agent.actions.service import NodeServiceAbstractAction
2024-07-01 11:31:27 +01:00
from primaite.interface.request import RequestFormat
2023-10-09 18:35:30 +01:00
2023-10-09 17:29:50 +01:00
_LOGGER = getLogger(__name__)
2023-10-02 17:21:43 +01:00
class NodeServiceScanAction(NodeServiceAbstractAction):
2023-10-19 01:56:40 +01:00
"""Action which scans a service."""
2023-10-09 18:35:30 +01:00
def __init__(self, manager: "ActionManager", num_nodes: int, num_services: int, **kwargs) -> None:
2023-10-06 20:32:52 +01:00
super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services)
2023-11-27 13:28:11 +00:00
self.verb: str = "scan"
2023-10-02 17:21:43 +01:00
2023-10-06 10:36:29 +01:00
2023-10-02 17:21:43 +01:00
class NodeServiceStopAction(NodeServiceAbstractAction):
2023-10-19 01:56:40 +01:00
"""Action which stops a service."""
2023-10-09 18:35:30 +01:00
def __init__(self, manager: "ActionManager", num_nodes: int, num_services: int, **kwargs) -> None:
2023-10-06 20:32:52 +01:00
super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services)
2023-11-27 13:28:11 +00:00
self.verb: str = "stop"
2023-10-02 17:21:43 +01:00
2023-10-06 10:36:29 +01:00
2023-10-02 17:21:43 +01:00
class NodeServiceStartAction(NodeServiceAbstractAction):
2023-10-19 01:56:40 +01:00
"""Action which starts a service."""
2023-10-09 18:35:30 +01:00
def __init__(self, manager: "ActionManager", num_nodes: int, num_services: int, **kwargs) -> None:
2023-10-06 20:32:52 +01:00
super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services)
2023-11-27 13:28:11 +00:00
self.verb: str = "start"
2023-10-02 17:21:43 +01:00
2023-10-06 10:36:29 +01:00
2023-10-02 17:21:43 +01:00
class NodeServicePauseAction(NodeServiceAbstractAction):
2023-10-19 01:56:40 +01:00
"""Action which pauses a service."""
2023-10-09 18:35:30 +01:00
def __init__(self, manager: "ActionManager", num_nodes: int, num_services: int, **kwargs) -> None:
2023-10-06 20:32:52 +01:00
super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services)
2023-11-27 13:28:11 +00:00
self.verb: str = "pause"
2023-10-02 17:21:43 +01:00
2023-10-06 10:36:29 +01:00
2023-10-02 17:21:43 +01:00
class NodeServiceResumeAction(NodeServiceAbstractAction):
2023-10-19 01:56:40 +01:00
"""Action which resumes a service."""
2023-10-09 18:35:30 +01:00
def __init__(self, manager: "ActionManager", num_nodes: int, num_services: int, **kwargs) -> None:
2023-10-06 20:32:52 +01:00
super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services)
2023-11-27 13:28:11 +00:00
self.verb: str = "resume"
2023-10-02 17:21:43 +01:00
2023-10-06 10:36:29 +01:00
2023-10-02 17:21:43 +01:00
class NodeServiceRestartAction(NodeServiceAbstractAction):
2023-10-19 01:56:40 +01:00
"""Action which restarts a service."""
2023-10-09 18:35:30 +01:00
def __init__(self, manager: "ActionManager", num_nodes: int, num_services: int, **kwargs) -> None:
2023-10-06 20:32:52 +01:00
super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services)
2023-11-27 13:28:11 +00:00
self.verb: str = "restart"
2023-10-02 17:21:43 +01:00
2023-10-06 10:36:29 +01:00
2023-10-02 17:21:43 +01:00
class NodeServiceDisableAction(NodeServiceAbstractAction):
2023-10-19 01:56:40 +01:00
"""Action which disables a service."""
2023-10-09 18:35:30 +01:00
def __init__(self, manager: "ActionManager", num_nodes: int, num_services: int, **kwargs) -> None:
2023-10-06 20:32:52 +01:00
super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services)
2023-11-27 13:28:11 +00:00
self.verb: str = "disable"
2023-10-02 17:21:43 +01:00
2023-10-06 10:36:29 +01:00
2023-10-02 17:21:43 +01:00
class NodeServiceEnableAction(NodeServiceAbstractAction):
2023-10-19 01:56:40 +01:00
"""Action which enables a service."""
2023-10-09 18:35:30 +01:00
def __init__(self, manager: "ActionManager", num_nodes: int, num_services: int, **kwargs) -> None:
2023-10-06 20:32:52 +01:00
super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services)
2023-11-27 13:28:11 +00:00
self.verb: str = "enable"
2023-10-02 17:21:43 +01:00
2024-03-26 10:51:33 +00:00
class NodeServiceFixAction(NodeServiceAbstractAction):
"""Action which fixes a service."""
2023-12-21 15:07:41 +00:00
def __init__(self, manager: "ActionManager", num_nodes: int, num_services: int, **kwargs) -> None:
super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services)
2024-03-26 10:51:33 +00:00
self.verb: str = "fix"
2023-12-21 15:07:41 +00:00
class NodeApplicationAbstractAction(AbstractAction):
"""
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
inherit from this base class.
"""
@abstractmethod
def __init__(self, manager: "ActionManager", num_nodes: int, num_applications: int, **kwargs) -> None:
super().__init__(manager=manager)
self.shape: Dict[str, int] = {"node_id": num_nodes, "application_id": num_applications}
2023-11-27 13:28:11 +00:00
self.verb: str # define but don't initialise: defends against children classes not defining this
def form_request(self, node_id: int, application_id: int) -> 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)
application_name = self.manager.get_application_name_by_idx(node_id, application_id)
if node_name is None or application_name is None:
return ["do_nothing"]
return ["network", "node", node_name, "application", application_name, self.verb]
class NodeApplicationExecuteAction(NodeApplicationAbstractAction):
"""Action which executes an application."""
def __init__(self, manager: "ActionManager", num_nodes: int, num_applications: int, **kwargs) -> None:
super().__init__(manager=manager, num_nodes=num_nodes, num_applications=num_applications)
2023-11-27 13:28:11 +00:00
self.verb: str = "execute"
class NodeApplicationScanAction(NodeApplicationAbstractAction):
"""Action which scans an application."""
def __init__(self, manager: "ActionManager", num_nodes: int, num_applications: int, **kwargs) -> None:
super().__init__(manager=manager, num_nodes=num_nodes, num_applications=num_applications)
self.verb: str = "scan"
class NodeApplicationCloseAction(NodeApplicationAbstractAction):
"""Action which closes an application."""
def __init__(self, manager: "ActionManager", num_nodes: int, num_applications: int, **kwargs) -> None:
super().__init__(manager=manager, num_nodes=num_nodes, num_applications=num_applications)
self.verb: str = "close"
class NodeApplicationFixAction(NodeApplicationAbstractAction):
"""Action which fixes an application."""
def __init__(self, manager: "ActionManager", num_nodes: int, num_applications: int, **kwargs) -> None:
super().__init__(manager=manager, num_nodes=num_nodes, num_applications=num_applications)
2024-03-26 10:51:33 +00:00
self.verb: str = "fix"
class NodeApplicationInstallAction(AbstractAction):
"""Action which installs an application."""
def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None:
super().__init__(manager=manager)
self.shape: Dict[str, int] = {"node_id": num_nodes}
def form_request(self, node_id: int, application_name: str) -> 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)
if node_name is None:
return ["do_nothing"]
return [
"network",
"node",
node_name,
"software_manager",
"application",
"install",
application_name,
]
2024-07-01 11:31:27 +01:00
class ConfigureDatabaseClientAction(AbstractAction):
"""Action which sets config parameters for a database client on a node."""
class _Opts(BaseModel):
"""Schema for options that can be passed to this action."""
2024-07-01 14:41:41 +01:00
model_config = ConfigDict(extra="forbid")
2024-07-01 11:31:27 +01:00
server_ip_address: Optional[str] = None
server_password: Optional[str] = None
def __init__(self, manager: "ActionManager", **kwargs) -> None:
super().__init__(manager=manager)
def form_request(self, node_id: int, config: Dict) -> RequestFormat:
2024-07-01 11:31:27 +01:00
"""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"]
ConfigureDatabaseClientAction._Opts.model_validate(config) # check that options adhere to schema
return ["network", "node", node_name, "application", "DatabaseClient", "configure", config]
2024-07-01 11:31:27 +01:00
2024-07-01 14:41:41 +01:00
class ConfigureRansomwareScriptAction(AbstractAction):
"""Action which sets config parameters for a ransomware script on a node."""
class _Opts(BaseModel, AbstractAction.ConfigSchema):
2024-07-01 14:41:41 +01:00
"""Schema for options that can be passed to this option."""
model_config = ConfigDict(extra="forbid")
server_ip_address: Optional[str] = None
server_password: Optional[str] = None
payload: Optional[str] = None
def __init__(self, manager: "ActionManager", **kwargs) -> None:
super().__init__(manager=manager)
def form_request(self, config: _Opts) -> RequestFormat:
2024-07-01 14:41:41 +01:00
"""Return the action formatted as a request that can be ingested by the simulation."""
if config.node_name is None:
2024-07-01 14:41:41 +01:00
return ["do_nothing"]
ConfigureRansomwareScriptAction._Opts.model_validate(config) # check that options adhere to schema
return [
"network",
"node",
config.node_name,
"application",
"RansomwareScript",
"configure",
config.model_config,
]
2024-07-01 14:41:41 +01:00
2024-07-01 15:25:20 +01:00
class ConfigureDoSBotAction(AbstractAction):
"""Action which sets config parameters for a DoS bot on a node."""
class _Opts(BaseModel):
2024-08-01 11:08:41 +01:00
"""Schema for options that can be passed to this action."""
2024-07-01 15:25:20 +01:00
model_config = ConfigDict(extra="forbid")
target_ip_address: Optional[str] = None
target_port: Optional[str] = None
payload: Optional[str] = None
repeat: Optional[bool] = None
port_scan_p_of_success: Optional[float] = None
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, node_id: int, config: Dict) -> RequestFormat:
2024-07-01 15:25:20 +01:00
"""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"]
self._Opts.model_validate(config) # check that options adhere to schema
return ["network", "node", node_name, "application", "DoSBot", "configure", config]
2024-07-01 15:25:20 +01:00
class NodeApplicationRemoveAction(AbstractAction):
"""Action which removes/uninstalls an application."""
def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None:
super().__init__(manager=manager)
self.shape: Dict[str, int] = {"node_id": num_nodes}
def form_request(self, node_id: int, application_name: str) -> 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)
if node_name is None:
return ["do_nothing"]
return ["network", "node", node_name, "software_manager", "application", "uninstall", application_name]
2023-10-02 17:21:43 +01:00
class NodeFolderAbstractAction(AbstractAction):
2023-10-19 01:56:40 +01:00
"""
Base class for folder actions.
Any action which applies to a folder and uses node_id and folder_id as its only two parameters can inherit from
this base class.
"""
2023-10-02 17:21:43 +01:00
@abstractmethod
2023-10-09 18:35:30 +01:00
def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, **kwargs) -> None:
2023-10-02 17:21:43 +01:00
super().__init__(manager=manager)
2023-10-06 10:36:29 +01:00
self.shape: Dict[str, int] = {"node_id": num_nodes, "folder_id": num_folders}
2023-11-27 13:28:11 +00:00
self.verb: str # define but don't initialise: defends against children classes not defining this
2023-10-02 17:21:43 +01:00
def form_request(self, node_id: int, folder_id: int) -> RequestFormat:
2023-10-19 01:56:40 +01:00
"""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)
folder_name = self.manager.get_folder_name_by_idx(node_idx=node_id, folder_idx=folder_id)
if node_name is None or folder_name is None:
2023-10-02 17:21:43 +01:00
return ["do_nothing"]
return ["network", "node", node_name, "file_system", "folder", folder_name, self.verb]
2023-10-06 10:36:29 +01:00
2023-10-02 17:21:43 +01:00
class NodeFolderScanAction(NodeFolderAbstractAction):
2023-10-19 01:56:40 +01:00
"""Action which scans a folder."""
def __init__(self, manager: "ActionManager", node_name: str, folder_name: str, **kwargs) -> None:
super().__init__(manager, node_name=node_name, folder_name=folder_name, **kwargs)
2023-10-06 10:36:29 +01:00
self.verb: str = "scan"
2023-10-02 17:21:43 +01:00
class NodeFolderCheckhashAction(NodeFolderAbstractAction):
2023-10-19 01:56:40 +01:00
"""Action which checks the hash of a folder."""
2023-10-09 18:35:30 +01:00
def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, **kwargs) -> None:
2023-10-06 20:32:52 +01:00
super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, **kwargs)
2023-10-06 10:36:29 +01:00
self.verb: str = "checkhash"
2023-10-02 17:21:43 +01:00
class NodeFolderRepairAction(NodeFolderAbstractAction):
2023-10-19 01:56:40 +01:00
"""Action which repairs a folder."""
2023-10-09 18:35:30 +01:00
def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, **kwargs) -> None:
2023-10-06 20:32:52 +01:00
super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, **kwargs)
2023-10-06 10:36:29 +01:00
self.verb: str = "repair"
2023-10-02 17:21:43 +01:00
class NodeFolderRestoreAction(NodeFolderAbstractAction):
2023-10-19 01:56:40 +01:00
"""Action which restores a folder."""
2023-10-09 18:35:30 +01:00
def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, **kwargs) -> None:
2023-10-06 20:32:52 +01:00
super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, **kwargs)
2023-10-06 10:36:29 +01:00
self.verb: str = "restore"
2023-10-02 17:21:43 +01:00
2024-05-17 17:05:14 +01:00
2024-05-20 13:10:21 +01:00
class NodeFileCreateAction(AbstractAction):
"""Action which creates a new file in a given 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)
self.verb: str = "create"
def form_request(
self, node_id: int, folder_name: str, file_name: str, force: Optional[bool] = False
) -> RequestFormat:
2024-05-20 13:10:21 +01:00
"""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)
if node_name is None or folder_name is None or file_name is None:
return ["do_nothing"]
return ["network", "node", node_name, "file_system", "create", "file", folder_name, file_name, force]
2024-05-20 13:10:21 +01:00
class NodeFolderCreateAction(AbstractAction):
"""Action which creates a new 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)
self.verb: str = "create"
def form_request(self, node_id: int, folder_name: str) -> RequestFormat:
2024-05-20 13:10:21 +01:00
"""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)
if node_name is None or folder_name is None:
return ["do_nothing"]
return ["network", "node", node_name, "file_system", "create", "folder", folder_name]
2023-10-02 17:21:43 +01:00
class NodeFileAbstractAction(AbstractAction):
2023-10-19 01:56:40 +01:00
"""Abstract base class for file actions.
Any action which applies to a file and uses node_id, folder_id, and file_id as its only three parameters can inherit
from this base class.
"""
2023-10-02 17:21:43 +01:00
@abstractmethod
2023-10-06 10:36:29 +01:00
def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None:
2023-10-02 17:21:43 +01:00
super().__init__(manager=manager)
2023-10-06 10:36:29 +01:00
self.shape: Dict[str, int] = {"node_id": num_nodes, "folder_id": num_folders, "file_id": num_files}
2023-11-27 13:28:11 +00:00
self.verb: str # define but don't initialise: defends against children classes not defining this
2023-10-02 17:21:43 +01:00
def form_request(self, node_id: int, folder_id: int, file_id: int) -> RequestFormat:
2023-10-19 01:56:40 +01:00
"""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)
folder_name = self.manager.get_folder_name_by_idx(node_idx=node_id, folder_idx=folder_id)
file_name = self.manager.get_file_name_by_idx(node_idx=node_id, folder_idx=folder_id, file_idx=file_id)
if node_name is None or folder_name is None or file_name is None:
2023-10-02 17:21:43 +01:00
return ["do_nothing"]
2024-01-31 13:29:46 +00:00
return ["network", "node", node_name, "file_system", "folder", folder_name, "file", file_name, self.verb]
2023-10-06 10:36:29 +01:00
2023-10-02 17:21:43 +01:00
class NodeFileScanAction(NodeFileAbstractAction):
2023-10-19 01:56:40 +01:00
"""Action which scans a file."""
2023-10-02 17:21:43 +01:00
def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None:
2023-10-06 20:32:52 +01:00
super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, num_files=num_files, **kwargs)
2023-11-27 13:28:11 +00:00
self.verb: str = "scan"
2023-10-02 17:21:43 +01:00
2023-10-06 10:36:29 +01:00
2023-10-02 17:21:43 +01:00
class NodeFileCheckhashAction(NodeFileAbstractAction):
2023-10-19 01:56:40 +01:00
"""Action which checks the hash of a file."""
2023-10-02 17:21:43 +01:00
def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None:
2023-10-06 20:32:52 +01:00
super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, num_files=num_files, **kwargs)
2023-11-27 13:28:11 +00:00
self.verb: str = "checkhash"
2023-10-02 17:21:43 +01:00
2023-10-06 10:36:29 +01:00
2023-10-02 17:21:43 +01:00
class NodeFileDeleteAction(NodeFileAbstractAction):
2023-10-19 01:56:40 +01:00
"""Action which deletes a file."""
2023-10-02 17:21:43 +01:00
def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None:
2023-10-06 20:32:52 +01:00
super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, num_files=num_files, **kwargs)
2023-11-27 13:28:11 +00:00
self.verb: str = "delete"
2023-10-02 17:21:43 +01:00
def form_request(self, node_id: int, folder_id: int, file_id: int) -> RequestFormat:
2024-01-15 09:48:14 +00:00
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
2024-01-31 10:48:40 +00:00
node_name = self.manager.get_node_name_by_idx(node_id)
folder_name = self.manager.get_folder_name_by_idx(node_idx=node_id, folder_idx=folder_id)
file_name = self.manager.get_file_name_by_idx(node_idx=node_id, folder_idx=folder_id, file_idx=file_id)
if node_name is None or folder_name is None or file_name is None:
2024-01-15 09:48:14 +00:00
return ["do_nothing"]
2024-01-31 13:50:20 +00:00
return ["network", "node", node_name, "file_system", "delete", "file", folder_name, file_name]
2024-01-15 09:48:14 +00:00
2023-10-06 10:36:29 +01:00
2023-10-02 17:21:43 +01:00
class NodeFileRepairAction(NodeFileAbstractAction):
2023-10-19 01:56:40 +01:00
"""Action which repairs a file."""
2023-10-02 17:21:43 +01:00
def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None:
2023-10-06 20:32:52 +01:00
super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, num_files=num_files, **kwargs)
2023-11-27 13:28:11 +00:00
self.verb: str = "repair"
2023-10-02 17:21:43 +01:00
2023-10-06 10:36:29 +01:00
2023-10-02 17:21:43 +01:00
class NodeFileRestoreAction(NodeFileAbstractAction):
2023-10-19 01:56:40 +01:00
"""Action which restores a file."""
2023-10-02 17:21:43 +01:00
def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None:
2023-10-06 20:32:52 +01:00
super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, num_files=num_files, **kwargs)
2023-11-27 13:28:11 +00:00
self.verb: str = "restore"
2023-10-02 17:21:43 +01:00
2023-10-06 10:36:29 +01:00
2023-10-06 20:32:52 +01:00
class NodeFileCorruptAction(NodeFileAbstractAction):
2023-10-19 01:56:40 +01:00
"""Action which corrupts a file."""
2023-10-06 20:32:52 +01:00
def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None:
super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, num_files=num_files, **kwargs)
2023-11-27 13:28:11 +00:00
self.verb: str = "corrupt"
2023-10-06 20:32:52 +01:00
2024-05-20 13:10:21 +01:00
class NodeFileAccessAction(AbstractAction):
"""Action which increases a file's access count."""
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)
2024-05-20 15:02:05 +01:00
self.verb: str = "access"
2024-05-20 13:10:21 +01:00
def form_request(self, node_id: int, folder_name: str, file_name: str) -> RequestFormat:
2024-05-20 13:10:21 +01:00
"""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)
if node_name is None or folder_name is None or file_name is None:
return ["do_nothing"]
return ["network", "node", node_name, "file_system", "access", folder_name, file_name]
2023-10-02 17:21:43 +01:00
class NodeAbstractAction(AbstractAction):
2023-10-19 01:56:40 +01:00
"""
Abstract base class for node actions.
Any action which applies to a node and uses node_id as its only parameter can inherit from this base class.
"""
2023-10-02 17:21:43 +01:00
@abstractmethod
def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None:
super().__init__(manager=manager)
2023-10-06 10:36:29 +01:00
self.shape: Dict[str, int] = {"node_id": num_nodes}
2023-11-27 13:28:11 +00:00
self.verb: str # define but don't initialise: defends against children classes not defining this
2023-10-02 17:21:43 +01:00
def form_request(self, node_id: int) -> RequestFormat:
2023-10-19 01:56:40 +01:00
"""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, self.verb]
2023-10-02 17:21:43 +01:00
2023-10-06 10:36:29 +01:00
2023-10-02 17:21:43 +01:00
class NodeOSScanAction(NodeAbstractAction):
2023-10-19 01:56:40 +01:00
"""Action which scans a node's OS."""
2023-10-02 17:21:43 +01:00
def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None:
2023-10-06 20:32:52 +01:00
super().__init__(manager=manager, num_nodes=num_nodes)
2023-11-27 13:28:11 +00:00
self.verb: str = "scan"
2023-10-06 10:36:29 +01:00
2023-10-02 17:21:43 +01:00
class NodeShutdownAction(NodeAbstractAction):
2023-10-19 01:56:40 +01:00
"""Action which shuts down a node."""
2023-10-02 17:21:43 +01:00
def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None:
2023-10-06 20:32:52 +01:00
super().__init__(manager=manager, num_nodes=num_nodes)
2023-11-27 13:28:11 +00:00
self.verb: str = "shutdown"
2023-10-06 10:36:29 +01:00
2023-10-02 17:21:43 +01:00
class NodeStartupAction(NodeAbstractAction):
2023-10-19 01:56:40 +01:00
"""Action which starts up a node."""
2023-10-02 17:21:43 +01:00
def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None:
2023-10-06 20:32:52 +01:00
super().__init__(manager=manager, num_nodes=num_nodes)
2023-11-27 13:28:11 +00:00
self.verb: str = "startup"
2023-10-06 10:36:29 +01:00
2023-10-02 17:21:43 +01:00
class NodeResetAction(NodeAbstractAction):
2023-10-19 01:56:40 +01:00
"""Action which resets a node."""
2023-10-02 17:21:43 +01:00
def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None:
2023-10-06 20:32:52 +01:00
super().__init__(manager=manager, num_nodes=num_nodes)
2023-11-27 13:28:11 +00:00
self.verb: str = "reset"
2023-10-06 10:36:29 +01:00
2023-10-02 17:21:43 +01:00
class RouterACLAddRuleAction(AbstractAction):
2023-10-19 01:56:40 +01:00
"""Action which adds a rule to a router's ACL."""
class ACLRuleOptions(BaseModel):
"""Validator for ACL_ADD_RULE options."""
target_router: str
"""On which router to add the rule, must be specified."""
position: int
"""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)
"""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)
"""Rule source port. By default, all source ports."""
dest_ip_id: int = Field(default=1, ge=1)
"""Rule destination IP address. By default, all ip addresses."""
dest_wildcard_id: int = Field(default=0, ge=0)
"""Rule destination IP wildcard. By default, use the wildcard at index 0 from action manager."""
dest_port_id: int = Field(default=1, ge=1)
"""Rule destination port. By default, all destination ports."""
protocol_id: int = Field(default=1, ge=1)
"""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_id",
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
2023-10-06 10:36:29 +01:00
def __init__(
self,
manager: "ActionManager",
max_acl_rules: int,
num_ips: int,
num_ports: int,
num_protocols: int,
**kwargs,
) -> None:
"""Init method for RouterACLAddRuleAction.
2023-10-19 01:56:40 +01:00
:param manager: Reference to the ActionManager which created this action.
:type manager: ActionManager
:param max_acl_rules: Maximum number of ACL rules that can be added to the router.
:type max_acl_rules: int
:param num_ips: Number of IP addresses in the simulation.
:type num_ips: int
:param num_ports: Number of ports in the simulation.
:type num_ports: int
:param num_protocols: Number of protocols in the simulation.
:type num_protocols: int
"""
2023-10-02 17:21:43 +01:00
super().__init__(manager=manager)
2023-10-09 17:29:50 +01:00
num_permissions = 3
2023-10-06 10:36:29 +01:00
self.shape: Dict[str, int] = {
"position": max_acl_rules,
"permission": num_permissions,
2023-10-09 17:29:50 +01:00
"source_ip_id": num_ips,
"dest_ip_id": num_ips,
"source_port_id": num_ports,
"dest_port_id": num_ports,
"protocol_id": num_protocols,
2023-10-06 10:36:29 +01:00
}
def form_request(
2023-10-19 01:56:40 +01:00
self,
target_router: str,
2023-10-19 01:56:40 +01:00
position: int,
permission: int,
source_ip_id: int,
2024-04-15 11:50:08 +01:00
source_wildcard_id: int,
2023-10-19 01:56:40 +01:00
dest_ip_id: int,
2024-04-15 11:50:08 +01:00
dest_wildcard_id: int,
2023-10-19 01:56:40 +01:00
source_port_id: int,
dest_port_id: int,
protocol_id: int,
2023-10-06 10:36:29 +01:00
) -> List[str]:
2023-10-19 01:56:40 +01:00
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
# Validate incoming data.
parsed_options = RouterACLAddRuleAction.ACLRuleOptions(
target_router=target_router,
position=position,
permission=permission,
source_ip_id=source_ip_id,
source_wildcard_id=source_wildcard_id,
dest_ip_id=dest_ip_id,
dest_wildcard_id=dest_wildcard_id,
source_port_id=source_port_id,
dest_port_id=dest_port_id,
protocol_id=protocol_id,
)
if parsed_options.permission == 1:
2024-03-29 16:30:39 +00:00
permission_str = "PERMIT"
elif parsed_options.permission == 2:
2023-10-09 17:29:50 +01:00
permission_str = "DENY"
else:
2023-11-30 21:11:35 +00:00
_LOGGER.warning(f"{self.__class__} received permission {permission}, expected 0 or 1.")
2023-10-09 17:29:50 +01:00
if parsed_options.protocol_id == 1:
2023-10-09 17:29:50 +01:00
protocol = "ALL"
else:
protocol = self.manager.get_internet_protocol_by_idx(parsed_options.protocol_id - 2)
2023-10-09 17:29:50 +01:00
# subtract 2 to account for UNUSED=0 and ALL=1.
if parsed_options.source_ip_id == 1:
2023-10-09 17:29:50 +01:00
src_ip = "ALL"
else:
src_ip = self.manager.get_ip_address_by_idx(parsed_options.source_ip_id - 2)
2023-10-09 17:29:50 +01:00
# subtract 2 to account for UNUSED=0, and ALL=1
src_wildcard = self.manager.get_wildcard_by_idx(parsed_options.source_wildcard_id)
if parsed_options.source_port_id == 1:
2023-10-09 17:29:50 +01:00
src_port = "ALL"
else:
src_port = self.manager.get_port_by_idx(parsed_options.source_port_id - 2)
2023-10-09 17:29:50 +01:00
# subtract 2 to account for UNUSED=0, and ALL=1
if parsed_options.dest_ip_id == 1:
2023-10-09 17:29:50 +01:00
dst_ip = "ALL"
else:
dst_ip = self.manager.get_ip_address_by_idx(parsed_options.dest_ip_id - 2)
2023-10-09 17:29:50 +01:00
# subtract 2 to account for UNUSED=0, and ALL=1
dst_wildcard = self.manager.get_wildcard_by_idx(parsed_options.dest_wildcard_id)
2023-10-09 17:29:50 +01:00
if parsed_options.dest_port_id == 1:
2023-10-09 17:29:50 +01:00
dst_port = "ALL"
else:
dst_port = self.manager.get_port_by_idx(parsed_options.dest_port_id - 2)
2023-10-09 17:29:50 +01:00
# subtract 2 to account for UNUSED=0, and ALL=1
2023-10-04 09:07:04 +01:00
return [
2023-10-06 10:36:29 +01:00
"network",
"node",
target_router,
2023-10-06 10:36:29 +01:00
"acl",
"add_rule",
2023-10-09 17:29:50 +01:00
permission_str,
2023-10-04 09:07:04 +01:00
protocol,
2024-03-08 15:57:43 +00:00
str(src_ip),
2024-04-15 11:50:08 +01:00
src_wildcard,
2023-10-04 09:07:04 +01:00
src_port,
2024-03-08 15:57:43 +00:00
str(dst_ip),
2024-04-15 11:50:08 +01:00
dst_wildcard,
2023-10-04 09:07:04 +01:00
dst_port,
2023-10-06 10:36:29 +01:00
position,
2023-10-04 09:07:04 +01:00
]
2023-10-02 17:21:43 +01:00
class RouterACLRemoveRuleAction(AbstractAction):
2023-10-19 01:56:40 +01:00
"""Action which removes a rule from a router's ACL."""
def __init__(self, manager: "ActionManager", max_acl_rules: int, **kwargs) -> None:
"""Init method for RouterACLRemoveRuleAction.
2023-10-19 01:56:40 +01:00
:param manager: Reference to the ActionManager which created this action.
:type manager: ActionManager
:param max_acl_rules: Maximum number of ACL rules that can be added to the router.
:type max_acl_rules: int
"""
2023-10-06 10:36:29 +01:00
super().__init__(manager=manager)
self.shape: Dict[str, int] = {"position": max_acl_rules}
def form_request(self, target_router: str, position: int) -> RequestFormat:
2023-10-19 01:56:40 +01:00
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
return ["network", "node", target_router, "acl", "remove_rule", position]
2023-10-06 10:36:29 +01:00
2024-03-29 16:30:39 +00:00
class FirewallACLAddRuleAction(AbstractAction):
"""Action which adds a rule to a firewall port's ACL."""
def __init__(
self,
manager: "ActionManager",
max_acl_rules: int,
num_ips: int,
num_ports: int,
num_protocols: int,
**kwargs,
) -> None:
"""Init method for FirewallACLAddRuleAction.
:param manager: Reference to the ActionManager which created this action.
:type manager: ActionManager
:param max_acl_rules: Maximum number of ACL rules that can be added to the router.
:type max_acl_rules: int
:param num_ips: Number of IP addresses in the simulation.
:type num_ips: int
:param num_ports: Number of ports in the simulation.
:type num_ports: int
:param num_protocols: Number of protocols in the simulation.
:type num_protocols: int
"""
super().__init__(manager=manager)
num_permissions = 3
self.shape: Dict[str, int] = {
"position": max_acl_rules,
"permission": num_permissions,
"source_ip_id": num_ips,
"dest_ip_id": num_ips,
"source_port_id": num_ports,
"dest_port_id": num_ports,
"protocol_id": num_protocols,
}
def form_request(
self,
target_firewall_nodename: str,
firewall_port_name: str,
firewall_port_direction: str,
position: int,
permission: int,
source_ip_id: int,
2024-04-15 11:50:08 +01:00
source_wildcard_id: int,
2024-03-29 16:30:39 +00:00
dest_ip_id: int,
2024-04-15 11:50:08 +01:00
dest_wildcard_id: int,
2024-03-29 16:30:39 +00:00
source_port_id: int,
dest_port_id: int,
protocol_id: int,
) -> List[str]:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
if permission == 0:
permission_str = "UNUSED"
return ["do_nothing"] # NOT SUPPORTED, JUST DO NOTHING IF WE COME ACROSS THIS
elif permission == 1:
permission_str = "PERMIT"
elif permission == 2:
permission_str = "DENY"
else:
_LOGGER.warning(f"{self.__class__} received permission {permission}, expected 0 or 1.")
if protocol_id == 0:
return ["do_nothing"] # NOT SUPPORTED, JUST DO NOTHING IF WE COME ACROSS THIS
if protocol_id == 1:
protocol = "ALL"
else:
protocol = self.manager.get_internet_protocol_by_idx(protocol_id - 2)
# subtract 2 to account for UNUSED=0 and ALL=1.
if source_ip_id == 0:
return ["do_nothing"] # invalid formulation
elif source_ip_id == 1:
src_ip = "ALL"
else:
src_ip = self.manager.get_ip_address_by_idx(source_ip_id - 2)
# subtract 2 to account for UNUSED=0, and ALL=1
if source_port_id == 0:
return ["do_nothing"] # invalid formulation
elif source_port_id == 1:
src_port = "ALL"
else:
src_port = self.manager.get_port_by_idx(source_port_id - 2)
# subtract 2 to account for UNUSED=0, and ALL=1
2024-04-15 11:50:08 +01:00
if dest_ip_id == 0:
2024-03-29 16:30:39 +00:00
return ["do_nothing"] # invalid formulation
elif dest_ip_id == 1:
dst_ip = "ALL"
else:
dst_ip = self.manager.get_ip_address_by_idx(dest_ip_id - 2)
# subtract 2 to account for UNUSED=0, and ALL=1
if dest_port_id == 0:
return ["do_nothing"] # invalid formulation
elif dest_port_id == 1:
dst_port = "ALL"
else:
dst_port = self.manager.get_port_by_idx(dest_port_id - 2)
# subtract 2 to account for UNUSED=0, and ALL=1
2024-04-15 11:50:08 +01:00
src_wildcard = self.manager.get_wildcard_by_idx(source_wildcard_id)
dst_wildcard = self.manager.get_wildcard_by_idx(dest_wildcard_id)
2024-03-29 16:30:39 +00:00
return [
"network",
"node",
target_firewall_nodename,
firewall_port_name,
firewall_port_direction,
"acl",
"add_rule",
permission_str,
protocol,
str(src_ip),
2024-04-15 11:50:08 +01:00
src_wildcard,
2024-03-29 16:30:39 +00:00
src_port,
str(dst_ip),
2024-04-15 11:50:08 +01:00
dst_wildcard,
2024-03-29 16:30:39 +00:00
dst_port,
position,
]
class FirewallACLRemoveRuleAction(AbstractAction):
"""Action which removes a rule from a firewall port's ACL."""
def __init__(self, manager: "ActionManager", max_acl_rules: int, **kwargs) -> None:
"""Init method for RouterACLRemoveRuleAction.
:param manager: Reference to the ActionManager which created this action.
:type manager: ActionManager
:param max_acl_rules: Maximum number of ACL rules that can be added to the router.
:type max_acl_rules: int
"""
super().__init__(manager=manager)
self.shape: Dict[str, int] = {"position": max_acl_rules}
def form_request(
self, target_firewall_nodename: str, firewall_port_name: str, firewall_port_direction: str, position: int
) -> List[str]:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
return [
"network",
"node",
target_firewall_nodename,
firewall_port_name,
firewall_port_direction,
"acl",
"remove_rule",
position,
]
class HostNICAbstractAction(AbstractAction):
2023-10-19 01:56:40 +01:00
"""
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.
"""
2023-10-06 10:36:29 +01:00
def __init__(self, manager: "ActionManager", num_nodes: int, max_nics_per_node: int, **kwargs) -> None:
"""Init method for HostNICAbstractAction.
2023-10-19 01:56:40 +01:00
:param manager: Reference to the ActionManager which created this action.
:type manager: ActionManager
:param num_nodes: Number of nodes in the simulation.
:type num_nodes: int
:param max_nics_per_node: Maximum number of NICs per node.
:type max_nics_per_node: int
"""
2023-10-06 10:36:29 +01:00
super().__init__(manager=manager)
self.shape: Dict[str, int] = {"node_id": num_nodes, "nic_id": max_nics_per_node}
2023-11-27 13:28:11 +00:00
self.verb: str # define but don't initialise: defends against children classes not defining this
2023-10-02 17:21:43 +01:00
def form_request(self, node_id: int, nic_id: int) -> RequestFormat:
2023-10-19 01:56:40 +01:00
"""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_idx=node_id)
nic_num = self.manager.get_nic_num_by_idx(node_idx=node_id, nic_idx=nic_id)
if node_name is None or nic_num is None:
2023-10-09 17:29:50 +01:00
return ["do_nothing"]
2024-02-08 15:27:02 +00:00
return ["network", "node", node_name, "network_interface", nic_num, self.verb]
2023-10-06 10:36:29 +01:00
class HostNICEnableAction(HostNICAbstractAction):
2023-10-19 01:56:40 +01:00
"""Action which enables a NIC."""
2023-10-06 10:36:29 +01:00
def __init__(self, manager: "ActionManager", num_nodes: int, max_nics_per_node: int, **kwargs) -> None:
2023-10-09 17:29:50 +01:00
super().__init__(manager=manager, num_nodes=num_nodes, max_nics_per_node=max_nics_per_node, **kwargs)
2023-11-27 13:28:11 +00:00
self.verb: str = "enable"
2023-10-06 10:36:29 +01:00
2023-10-09 17:29:50 +01:00
class HostNICDisableAction(HostNICAbstractAction):
2023-10-19 01:56:40 +01:00
"""Action which disables a NIC."""
2023-10-09 17:29:50 +01:00
def __init__(self, manager: "ActionManager", num_nodes: int, max_nics_per_node: int, **kwargs) -> None:
super().__init__(manager=manager, num_nodes=num_nodes, max_nics_per_node=max_nics_per_node, **kwargs)
2023-11-27 13:28:11 +00:00
self.verb: str = "disable"
2023-10-09 17:29:50 +01:00
2023-10-09 18:35:30 +01:00
class NetworkPortEnableAction(AbstractAction):
"""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.
: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}
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):
"""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.
: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}
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"]
class NodeNMAPPingScanAction(AbstractAction):
"""Action which performs an NMAP ping scan."""
def __init__(self, manager: "ActionManager", **kwargs) -> None:
super().__init__(manager=manager)
def form_request(
self, source_node: str, target_ip_address: Union[str, List[str]], show: Optional[bool] = False
) -> List[str]: # noqa
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
return [
"network",
"node",
source_node,
"application",
"NMAP",
"ping_scan",
{"target_ip_address": target_ip_address, "show": show},
]
class NodeNMAPPortScanAction(AbstractAction):
"""Action which performs an NMAP port scan."""
def __init__(self, manager: "ActionManager", **kwargs) -> None:
super().__init__(manager=manager)
def form_request(
self,
source_node: str,
target_ip_address: Union[str, List[str]],
target_protocol: Optional[Union[str, List[str]]] = None,
target_port: Optional[Union[str, List[str]]] = None,
show: Optional[bool] = False,
) -> List[str]: # noqa
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
return [
"network",
"node",
source_node,
"application",
"NMAP",
"port_scan",
{
"target_ip_address": target_ip_address,
"target_port": target_port,
"target_protocol": target_protocol,
"show": show,
},
]
class NodeNetworkServiceReconAction(AbstractAction):
"""Action which performs an NMAP network service recon (ping scan followed by port scan)."""
def __init__(self, manager: "ActionManager", **kwargs) -> None:
super().__init__(manager=manager)
def form_request(
self,
source_node: str,
target_ip_address: Union[str, List[str]],
target_protocol: Optional[Union[str, List[str]]] = None,
target_port: Optional[Union[str, List[str]]] = None,
show: Optional[bool] = False,
) -> List[str]: # noqa
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
return [
"network",
"node",
source_node,
"application",
"NMAP",
"network_service_recon",
{
"target_ip_address": target_ip_address,
"target_port": target_port,
"target_protocol": target_protocol,
"show": show,
},
]
class ConfigureC2BeaconAction(AbstractAction):
"""Action which configures a C2 Beacon based on the parameters given."""
class _Opts(BaseModel):
"""Schema for options that can be passed to this action."""
c2_server_ip_address: str
keep_alive_frequency: int = Field(default=5, ge=1)
masquerade_protocol: str = Field(default="TCP")
masquerade_port: str = Field(default="HTTP")
@field_validator(
"c2_server_ip_address",
"keep_alive_frequency",
"masquerade_protocol",
"masquerade_port",
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
def __init__(self, manager: "ActionManager", **kwargs) -> None:
super().__init__(manager=manager)
def form_request(self, node_id: int, config: 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"]
config = ConfigureC2BeaconAction._Opts(
c2_server_ip_address=config["c2_server_ip_address"],
keep_alive_frequency=config["keep_alive_frequency"],
masquerade_port=config["masquerade_port"],
masquerade_protocol=config["masquerade_protocol"],
)
ConfigureC2BeaconAction._Opts.model_validate(config) # check that options adhere to schema
return ["network", "node", node_name, "application", "C2Beacon", "configure", config.__dict__]
class NodeAccountsChangePasswordAction(AbstractAction):
"""Action which changes the password for a user."""
def __init__(self, manager: "ActionManager", **kwargs) -> None:
super().__init__(manager=manager)
def form_request(self, node_id: str, username: str, current_password: str, new_password: str) -> 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",
"UserManager",
"change_password",
username,
current_password,
new_password,
]
class NodeSessionsRemoteLoginAction(AbstractAction):
"""Action which performs a remote session login."""
def __init__(self, manager: "ActionManager", **kwargs) -> None:
super().__init__(manager=manager)
def form_request(self, node_id: str, username: str, password: str, remote_ip: str) -> 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",
"ssh_to_remote",
username,
password,
remote_ip,
]
class NodeSessionsRemoteLogoutAction(AbstractAction):
"""Action which performs a remote session logout."""
def __init__(self, manager: "ActionManager", **kwargs) -> None:
super().__init__(manager=manager)
2024-08-20 11:15:04 +01:00
def form_request(self, node_id: str, remote_ip: str) -> 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)
2024-08-20 11:15:04 +01:00
return ["network", "node", node_name, "service", "Terminal", "remote_logoff", remote_ip]
class RansomwareConfigureC2ServerAction(AbstractAction):
"""Action which sends a command from the C2 Server to the C2 Beacon which configures a local RansomwareScript."""
def __init__(self, manager: "ActionManager", **kwargs) -> None:
super().__init__(manager=manager)
def form_request(self, node_id: int, config: 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"]
# Using the ransomware scripts model to validate.
ConfigureRansomwareScriptAction._Opts.model_validate(config) # check that options adhere to schema
return ["network", "node", node_name, "application", "C2Server", "ransomware_configure", config]
class RansomwareLaunchC2ServerAction(AbstractAction):
"""Action which causes the C2 Server to send a command to the C2 Beacon to launch the RansomwareScript."""
def __init__(self, manager: "ActionManager", **kwargs) -> None:
super().__init__(manager=manager)
def form_request(self, node_id: int) -> 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"]
# This action currently doesn't require any further configuration options.
return ["network", "node", node_name, "application", "C2Server", "ransomware_launch"]
class ExfiltrationC2ServerAction(AbstractAction):
"""Action which exfiltrates a target file from a certain node onto the C2 beacon and then the C2 Server."""
class _Opts(BaseModel):
"""Schema for options that can be passed to this action."""
username: Optional[str]
password: Optional[str]
target_ip_address: str
target_file_name: str
target_folder_name: str
exfiltration_folder_name: Optional[str]
def __init__(self, manager: "ActionManager", **kwargs) -> None:
super().__init__(manager=manager)
def form_request(
self,
node_id: int,
account: dict,
target_ip_address: str,
target_file_name: str,
target_folder_name: str,
exfiltration_folder_name: Optional[str],
) -> 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 = {
"target_file_name": target_file_name,
"target_folder_name": target_folder_name,
"exfiltration_folder_name": exfiltration_folder_name,
"target_ip_address": target_ip_address,
"username": account["username"],
"password": account["password"],
}
ExfiltrationC2ServerAction._Opts.model_validate(command_model)
return ["network", "node", node_name, "application", "C2Server", "exfiltrate", command_model]
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."""
def __init__(self, manager: "ActionManager", **kwargs) -> None:
super().__init__(manager=manager)
def form_request(self, node_id: int) -> 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"]
# This action currently doesn't require any further configuration options.
return ["network", "node", node_name, "application", "C2Server", "ransomware_launch"]