#2912 - End of day commit
This commit is contained in:
@@ -5,6 +5,7 @@ from pydantic import BaseModel, Field, field_validator, ValidationInfo
|
||||
|
||||
from primaite.game.agent.actions.manager import AbstractAction
|
||||
from primaite.game.game import _LOGGER
|
||||
from primaite.interface.request import RequestFormat
|
||||
|
||||
|
||||
class RouterACLAddRuleAction(AbstractAction, identifier="router_acl_add_rule"):
|
||||
@@ -168,3 +169,173 @@ class RouterACLAddRuleAction(AbstractAction, identifier="router_acl_add_rule"):
|
||||
dst_port,
|
||||
position,
|
||||
]
|
||||
|
||||
|
||||
class RouterACLRemoveRuleAction(AbstractAction):
|
||||
"""Action which removes a rule from a router's ACL."""
|
||||
|
||||
class ConfigSchema(AbstractAction.ConfigSchema):
|
||||
"""Configuration schema for RouterACLRemoveRuleAction."""
|
||||
|
||||
target_router: str
|
||||
position: str
|
||||
|
||||
@classmethod
|
||||
def form_request(cls, config: ConfigSchema) -> RequestFormat:
|
||||
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
|
||||
return ["network", "node", config.target_router, "acl", "remove_rule", config.position]
|
||||
|
||||
|
||||
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,
|
||||
source_wildcard_id: int,
|
||||
dest_ip_id: int,
|
||||
dest_wildcard_id: int,
|
||||
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
|
||||
|
||||
if dest_ip_id == 0:
|
||||
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
|
||||
src_wildcard = self.manager.get_wildcard_by_idx(source_wildcard_id)
|
||||
dst_wildcard = self.manager.get_wildcard_by_idx(dest_wildcard_id)
|
||||
|
||||
return [
|
||||
"network",
|
||||
"node",
|
||||
target_firewall_nodename,
|
||||
firewall_port_name,
|
||||
firewall_port_direction,
|
||||
"acl",
|
||||
"add_rule",
|
||||
permission_str,
|
||||
protocol,
|
||||
str(src_ip),
|
||||
src_wildcard,
|
||||
src_port,
|
||||
str(dst_ip),
|
||||
dst_wildcard,
|
||||
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,
|
||||
]
|
||||
@@ -1,7 +1,5 @@
|
||||
# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
|
||||
# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
|
||||
from abc import abstractmethod
|
||||
from typing import ClassVar, Dict
|
||||
from typing import ClassVar
|
||||
|
||||
from primaite.game.agent.actions.manager import AbstractAction
|
||||
from primaite.interface.request import RequestFormat
|
||||
@@ -16,6 +14,8 @@ class NodeApplicationAbstractAction(AbstractAction):
|
||||
"""
|
||||
|
||||
class ConfigSchema(AbstractAction.ConfigSchema):
|
||||
"""Base Configuration schema for Node Application actions."""
|
||||
|
||||
node_name: str
|
||||
application_name: str
|
||||
|
||||
@@ -33,6 +33,8 @@ class NodeApplicationExecuteAction(NodeApplicationAbstractAction, identifier="no
|
||||
"""Action which executes an application."""
|
||||
|
||||
class ConfigSchema(NodeApplicationAbstractAction.ConfigSchema):
|
||||
"""Configuration schema for NodeApplicationExecuteAction."""
|
||||
|
||||
verb: str = "execute"
|
||||
|
||||
|
||||
@@ -40,6 +42,8 @@ class NodeApplicationScanAction(NodeApplicationAbstractAction, identifier="node_
|
||||
"""Action which scans an application."""
|
||||
|
||||
class ConfigSchema(NodeApplicationAbstractAction.ConfigSchema):
|
||||
"""Configuration schema for NodeApplicationScanAction."""
|
||||
|
||||
verb: str = "scan"
|
||||
|
||||
|
||||
@@ -47,6 +51,8 @@ class NodeApplicationCloseAction(NodeApplicationAbstractAction, identifier="node
|
||||
"""Action which closes an application."""
|
||||
|
||||
class ConfigSchema(NodeApplicationAbstractAction.ConfigSchema):
|
||||
"""Configuration schema for NodeApplicationCloseAction."""
|
||||
|
||||
verb: str = "close"
|
||||
|
||||
|
||||
@@ -54,11 +60,28 @@ class NodeApplicationFixAction(NodeApplicationAbstractAction, identifier="node_a
|
||||
"""Action which fixes an application."""
|
||||
|
||||
class ConfigSchema(NodeApplicationAbstractAction.ConfigSchema):
|
||||
"""Configuration schema for NodeApplicationFixAction."""
|
||||
|
||||
verb: str = "fix"
|
||||
|
||||
|
||||
class NodeApplicationInstallAction(AbstractAction):
|
||||
class NodeApplicationInstallAction(NodeApplicationAbstractAction):
|
||||
"""Action which installs an application."""
|
||||
|
||||
class ConfigSchema(NodeApplicationAbstractAction.ConfigSchema):
|
||||
"""Configuration schema for NodeApplicationInstallAction."""
|
||||
|
||||
verb: str = "install"
|
||||
|
||||
# TODO: Either changes to application form_request bits, or add that here.
|
||||
|
||||
class NodeApplicationRemoveAction(NodeApplicationAbstractAction):
|
||||
"""Action which removes/uninstalls an application"""
|
||||
|
||||
class ConfigSchema(NodeApplicationAbstractAction.ConfigSchema):
|
||||
"""Configuration schema for NodeApplicationRemoveAction."""
|
||||
|
||||
verb: str = "uninstall"
|
||||
|
||||
# TODO: Either changes to application form_request bits, or add that here.
|
||||
|
||||
|
||||
107
src/primaite/game/agent/actions/config.py
Normal file
107
src/primaite/game/agent/actions/config.py
Normal file
@@ -0,0 +1,107 @@
|
||||
# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
|
||||
|
||||
from typing import Dict, Optional
|
||||
from pydantic import BaseModel, ConfigDict, Field, ValidationInfo, field_validator
|
||||
from primaite.game.agent.actions.manager import AbstractAction
|
||||
from primaite.interface.request import RequestFormat
|
||||
|
||||
|
||||
class ConfigureRansomwareScriptAction(AbstractAction):
|
||||
"""Action which sets config parameters for a ransomware script on a node."""
|
||||
|
||||
class _Opts(BaseModel):
|
||||
"""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, 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"]
|
||||
ConfigureRansomwareScriptAction._Opts.model_validate(config) # check that options adhere to schema
|
||||
return ["network", "node", node_name, "application", "RansomwareScript", "configure", config]
|
||||
|
||||
class ConfigureDoSBotAction(AbstractAction):
|
||||
"""Action which sets config parameters for a DoS bot on a node."""
|
||||
|
||||
class _Opts(BaseModel):
|
||||
"""Schema for options that can be passed to this action."""
|
||||
|
||||
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:
|
||||
"""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]
|
||||
|
||||
|
||||
class ConfigureC2BeaconAction(AbstractAction):
|
||||
"""Action which configures a C2 Beacon based on the parameters given."""
|
||||
|
||||
class ConfigSchema(AbstractAction.ConfigSchema):
|
||||
"""Configuration schema for ConfigureC2BeaconAction."""
|
||||
|
||||
node_name: str
|
||||
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")
|
||||
|
||||
|
||||
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
|
||||
|
||||
@classmethod
|
||||
def form_request(self, config: ConfigSchema) -> RequestFormat:
|
||||
"""Return the action formatted as a request that can be ingested by the simulation."""
|
||||
if config.node_name is None:
|
||||
return ["do_nothing"]
|
||||
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", config.node_name, "application", "C2Beacon", "configure", config.__dict__]
|
||||
@@ -8,11 +8,13 @@ from primaite.interface.request import RequestFormat
|
||||
class NodeFileAbstractAction(AbstractAction):
|
||||
"""Abstract base class for file actions.
|
||||
|
||||
Any action which applies to a file and uses node_name, folder_name, and file_name as its only three parameters can inherit
|
||||
from this base class.
|
||||
Any action which applies to a file and uses node_name, folder_name, and file_name as its
|
||||
only three parameters can inherit from this base class.
|
||||
"""
|
||||
|
||||
class ConfigSchema(AbstractAction.ConfigSchema):
|
||||
"""Configuration Schema for NodeFileAbstractAction."""
|
||||
|
||||
node_name: str
|
||||
folder_name: str
|
||||
file_name: str
|
||||
@@ -41,6 +43,8 @@ class NodeFileCreateAction(NodeFileAbstractAction, identifier="node_file_create"
|
||||
"""Action which creates a new file in a given folder."""
|
||||
|
||||
class ConfigSchema(NodeFileAbstractAction.ConfigSchema):
|
||||
"""Configuration schema for NodeFileCreateAction."""
|
||||
|
||||
verb: str = "create"
|
||||
|
||||
|
||||
@@ -48,6 +52,8 @@ class NodeFileScanAction(NodeFileAbstractAction, identifier="node_file_scan"):
|
||||
"""Action which scans a file."""
|
||||
|
||||
class ConfigSchema(NodeFileAbstractAction.ConfigSchema):
|
||||
"""Configuration schema for NodeFileScanAction."""
|
||||
|
||||
verb: str = "scan"
|
||||
|
||||
|
||||
@@ -55,6 +61,8 @@ class NodeFileDeleteAction(NodeFileAbstractAction, identifier="node_file_delete"
|
||||
"""Action which deletes a file."""
|
||||
|
||||
class ConfigSchema(NodeFileAbstractAction.ConfigSchema):
|
||||
"""Configuration schema for NodeFileDeleteAction."""
|
||||
|
||||
verb: str = "delete"
|
||||
|
||||
|
||||
@@ -62,6 +70,8 @@ class NodeFileRestoreAction(NodeFileAbstractAction, identifier="node_file_restor
|
||||
"""Action which restores a file."""
|
||||
|
||||
class ConfigSchema(NodeFileAbstractAction.ConfigSchema):
|
||||
"""Configuration schema for NodeFileRestoreAction."""
|
||||
|
||||
verb: str = "restore"
|
||||
|
||||
|
||||
@@ -69,6 +79,8 @@ class NodeFileCorruptAction(NodeFileAbstractAction, identifier="node_file_corrup
|
||||
"""Action which corrupts a file."""
|
||||
|
||||
class ConfigSchema(NodeFileAbstractAction.ConfigSchema):
|
||||
"""Configuration schema for NodeFileCorruptAction."""
|
||||
|
||||
verb: str = "corrupt"
|
||||
|
||||
|
||||
@@ -76,4 +88,6 @@ class NodeFileAccessAction(NodeFileAbstractAction, identifier="node_file_access"
|
||||
"""Action which increases a file's access count."""
|
||||
|
||||
class ConfigSchema(NodeFileAbstractAction.ConfigSchema):
|
||||
"""Configuration schema for NodeFileAccessAction."""
|
||||
|
||||
verb: str = "access"
|
||||
|
||||
@@ -15,6 +15,8 @@ class NodeFolderAbstractAction(AbstractAction):
|
||||
"""
|
||||
|
||||
class ConfigSchema(AbstractAction.ConfigSchema):
|
||||
"""Base configuration schema for NodeFolder actions."""
|
||||
|
||||
node_name: str
|
||||
folder_name: str
|
||||
|
||||
@@ -34,6 +36,8 @@ class NodeFolderScanAction(NodeFolderAbstractAction, identifier="node_folder_sca
|
||||
"""Action which scans a folder."""
|
||||
|
||||
class ConfigSchema(NodeFolderAbstractAction.ConfigSchema):
|
||||
"""Configuration schema for NodeFolderScanAction."""
|
||||
|
||||
verb: str = "scan"
|
||||
|
||||
|
||||
@@ -41,6 +45,8 @@ class NodeFolderCheckhashAction(NodeFolderAbstractAction, identifier="node_folde
|
||||
"""Action which checks the hash of a folder."""
|
||||
|
||||
class ConfigSchema(NodeFolderAbstractAction.ConfigSchema):
|
||||
"""Configuration schema for NodeFolderCheckhashAction."""
|
||||
|
||||
verb: str = "checkhash"
|
||||
|
||||
|
||||
@@ -48,6 +54,8 @@ class NodeFolderRepairAction(NodeFolderAbstractAction, identifier="node_folder_r
|
||||
"""Action which repairs a folder."""
|
||||
|
||||
class ConfigSchema(NodeFolderAbstractAction.ConfigSchema):
|
||||
"""Configuration schema for NodeFolderRepairAction."""
|
||||
|
||||
verb: str = "repair"
|
||||
|
||||
|
||||
@@ -55,6 +63,8 @@ class NodeFolderRestoreAction(NodeFolderAbstractAction, identifier="node_folder_
|
||||
"""Action which restores a folder."""
|
||||
|
||||
class ConfigSchema(NodeFolderAbstractAction.ConfigSchema):
|
||||
"""Configuration schema for NodeFolderRestoreAction."""
|
||||
|
||||
verb: str = "restore"
|
||||
|
||||
|
||||
@@ -62,4 +72,6 @@ class NodeFolderCreateAction(AbstractAction, identifier="node_folder_create"):
|
||||
"""Action which creates a new folder."""
|
||||
|
||||
class ConfigSchema(NodeFolderAbstractAction.ConfigSchema):
|
||||
"""Configuration schema for NodeFolderCreateAction."""
|
||||
|
||||
verb: str = "create"
|
||||
|
||||
57
src/primaite/game/agent/actions/host_nic.py
Normal file
57
src/primaite/game/agent/actions/host_nic.py
Normal file
@@ -0,0 +1,57 @@
|
||||
# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
|
||||
|
||||
from typing import Dict, Optional
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
from primaite.game.agent.actions.manager import AbstractAction
|
||||
from primaite.interface.request import RequestFormat
|
||||
|
||||
class HostNICAbstractAction(AbstractAction):
|
||||
"""
|
||||
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.
|
||||
"""
|
||||
|
||||
class ConfigSchema(AbstractAction.ConfigSchema):
|
||||
"""Base Configuration schema for HostNIC actions."""
|
||||
num_nodes: str
|
||||
max_nics_per_node: str
|
||||
node_name: 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
|
||||
def form_request(cls, config: ConfigSchema) -> RequestFormat:
|
||||
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
|
||||
if config.node_name is None or config.nic_num is None:
|
||||
return ["do_nothing"]
|
||||
return ["network", "node", config.node_name, "network_interface", config.nic_num, cls.verb]
|
||||
|
||||
class HostNICEnableAction(HostNICAbstractAction):
|
||||
"""Action which enables a NIC."""
|
||||
|
||||
class ConfigSchema(HostNICAbstractAction.ConfigSchema):
|
||||
"""Configuration schema for HostNICEnableAction."""
|
||||
verb: str = "enable"
|
||||
|
||||
|
||||
class HostNICDisableAction(HostNICAbstractAction):
|
||||
"""Action which disables a NIC."""
|
||||
|
||||
class ConfigSchema(HostNICAbstractAction.ConfigSchema):
|
||||
"""Configuration schema for HostNICDisableAction."""
|
||||
verb: str = "disable"
|
||||
@@ -14,13 +14,13 @@ agents:
|
||||
from __future__ import annotations
|
||||
|
||||
import itertools
|
||||
from abc import ABC, abstractmethod
|
||||
from abc import ABC
|
||||
from typing import Any, ClassVar, Dict, List, Literal, Optional, Tuple, Type
|
||||
|
||||
from gymnasium import spaces
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
|
||||
from primaite.game.game import _LOGGER, PrimaiteGame
|
||||
from primaite.game.game import PrimaiteGame
|
||||
from primaite.interface.request import RequestFormat
|
||||
|
||||
|
||||
@@ -57,6 +57,7 @@ class DoNothingAction(AbstractAction):
|
||||
type: Literal["do_nothing"] = "do_nothing"
|
||||
|
||||
def form_request(self, options: ConfigSchema) -> RequestFormat:
|
||||
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
|
||||
return ["do_nothing"]
|
||||
|
||||
|
||||
|
||||
43
src/primaite/game/agent/actions/network.py
Normal file
43
src/primaite/game/agent/actions/network.py
Normal file
@@ -0,0 +1,43 @@
|
||||
# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
|
||||
|
||||
from typing import Dict, Optional
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
from primaite.game.agent.actions.manager import AbstractAction
|
||||
from primaite.interface.request import RequestFormat
|
||||
|
||||
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"]
|
||||
@@ -14,6 +14,8 @@ class NodeAbstractAction(AbstractAction):
|
||||
"""
|
||||
|
||||
class ConfigSchema(AbstractAction.ConfigSchema):
|
||||
"""Base Configuration schema for Node actions."""
|
||||
|
||||
node_name: str
|
||||
|
||||
verb: ClassVar[str]
|
||||
@@ -28,6 +30,8 @@ class NodeOSScanAction(NodeAbstractAction, identifier="node_os_scan"):
|
||||
"""Action which scans a node's OS."""
|
||||
|
||||
class ConfigSchema(NodeAbstractAction.ConfigSchema):
|
||||
"""Configuration schema for NodeOSScanAction."""
|
||||
|
||||
verb: str = "scan"
|
||||
|
||||
|
||||
@@ -35,6 +39,8 @@ class NodeShutdownAction(NodeAbstractAction, identifier="node_shutdown"):
|
||||
"""Action which shuts down a node."""
|
||||
|
||||
class ConfigSchema(NodeAbstractAction.ConfigSchema):
|
||||
"""Configuration schema for NodeShutdownAction."""
|
||||
|
||||
verb: str = "shutdown"
|
||||
|
||||
|
||||
@@ -42,6 +48,8 @@ class NodeStartupAction(NodeAbstractAction, identifier="node_startup"):
|
||||
"""Action which starts up a node."""
|
||||
|
||||
class ConfigSchema(NodeAbstractAction.ConfigSchema):
|
||||
"""Configuration schema for NodeStartupAction."""
|
||||
|
||||
verb: str = "startup"
|
||||
|
||||
|
||||
@@ -49,4 +57,6 @@ class NodeResetAction(NodeAbstractAction, identifier="node_reset"):
|
||||
"""Action which resets a node."""
|
||||
|
||||
class ConfigSchema(NodeAbstractAction.ConfigSchema):
|
||||
"""Configuration schema for NodeResetAction."""
|
||||
|
||||
verb: str = "reset"
|
||||
|
||||
Reference in New Issue
Block a user