2024-06-05 09:11:37 +01:00
|
|
|
# © 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
|
2024-10-17 12:22:30 +01:00
|
|
|
from typing import Dict, List, Literal, Optional, Union
|
2023-09-21 10:13:01 +01:00
|
|
|
|
2024-07-01 14:41:41 +01:00
|
|
|
from pydantic import BaseModel, ConfigDict, Field, field_validator, ValidationInfo
|
2023-09-21 10:13:01 +01:00
|
|
|
|
2023-10-09 17:29:50 +01:00
|
|
|
from primaite import getLogger
|
2024-10-17 12:22:30 +01:00
|
|
|
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
|
|
|
|
|
|
|
|
|
2023-11-23 16:06:19 +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
|
2023-11-23 16:06:19 +00:00
|
|
|
|
2024-07-09 13:13:13 +01:00
|
|
|
def form_request(self, node_id: int, application_id: int) -> RequestFormat:
|
2023-11-23 16:06:19 +00:00
|
|
|
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
|
2023-12-18 14:03:47 +00:00
|
|
|
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:
|
2023-11-23 16:06:19 +00:00
|
|
|
return ["do_nothing"]
|
2023-12-18 14:03:47 +00:00
|
|
|
return ["network", "node", node_name, "application", application_name, self.verb]
|
2023-11-23 16:06:19 +00:00
|
|
|
|
|
|
|
|
|
|
|
|
|
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"
|
2023-11-23 16:06:19 +00:00
|
|
|
|
|
|
|
|
|
2024-03-25 16:58:27 +00:00
|
|
|
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"
|
2024-03-25 16:58:27 +00:00
|
|
|
|
|
|
|
|
|
2024-03-27 17:07:12 +00:00
|
|
|
class NodeApplicationInstallAction(AbstractAction):
|
|
|
|
|
"""Action which installs an application."""
|
|
|
|
|
|
2024-03-28 14:08:08 +00:00
|
|
|
def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None:
|
2024-03-27 17:07:12 +00:00
|
|
|
super().__init__(manager=manager)
|
|
|
|
|
self.shape: Dict[str, int] = {"node_id": num_nodes}
|
|
|
|
|
|
2024-07-09 13:13:13 +01:00
|
|
|
def form_request(self, node_id: int, application_name: str) -> RequestFormat:
|
2024-03-27 17:07:12 +00: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:
|
|
|
|
|
return ["do_nothing"]
|
|
|
|
|
return [
|
|
|
|
|
"network",
|
|
|
|
|
"node",
|
|
|
|
|
node_name,
|
|
|
|
|
"software_manager",
|
|
|
|
|
"application",
|
|
|
|
|
"install",
|
2024-03-28 14:08:08 +00:00
|
|
|
application_name,
|
2024-03-27 17:07:12 +00:00
|
|
|
]
|
|
|
|
|
|
|
|
|
|
|
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)
|
|
|
|
|
|
2024-07-02 11:10:19 +01:00
|
|
|
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"]
|
2024-07-02 11:10:19 +01:00
|
|
|
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."""
|
|
|
|
|
|
2024-11-12 14:49:44 +00:00
|
|
|
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)
|
|
|
|
|
|
2024-11-12 14:49:44 +00:00
|
|
|
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."""
|
2024-11-12 14:49:44 +00:00
|
|
|
if config.node_name is None:
|
2024-07-01 14:41:41 +01:00
|
|
|
return ["do_nothing"]
|
2024-07-02 11:10:19 +01:00
|
|
|
ConfigureRansomwareScriptAction._Opts.model_validate(config) # check that options adhere to schema
|
2024-11-12 14:49:44 +00:00
|
|
|
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)
|
|
|
|
|
|
2024-07-02 11:10:19 +01:00
|
|
|
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"]
|
2024-07-02 11:10:19 +01:00
|
|
|
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
|
|
|
|
|
|
|
|
|
2024-03-27 17:07:12 +00:00
|
|
|
class NodeApplicationRemoveAction(AbstractAction):
|
|
|
|
|
"""Action which removes/uninstalls an application."""
|
|
|
|
|
|
2024-03-28 14:08:08 +00:00
|
|
|
def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None:
|
2024-03-27 17:07:12 +00:00
|
|
|
super().__init__(manager=manager)
|
|
|
|
|
self.shape: Dict[str, int] = {"node_id": num_nodes}
|
|
|
|
|
|
2024-07-09 13:13:13 +01:00
|
|
|
def form_request(self, node_id: int, application_name: str) -> RequestFormat:
|
2024-03-27 17:07:12 +00: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:
|
|
|
|
|
return ["do_nothing"]
|
2024-03-28 14:08:08 +00:00
|
|
|
return ["network", "node", node_name, "software_manager", "application", "uninstall", application_name]
|
2024-03-27 17:07:12 +00:00
|
|
|
|
|
|
|
|
|
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
|
|
|
|
2024-07-09 13:13:13 +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."""
|
2023-12-18 14:03:47 +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)
|
|
|
|
|
if node_name is None or folder_name is None:
|
2023-10-02 17:21:43 +01:00
|
|
|
return ["do_nothing"]
|
2023-12-18 14:03:47 +00:00
|
|
|
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."""
|
|
|
|
|
|
2024-11-13 10:40:51 +00:00
|
|
|
def __init__(self, manager: "ActionManager", node_name: str, folder_name: str, **kwargs) -> None:
|
2024-11-12 14:49:44 +00:00
|
|
|
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):
|
2024-05-17 09:28:03 +01:00
|
|
|
"""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"
|
|
|
|
|
|
2024-07-09 13:13:13 +01:00
|
|
|
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"]
|
2024-06-25 11:04:52 +01:00
|
|
|
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"
|
|
|
|
|
|
2024-07-09 13:13:13 +01:00
|
|
|
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
|
|
|
|
2024-07-09 13:13:13 +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."""
|
2023-12-18 14:03:47 +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:
|
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
|
|
|
|
2024-07-09 13:13:13 +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
|
|
|
|
2024-07-09 13:13:13 +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
|
|
|
|
2024-07-09 13:13:13 +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."""
|
2023-12-18 14:03:47 +00:00
|
|
|
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
|
|
|
|
2024-03-29 11:34:43 +00:00
|
|
|
class RouterACLAddRuleAction(AbstractAction):
|
2023-10-19 01:56:40 +01:00
|
|
|
"""Action which adds a rule to a router's ACL."""
|
|
|
|
|
|
2024-06-25 11:04:52 +01:00
|
|
|
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:
|
2024-03-29 11:34:43 +00:00
|
|
|
"""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,
|
2024-06-25 11:04:52 +01:00
|
|
|
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."""
|
2024-06-25 11:04:52 +01:00
|
|
|
# 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"
|
2024-06-25 11:04:52 +01:00
|
|
|
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
|
|
|
|
2024-06-25 11:04:52 +01:00
|
|
|
if parsed_options.protocol_id == 1:
|
2023-10-09 17:29:50 +01:00
|
|
|
protocol = "ALL"
|
|
|
|
|
else:
|
2024-06-25 11:04:52 +01:00
|
|
|
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.
|
|
|
|
|
|
2024-06-25 11:04:52 +01:00
|
|
|
if parsed_options.source_ip_id == 1:
|
2023-10-09 17:29:50 +01:00
|
|
|
src_ip = "ALL"
|
|
|
|
|
else:
|
2024-06-25 11:04:52 +01:00
|
|
|
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
|
2024-06-25 11:04:52 +01:00
|
|
|
|
|
|
|
|
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:
|
2024-06-25 11:04:52 +01:00
|
|
|
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
|
|
|
|
|
|
2024-06-25 11:04:52 +01:00
|
|
|
if parsed_options.dest_ip_id == 1:
|
2023-10-09 17:29:50 +01:00
|
|
|
dst_ip = "ALL"
|
|
|
|
|
else:
|
2024-06-25 11:04:52 +01:00
|
|
|
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
|
2024-06-25 11:04:52 +01:00
|
|
|
dst_wildcard = self.manager.get_wildcard_by_idx(parsed_options.dest_wildcard_id)
|
2023-10-09 17:29:50 +01:00
|
|
|
|
2024-06-25 11:04:52 +01:00
|
|
|
if parsed_options.dest_port_id == 1:
|
2023-10-09 17:29:50 +01:00
|
|
|
dst_port = "ALL"
|
|
|
|
|
else:
|
2024-06-25 11:04:52 +01:00
|
|
|
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",
|
2024-06-25 11:04:52 +01:00
|
|
|
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
|
|
|
|
|
|
|
|
|
2024-03-29 11:34:43 +00:00
|
|
|
class RouterACLRemoveRuleAction(AbstractAction):
|
2023-10-19 01:56:40 +01:00
|
|
|
"""Action which removes a rule from a router's ACL."""
|
|
|
|
|
|
2024-03-29 11:34:43 +00:00
|
|
|
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}
|
|
|
|
|
|
2024-07-09 13:13:13 +01:00
|
|
|
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."""
|
2024-06-25 11:04:52 +01:00
|
|
|
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,
|
|
|
|
|
]
|
|
|
|
|
|
|
|
|
|
|
2024-03-31 11:59:31 +00:00
|
|
|
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:
|
2024-03-31 11:59:31 +00:00
|
|
|
"""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
|
|
|
|
2024-07-09 13:13:13 +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."""
|
2023-12-18 14:03:47 +00:00
|
|
|
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
|
|
|
|
|
|
|
|
|
2024-03-31 11:59:31 +00: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
|
|
|
|
2024-03-31 11:59:31 +00: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
|
|
|
|
2024-03-31 11:59:31 +00:00
|
|
|
class NetworkPortEnableAction(AbstractAction):
|
|
|
|
|
"""Action which enables are port on a router or a firewall."""
|
2024-03-22 16:35:53 +00:00
|
|
|
|
2024-03-31 11:59:31 +00:00
|
|
|
def __init__(self, manager: "ActionManager", max_nics_per_node: int, **kwargs) -> None:
|
|
|
|
|
"""Init method for NetworkPortEnableAction.
|
2024-03-22 16:35:53 +00:00
|
|
|
|
|
|
|
|
:param max_nics_per_node: Maximum number of NICs per node.
|
|
|
|
|
:type max_nics_per_node: int
|
|
|
|
|
"""
|
|
|
|
|
super().__init__(manager=manager)
|
2024-03-31 11:59:31 +00:00
|
|
|
self.shape: Dict[str, int] = {"port_id": max_nics_per_node}
|
2024-03-22 16:35:53 +00:00
|
|
|
|
2024-07-09 13:13:13 +01:00
|
|
|
def form_request(self, target_nodename: str, port_id: int) -> RequestFormat:
|
2024-03-22 16:35:53 +00:00
|
|
|
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
|
2024-03-31 11:59:31 +00:00
|
|
|
if target_nodename is None or port_id is None:
|
2024-03-22 16:35:53 +00:00
|
|
|
return ["do_nothing"]
|
2024-03-31 11:59:31 +00:00
|
|
|
return ["network", "node", target_nodename, "network_interface", port_id, "enable"]
|
2024-03-22 16:35:53 +00:00
|
|
|
|
|
|
|
|
|
2024-03-31 11:59:31 +00:00
|
|
|
class NetworkPortDisableAction(AbstractAction):
|
|
|
|
|
"""Action which disables are port on a router or a firewall."""
|
2024-03-22 16:35:53 +00:00
|
|
|
|
2024-03-31 11:59:31 +00:00
|
|
|
def __init__(self, manager: "ActionManager", max_nics_per_node: int, **kwargs) -> None:
|
|
|
|
|
"""Init method for NetworkPortDisableAction.
|
2024-03-22 16:35:53 +00:00
|
|
|
|
2024-03-31 11:59:31 +00:00
|
|
|
: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}
|
2024-03-22 16:35:53 +00:00
|
|
|
|
2024-07-09 13:13:13 +01:00
|
|
|
def form_request(self, target_nodename: str, port_id: int) -> RequestFormat:
|
2024-03-31 11:59:31 +00:00
|
|
|
"""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"]
|
2024-03-22 16:35:53 +00:00
|
|
|
|
|
|
|
|
|
2024-05-31 13:53:18 +01:00
|
|
|
class NodeNMAPPingScanAction(AbstractAction):
|
|
|
|
|
"""Action which performs an NMAP ping scan."""
|
|
|
|
|
|
|
|
|
|
def __init__(self, manager: "ActionManager", **kwargs) -> None:
|
|
|
|
|
super().__init__(manager=manager)
|
|
|
|
|
|
2024-06-25 11:04:52 +01:00
|
|
|
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."""
|
2024-05-31 13:53:18 +01:00
|
|
|
return [
|
|
|
|
|
"network",
|
|
|
|
|
"node",
|
|
|
|
|
source_node,
|
|
|
|
|
"application",
|
|
|
|
|
"NMAP",
|
|
|
|
|
"ping_scan",
|
2024-06-25 11:04:52 +01:00
|
|
|
{"target_ip_address": target_ip_address, "show": show},
|
2024-05-31 13:53:18 +01:00
|
|
|
]
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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,
|
2024-06-25 11:04:52 +01:00
|
|
|
show: Optional[bool] = False,
|
2024-05-31 13:53:18 +01:00
|
|
|
) -> 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",
|
2024-06-25 11:04:52 +01:00
|
|
|
{
|
|
|
|
|
"target_ip_address": target_ip_address,
|
|
|
|
|
"target_port": target_port,
|
|
|
|
|
"target_protocol": target_protocol,
|
|
|
|
|
"show": show,
|
|
|
|
|
},
|
2024-05-31 13:53:18 +01:00
|
|
|
]
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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,
|
2024-06-25 11:04:52 +01:00
|
|
|
show: Optional[bool] = False,
|
2024-05-31 13:53:18 +01:00
|
|
|
) -> 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",
|
2024-06-25 11:04:52 +01:00
|
|
|
{
|
|
|
|
|
"target_ip_address": target_ip_address,
|
|
|
|
|
"target_port": target_port,
|
|
|
|
|
"target_protocol": target_protocol,
|
|
|
|
|
"show": show,
|
|
|
|
|
},
|
2024-05-31 13:53:18 +01:00
|
|
|
]
|
|
|
|
|
|
|
|
|
|
|
2024-08-06 17:05:00 +01:00
|
|
|
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"],
|
2024-08-09 12:45:15 +01:00
|
|
|
masquerade_port=config["masquerade_port"],
|
|
|
|
|
masquerade_protocol=config["masquerade_protocol"],
|
2024-08-06 17:05:00 +01:00
|
|
|
)
|
|
|
|
|
|
|
|
|
|
ConfigureC2BeaconAction._Opts.model_validate(config) # check that options adhere to schema
|
|
|
|
|
|
|
|
|
|
return ["network", "node", node_name, "application", "C2Beacon", "configure", config.__dict__]
|
|
|
|
|
|
|
|
|
|
|
2024-08-20 10:51:29 +01:00
|
|
|
class NodeAccountsChangePasswordAction(AbstractAction):
|
|
|
|
|
"""Action which changes the password for a user."""
|
|
|
|
|
|
|
|
|
|
def __init__(self, manager: "ActionManager", **kwargs) -> None:
|
|
|
|
|
super().__init__(manager=manager)
|
2024-08-06 17:05:00 +01:00
|
|
|
|
2024-07-30 11:33:52 +01:00
|
|
|
def form_request(self, node_id: str, username: str, current_password: str, new_password: str) -> RequestFormat:
|
2024-07-24 16:38:06 +01:00
|
|
|
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
|
2024-08-05 16:27:53 +01:00
|
|
|
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,
|
|
|
|
|
]
|
2024-07-24 16:38:06 +01:00
|
|
|
|
|
|
|
|
|
|
|
|
|
class NodeSessionsRemoteLoginAction(AbstractAction):
|
|
|
|
|
"""Action which performs a remote session login."""
|
|
|
|
|
|
2024-08-06 17:05:00 +01:00
|
|
|
def __init__(self, manager: "ActionManager", **kwargs) -> None:
|
|
|
|
|
super().__init__(manager=manager)
|
|
|
|
|
|
2024-08-05 16:27:53 +01:00
|
|
|
def form_request(self, node_id: str, username: str, password: str, remote_ip: str) -> RequestFormat:
|
2024-07-24 16:38:06 +01:00
|
|
|
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
|
2024-08-05 16:27:53 +01:00
|
|
|
node_name = self.manager.get_node_name_by_idx(node_id)
|
|
|
|
|
return [
|
|
|
|
|
"network",
|
|
|
|
|
"node",
|
|
|
|
|
node_name,
|
|
|
|
|
"service",
|
2024-08-11 23:24:29 +01:00
|
|
|
"Terminal",
|
|
|
|
|
"ssh_to_remote",
|
2024-08-05 16:27:53 +01:00
|
|
|
username,
|
|
|
|
|
password,
|
|
|
|
|
remote_ip,
|
|
|
|
|
]
|
2024-07-24 16:38:06 +01:00
|
|
|
|
|
|
|
|
|
|
|
|
|
class NodeSessionsRemoteLogoutAction(AbstractAction):
|
|
|
|
|
"""Action which performs a remote session logout."""
|
|
|
|
|
|
2024-08-20 10:51:29 +01:00
|
|
|
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:
|
2024-08-20 10:51:29 +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)
|
2024-08-20 11:15:04 +01:00
|
|
|
return ["network", "node", node_name, "service", "Terminal", "remote_logoff", remote_ip]
|
2024-08-20 10:51:29 +01:00
|
|
|
|
|
|
|
|
|
|
|
|
|
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."""
|
|
|
|
|
|
2024-08-06 17:05:00 +01:00
|
|
|
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"]
|
2024-08-15 17:47:33 +01:00
|
|
|
# This action currently doesn't require any further configuration options.
|
2024-08-06 17:05:00 +01:00
|
|
|
return ["network", "node", node_name, "application", "C2Server", "ransomware_launch"]
|
|
|
|
|
|
|
|
|
|
|
2024-08-14 19:49:58 +01:00
|
|
|
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]
|
|
|
|
|
|
|
|
|
|
|
2024-08-20 10:51:29 +01:00
|
|
|
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},
|
|
|
|
|
]
|
|
|
|
|
|
|
|
|
|
|
2024-08-06 17:05:00 +01:00
|
|
|
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."""
|
|
|
|
|
|
2024-08-15 11:36:55 +01:00
|
|
|
commands: Union[List[RequestFormat], RequestFormat]
|
2024-08-09 09:58:44 +01:00
|
|
|
ip_address: Optional[str]
|
|
|
|
|
username: Optional[str]
|
|
|
|
|
password: Optional[str]
|
2024-08-06 17:05:00 +01:00
|
|
|
|
|
|
|
|
def __init__(self, manager: "ActionManager", **kwargs) -> None:
|
|
|
|
|
super().__init__(manager=manager)
|
|
|
|
|
|
2024-08-09 09:58:44 +01:00
|
|
|
def form_request(self, node_id: int, commands: List, ip_address: Optional[str], account: dict) -> RequestFormat:
|
2024-08-06 17:05:00 +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"]
|
2024-08-09 09:58:44 +01:00
|
|
|
|
|
|
|
|
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]
|
2024-08-06 17:05:00 +01:00
|
|
|
|
2024-08-20 10:51:29 +01:00
|
|
|
|
|
|
|
|
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."""
|
2024-08-11 23:24:29 +01:00
|
|
|
node_name = self.manager.get_node_name_by_idx(node_id)
|
2024-08-20 10:51:29 +01:00
|
|
|
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"]
|