oxbrbtx/PrimAITE
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge agents and actions branches + fix import / subclass errors

Browse Source
This commit is contained in:
Marek Wolan
2025-01-14 11:34:01 +00:00
parent 9f5e16dd85 1ac562ebc9
commit fb5242bd78
57 changed files with 2297 additions and 2381 deletions
Download Patch File Download Diff File Expand all files Collapse all files

33
src/primaite/game/agent/actions/__init__.py Normal file
View File

@@ -0,0 +1,33 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from primaite.game.agent.actions import (
abstract,
acl,
application,
file,
folder,
host_nic,
manager,
network,
node,
service,
session,
software,
)
from primaite.game.agent.actions.manager import ActionManager
__all__ = (
"abstract",
"acl",
"application",
"software",
"file",
"folder",
"host_nic",
"manager",
"network",
"node",
"service",
"session",
"ActionManager",
)

36
src/primaite/game/agent/actions/abstract.py Normal file
View File

@@ -0,0 +1,36 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from __future__ import annotations
from abc import ABC
from typing import Any, ClassVar, Dict, Optional, Type
from pydantic import BaseModel, ConfigDict
from primaite.interface.request import RequestFormat
class AbstractAction(BaseModel, ABC):
"""Base class for actions."""
config: "AbstractAction.ConfigSchema"
class ConfigSchema(BaseModel, ABC):
"""Base configuration schema for Actions."""
model_config = ConfigDict(extra="forbid")
type: str
_registry: ClassVar[Dict[str, Type[AbstractAction]]] = {}
def __init_subclass__(cls, identifier: Optional[str] = None, **kwargs: Any) -> None:
super().__init_subclass__(**kwargs)
if identifier is None:
return
if identifier in cls._registry:
raise ValueError(f"Cannot create new action under reserved name {identifier}")
cls._registry[identifier] = cls
@classmethod
def form_request(cls, config: ConfigSchema) -> RequestFormat:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
pass

188
src/primaite/game/agent/actions/acl.py Normal file
View File

@@ -0,0 +1,188 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from __future__ import annotations
from abc import ABC
from typing import List
from pydantic import field_validator
from primaite.game.agent.actions.manager import AbstractAction
from primaite.interface.request import RequestFormat
from primaite.utils.validation.ip_protocol import protocol_validator
from primaite.utils.validation.ipv4_address import ipv4_validator, IPV4Address
from primaite.utils.validation.port import port_validator
__all__ = (
"RouterACLAddRuleAction",
"RouterACLRemoveRuleAction",
"FirewallACLAddRuleAction",
"FirewallACLRemoveRuleAction",
)
class ACLAddRuleAbstractAction(AbstractAction, ABC):
"""Base abstract class for ACL add rule actions."""
config: ConfigSchema = "ACLAddRuleAbstractAction.ConfigSchema"
class ConfigSchema(AbstractAction.ConfigSchema):
"""Configuration Schema base for ACL add rule abstract actions."""
src_ip: IPV4Address
protocol_name: str
permission: str
position: int
dst_ip: IPV4Address
src_port: int
dst_port: int
src_wildcard: int
dst_wildcard: int
@field_validator(
"src_port",
"dst_port",
mode="before",
)
@classmethod
def valid_port(cls, v: str) -> int:
"""Check that inputs are valid."""
return port_validator(v)
@field_validator(
"src_ip",
"dst_ip",
mode="before",
)
@classmethod
def valid_ip(cls, v: str) -> str:
"""Check that a valid IP has been provided for src and dst."""
return ipv4_validator(v)
@field_validator(
"protocol_name",
mode="before",
)
@classmethod
def is_valid_protocol(cls, v: str) -> bool:
"""Check that we are using a valid protocol."""
return protocol_validator(v)
class ACLRemoveRuleAbstractAction(AbstractAction, identifier="acl_remove_rule_abstract_action"):
"""Base abstract class for ACL remove rule actions."""
config: ConfigSchema = "ACLRemoveRuleAbstractAction.ConfigSchema"
class ConfigSchema(AbstractAction.ConfigSchema):
"""Configuration Schema base for ACL remove rule abstract actions."""
position: int
class RouterACLAddRuleAction(ACLAddRuleAbstractAction, identifier="router_acl_add_rule"):
"""Action which adds a rule to a router's ACL."""
config: "RouterACLAddRuleAction.ConfigSchema"
class ConfigSchema(ACLAddRuleAbstractAction.ConfigSchema):
"""Configuration Schema for RouterACLAddRuleAction."""
target_router: str
@classmethod
def form_request(cls, config: ConfigSchema) -> List[str]:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
return [
"network",
"node",
config.target_router,
"acl",
"add_rule",
config.permission,
config.protocol_name,
config.src_ip,
config.src_wildcard,
config.src_port,
config.dst_ip,
config.dst_wildcard,
config.dst_port,
config.position,
]
class RouterACLRemoveRuleAction(ACLRemoveRuleAbstractAction, identifier="router_acl_remove_rule"):
"""Action which removes a rule from a router's ACL."""
config: "RouterACLRemoveRuleAction.ConfigSchema"
class ConfigSchema(ACLRemoveRuleAbstractAction.ConfigSchema):
"""Configuration schema for RouterACLRemoveRuleAction."""
target_router: 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(ACLAddRuleAbstractAction, identifier="firewall_acl_add_rule"):
"""Action which adds a rule to a firewall port's ACL."""
config: "FirewallACLAddRuleAction.ConfigSchema"
class ConfigSchema(ACLAddRuleAbstractAction.ConfigSchema):
"""Configuration schema for FirewallACLAddRuleAction."""
target_firewall_nodename: str
firewall_port_name: str
firewall_port_direction: str
@classmethod
def form_request(cls, config: ConfigSchema) -> List[str]:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
return [
"network",
"node",
config.target_firewall_nodename,
config.firewall_port_name,
config.firewall_port_direction,
"acl",
"add_rule",
config.permission,
config.protocol_name,
config.src_ip,
config.src_wildcard,
config.src_port,
config.dst_ip,
config.dst_wildcard,
config.dst_port,
config.position,
]
class FirewallACLRemoveRuleAction(ACLRemoveRuleAbstractAction, identifier="firewall_acl_remove_rule"):
"""Action which removes a rule from a firewall port's ACL."""
config: "FirewallACLRemoveRuleAction.ConfigSchema"
class ConfigSchema(ACLRemoveRuleAbstractAction.ConfigSchema):
"""Configuration schema for FirewallACLRemoveRuleAction."""
target_firewall_nodename: str
firewall_port_name: str
firewall_port_direction: str
@classmethod
def form_request(cls, config: ConfigSchema) -> List[str]:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
return [
"network",
"node",
config.target_firewall_nodename,
config.firewall_port_name,
config.firewall_port_direction,
"acl",
"remove_rule",
config.position,
]

137
src/primaite/game/agent/actions/application.py Normal file
View File

@@ -0,0 +1,137 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from abc import ABC
from typing import ClassVar
from primaite.game.agent.actions.abstract import AbstractAction
from primaite.interface.request import RequestFormat
__all__ = (
"NodeApplicationExecuteAction",
"NodeApplicationScanAction",
"NodeApplicationCloseAction",
"NodeApplicationFixAction",
"NodeApplicationInstallAction",
"NodeApplicationRemoveAction",
)
class NodeApplicationAbstractAction(AbstractAction, ABC):
"""
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.
"""
config: "NodeApplicationAbstractAction.ConfigSchema"
class ConfigSchema(AbstractAction.ConfigSchema):
"""Base Configuration schema for Node Application actions."""
node_name: str
application_name: str
verb: ClassVar[str]
@classmethod
def form_request(cls, config: ConfigSchema) -> RequestFormat:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
return [
"network",
"node",
config.node_name,
"application",
config.application_name,
config.verb,
]
class NodeApplicationExecuteAction(NodeApplicationAbstractAction, identifier="node_application_execute"):
"""Action which executes an application."""
config: "NodeApplicationExecuteAction.ConfigSchema"
class ConfigSchema(NodeApplicationAbstractAction.ConfigSchema):
"""Configuration schema for NodeApplicationExecuteAction."""
verb: str = "execute"
class NodeApplicationScanAction(NodeApplicationAbstractAction, identifier="node_application_scan"):
"""Action which scans an application."""
config: "NodeApplicationScanAction.ConfigSchema"
class ConfigSchema(NodeApplicationAbstractAction.ConfigSchema):
"""Configuration schema for NodeApplicationScanAction."""
verb: str = "scan"
class NodeApplicationCloseAction(NodeApplicationAbstractAction, identifier="node_application_close"):
"""Action which closes an application."""
config: "NodeApplicationCloseAction.ConfigSchema"
class ConfigSchema(NodeApplicationAbstractAction.ConfigSchema):
"""Configuration schema for NodeApplicationCloseAction."""
verb: str = "close"
class NodeApplicationFixAction(NodeApplicationAbstractAction, identifier="node_application_fix"):
"""Action which fixes an application."""
config: "NodeApplicationFixAction.ConfigSchema"
class ConfigSchema(NodeApplicationAbstractAction.ConfigSchema):
"""Configuration schema for NodeApplicationFixAction."""
verb: str = "fix"
class NodeApplicationInstallAction(NodeApplicationAbstractAction, identifier="node_application_install"):
"""Action which installs an application."""
config: "NodeApplicationInstallAction.ConfigSchema"
class ConfigSchema(NodeApplicationAbstractAction.ConfigSchema):
"""Configuration schema for NodeApplicationInstallAction."""
verb: str = "install"
@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.node_name,
"software_manager",
"application",
config.verb,
config.application_name,
]
class NodeApplicationRemoveAction(NodeApplicationAbstractAction, identifier="node_application_remove"):
"""Action which removes/uninstalls an application."""
config: "NodeApplicationRemoveAction.ConfigSchema"
class ConfigSchema(NodeApplicationAbstractAction.ConfigSchema):
"""Configuration schema for NodeApplicationRemoveAction."""
verb: str = "uninstall"
@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.node_name,
"software_manager",
"application",
config.verb,
config.application_name,
]

189
src/primaite/game/agent/actions/file.py Normal file
View File

@@ -0,0 +1,189 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from abc import ABC
from typing import ClassVar
from primaite.game.agent.actions.manager import AbstractAction
from primaite.interface.request import RequestFormat
__all__ = (
"NodeFileCreateAction",
"NodeFileScanAction",
"NodeFileDeleteAction",
"NodeFileRestoreAction",
"NodeFileCorruptAction",
"NodeFileAccessAction",
"NodeFileCheckhashAction",
"NodeFileRepairAction",
)
class NodeFileAbstractAction(AbstractAction, ABC):
"""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.
"""
config: "NodeFileAbstractAction.ConfigSchema"
class ConfigSchema(AbstractAction.ConfigSchema):
"""Configuration Schema for NodeFileAbstractAction."""
node_name: str
folder_name: str
file_name: str
verb: ClassVar[str]
@classmethod
def form_request(cls, config: ConfigSchema) -> RequestFormat:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
if config.node_name is None or config.folder_name is None or config.file_name is None:
return ["do_nothing"]
return [
"network",
"node",
config.node_name,
"file_system",
"folder",
config.folder_name,
"file",
config.file_name,
config.verb,
]
class NodeFileCreateAction(NodeFileAbstractAction, identifier="node_file_create"):
"""Action which creates a new file in a given folder."""
config: "NodeFileCreateAction.ConfigSchema"
class ConfigSchema(NodeFileAbstractAction.ConfigSchema):
"""Configuration schema for NodeFileCreateAction."""
verb: ClassVar[str] = "create"
force: bool = False
@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.folder_name is None or config.file_name is None:
return ["do_nothing"]
return [
"network",
"node",
config.node_name,
"file_system",
config.verb,
"file",
config.folder_name,
config.file_name,
config.verb,
]
class NodeFileScanAction(NodeFileAbstractAction, identifier="node_file_scan"):
"""Action which scans a file."""
config: "NodeFileScanAction.ConfigSchema"
class ConfigSchema(NodeFileAbstractAction.ConfigSchema):
"""Configuration schema for NodeFileScanAction."""
verb: ClassVar[str] = "scan"
class NodeFileDeleteAction(NodeFileAbstractAction, identifier="node_file_delete"):
"""Action which deletes a file."""
config: "NodeFileDeleteAction.ConfigSchema"
class ConfigSchema(NodeFileAbstractAction.ConfigSchema):
"""Configuration schema for NodeFileDeleteAction."""
verb: ClassVar[str] = "delete"
@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.folder_name is None or config.file_name is None:
return ["do_nothing"]
return [
"network",
"node",
config.node_name,
"file_system",
config.verb,
"file",
config.folder_name,
config.file_name,
]
class NodeFileRestoreAction(NodeFileAbstractAction, identifier="node_file_restore"):
"""Action which restores a file."""
config: "NodeFileRestoreAction.ConfigSchema"
class ConfigSchema(NodeFileAbstractAction.ConfigSchema):
"""Configuration schema for NodeFileRestoreAction."""
verb: ClassVar[str] = "restore"
class NodeFileCorruptAction(NodeFileAbstractAction, identifier="node_file_corrupt"):
"""Action which corrupts a file."""
config: "NodeFileCorruptAction.ConfigSchema"
class ConfigSchema(NodeFileAbstractAction.ConfigSchema):
"""Configuration schema for NodeFileCorruptAction."""
verb: ClassVar[str] = "corrupt"
class NodeFileAccessAction(NodeFileAbstractAction, identifier="node_file_access"):
"""Action which increases a file's access count."""
config: "NodeFileAccessAction.ConfigSchema"
class ConfigSchema(NodeFileAbstractAction.ConfigSchema):
"""Configuration schema for NodeFileAccessAction."""
verb: ClassVar[str] = "access"
@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.folder_name is None or config.file_name is None:
return ["do_nothing"]
return [
"network",
"node",
config.node_name,
"file_system",
config.verb,
config.folder_name,
config.file_name,
]
class NodeFileCheckhashAction(NodeFileAbstractAction, identifier="node_file_checkhash"):
"""Action which checks the hash of a file."""
config: "NodeFileCheckhashAction.ConfigSchema"
class ConfigSchema(NodeFileAbstractAction.ConfigSchema):
"""Configuration schema for NodeFileCheckhashAction."""
verb: ClassVar[str] = "checkhash"
class NodeFileRepairAction(NodeFileAbstractAction, identifier="node_file_repair"):
"""Action which repairs a file."""
config: "NodeFileRepairAction.ConfigSchema"
class ConfigSchema(NodeFileAbstractAction.ConfigSchema):
"""Configuration Schema for NodeFileRepairAction."""
verb: ClassVar[str] = "repair"

117
src/primaite/game/agent/actions/folder.py Normal file
View File

@@ -0,0 +1,117 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from abc import ABC
from typing import ClassVar
from primaite.game.agent.actions.manager import AbstractAction
from primaite.interface.request import RequestFormat
__all__ = (
"NodeFolderScanAction",
"NodeFolderCheckhashAction",
"NodeFolderRepairAction",
"NodeFolderRestoreAction",
"NodeFolderCreateAction",
)
class NodeFolderAbstractAction(AbstractAction, ABC):
"""
Base class for folder actions.
Any action which applies to a folder and uses node_name and folder_name as its only two parameters can inherit from
this base class.
"""
config: "NodeFolderAbstractAction.ConfigSchema"
class ConfigSchema(AbstractAction.ConfigSchema):
"""Base configuration schema for NodeFolder actions."""
node_name: str
folder_name: str
verb: ClassVar[str]
@classmethod
def form_request(cls, config: ConfigSchema) -> RequestFormat:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
if config.node_name is None or config.folder_name is None:
return ["do_nothing"]
return [
"network",
"node",
config.node_name,
"file_system",
"folder",
config.folder_name,
config.verb,
]
class NodeFolderScanAction(NodeFolderAbstractAction, identifier="node_folder_scan"):
"""Action which scans a folder."""
config: "NodeFolderScanAction.ConfigSchema"
class ConfigSchema(NodeFolderAbstractAction.ConfigSchema):
"""Configuration schema for NodeFolderScanAction."""
verb: ClassVar[str] = "scan"
class NodeFolderCheckhashAction(NodeFolderAbstractAction, identifier="node_folder_checkhash"):
"""Action which checks the hash of a folder."""
config: "NodeFolderCheckhashAction.ConfigSchema"
class ConfigSchema(NodeFolderAbstractAction.ConfigSchema):
"""Configuration schema for NodeFolderCheckhashAction."""
verb: ClassVar[str] = "checkhash"
class NodeFolderRepairAction(NodeFolderAbstractAction, identifier="node_folder_repair"):
"""Action which repairs a folder."""
config: "NodeFolderRepairAction.ConfigSchema"
class ConfigSchema(NodeFolderAbstractAction.ConfigSchema):
"""Configuration schema for NodeFolderRepairAction."""
verb: ClassVar[str] = "repair"
class NodeFolderRestoreAction(NodeFolderAbstractAction, identifier="node_folder_restore"):
"""Action which restores a folder."""
config: "NodeFolderRestoreAction.ConfigSchema"
class ConfigSchema(NodeFolderAbstractAction.ConfigSchema):
"""Configuration schema for NodeFolderRestoreAction."""
verb: ClassVar[str] = "restore"
class NodeFolderCreateAction(NodeFolderAbstractAction, identifier="node_folder_create"):
"""Action which creates a new folder."""
config: "NodeFolderCreateAction.ConfigSchema"
class ConfigSchema(NodeFolderAbstractAction.ConfigSchema):
"""Configuration schema for NodeFolderCreateAction."""
verb: ClassVar[str] = "create"
@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.folder_name is None:
return ["do_nothing"]
return [
"network",
"node",
config.node_name,
"file_system",
config.verb,
"folder",
config.folder_name,
]

62
src/primaite/game/agent/actions/host_nic.py Normal file
View File

@@ -0,0 +1,62 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from abc import ABC
from typing import ClassVar
from primaite.game.agent.actions.manager import AbstractAction
from primaite.interface.request import RequestFormat
__all__ = ("HostNICEnableAction", "HostNICDisableAction")
class HostNICAbstractAction(AbstractAction, ABC):
"""
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.
"""
config: "HostNICAbstractAction.ConfigSchema"
class ConfigSchema(AbstractAction.ConfigSchema):
"""Base Configuration schema for HostNIC actions."""
node_name: str
nic_num: int
verb: ClassVar[str]
@classmethod
def form_request(cls, config: ConfigSchema) -> RequestFormat:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
if config.node_name is None or config.nic_num is None:
return ["do_nothing"]
return [
"network",
"node",
config.node_name,
"network_interface",
config.nic_num,
config.verb,
]
class HostNICEnableAction(HostNICAbstractAction, identifier="host_nic_enable"):
"""Action which enables a NIC."""
config: "HostNICEnableAction.ConfigSchema"
class ConfigSchema(HostNICAbstractAction.ConfigSchema):
"""Configuration schema for HostNICEnableAction."""
verb: ClassVar[str] = "enable"
class HostNICDisableAction(HostNICAbstractAction, identifier="host_nic_disable"):
"""Action which disables a NIC."""
config: "HostNICDisableAction.ConfigSchema"
class ConfigSchema(HostNICAbstractAction.ConfigSchema):
"""Configuration schema for HostNICDisableAction."""
verb: ClassVar[str] = "disable"

138
src/primaite/game/agent/actions/manager.py Normal file
View File

@@ -0,0 +1,138 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
"""yaml example.
agents:
- name: agent_1
action_space:
actions:
- do_nothing
- node_service_start
- node_service_stop
action_map:
"""
from __future__ import annotations
from typing import Dict, List, Optional, Tuple
from gymnasium import spaces
# from primaite.game.game import PrimaiteGame # TODO: Breaks things
from primaite.game.agent.actions.abstract import AbstractAction
from primaite.interface.request import RequestFormat
__all__ = ("DoNothingAction", "ActionManager")
class DoNothingAction(AbstractAction, identifier="do_nothing"):
"""Do Nothing Action."""
class ConfigSchema(AbstractAction.ConfigSchema):
"""Configuration Schema for do_nothingAction."""
type: str = "do_nothing"
@classmethod
def form_request(cls, config: ConfigSchema) -> RequestFormat:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
return ["do_nothing"]
class ActionManager:
"""Class which manages the action space for an agent."""
def __init__(
self,
actions: List[Dict], # stores list of actions available to agent
nodes: List[Dict], # extra configuration for each node
act_map: Optional[
Dict[int, Dict]
] = None, # allows restricting set of possible actions - TODO: Refactor to be a list?
*args,
**kwargs,
) -> None:
"""Init method for ActionManager.
:param game: Reference to the game to which the agent belongs.
:type game: PrimaiteGame
:param actions: List of action specs which should be made available to the agent. The keys of each spec are:
'type' and 'options' for passing any options to the action class's init method
:type actions: List[dict]
:param act_map: Action map which maps integers to actions. Used for restricting the set of possible actions.
:type act_map: Optional[Dict[int, Dict]]
"""
self.actions: Dict[str, AbstractAction] = {}
for act_spec in actions:
act_type = act_spec.get("type")
self.actions[act_type] = AbstractAction._registry[act_type]
self.action_map: Dict[int, Tuple[str, Dict]] = {}
"""
Action mapping that converts an integer to a specific action and parameter choice.
For example :
{0: ("node_service_scan", {node_name:"client_1", service_name:"WebBrowser"})}
"""
if act_map is None:
# raise RuntimeError("Action map must be specified in the config file.")
pass
else:
self.action_map = {i: (a["action"], a["options"]) for i, a in act_map.items()}
# make sure all numbers between 0 and N are represented as dict keys in action map
assert all([i in self.action_map.keys() for i in range(len(self.action_map))])
def get_action(self, action: int) -> Tuple[str, Dict]:
"""Produce action in CAOS format."""
"""the agent chooses an action (as an integer), this is converted into an action in CAOS format"""
"""The CAOS format is basically a action identifier, followed by parameters stored in a dictionary"""
act_identifier, act_options = self.action_map[action]
return act_identifier, act_options
def form_request(self, action_identifier: str, action_options: Dict) -> RequestFormat:
"""Take action in CAOS format and use the execution definition to change it into PrimAITE request format."""
act_class = AbstractAction._registry[action_identifier]
config = act_class.ConfigSchema(**action_options)
return act_class.form_request(config=config)
@property
def space(self) -> spaces.Space:
"""Return the gymnasium action space for this agent."""
return spaces.Discrete(len(self.action_map))
@classmethod
def from_config(cls, game: "PrimaiteGame", cfg: Dict) -> "ActionManager": # noqa: F821
"""
Construct an ActionManager from a config definition.
The action space config supports the following three sections:
1. ``action_list``
``action_list`` contains a list action components which need to be included in the action space.
Each action component has a ``type`` which maps to a subclass of AbstractAction, and additional options
which will be passed to the action class's __init__ method during initialisation.
2. ``action_map``
Since the agent uses a discrete action space which acts as a flattened version of the component-based
action space, action_map provides a mapping between an integer (chosen by the agent) and a meaningful
action and values of parameters. For example action 0 can correspond to do nothing, action 1 can
correspond to "node_service_scan" with ``node_name="server"`` and
``service_name="WebBrowser"``, action 2 can be "
3. ``options``
``options`` contains a dictionary of options which are passed to the ActionManager's __init__ method.
These options are used to calculate the shape of the action space, and to provide additional information
to the ActionManager which is required to convert the agent's action choice into a CAOS request.
:param game: The Primaite Game to which the agent belongs.
:type game: PrimaiteGame
:param cfg: The action space config.
:type cfg: Dict
:return: The constructed ActionManager.
:rtype: ActionManager
"""
obj = cls(
actions=cfg["action_list"],
**cfg["options"],
protocols=game.options.protocols,
ports=game.options.ports,
act_map=cfg.get("action_map"),
)
return obj

57
src/primaite/game/agent/actions/network.py Normal file
View File

@@ -0,0 +1,57 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from typing import ClassVar
from primaite.game.agent.actions.manager import AbstractAction
from primaite.interface.request import RequestFormat
__all__ = ("NetworkPortEnableAction", "NetworkPortDisableAction")
class NetworkPortAbstractAction(AbstractAction, identifier="network_port_abstract"):
"""Base class for Network port actions."""
config: "NetworkPortAbstractAction.ConfigSchema"
class ConfigSchema(AbstractAction.ConfigSchema):
"""Base configuration schema for NetworkPort actions."""
target_nodename: str
port_id: int
verb: ClassVar[str]
@classmethod
def form_request(cls, config: ConfigSchema) -> RequestFormat:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
if config.target_nodename is None or config.port_id is None:
return ["do_nothing"]
return [
"network",
"node",
config.target_nodename,
"network_interface",
config.port_id,
config.verb,
]
class NetworkPortEnableAction(NetworkPortAbstractAction, identifier="network_port_enable"):
"""Action which enables are port on a router or a firewall."""
config: "NetworkPortEnableAction.ConfigSchema"
class ConfigSchema(NetworkPortAbstractAction.ConfigSchema):
"""Configuration schema for NetworkPortEnableAction."""
verb: ClassVar[str] = "enable"
class NetworkPortDisableAction(NetworkPortAbstractAction, identifier="network_port_disable"):
"""Action which disables are port on a router or a firewall."""
config: "NetworkPortDisableAction.ConfigSchema"
class ConfigSchema(NetworkPortAbstractAction.ConfigSchema):
"""Configuration schema for NetworkPortDisableAction."""
verb: ClassVar[str] = "disable"

195
src/primaite/game/agent/actions/node.py Normal file
View File

@@ -0,0 +1,195 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from abc import abstractmethod
from typing import ClassVar, List, Optional, Union
from primaite.game.agent.actions.manager import AbstractAction
from primaite.interface.request import RequestFormat
__all__ = (
"NodeOSScanAction",
"NodeShutdownAction",
"NodeStartupAction",
"NodeResetAction",
"NodeNMAPPingScanAction",
"NodeNMAPPortScanAction",
"NodeNetworkServiceReconAction",
)
class NodeAbstractAction(AbstractAction, identifier="node_abstract"):
"""
Abstract base class for node actions.
Any action which applies to a node and uses node_name as its only parameter can inherit from this base class.
"""
config: "NodeAbstractAction.ConfigSchema"
class ConfigSchema(AbstractAction.ConfigSchema):
"""Base Configuration schema for Node actions."""
node_name: str
verb: ClassVar[str]
@classmethod
def form_request(cls, config: ConfigSchema) -> RequestFormat:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
print(config)
return ["network", "node", config.node_name, config.verb]
class NodeOSScanAction(NodeAbstractAction, identifier="node_os_scan"):
"""Action which scans a node's OS."""
config: "NodeOSScanAction.ConfigSchema"
class ConfigSchema(NodeAbstractAction.ConfigSchema):
"""Configuration schema for NodeOSScanAction."""
verb: ClassVar[str] = "scan"
class NodeShutdownAction(NodeAbstractAction, identifier="node_shutdown"):
"""Action which shuts down a node."""
config: "NodeShutdownAction.ConfigSchema"
class ConfigSchema(NodeAbstractAction.ConfigSchema):
"""Configuration schema for NodeShutdownAction."""
verb: ClassVar[str] = "shutdown"
class NodeStartupAction(NodeAbstractAction, identifier="node_startup"):
"""Action which starts up a node."""
config: "NodeStartupAction.ConfigSchema"
class ConfigSchema(NodeAbstractAction.ConfigSchema):
"""Configuration schema for NodeStartupAction."""
verb: ClassVar[str] = "startup"
class NodeResetAction(NodeAbstractAction, identifier="node_reset"):
"""Action which resets a node."""
config: "NodeResetAction.ConfigSchema"
class ConfigSchema(NodeAbstractAction.ConfigSchema):
"""Configuration schema for NodeResetAction."""
verb: ClassVar[str] = "reset"
class NodeNMAPAbstractAction(AbstractAction, identifier="node_nmap_abstract_action"):
"""Base class for NodeNMAP actions."""
config: "NodeNMAPAbstractAction.ConfigSchema"
class ConfigSchema(AbstractAction.ConfigSchema):
"""Base Configuration Schema for NodeNMAP actions."""
target_ip_address: Union[str, List[str]]
show: bool = False
node_name: str
@classmethod
@abstractmethod
def form_request(cls, config: ConfigSchema) -> RequestFormat:
# NMAP action requests don't share a common format for their requests
# This is just a placeholder to ensure the method is defined.
pass
class NodeNMAPPingScanAction(NodeNMAPAbstractAction, identifier="node_nmap_ping_scan"):
"""Action which performs an NMAP ping scan."""
config: "NodeNMAPPingScanAction.ConfigSchema"
class ConfigSchema(NodeNMAPAbstractAction.ConfigSchema):
"""Configuration schema for NodeNMAPPingScanAction."""
pass
@classmethod
def form_request(cls, config: ConfigSchema) -> List[str]: # noqa
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
return [
"network",
"node",
config.node_name,
"application",
"NMAP",
"ping_scan",
{"target_ip_address": config.target_ip_address, "show": config.show},
]
class NodeNMAPPortScanAction(NodeNMAPAbstractAction, identifier="node_nmap_port_scan"):
"""Action which performs an NMAP port scan."""
config: "NodeNMAPPortScanAction.ConfigSchema"
class ConfigSchema(NodeNMAPAbstractAction.ConfigSchema):
"""Configuration Schema for NodeNMAPPortScanAction."""
source_node: str
target_protocol: Optional[Union[str, List[str]]] = (None,)
target_port: Optional[Union[str, List[str]]] = (None,)
show: Optional[bool] = (False,)
@classmethod
def form_request(
cls,
config: ConfigSchema,
) -> List[str]: # noqa
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
return [
"network",
"node",
config.source_node,
"application",
"NMAP",
"port_scan",
{
"target_ip_address": config.target_ip_address,
"target_port": config.target_port,
"target_protocol": config.target_protocol,
"show": config.show,
},
]
class NodeNetworkServiceReconAction(NodeNMAPAbstractAction, identifier="node_network_service_recon"):
"""Action which performs an NMAP network service recon (ping scan followed by port scan)."""
config: "NodeNetworkServiceReconAction.ConfigSchema"
class ConfigSchema(AbstractAction.ConfigSchema):
"""Configuration schema for NodeNetworkServiceReconAction."""
target_protocol: Optional[Union[str, List[str]]] = (None,)
target_port: Optional[Union[str, List[str]]] = (None,)
show: Optional[bool] = (False,)
@classmethod
def form_request(
cls,
config: ConfigSchema,
) -> List[str]: # noqa
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
return [
"network",
"node",
config.source_node,
"application",
"NMAP",
"network_service_recon",
{
"target_ip_address": config.target_ip_address,
"target_port": config.target_port,
"target_protocol": config.target_protocol,
"show": config.show,
},
]

135
src/primaite/game/agent/actions/service.py Normal file
View File

@@ -0,0 +1,135 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from typing import ClassVar
from primaite.game.agent.actions.manager import AbstractAction
from primaite.interface.request import RequestFormat
__all__ = (
"NodeServiceScanAction",
"NodeServiceStopAction",
"NodeServiceStartAction",
"NodeServicePauseAction",
"NodeServiceResumeAction",
"NodeServiceRestartAction",
"NodeServiceDisableAction",
"NodeServiceEnableAction",
"NodeServiceFixAction",
)
class NodeServiceAbstractAction(AbstractAction, identifier="node_service_abstract"):
"""Abstract Action for Node Service related actions.
Any actions which use node_name and service_name can inherit from this class.
"""
config: "NodeServiceAbstractAction.ConfigSchema"
class ConfigSchema(AbstractAction.ConfigSchema):
node_name: str
service_name: str
verb: ClassVar[str]
@classmethod
def form_request(cls, config: ConfigSchema) -> RequestFormat:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
return ["network", "node", config.node_name, "service", config.service_name, config.verb]
class NodeServiceScanAction(NodeServiceAbstractAction, identifier="node_service_scan"):
"""Action which scans a service."""
config: "NodeServiceScanAction.ConfigSchema"
class ConfigSchema(NodeServiceAbstractAction.ConfigSchema):
"""Configuration Schema for NodeServiceScanAction."""
verb: ClassVar[str] = "scan"
class NodeServiceStopAction(NodeServiceAbstractAction, identifier="node_service_stop"):
"""Action which stops a service."""
config: "NodeServiceStopAction.ConfigSchema"
class ConfigSchema(NodeServiceAbstractAction.ConfigSchema):
"""Configuration Schema for NodeServiceStopAction."""
verb: ClassVar[str] = "stop"
class NodeServiceStartAction(NodeServiceAbstractAction, identifier="node_service_start"):
"""Action which starts a service."""
config: "NodeServiceStartAction.ConfigSchema"
class ConfigSchema(NodeServiceAbstractAction.ConfigSchema):
"""Configuration Schema for NodeServiceStartAction."""
verb: ClassVar[str] = "start"
class NodeServicePauseAction(NodeServiceAbstractAction, identifier="node_service_pause"):
"""Action which pauses a service."""
config: "NodeServicePauseAction.ConfigSchema"
class ConfigSchema(NodeServiceAbstractAction.ConfigSchema):
"""Configuration Schema for NodeServicePauseAction."""
verb: ClassVar[str] = "pause"
class NodeServiceResumeAction(NodeServiceAbstractAction, identifier="node_service_resume"):
"""Action which resumes a service."""
config: "NodeServiceResumeAction.ConfigSchema"
class ConfigSchema(NodeServiceAbstractAction.ConfigSchema):
"""Configuration Schema for NodeServiceResumeAction."""
verb: ClassVar[str] = "resume"
class NodeServiceRestartAction(NodeServiceAbstractAction, identifier="node_service_restart"):
"""Action which restarts a service."""
config: "NodeServiceRestartAction.ConfigSchema"
class ConfigSchema(NodeServiceAbstractAction.ConfigSchema):
"""Configuration Schema for NodeServiceRestartAction."""
verb: ClassVar[str] = "restart"
class NodeServiceDisableAction(NodeServiceAbstractAction, identifier="node_service_disable"):
"""Action which disables a service."""
config: "NodeServiceDisableAction.ConfigSchema"
class ConfigSchema(NodeServiceAbstractAction.ConfigSchema):
"""Configuration Schema for NodeServiceDisableAction."""
verb: ClassVar[str] = "disable"
class NodeServiceEnableAction(NodeServiceAbstractAction, identifier="node_service_enable"):
"""Action which enables a service."""
config: "NodeServiceEnableAction.ConfigSchema"
class ConfigSchema(NodeServiceAbstractAction.ConfigSchema):
"""Configuration Schema for NodeServiceEnableAction."""
verb: ClassVar[str] = "enable"
class NodeServiceFixAction(NodeServiceAbstractAction, identifier="node_service_fix"):
"""Action which fixes a service."""
config: "NodeServiceFixAction.ConfigSchema"
class ConfigSchema(NodeServiceAbstractAction.ConfigSchema):
"""Configuration Schema for NodeServiceFixAction."""
verb: ClassVar[str] = "fix"

108
src/primaite/game/agent/actions/session.py Normal file
View File

@@ -0,0 +1,108 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from abc import abstractmethod
from primaite.game.agent.actions.manager import AbstractAction
from primaite.interface.request import RequestFormat
__all__ = (
"NodeSessionsRemoteLoginAction",
"NodeSessionsRemoteLogoutAction",
"NodeAccountChangePasswordAction",
)
class NodeSessionAbstractAction(AbstractAction, identifier="node_session_abstract"):
"""Base class for NodeSession actions."""
config: "NodeSessionAbstractAction.ConfigSchema"
class ConfigSchema(AbstractAction.ConfigSchema):
"""Base configuration schema for NodeSessionAbstractActions."""
node_name: str
remote_ip: str
@classmethod
@abstractmethod
def form_request(cls, config: ConfigSchema) -> RequestFormat:
"""
Abstract method for request forming.
Should return the action formatted as a request which can be ingested by the PrimAITE simulation.
"""
pass
class NodeSessionsRemoteLoginAction(NodeSessionAbstractAction, identifier="node_session_remote_login"):
"""Action which performs a remote session login."""
config: "NodeSessionsRemoteLoginAction.ConfigSchema"
class ConfigSchema(NodeSessionAbstractAction.ConfigSchema):
"""Configuration schema for NodeSessionsRemoteLoginAction."""
username: str
password: str
@classmethod
def form_request(cls, config: ConfigSchema) -> RequestFormat:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
if config.node_name is None or config.remote_ip is None:
return ["do_nothing"]
return [
"network",
"node",
config.node_name,
"service",
"Terminal",
"ssh_to_remote",
config.username,
config.password,
config.remote_ip,
]
class NodeSessionsRemoteLogoutAction(NodeSessionAbstractAction, identifier="node_session_remote_logoff"):
"""Action which performs a remote session logout."""
config: "NodeSessionsRemoteLogoutAction.ConfigSchema"
class ConfigSchema(NodeSessionAbstractAction.ConfigSchema):
"""Configuration schema for NodeSessionsRemoteLogoutAction."""
verb: str = "remote_logoff"
@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.remote_ip is None:
return ["do_nothing"]
return ["network", "node", config.node_name, "service", "Terminal", config.verb, config.remote_ip]
class NodeAccountChangePasswordAction(NodeSessionAbstractAction, identifier="node_account_change_password"):
"""Action which changes the password for a user."""
config: "NodeAccountChangePasswordAction.ConfigSchema"
class ConfigSchema(NodeSessionAbstractAction.ConfigSchema):
"""Configuration schema for NodeAccountsChangePasswordAction."""
username: str
current_password: str
new_password: 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.node_name,
"service",
"UserManager",
"change_password",
config.username,
config.current_password,
config.new_password,
]

238
src/primaite/game/agent/actions/software.py Normal file
View File

@@ -0,0 +1,238 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from typing import List, Optional, Union
from pydantic import ConfigDict, Field, field_validator, ValidationInfo
from primaite.game.agent.actions.manager import AbstractAction, ActionManager
from primaite.interface.request import RequestFormat
__all__ = (
"ConfigureRansomwareScriptAction",
"ConfigureDoSBotAction",
"ConfigureC2BeaconAction",
"NodeSendRemoteCommandAction",
"TerminalC2ServerAction",
"RansomwareLaunchC2ServerAction",
"ExfiltrationC2ServerAction",
"ConfigureDatabaseClientAction",
)
class ConfigureRansomwareScriptAction(AbstractAction, identifier="c2_server_ransomware_configure"):
"""Action which sets config parameters for a ransomware script on a node."""
config: "ConfigureRansomwareScriptAction.ConfigSchema"
class ConfigSchema(AbstractAction.ConfigSchema):
"""Configuration schema for ConfigureRansomwareScriptAction."""
node_name: str
server_ip_address: Optional[str]
server_password: Optional[str]
payload: Optional[str]
@classmethod
def form_request(cls, config: ConfigSchema) -> RequestFormat:
"""Return the action formatted as a request that can be ingested by the simulation."""
if config.node_name is None:
return ["do_nothing"]
return [
"network",
"node",
config.node_name,
"application",
"RansomwareScript",
"configure",
config.model_config,
]
class ConfigureDoSBotAction(AbstractAction, identifier="configure_dos_bot"):
"""Action which sets config parameters for a DoS bot on a node."""
config: "ConfigureDoSBotAction.ConfigSchema"
class ConfigSchema(AbstractAction.ConfigSchema):
"""Schema for options that can be passed to this action."""
node_name: str
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, 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"]
self.ConfigSchema.model_validate(config) # check that options adhere to schema
return ["network", "node", config.node_name, "application", "DoSBot", "configure", config]
class ConfigureC2BeaconAction(AbstractAction, identifier="configure_c2_beacon"):
"""Action which configures a C2 Beacon based on the parameters given."""
config: "ConfigureC2BeaconAction.ConfigSchema"
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")
@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."""
return ["network", "node", config.node_name, "application", "C2Beacon", "configure", config]
class NodeSendRemoteCommandAction(AbstractAction, identifier="node_send_remote_command"):
"""Action which sends a terminal command to a remote node via SSH."""
config: "NodeSendRemoteCommandAction.ConfigSchema"
class ConfigSchema(AbstractAction.ConfigSchema):
"""Configuration schema for NodeSendRemoteCommandAction."""
node_name: str
remote_ip: str
command: RequestFormat
@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.node_name,
"service",
"Terminal",
"send_remote_command",
config.remote_ip,
{"command": config.command},
]
class TerminalC2ServerAction(AbstractAction, identifier="c2_server_terminal_command"):
"""Action which causes the C2 Server to send a command to the C2 Beacon to execute the terminal command passed."""
config: "TerminalC2ServerAction.ConfigSchema"
class ConfigSchema(AbstractAction.ConfigSchema):
"""Schema for options that can be passed to this action."""
node_name: str
commands: Union[List[RequestFormat], RequestFormat]
ip_address: Optional[str]
username: Optional[str]
password: Optional[str]
@classmethod
def form_request(cls, config: ConfigSchema) -> RequestFormat:
"""Return the action formatted as a request that can be ingested by the simulation."""
if config.node_name is None:
return ["do_nothing"]
command_model = {
"commands": config.commands,
"ip_address": config.ip_address,
"username": config.username,
"password": config.password,
}
return ["network", "node", config.node_name, "application", "C2Server", "terminal_command", command_model]
class RansomwareLaunchC2ServerAction(AbstractAction, identifier="c2_server_ransomware_launch"):
"""Action which causes the C2 Server to send a command to the C2 Beacon to launch the RansomwareScript."""
config: "RansomwareLaunchC2ServerAction.ConfigSchema"
class ConfigSchema(AbstractAction.ConfigSchema):
"""Configuration schema for RansomwareLaunchC2ServerAction."""
node_name: str
@classmethod
def form_request(cls, config: ConfigSchema) -> RequestFormat:
"""Return the action formatted as a request that can be ingested by the simulation."""
if config.node_name is None:
return ["do_nothing"]
# This action currently doesn't require any further configuration options.
return ["network", "node", config.node_name, "application", "C2Server", "ransomware_launch"]
class ExfiltrationC2ServerAction(AbstractAction, identifier="c2_server_data_exfiltrate"):
"""Action which exfiltrates a target file from a certain node onto the C2 beacon and then the C2 Server."""
config: "ExfiltrationC2ServerAction.ConfigSchema"
class ConfigSchema(AbstractAction.ConfigSchema):
"""Schema for options that can be passed to this action."""
node_name: str
username: Optional[str]
password: Optional[str]
target_ip_address: str
target_file_name: str
target_folder_name: str
exfiltration_folder_name: Optional[str]
@classmethod
def form_request(cls, config: ConfigSchema) -> RequestFormat:
"""Return the action formatted as a request that can be ingested by the simulation."""
if config.node_name is None:
return ["do_nothing"]
command_model = {
"target_file_name": config.target_file_name,
"target_folder_name": config.target_folder_name,
"exfiltration_folder_name": config.exfiltration_folder_name,
"target_ip_address": config.target_ip_address,
"username": config.username,
"password": config.password,
}
return ["network", "node", config.node_name, "application", "C2Server", "exfiltrate", command_model]
class ConfigureDatabaseClientAction(AbstractAction, identifier="configure_database_client"):
"""Action which sets config parameters for a database client on a node."""
config: "ConfigureDatabaseClientAction.ConfigSchema"
class ConfigSchema(AbstractAction.ConfigSchema):
"""Schema for options that can be passed to this action."""
node_name: str
model_config = ConfigDict(extra="forbid")
@classmethod
def form_request(cls, config: ConfigSchema) -> RequestFormat:
"""Return the action formatted as a request that can be ingested by the simulation."""
if config.node_name is None:
return ["do_nothing"]
return ["network", "node", config.node_name, "application", "DatabaseClient", "configure", config.model_config]