#2912 - Updates to remaining action refactoring
This commit is contained in:
@@ -25,3 +25,16 @@ __all__ = (
|
|||||||
"NodeServiceStopAction",
|
"NodeServiceStopAction",
|
||||||
"ActionManager",
|
"ActionManager",
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# __all__ = (
|
||||||
|
# "acl",
|
||||||
|
# "application",
|
||||||
|
# "config",
|
||||||
|
# "file",
|
||||||
|
# "folder",
|
||||||
|
# "host_nic",
|
||||||
|
# "manager",
|
||||||
|
# "network",
|
||||||
|
# "node",
|
||||||
|
# "service",
|
||||||
|
# )
|
||||||
|
|||||||
@@ -171,7 +171,7 @@ class RouterACLAddRuleAction(AbstractAction, identifier="router_acl_add_rule"):
|
|||||||
]
|
]
|
||||||
|
|
||||||
|
|
||||||
class RouterACLRemoveRuleAction(AbstractAction):
|
class RouterACLRemoveRuleAction(AbstractAction, identifier="router_acl_remove_rule"):
|
||||||
"""Action which removes a rule from a router's ACL."""
|
"""Action which removes a rule from a router's ACL."""
|
||||||
|
|
||||||
class ConfigSchema(AbstractAction.ConfigSchema):
|
class ConfigSchema(AbstractAction.ConfigSchema):
|
||||||
@@ -186,7 +186,7 @@ class RouterACLRemoveRuleAction(AbstractAction):
|
|||||||
return ["network", "node", config.target_router, "acl", "remove_rule", config.position]
|
return ["network", "node", config.target_router, "acl", "remove_rule", config.position]
|
||||||
|
|
||||||
|
|
||||||
class FirewallACLAddRuleAction(AbstractAction):
|
class FirewallACLAddRuleAction(AbstractAction, identifier="firewall_acl_add_rule"):
|
||||||
"""Action which adds a rule to a firewall port's ACL."""
|
"""Action which adds a rule to a firewall port's ACL."""
|
||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
@@ -310,8 +310,9 @@ class FirewallACLAddRuleAction(AbstractAction):
|
|||||||
dst_port,
|
dst_port,
|
||||||
position,
|
position,
|
||||||
]
|
]
|
||||||
|
|
||||||
class FirewallACLRemoveRuleAction(AbstractAction):
|
|
||||||
|
class FirewallACLRemoveRuleAction(AbstractAction, identifier="firewall_acl_remove_rule"):
|
||||||
"""Action which removes a rule from a firewall port's ACL."""
|
"""Action which removes a rule from a firewall port's ACL."""
|
||||||
|
|
||||||
def __init__(self, manager: "ActionManager", max_acl_rules: int, **kwargs) -> None:
|
def __init__(self, manager: "ActionManager", max_acl_rules: int, **kwargs) -> None:
|
||||||
@@ -325,8 +326,9 @@ class FirewallACLRemoveRuleAction(AbstractAction):
|
|||||||
super().__init__(manager=manager)
|
super().__init__(manager=manager)
|
||||||
self.shape: Dict[str, int] = {"position": max_acl_rules}
|
self.shape: Dict[str, int] = {"position": max_acl_rules}
|
||||||
|
|
||||||
|
@classmethod
|
||||||
def form_request(
|
def form_request(
|
||||||
self, target_firewall_nodename: str, firewall_port_name: str, firewall_port_direction: str, position: int
|
cls, target_firewall_nodename: str, firewall_port_name: str, firewall_port_direction: str, position: int
|
||||||
) -> List[str]:
|
) -> List[str]:
|
||||||
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
|
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
|
||||||
return [
|
return [
|
||||||
@@ -338,4 +340,4 @@ class FirewallACLRemoveRuleAction(AbstractAction):
|
|||||||
"acl",
|
"acl",
|
||||||
"remove_rule",
|
"remove_rule",
|
||||||
position,
|
position,
|
||||||
]
|
]
|
||||||
|
|||||||
@@ -75,6 +75,7 @@ class NodeApplicationInstallAction(NodeApplicationAbstractAction):
|
|||||||
|
|
||||||
# TODO: Either changes to application form_request bits, or add that here.
|
# TODO: Either changes to application form_request bits, or add that here.
|
||||||
|
|
||||||
|
|
||||||
class NodeApplicationRemoveAction(NodeApplicationAbstractAction):
|
class NodeApplicationRemoveAction(NodeApplicationAbstractAction):
|
||||||
"""Action which removes/uninstalls an application"""
|
"""Action which removes/uninstalls an application"""
|
||||||
|
|
||||||
@@ -84,4 +85,3 @@ class NodeApplicationRemoveAction(NodeApplicationAbstractAction):
|
|||||||
verb: str = "uninstall"
|
verb: str = "uninstall"
|
||||||
|
|
||||||
# TODO: Either changes to application form_request bits, or add that here.
|
# TODO: Either changes to application form_request bits, or add that here.
|
||||||
|
|
||||||
|
|||||||
@@ -1,7 +1,9 @@
|
|||||||
# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
|
# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
|
||||||
|
|
||||||
from typing import Dict, Optional
|
from typing import Dict, Optional
|
||||||
from pydantic import BaseModel, ConfigDict, Field, ValidationInfo, field_validator
|
|
||||||
|
from pydantic import BaseModel, ConfigDict, Field, field_validator, ValidationInfo
|
||||||
|
|
||||||
from primaite.game.agent.actions.manager import AbstractAction
|
from primaite.game.agent.actions.manager import AbstractAction
|
||||||
from primaite.interface.request import RequestFormat
|
from primaite.interface.request import RequestFormat
|
||||||
|
|
||||||
@@ -28,6 +30,7 @@ class ConfigureRansomwareScriptAction(AbstractAction):
|
|||||||
ConfigureRansomwareScriptAction._Opts.model_validate(config) # check that options adhere to schema
|
ConfigureRansomwareScriptAction._Opts.model_validate(config) # check that options adhere to schema
|
||||||
return ["network", "node", node_name, "application", "RansomwareScript", "configure", config]
|
return ["network", "node", node_name, "application", "RansomwareScript", "configure", config]
|
||||||
|
|
||||||
|
|
||||||
class ConfigureDoSBotAction(AbstractAction):
|
class ConfigureDoSBotAction(AbstractAction):
|
||||||
"""Action which sets config parameters for a DoS bot on a node."""
|
"""Action which sets config parameters for a DoS bot on a node."""
|
||||||
|
|
||||||
@@ -53,7 +56,7 @@ class ConfigureDoSBotAction(AbstractAction):
|
|||||||
return ["do_nothing"]
|
return ["do_nothing"]
|
||||||
self._Opts.model_validate(config) # check that options adhere to schema
|
self._Opts.model_validate(config) # check that options adhere to schema
|
||||||
return ["network", "node", node_name, "application", "DoSBot", "configure", config]
|
return ["network", "node", node_name, "application", "DoSBot", "configure", config]
|
||||||
|
|
||||||
|
|
||||||
class ConfigureC2BeaconAction(AbstractAction):
|
class ConfigureC2BeaconAction(AbstractAction):
|
||||||
"""Action which configures a C2 Beacon based on the parameters given."""
|
"""Action which configures a C2 Beacon based on the parameters given."""
|
||||||
@@ -67,7 +70,6 @@ class ConfigureC2BeaconAction(AbstractAction):
|
|||||||
masquerade_protocol: str = Field(default="TCP")
|
masquerade_protocol: str = Field(default="TCP")
|
||||||
masquerade_port: str = Field(default="HTTP")
|
masquerade_port: str = Field(default="HTTP")
|
||||||
|
|
||||||
|
|
||||||
class _Opts(BaseModel):
|
class _Opts(BaseModel):
|
||||||
"""Schema for options that can be passed to this action."""
|
"""Schema for options that can be passed to this action."""
|
||||||
|
|
||||||
@@ -104,4 +106,4 @@ class ConfigureC2BeaconAction(AbstractAction):
|
|||||||
|
|
||||||
ConfigureC2BeaconAction._Opts.model_validate(config) # check that options adhere to schema
|
ConfigureC2BeaconAction._Opts.model_validate(config) # check that options adhere to schema
|
||||||
|
|
||||||
return ["network", "node", config.node_name, "application", "C2Beacon", "configure", config.__dict__]
|
return ["network", "node", config.node_name, "application", "C2Beacon", "configure", config.__dict__]
|
||||||
|
|||||||
@@ -23,13 +23,11 @@ class NodeFolderAbstractAction(AbstractAction):
|
|||||||
verb: ClassVar[str]
|
verb: ClassVar[str]
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def form_request(cls, node_id: int, folder_id: int) -> RequestFormat:
|
def form_request(cls, config: ConfigSchema) -> RequestFormat:
|
||||||
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
|
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
|
||||||
node_name = cls.manager.get_node_name_by_idx(node_id)
|
if config.node_name is None or config.folder_name is None:
|
||||||
folder_name = cls.manager.get_folder_name_by_idx(node_idx=node_id, folder_idx=folder_id)
|
|
||||||
if node_name is None or folder_name is None:
|
|
||||||
return ["do_nothing"]
|
return ["do_nothing"]
|
||||||
return ["network", "node", node_name, "file_system", "folder", folder_name, cls.verb]
|
return ["network", "node", config.node_name, "file_system", "folder", config.folder_name, cls.verb]
|
||||||
|
|
||||||
|
|
||||||
class NodeFolderScanAction(NodeFolderAbstractAction, identifier="node_folder_scan"):
|
class NodeFolderScanAction(NodeFolderAbstractAction, identifier="node_folder_scan"):
|
||||||
|
|||||||
@@ -1,10 +1,13 @@
|
|||||||
# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
|
# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
|
||||||
|
|
||||||
from typing import Dict, Optional
|
from typing import Dict, Optional
|
||||||
|
|
||||||
from pydantic import BaseModel, ConfigDict
|
from pydantic import BaseModel, ConfigDict
|
||||||
|
|
||||||
from primaite.game.agent.actions.manager import AbstractAction
|
from primaite.game.agent.actions.manager import AbstractAction
|
||||||
from primaite.interface.request import RequestFormat
|
from primaite.interface.request import RequestFormat
|
||||||
|
|
||||||
|
|
||||||
class HostNICAbstractAction(AbstractAction):
|
class HostNICAbstractAction(AbstractAction):
|
||||||
"""
|
"""
|
||||||
Abstract base class for NIC actions.
|
Abstract base class for NIC actions.
|
||||||
@@ -15,25 +18,12 @@ class HostNICAbstractAction(AbstractAction):
|
|||||||
|
|
||||||
class ConfigSchema(AbstractAction.ConfigSchema):
|
class ConfigSchema(AbstractAction.ConfigSchema):
|
||||||
"""Base Configuration schema for HostNIC actions."""
|
"""Base Configuration schema for HostNIC actions."""
|
||||||
|
|
||||||
num_nodes: str
|
num_nodes: str
|
||||||
max_nics_per_node: str
|
max_nics_per_node: str
|
||||||
node_name: str
|
node_name: str
|
||||||
nic_num: str
|
nic_num: str
|
||||||
|
|
||||||
def __init__(self, manager: "ActionManager", num_nodes: int, max_nics_per_node: int, **kwargs) -> None:
|
|
||||||
"""Init method for HostNICAbstractAction.
|
|
||||||
|
|
||||||
: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
|
|
||||||
"""
|
|
||||||
super().__init__(manager=manager)
|
|
||||||
self.shape: Dict[str, int] = {"node_id": num_nodes, "nic_id": max_nics_per_node}
|
|
||||||
self.verb: str # define but don't initialise: defends against children classes not defining this
|
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def form_request(cls, config: ConfigSchema) -> RequestFormat:
|
def form_request(cls, config: ConfigSchema) -> RequestFormat:
|
||||||
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
|
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
|
||||||
@@ -41,11 +31,13 @@ class HostNICAbstractAction(AbstractAction):
|
|||||||
return ["do_nothing"]
|
return ["do_nothing"]
|
||||||
return ["network", "node", config.node_name, "network_interface", config.nic_num, cls.verb]
|
return ["network", "node", config.node_name, "network_interface", config.nic_num, cls.verb]
|
||||||
|
|
||||||
|
|
||||||
class HostNICEnableAction(HostNICAbstractAction):
|
class HostNICEnableAction(HostNICAbstractAction):
|
||||||
"""Action which enables a NIC."""
|
"""Action which enables a NIC."""
|
||||||
|
|
||||||
class ConfigSchema(HostNICAbstractAction.ConfigSchema):
|
class ConfigSchema(HostNICAbstractAction.ConfigSchema):
|
||||||
"""Configuration schema for HostNICEnableAction."""
|
"""Configuration schema for HostNICEnableAction."""
|
||||||
|
|
||||||
verb: str = "enable"
|
verb: str = "enable"
|
||||||
|
|
||||||
|
|
||||||
@@ -54,4 +46,5 @@ class HostNICDisableAction(HostNICAbstractAction):
|
|||||||
|
|
||||||
class ConfigSchema(HostNICAbstractAction.ConfigSchema):
|
class ConfigSchema(HostNICAbstractAction.ConfigSchema):
|
||||||
"""Configuration schema for HostNICDisableAction."""
|
"""Configuration schema for HostNICDisableAction."""
|
||||||
verb: str = "disable"
|
|
||||||
|
verb: str = "disable"
|
||||||
|
|||||||
@@ -23,6 +23,8 @@ from pydantic import BaseModel, ConfigDict
|
|||||||
from primaite.game.game import PrimaiteGame
|
from primaite.game.game import PrimaiteGame
|
||||||
from primaite.interface.request import RequestFormat
|
from primaite.interface.request import RequestFormat
|
||||||
|
|
||||||
|
# TODO: Make sure that actions are backwards compatible where the old YAML format is used.
|
||||||
|
|
||||||
|
|
||||||
class AbstractAction(BaseModel):
|
class AbstractAction(BaseModel):
|
||||||
"""Base class for actions."""
|
"""Base class for actions."""
|
||||||
|
|||||||
@@ -1,10 +1,13 @@
|
|||||||
# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
|
# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
|
||||||
|
|
||||||
from typing import Dict, Optional
|
from typing import Dict, Optional
|
||||||
|
|
||||||
from pydantic import BaseModel, ConfigDict
|
from pydantic import BaseModel, ConfigDict
|
||||||
|
|
||||||
from primaite.game.agent.actions.manager import AbstractAction
|
from primaite.game.agent.actions.manager import AbstractAction
|
||||||
from primaite.interface.request import RequestFormat
|
from primaite.interface.request import RequestFormat
|
||||||
|
|
||||||
|
|
||||||
class NetworkPortEnableAction(AbstractAction):
|
class NetworkPortEnableAction(AbstractAction):
|
||||||
"""Action which enables are port on a router or a firewall."""
|
"""Action which enables are port on a router or a firewall."""
|
||||||
|
|
||||||
@@ -40,4 +43,4 @@ class NetworkPortDisableAction(AbstractAction):
|
|||||||
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
|
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
|
||||||
if target_nodename is None or port_id is None:
|
if target_nodename is None or port_id is None:
|
||||||
return ["do_nothing"]
|
return ["do_nothing"]
|
||||||
return ["network", "node", target_nodename, "network_interface", port_id, "disable"]
|
return ["network", "node", target_nodename, "network_interface", port_id, "disable"]
|
||||||
|
|||||||
Reference in New Issue
Block a user