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

#3062 - rename identifier to discriminator

Browse Source
This commit is contained in:
Marek Wolan
2025-01-31 16:00:32 +00:00
parent 8feb2db954
commit 055c853b0f
78 changed files with 223 additions and 213 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
docs/source/configuration/agents.rst
View File

@@ -103,7 +103,7 @@ Similar to action space, this is defined as a list of components from the :py:mo
``reward_components`` ``reward_components``
^^^^^^^^^^^^^^^^^^^^^ ^^^^^^^^^^^^^^^^^^^^^
TODO: update description
A list of reward types from :py:mod:`primaite.game.agent.rewards.RewardFunction.rew_class_identifiers` A list of reward types from :py:mod:`primaite.game.agent.rewards.RewardFunction.rew_class_identifiers`
e.g. e.g.

2
docs/source/configuration/simulation.rst
View File

@@ -6,7 +6,7 @@
``simulation`` ``simulation``
============== ==============
In this section the network layout is defined. This part of the config follows a hierarchical structure. Almost every component defines a ``ref`` field which acts as a human-readable unique identifier, used by other parts of the config, such as agents. In this section the network layout is defined. This part of the config follows a hierarchical structure. Almost every component defines a ``ref`` field which acts as a human-readable unique identifier, used by other parts of the config, such as agents.
# TODO: ref field is no longer real
At the top level of the network are ``nodes``, ``links`` and ``airspace``. At the top level of the network are ``nodes``, ``links`` and ``airspace``.
e.g. e.g.

10
docs/source/how_to_guides/extensible_actions.rst
View File

@@ -20,7 +20,7 @@ Custom actions within PrimAITE must be a sub-class of `AbstractAction`, and cont
#. ConfigSchema class #. ConfigSchema class
#. Unique Identifier #. Unique discriminator
#. `form_request` method. #. `form_request` method.
@@ -31,14 +31,14 @@ ConfigSchema
The ConfigSchema sub-class of the action must contain all `configurable` variables within the action, that would be specified within the environments configuration YAML file. The ConfigSchema sub-class of the action must contain all `configurable` variables within the action, that would be specified within the environments configuration YAML file.
Unique Identifier Unique discriminator
################# #################
When declaring a custom class, it must have a unique identifier string, that allows PrimAITE to generate the correct action when needed. When declaring a custom class, it must have a unique discriminator string, that allows PrimAITE to generate the correct action when needed.
.. code:: Python .. code:: Python
class CreateDirectoryAction(AbstractAction, identifier="node_folder_create") class CreateDirectoryAction(AbstractAction, discriminator="node_folder_create")
config: CreateDirectoryAction.ConfigSchema config: CreateDirectoryAction.ConfigSchema
@@ -58,7 +58,7 @@ When declaring a custom class, it must have a unique identifier string, that all
config.directory_name, config.directory_name,
] ]
The above action would fail pydantic validation as the identifier "node_folder_create" is already used by the `NodeFolderCreateAction`, and would create a duplicate listing within `AbstractAction._registry`. The above action would fail pydantic validation as the discriminator "node_folder_create" is already used by the `NodeFolderCreateAction`, and would create a duplicate listing within `AbstractAction._registry`.
form_request method form_request method

6
docs/source/how_to_guides/extensible_agents.rst
View File

@@ -25,7 +25,7 @@ The core features that should be implemented in any new agent are detailed below
.. code-block:: python .. code-block:: python
class ExampleAgent(AbstractAgent, identifier = "ExampleAgent"): class ExampleAgent(AbstractAgent, discriminator = "ExampleAgent"):
"""An example agent for demonstration purposes.""" """An example agent for demonstration purposes."""
config: "ExampleAgent.ConfigSchema" = Field(default_factory= lambda: ExampleAgent.ConfigSchema()) config: "ExampleAgent.ConfigSchema" = Field(default_factory= lambda: ExampleAgent.ConfigSchema())
@@ -64,9 +64,9 @@ The core features that should be implemented in any new agent are detailed below
starting_host: "Server_1" starting_host: "Server_1"
#. **Identifiers**: #. **discriminators**:
All agent classes should have an ``identifier`` attribute, a unique kebab-case string, for when they are added to the base ``AbstractAgent`` registry. This is then specified in your configuration YAML, and used by PrimAITE to generate the correct Agent. All agent classes should have an ``discriminator`` attribute, a unique kebab-case string, for when they are added to the base ``AbstractAgent`` registry. This is then specified in your configuration YAML, and used by PrimAITE to generate the correct Agent.
Changes to YAML file Changes to YAML file
==================== ====================

4
docs/source/how_to_guides/extensible_rewards.rst
View File

@@ -17,7 +17,7 @@ Reward classes are inherited from AbstractReward (a sub-class of Pydantic's Base
Within the reward class there is a ConfigSchema class responsible for ensuring the config file data Within the reward class there is a ConfigSchema class responsible for ensuring the config file data
is in the correct format. This also means there is little (if no) requirement for and `__init__` is in the correct format. This also means there is little (if no) requirement for and `__init__`
method. The `.from_config` method is no longer required as it's inherited from `AbstractReward`. method. The `.from_config` method is no longer required as it's inherited from `AbstractReward`.
Each class requires an identifier string which is used by the ConfigSchema class to verify that it Each class requires an discriminator string which is used by the ConfigSchema class to verify that it
hasn't previously been added to the registry. hasn't previously been added to the registry.
Inheriting from `BaseModel` removes the need for an `__init__` method but means that object Inheriting from `BaseModel` removes the need for an `__init__` method but means that object
@@ -28,7 +28,7 @@ To add a new reward class follow the example below. Note that the type attribute
.. code-block:: Python .. code-block:: Python
class DatabaseFileIntegrity(AbstractReward, identifier="DATABASE_FILE_INTEGRITY"): class DatabaseFileIntegrity(AbstractReward, discriminator="DATABASE_FILE_INTEGRITY"):
"""Reward function component which rewards the agent for maintaining the integrity of a database file.""" """Reward function component which rewards the agent for maintaining the integrity of a database file."""
config: "DatabaseFileIntegrity.ConfigSchema" config: "DatabaseFileIntegrity.ConfigSchema"

2
docs/source/node_sets.rst
View File

@@ -82,7 +82,7 @@ Here is an example of creating a custom node adder, DataCenterAdder:
.. code-block:: python .. code-block:: python
class DataCenterAdder(NetworkNodeAdder, identifier="data_center"): class DataCenterAdder(NetworkNodeAdder, discriminator="data_center"):
class ConfigSchema(NetworkNodeAdder.ConfigSchema): class ConfigSchema(NetworkNodeAdder.ConfigSchema):
type: Literal["data_center"] = "data_center" type: Literal["data_center"] = "data_center"
num_servers: int num_servers: int

17
src/primaite/game/agent/actions/abstract.py
View File

@@ -22,13 +22,20 @@ class AbstractAction(BaseModel, ABC):
_registry: ClassVar[Dict[str, Type[AbstractAction]]] = {} _registry: ClassVar[Dict[str, Type[AbstractAction]]] = {}
def __init_subclass__(cls, identifier: Optional[str] = None, **kwargs: Any) -> None: def __init_subclass__(cls, discriminator: Optional[str] = None, **kwargs: Any) -> None:
"""
Register an action type.
:param discriminator: discriminator used to uniquely specify action types.
:type discriminator: str
:raises ValueError: When attempting to create an action with a name that is already in use.
"""
super().__init_subclass__(**kwargs) super().__init_subclass__(**kwargs)
if identifier is None: if discriminator is None:
return return
if identifier in cls._registry: if discriminator in cls._registry:
raise ValueError(f"Cannot create new action under reserved name {identifier}") raise ValueError(f"Cannot create new action under reserved name {discriminator}")
cls._registry[identifier] = cls cls._registry[discriminator] = cls
@classmethod @classmethod
def form_request(cls, config: ConfigSchema) -> RequestFormat: def form_request(cls, config: ConfigSchema) -> RequestFormat:

10
src/primaite/game/agent/actions/acl.py
View File

@@ -37,7 +37,7 @@ class ACLAddRuleAbstractAction(AbstractAction, ABC):
dst_wildcard: Union[IPV4Address, Literal["NONE"]] dst_wildcard: Union[IPV4Address, Literal["NONE"]]
class ACLRemoveRuleAbstractAction(AbstractAction, identifier="acl_remove_rule_abstract_action"): class ACLRemoveRuleAbstractAction(AbstractAction, discriminator="acl_remove_rule_abstract_action"):
"""Base abstract class for ACL remove rule actions.""" """Base abstract class for ACL remove rule actions."""
config: ConfigSchema = "ACLRemoveRuleAbstractAction.ConfigSchema" config: ConfigSchema = "ACLRemoveRuleAbstractAction.ConfigSchema"
@@ -48,7 +48,7 @@ class ACLRemoveRuleAbstractAction(AbstractAction, identifier="acl_remove_rule_ab
position: int position: int
class RouterACLAddRuleAction(ACLAddRuleAbstractAction, identifier="router_acl_add_rule"): class RouterACLAddRuleAction(ACLAddRuleAbstractAction, discriminator="router_acl_add_rule"):
"""Action which adds a rule to a router's ACL.""" """Action which adds a rule to a router's ACL."""
config: "RouterACLAddRuleAction.ConfigSchema" config: "RouterACLAddRuleAction.ConfigSchema"
@@ -79,7 +79,7 @@ class RouterACLAddRuleAction(ACLAddRuleAbstractAction, identifier="router_acl_ad
] ]
class RouterACLRemoveRuleAction(ACLRemoveRuleAbstractAction, identifier="router_acl_remove_rule"): class RouterACLRemoveRuleAction(ACLRemoveRuleAbstractAction, discriminator="router_acl_remove_rule"):
"""Action which removes a rule from a router's ACL.""" """Action which removes a rule from a router's ACL."""
config: "RouterACLRemoveRuleAction.ConfigSchema" config: "RouterACLRemoveRuleAction.ConfigSchema"
@@ -95,7 +95,7 @@ class RouterACLRemoveRuleAction(ACLRemoveRuleAbstractAction, identifier="router_
return ["network", "node", config.target_router, "acl", "remove_rule", config.position] return ["network", "node", config.target_router, "acl", "remove_rule", config.position]
class FirewallACLAddRuleAction(ACLAddRuleAbstractAction, identifier="firewall_acl_add_rule"): class FirewallACLAddRuleAction(ACLAddRuleAbstractAction, discriminator="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."""
config: "FirewallACLAddRuleAction.ConfigSchema" config: "FirewallACLAddRuleAction.ConfigSchema"
@@ -130,7 +130,7 @@ class FirewallACLAddRuleAction(ACLAddRuleAbstractAction, identifier="firewall_ac
] ]
class FirewallACLRemoveRuleAction(ACLRemoveRuleAbstractAction, identifier="firewall_acl_remove_rule"): class FirewallACLRemoveRuleAction(ACLRemoveRuleAbstractAction, discriminator="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."""
config: "FirewallACLRemoveRuleAction.ConfigSchema" config: "FirewallACLRemoveRuleAction.ConfigSchema"

12
src/primaite/game/agent/actions/application.py
View File

@@ -45,7 +45,7 @@ class NodeApplicationAbstractAction(AbstractAction, ABC):
] ]
class NodeApplicationExecuteAction(NodeApplicationAbstractAction, identifier="node_application_execute"): class NodeApplicationExecuteAction(NodeApplicationAbstractAction, discriminator="node_application_execute"):
"""Action which executes an application.""" """Action which executes an application."""
config: "NodeApplicationExecuteAction.ConfigSchema" config: "NodeApplicationExecuteAction.ConfigSchema"
@@ -56,7 +56,7 @@ class NodeApplicationExecuteAction(NodeApplicationAbstractAction, identifier="no
verb: str = "execute" verb: str = "execute"
class NodeApplicationScanAction(NodeApplicationAbstractAction, identifier="node_application_scan"): class NodeApplicationScanAction(NodeApplicationAbstractAction, discriminator="node_application_scan"):
"""Action which scans an application.""" """Action which scans an application."""
config: "NodeApplicationScanAction.ConfigSchema" config: "NodeApplicationScanAction.ConfigSchema"
@@ -67,7 +67,7 @@ class NodeApplicationScanAction(NodeApplicationAbstractAction, identifier="node_
verb: str = "scan" verb: str = "scan"
class NodeApplicationCloseAction(NodeApplicationAbstractAction, identifier="node_application_close"): class NodeApplicationCloseAction(NodeApplicationAbstractAction, discriminator="node_application_close"):
"""Action which closes an application.""" """Action which closes an application."""
config: "NodeApplicationCloseAction.ConfigSchema" config: "NodeApplicationCloseAction.ConfigSchema"
@@ -78,7 +78,7 @@ class NodeApplicationCloseAction(NodeApplicationAbstractAction, identifier="node
verb: str = "close" verb: str = "close"
class NodeApplicationFixAction(NodeApplicationAbstractAction, identifier="node_application_fix"): class NodeApplicationFixAction(NodeApplicationAbstractAction, discriminator="node_application_fix"):
"""Action which fixes an application.""" """Action which fixes an application."""
config: "NodeApplicationFixAction.ConfigSchema" config: "NodeApplicationFixAction.ConfigSchema"
@@ -89,7 +89,7 @@ class NodeApplicationFixAction(NodeApplicationAbstractAction, identifier="node_a
verb: str = "fix" verb: str = "fix"
class NodeApplicationInstallAction(NodeApplicationAbstractAction, identifier="node_application_install"): class NodeApplicationInstallAction(NodeApplicationAbstractAction, discriminator="node_application_install"):
"""Action which installs an application.""" """Action which installs an application."""
config: "NodeApplicationInstallAction.ConfigSchema" config: "NodeApplicationInstallAction.ConfigSchema"
@@ -113,7 +113,7 @@ class NodeApplicationInstallAction(NodeApplicationAbstractAction, identifier="no
] ]
class NodeApplicationRemoveAction(NodeApplicationAbstractAction, identifier="node_application_remove"): class NodeApplicationRemoveAction(NodeApplicationAbstractAction, discriminator="node_application_remove"):
"""Action which removes/uninstalls an application.""" """Action which removes/uninstalls an application."""
config: "NodeApplicationRemoveAction.ConfigSchema" config: "NodeApplicationRemoveAction.ConfigSchema"

16
src/primaite/game/agent/actions/file.py
View File

@@ -52,7 +52,7 @@ class NodeFileAbstractAction(AbstractAction, ABC):
] ]
class NodeFileCreateAction(NodeFileAbstractAction, identifier="node_file_create"): class NodeFileCreateAction(NodeFileAbstractAction, discriminator="node_file_create"):
"""Action which creates a new file in a given folder.""" """Action which creates a new file in a given folder."""
config: "NodeFileCreateAction.ConfigSchema" config: "NodeFileCreateAction.ConfigSchema"
@@ -81,7 +81,7 @@ class NodeFileCreateAction(NodeFileAbstractAction, identifier="node_file_create"
] ]
class NodeFileScanAction(NodeFileAbstractAction, identifier="node_file_scan"): class NodeFileScanAction(NodeFileAbstractAction, discriminator="node_file_scan"):
"""Action which scans a file.""" """Action which scans a file."""
config: "NodeFileScanAction.ConfigSchema" config: "NodeFileScanAction.ConfigSchema"
@@ -92,7 +92,7 @@ class NodeFileScanAction(NodeFileAbstractAction, identifier="node_file_scan"):
verb: ClassVar[str] = "scan" verb: ClassVar[str] = "scan"
class NodeFileDeleteAction(NodeFileAbstractAction, identifier="node_file_delete"): class NodeFileDeleteAction(NodeFileAbstractAction, discriminator="node_file_delete"):
"""Action which deletes a file.""" """Action which deletes a file."""
config: "NodeFileDeleteAction.ConfigSchema" config: "NodeFileDeleteAction.ConfigSchema"
@@ -119,7 +119,7 @@ class NodeFileDeleteAction(NodeFileAbstractAction, identifier="node_file_delete"
] ]
class NodeFileRestoreAction(NodeFileAbstractAction, identifier="node_file_restore"): class NodeFileRestoreAction(NodeFileAbstractAction, discriminator="node_file_restore"):
"""Action which restores a file.""" """Action which restores a file."""
config: "NodeFileRestoreAction.ConfigSchema" config: "NodeFileRestoreAction.ConfigSchema"
@@ -130,7 +130,7 @@ class NodeFileRestoreAction(NodeFileAbstractAction, identifier="node_file_restor
verb: ClassVar[str] = "restore" verb: ClassVar[str] = "restore"
class NodeFileCorruptAction(NodeFileAbstractAction, identifier="node_file_corrupt"): class NodeFileCorruptAction(NodeFileAbstractAction, discriminator="node_file_corrupt"):
"""Action which corrupts a file.""" """Action which corrupts a file."""
config: "NodeFileCorruptAction.ConfigSchema" config: "NodeFileCorruptAction.ConfigSchema"
@@ -141,7 +141,7 @@ class NodeFileCorruptAction(NodeFileAbstractAction, identifier="node_file_corrup
verb: ClassVar[str] = "corrupt" verb: ClassVar[str] = "corrupt"
class NodeFileAccessAction(NodeFileAbstractAction, identifier="node_file_access"): class NodeFileAccessAction(NodeFileAbstractAction, discriminator="node_file_access"):
"""Action which increases a file's access count.""" """Action which increases a file's access count."""
config: "NodeFileAccessAction.ConfigSchema" config: "NodeFileAccessAction.ConfigSchema"
@@ -167,7 +167,7 @@ class NodeFileAccessAction(NodeFileAbstractAction, identifier="node_file_access"
] ]
class NodeFileCheckhashAction(NodeFileAbstractAction, identifier="node_file_checkhash"): class NodeFileCheckhashAction(NodeFileAbstractAction, discriminator="node_file_checkhash"):
"""Action which checks the hash of a file.""" """Action which checks the hash of a file."""
config: "NodeFileCheckhashAction.ConfigSchema" config: "NodeFileCheckhashAction.ConfigSchema"
@@ -178,7 +178,7 @@ class NodeFileCheckhashAction(NodeFileAbstractAction, identifier="node_file_chec
verb: ClassVar[str] = "checkhash" verb: ClassVar[str] = "checkhash"
class NodeFileRepairAction(NodeFileAbstractAction, identifier="node_file_repair"): class NodeFileRepairAction(NodeFileAbstractAction, discriminator="node_file_repair"):
"""Action which repairs a file.""" """Action which repairs a file."""
config: "NodeFileRepairAction.ConfigSchema" config: "NodeFileRepairAction.ConfigSchema"

10
src/primaite/game/agent/actions/folder.py
View File

@@ -47,7 +47,7 @@ class NodeFolderAbstractAction(AbstractAction, ABC):
] ]
class NodeFolderScanAction(NodeFolderAbstractAction, identifier="node_folder_scan"): class NodeFolderScanAction(NodeFolderAbstractAction, discriminator="node_folder_scan"):
"""Action which scans a folder.""" """Action which scans a folder."""
config: "NodeFolderScanAction.ConfigSchema" config: "NodeFolderScanAction.ConfigSchema"
@@ -58,7 +58,7 @@ class NodeFolderScanAction(NodeFolderAbstractAction, identifier="node_folder_sca
verb: ClassVar[str] = "scan" verb: ClassVar[str] = "scan"
class NodeFolderCheckhashAction(NodeFolderAbstractAction, identifier="node_folder_checkhash"): class NodeFolderCheckhashAction(NodeFolderAbstractAction, discriminator="node_folder_checkhash"):
"""Action which checks the hash of a folder.""" """Action which checks the hash of a folder."""
config: "NodeFolderCheckhashAction.ConfigSchema" config: "NodeFolderCheckhashAction.ConfigSchema"
@@ -69,7 +69,7 @@ class NodeFolderCheckhashAction(NodeFolderAbstractAction, identifier="node_folde
verb: ClassVar[str] = "checkhash" verb: ClassVar[str] = "checkhash"
class NodeFolderRepairAction(NodeFolderAbstractAction, identifier="node_folder_repair"): class NodeFolderRepairAction(NodeFolderAbstractAction, discriminator="node_folder_repair"):
"""Action which repairs a folder.""" """Action which repairs a folder."""
config: "NodeFolderRepairAction.ConfigSchema" config: "NodeFolderRepairAction.ConfigSchema"
@@ -80,7 +80,7 @@ class NodeFolderRepairAction(NodeFolderAbstractAction, identifier="node_folder_r
verb: ClassVar[str] = "repair" verb: ClassVar[str] = "repair"
class NodeFolderRestoreAction(NodeFolderAbstractAction, identifier="node_folder_restore"): class NodeFolderRestoreAction(NodeFolderAbstractAction, discriminator="node_folder_restore"):
"""Action which restores a folder.""" """Action which restores a folder."""
config: "NodeFolderRestoreAction.ConfigSchema" config: "NodeFolderRestoreAction.ConfigSchema"
@@ -91,7 +91,7 @@ class NodeFolderRestoreAction(NodeFolderAbstractAction, identifier="node_folder_
verb: ClassVar[str] = "restore" verb: ClassVar[str] = "restore"
class NodeFolderCreateAction(NodeFolderAbstractAction, identifier="node_folder_create"): class NodeFolderCreateAction(NodeFolderAbstractAction, discriminator="node_folder_create"):
"""Action which creates a new folder.""" """Action which creates a new folder."""
config: "NodeFolderCreateAction.ConfigSchema" config: "NodeFolderCreateAction.ConfigSchema"

4
src/primaite/game/agent/actions/host_nic.py
View File

@@ -40,7 +40,7 @@ class HostNICAbstractAction(AbstractAction, ABC):
] ]
class HostNICEnableAction(HostNICAbstractAction, identifier="host_nic_enable"): class HostNICEnableAction(HostNICAbstractAction, discriminator="host_nic_enable"):
"""Action which enables a NIC.""" """Action which enables a NIC."""
config: "HostNICEnableAction.ConfigSchema" config: "HostNICEnableAction.ConfigSchema"
@@ -51,7 +51,7 @@ class HostNICEnableAction(HostNICAbstractAction, identifier="host_nic_enable"):
verb: ClassVar[str] = "enable" verb: ClassVar[str] = "enable"
class HostNICDisableAction(HostNICAbstractAction, identifier="host_nic_disable"): class HostNICDisableAction(HostNICAbstractAction, discriminator="host_nic_disable"):
"""Action which disables a NIC.""" """Action which disables a NIC."""
config: "HostNICDisableAction.ConfigSchema" config: "HostNICDisableAction.ConfigSchema"

2
src/primaite/game/agent/actions/manager.py
View File

@@ -24,7 +24,7 @@ from primaite.interface.request import RequestFormat
__all__ = ("DoNothingAction", "ActionManager") __all__ = ("DoNothingAction", "ActionManager")
class DoNothingAction(AbstractAction, identifier="do_nothing"): class DoNothingAction(AbstractAction, discriminator="do_nothing"):
"""Do Nothing Action.""" """Do Nothing Action."""
class ConfigSchema(AbstractAction.ConfigSchema): class ConfigSchema(AbstractAction.ConfigSchema):

6
src/primaite/game/agent/actions/network.py
View File

@@ -8,7 +8,7 @@ from primaite.interface.request import RequestFormat
__all__ = ("NetworkPortEnableAction", "NetworkPortDisableAction") __all__ = ("NetworkPortEnableAction", "NetworkPortDisableAction")
class NetworkPortAbstractAction(AbstractAction, identifier="network_port_abstract"): class NetworkPortAbstractAction(AbstractAction, discriminator="network_port_abstract"):
"""Base class for Network port actions.""" """Base class for Network port actions."""
config: "NetworkPortAbstractAction.ConfigSchema" config: "NetworkPortAbstractAction.ConfigSchema"
@@ -35,7 +35,7 @@ class NetworkPortAbstractAction(AbstractAction, identifier="network_port_abstrac
] ]
class NetworkPortEnableAction(NetworkPortAbstractAction, identifier="network_port_enable"): class NetworkPortEnableAction(NetworkPortAbstractAction, discriminator="network_port_enable"):
"""Action which enables are port on a router or a firewall.""" """Action which enables are port on a router or a firewall."""
config: "NetworkPortEnableAction.ConfigSchema" config: "NetworkPortEnableAction.ConfigSchema"
@@ -46,7 +46,7 @@ class NetworkPortEnableAction(NetworkPortAbstractAction, identifier="network_por
verb: ClassVar[str] = "enable" verb: ClassVar[str] = "enable"
class NetworkPortDisableAction(NetworkPortAbstractAction, identifier="network_port_disable"): class NetworkPortDisableAction(NetworkPortAbstractAction, discriminator="network_port_disable"):
"""Action which disables are port on a router or a firewall.""" """Action which disables are port on a router or a firewall."""
config: "NetworkPortDisableAction.ConfigSchema" config: "NetworkPortDisableAction.ConfigSchema"

18
src/primaite/game/agent/actions/node.py
View File

@@ -18,7 +18,7 @@ __all__ = (
) )
class NodeAbstractAction(AbstractAction, identifier="node_abstract"): class NodeAbstractAction(AbstractAction, discriminator="node_abstract"):
""" """
Abstract base class for node actions. Abstract base class for node actions.
@@ -39,7 +39,7 @@ class NodeAbstractAction(AbstractAction, identifier="node_abstract"):
return ["network", "node", config.node_name, config.verb] return ["network", "node", config.node_name, config.verb]
class NodeOSScanAction(NodeAbstractAction, identifier="node_os_scan"): class NodeOSScanAction(NodeAbstractAction, discriminator="node_os_scan"):
"""Action which scans a node's OS.""" """Action which scans a node's OS."""
config: "NodeOSScanAction.ConfigSchema" config: "NodeOSScanAction.ConfigSchema"
@@ -50,7 +50,7 @@ class NodeOSScanAction(NodeAbstractAction, identifier="node_os_scan"):
verb: ClassVar[str] = "scan" verb: ClassVar[str] = "scan"
class NodeShutdownAction(NodeAbstractAction, identifier="node_shutdown"): class NodeShutdownAction(NodeAbstractAction, discriminator="node_shutdown"):
"""Action which shuts down a node.""" """Action which shuts down a node."""
config: "NodeShutdownAction.ConfigSchema" config: "NodeShutdownAction.ConfigSchema"
@@ -61,7 +61,7 @@ class NodeShutdownAction(NodeAbstractAction, identifier="node_shutdown"):
verb: ClassVar[str] = "shutdown" verb: ClassVar[str] = "shutdown"
class NodeStartupAction(NodeAbstractAction, identifier="node_startup"): class NodeStartupAction(NodeAbstractAction, discriminator="node_startup"):
"""Action which starts up a node.""" """Action which starts up a node."""
config: "NodeStartupAction.ConfigSchema" config: "NodeStartupAction.ConfigSchema"
@@ -72,7 +72,7 @@ class NodeStartupAction(NodeAbstractAction, identifier="node_startup"):
verb: ClassVar[str] = "startup" verb: ClassVar[str] = "startup"
class NodeResetAction(NodeAbstractAction, identifier="node_reset"): class NodeResetAction(NodeAbstractAction, discriminator="node_reset"):
"""Action which resets a node.""" """Action which resets a node."""
config: "NodeResetAction.ConfigSchema" config: "NodeResetAction.ConfigSchema"
@@ -83,7 +83,7 @@ class NodeResetAction(NodeAbstractAction, identifier="node_reset"):
verb: ClassVar[str] = "reset" verb: ClassVar[str] = "reset"
class NodeNMAPAbstractAction(AbstractAction, identifier="node_nmap_abstract_action"): class NodeNMAPAbstractAction(AbstractAction, discriminator="node_nmap_abstract_action"):
"""Base class for NodeNMAP actions.""" """Base class for NodeNMAP actions."""
config: "NodeNMAPAbstractAction.ConfigSchema" config: "NodeNMAPAbstractAction.ConfigSchema"
@@ -103,7 +103,7 @@ class NodeNMAPAbstractAction(AbstractAction, identifier="node_nmap_abstract_acti
pass pass
class NodeNMAPPingScanAction(NodeNMAPAbstractAction, identifier="node_nmap_ping_scan"): class NodeNMAPPingScanAction(NodeNMAPAbstractAction, discriminator="node_nmap_ping_scan"):
"""Action which performs an NMAP ping scan.""" """Action which performs an NMAP ping scan."""
config: "NodeNMAPPingScanAction.ConfigSchema" config: "NodeNMAPPingScanAction.ConfigSchema"
@@ -122,7 +122,7 @@ class NodeNMAPPingScanAction(NodeNMAPAbstractAction, identifier="node_nmap_ping_
] ]
class NodeNMAPPortScanAction(NodeNMAPAbstractAction, identifier="node_nmap_port_scan"): class NodeNMAPPortScanAction(NodeNMAPAbstractAction, discriminator="node_nmap_port_scan"):
"""Action which performs an NMAP port scan.""" """Action which performs an NMAP port scan."""
config: "NodeNMAPPortScanAction.ConfigSchema" config: "NodeNMAPPortScanAction.ConfigSchema"
@@ -154,7 +154,7 @@ class NodeNMAPPortScanAction(NodeNMAPAbstractAction, identifier="node_nmap_port_
] ]
class NodeNetworkServiceReconAction(NodeNMAPAbstractAction, identifier="node_network_service_recon"): class NodeNetworkServiceReconAction(NodeNMAPAbstractAction, discriminator="node_network_service_recon"):
"""Action which performs an NMAP network service recon (ping scan followed by port scan).""" """Action which performs an NMAP network service recon (ping scan followed by port scan)."""
config: "NodeNetworkServiceReconAction.ConfigSchema" config: "NodeNetworkServiceReconAction.ConfigSchema"

20
src/primaite/game/agent/actions/service.py
View File

@@ -17,7 +17,7 @@ __all__ = (
) )
class NodeServiceAbstractAction(AbstractAction, identifier="node_service_abstract"): class NodeServiceAbstractAction(AbstractAction, discriminator="node_service_abstract"):
"""Abstract Action for Node Service related actions. """Abstract Action for Node Service related actions.
Any actions which use node_name and service_name can inherit from this class. Any actions which use node_name and service_name can inherit from this class.
@@ -36,7 +36,7 @@ class NodeServiceAbstractAction(AbstractAction, identifier="node_service_abstrac
return ["network", "node", config.node_name, "service", config.service_name, config.verb] return ["network", "node", config.node_name, "service", config.service_name, config.verb]
class NodeServiceScanAction(NodeServiceAbstractAction, identifier="node_service_scan"): class NodeServiceScanAction(NodeServiceAbstractAction, discriminator="node_service_scan"):
"""Action which scans a service.""" """Action which scans a service."""
config: "NodeServiceScanAction.ConfigSchema" config: "NodeServiceScanAction.ConfigSchema"
@@ -47,7 +47,7 @@ class NodeServiceScanAction(NodeServiceAbstractAction, identifier="node_service_
verb: ClassVar[str] = "scan" verb: ClassVar[str] = "scan"
class NodeServiceStopAction(NodeServiceAbstractAction, identifier="node_service_stop"): class NodeServiceStopAction(NodeServiceAbstractAction, discriminator="node_service_stop"):
"""Action which stops a service.""" """Action which stops a service."""
config: "NodeServiceStopAction.ConfigSchema" config: "NodeServiceStopAction.ConfigSchema"
@@ -58,7 +58,7 @@ class NodeServiceStopAction(NodeServiceAbstractAction, identifier="node_service_
verb: ClassVar[str] = "stop" verb: ClassVar[str] = "stop"
class NodeServiceStartAction(NodeServiceAbstractAction, identifier="node_service_start"): class NodeServiceStartAction(NodeServiceAbstractAction, discriminator="node_service_start"):
"""Action which starts a service.""" """Action which starts a service."""
config: "NodeServiceStartAction.ConfigSchema" config: "NodeServiceStartAction.ConfigSchema"
@@ -69,7 +69,7 @@ class NodeServiceStartAction(NodeServiceAbstractAction, identifier="node_service
verb: ClassVar[str] = "start" verb: ClassVar[str] = "start"
class NodeServicePauseAction(NodeServiceAbstractAction, identifier="node_service_pause"): class NodeServicePauseAction(NodeServiceAbstractAction, discriminator="node_service_pause"):
"""Action which pauses a service.""" """Action which pauses a service."""
config: "NodeServicePauseAction.ConfigSchema" config: "NodeServicePauseAction.ConfigSchema"
@@ -80,7 +80,7 @@ class NodeServicePauseAction(NodeServiceAbstractAction, identifier="node_service
verb: ClassVar[str] = "pause" verb: ClassVar[str] = "pause"
class NodeServiceResumeAction(NodeServiceAbstractAction, identifier="node_service_resume"): class NodeServiceResumeAction(NodeServiceAbstractAction, discriminator="node_service_resume"):
"""Action which resumes a service.""" """Action which resumes a service."""
config: "NodeServiceResumeAction.ConfigSchema" config: "NodeServiceResumeAction.ConfigSchema"
@@ -91,7 +91,7 @@ class NodeServiceResumeAction(NodeServiceAbstractAction, identifier="node_servic
verb: ClassVar[str] = "resume" verb: ClassVar[str] = "resume"
class NodeServiceRestartAction(NodeServiceAbstractAction, identifier="node_service_restart"): class NodeServiceRestartAction(NodeServiceAbstractAction, discriminator="node_service_restart"):
"""Action which restarts a service.""" """Action which restarts a service."""
config: "NodeServiceRestartAction.ConfigSchema" config: "NodeServiceRestartAction.ConfigSchema"
@@ -102,7 +102,7 @@ class NodeServiceRestartAction(NodeServiceAbstractAction, identifier="node_servi
verb: ClassVar[str] = "restart" verb: ClassVar[str] = "restart"
class NodeServiceDisableAction(NodeServiceAbstractAction, identifier="node_service_disable"): class NodeServiceDisableAction(NodeServiceAbstractAction, discriminator="node_service_disable"):
"""Action which disables a service.""" """Action which disables a service."""
config: "NodeServiceDisableAction.ConfigSchema" config: "NodeServiceDisableAction.ConfigSchema"
@@ -113,7 +113,7 @@ class NodeServiceDisableAction(NodeServiceAbstractAction, identifier="node_servi
verb: ClassVar[str] = "disable" verb: ClassVar[str] = "disable"
class NodeServiceEnableAction(NodeServiceAbstractAction, identifier="node_service_enable"): class NodeServiceEnableAction(NodeServiceAbstractAction, discriminator="node_service_enable"):
"""Action which enables a service.""" """Action which enables a service."""
config: "NodeServiceEnableAction.ConfigSchema" config: "NodeServiceEnableAction.ConfigSchema"
@@ -124,7 +124,7 @@ class NodeServiceEnableAction(NodeServiceAbstractAction, identifier="node_servic
verb: ClassVar[str] = "enable" verb: ClassVar[str] = "enable"
class NodeServiceFixAction(NodeServiceAbstractAction, identifier="node_service_fix"): class NodeServiceFixAction(NodeServiceAbstractAction, discriminator="node_service_fix"):
"""Action which fixes a service.""" """Action which fixes a service."""
config: "NodeServiceFixAction.ConfigSchema" config: "NodeServiceFixAction.ConfigSchema"

8
src/primaite/game/agent/actions/session.py
View File

@@ -11,7 +11,7 @@ __all__ = (
) )
class NodeSessionAbstractAction(AbstractAction, identifier="node_session_abstract"): class NodeSessionAbstractAction(AbstractAction, discriminator="node_session_abstract"):
"""Base class for NodeSession actions.""" """Base class for NodeSession actions."""
config: "NodeSessionAbstractAction.ConfigSchema" config: "NodeSessionAbstractAction.ConfigSchema"
@@ -33,7 +33,7 @@ class NodeSessionAbstractAction(AbstractAction, identifier="node_session_abstrac
pass pass
class NodeSessionsRemoteLoginAction(NodeSessionAbstractAction, identifier="node_session_remote_login"): class NodeSessionsRemoteLoginAction(NodeSessionAbstractAction, discriminator="node_session_remote_login"):
"""Action which performs a remote session login.""" """Action which performs a remote session login."""
config: "NodeSessionsRemoteLoginAction.ConfigSchema" config: "NodeSessionsRemoteLoginAction.ConfigSchema"
@@ -62,7 +62,7 @@ class NodeSessionsRemoteLoginAction(NodeSessionAbstractAction, identifier="node_
] ]
class NodeSessionsRemoteLogoutAction(NodeSessionAbstractAction, identifier="node_session_remote_logoff"): class NodeSessionsRemoteLogoutAction(NodeSessionAbstractAction, discriminator="node_session_remote_logoff"):
"""Action which performs a remote session logout.""" """Action which performs a remote session logout."""
config: "NodeSessionsRemoteLogoutAction.ConfigSchema" config: "NodeSessionsRemoteLogoutAction.ConfigSchema"
@@ -80,7 +80,7 @@ class NodeSessionsRemoteLogoutAction(NodeSessionAbstractAction, identifier="node
return ["network", "node", config.node_name, "service", "Terminal", config.verb, config.remote_ip] return ["network", "node", config.node_name, "service", "Terminal", config.verb, config.remote_ip]
class NodeAccountChangePasswordAction(NodeSessionAbstractAction, identifier="node_account_change_password"): class NodeAccountChangePasswordAction(NodeSessionAbstractAction, discriminator="node_account_change_password"):
"""Action which changes the password for a user.""" """Action which changes the password for a user."""
config: "NodeAccountChangePasswordAction.ConfigSchema" config: "NodeAccountChangePasswordAction.ConfigSchema"

20
src/primaite/game/agent/actions/software.py
View File

@@ -22,7 +22,7 @@ __all__ = (
) )
class ConfigureRansomwareScriptAction(AbstractAction, identifier="configure_ransomware_script"): class ConfigureRansomwareScriptAction(AbstractAction, discriminator="configure_ransomware_script"):
"""Action which sets config parameters for a ransomware script on a node.""" """Action which sets config parameters for a ransomware script on a node."""
config: "ConfigureRansomwareScriptAction.ConfigSchema" config: "ConfigureRansomwareScriptAction.ConfigSchema"
@@ -48,7 +48,9 @@ class ConfigureRansomwareScriptAction(AbstractAction, identifier="configure_rans
return ["network", "node", config.node_name, "application", "RansomwareScript", "configure", data] return ["network", "node", config.node_name, "application", "RansomwareScript", "configure", data]
class RansomwareConfigureC2ServerAction(ConfigureRansomwareScriptAction, identifier="c2_server_ransomware_configure"): class RansomwareConfigureC2ServerAction(
ConfigureRansomwareScriptAction, discriminator="c2_server_ransomware_configure"
):
"""Action which causes a C2 server to send a command to set options on a ransomware script remotely.""" """Action which causes a C2 server to send a command to set options on a ransomware script remotely."""
@classmethod @classmethod
@@ -59,7 +61,7 @@ class RansomwareConfigureC2ServerAction(ConfigureRansomwareScriptAction, identif
return ["network", "node", config.node_name, "application", "C2Server", "ransomware_configure", data] return ["network", "node", config.node_name, "application", "C2Server", "ransomware_configure", data]
class ConfigureDoSBotAction(AbstractAction, identifier="configure_dos_bot"): class ConfigureDoSBotAction(AbstractAction, discriminator="configure_dos_bot"):
"""Action which sets config parameters for a DoS bot on a node.""" """Action which sets config parameters for a DoS bot on a node."""
class ConfigSchema(AbstractAction.ConfigSchema): class ConfigSchema(AbstractAction.ConfigSchema):
@@ -91,7 +93,7 @@ class ConfigureDoSBotAction(AbstractAction, identifier="configure_dos_bot"):
return ["network", "node", config.node_name, "application", "DoSBot", "configure", data] return ["network", "node", config.node_name, "application", "DoSBot", "configure", data]
class ConfigureC2BeaconAction(AbstractAction, identifier="configure_c2_beacon"): class ConfigureC2BeaconAction(AbstractAction, discriminator="configure_c2_beacon"):
"""Action which configures a C2 Beacon based on the parameters given.""" """Action which configures a C2 Beacon based on the parameters given."""
class ConfigSchema(AbstractAction.ConfigSchema): class ConfigSchema(AbstractAction.ConfigSchema):
@@ -115,7 +117,7 @@ class ConfigureC2BeaconAction(AbstractAction, identifier="configure_c2_beacon"):
return ["network", "node", config.node_name, "application", "C2Beacon", "configure", data] return ["network", "node", config.node_name, "application", "C2Beacon", "configure", data]
class NodeSendRemoteCommandAction(AbstractAction, identifier="node_send_remote_command"): class NodeSendRemoteCommandAction(AbstractAction, discriminator="node_send_remote_command"):
"""Action which sends a terminal command to a remote node via SSH.""" """Action which sends a terminal command to a remote node via SSH."""
config: "NodeSendRemoteCommandAction.ConfigSchema" config: "NodeSendRemoteCommandAction.ConfigSchema"
@@ -142,7 +144,7 @@ class NodeSendRemoteCommandAction(AbstractAction, identifier="node_send_remote_c
] ]
class TerminalC2ServerAction(AbstractAction, identifier="c2_server_terminal_command"): class TerminalC2ServerAction(AbstractAction, discriminator="c2_server_terminal_command"):
"""Action which causes the C2 Server to send a command to the C2 Beacon to execute the terminal command passed.""" """Action which causes the C2 Server to send a command to the C2 Beacon to execute the terminal command passed."""
config: "TerminalC2ServerAction.ConfigSchema" config: "TerminalC2ServerAction.ConfigSchema"
@@ -171,7 +173,7 @@ class TerminalC2ServerAction(AbstractAction, identifier="c2_server_terminal_comm
return ["network", "node", config.node_name, "application", "C2Server", "terminal_command", command_model] return ["network", "node", config.node_name, "application", "C2Server", "terminal_command", command_model]
class RansomwareLaunchC2ServerAction(AbstractAction, identifier="c2_server_ransomware_launch"): class RansomwareLaunchC2ServerAction(AbstractAction, discriminator="c2_server_ransomware_launch"):
"""Action which causes the C2 Server to send a command to the C2 Beacon to launch the RansomwareScript.""" """Action which causes the C2 Server to send a command to the C2 Beacon to launch the RansomwareScript."""
config: "RansomwareLaunchC2ServerAction.ConfigSchema" config: "RansomwareLaunchC2ServerAction.ConfigSchema"
@@ -190,7 +192,7 @@ class RansomwareLaunchC2ServerAction(AbstractAction, identifier="c2_server_ranso
return ["network", "node", config.node_name, "application", "C2Server", "ransomware_launch"] return ["network", "node", config.node_name, "application", "C2Server", "ransomware_launch"]
class ExfiltrationC2ServerAction(AbstractAction, identifier="c2_server_data_exfiltrate"): class ExfiltrationC2ServerAction(AbstractAction, discriminator="c2_server_data_exfiltrate"):
"""Action which exfiltrates a target file from a certain node onto the C2 beacon and then the C2 Server.""" """Action which exfiltrates a target file from a certain node onto the C2 beacon and then the C2 Server."""
config: "ExfiltrationC2ServerAction.ConfigSchema" config: "ExfiltrationC2ServerAction.ConfigSchema"
@@ -223,7 +225,7 @@ class ExfiltrationC2ServerAction(AbstractAction, identifier="c2_server_data_exfi
return ["network", "node", config.node_name, "application", "C2Server", "exfiltrate", command_model] return ["network", "node", config.node_name, "application", "C2Server", "exfiltrate", command_model]
class ConfigureDatabaseClientAction(AbstractAction, identifier="configure_database_client"): class ConfigureDatabaseClientAction(AbstractAction, discriminator="configure_database_client"):
"""Action which sets config parameters for a database client on a node.""" """Action which sets config parameters for a database client on a node."""
config: "ConfigureDatabaseClientAction.ConfigSchema" config: "ConfigureDatabaseClientAction.ConfigSchema"

14
src/primaite/game/agent/interface.py
View File

@@ -79,13 +79,13 @@ class AbstractAgent(BaseModel, ABC):
_registry: ClassVar[Dict[str, Type[AbstractAgent]]] = {} _registry: ClassVar[Dict[str, Type[AbstractAgent]]] = {}
def __init_subclass__(cls, identifier: Optional[str] = None, **kwargs: Any) -> None: def __init_subclass__(cls, discriminator: Optional[str] = None, **kwargs: Any) -> None:
super().__init_subclass__(**kwargs) super().__init_subclass__(**kwargs)
if identifier is None: if discriminator is None:
return return
if identifier in cls._registry: if discriminator in cls._registry:
raise ValueError(f"Cannot create a new agent under reserved name {identifier}") raise ValueError(f"Cannot create a new agent under reserved name {discriminator}")
cls._registry[identifier] = cls cls._registry[discriminator] = cls
def model_post_init(self, __context: Any) -> None: def model_post_init(self, __context: Any) -> None:
"""Overwrite the default empty action, observation, and rewards with ones defined through the config.""" """Overwrite the default empty action, observation, and rewards with ones defined through the config."""
@@ -161,7 +161,7 @@ class AbstractAgent(BaseModel, ABC):
return agent_class(config=config) return agent_class(config=config)
class AbstractScriptedAgent(AbstractAgent, identifier="AbstractScriptedAgent"): class AbstractScriptedAgent(AbstractAgent, discriminator="AbstractScriptedAgent"):
"""Base class for actors which generate their own behaviour.""" """Base class for actors which generate their own behaviour."""
config: "AbstractScriptedAgent.ConfigSchema" = Field(default_factory=lambda: AbstractScriptedAgent.ConfigSchema()) config: "AbstractScriptedAgent.ConfigSchema" = Field(default_factory=lambda: AbstractScriptedAgent.ConfigSchema())
@@ -177,7 +177,7 @@ class AbstractScriptedAgent(AbstractAgent, identifier="AbstractScriptedAgent"):
return super().get_action(obs=obs, timestep=timestep) return super().get_action(obs=obs, timestep=timestep)
class ProxyAgent(AbstractAgent, identifier="ProxyAgent"): class ProxyAgent(AbstractAgent, discriminator="ProxyAgent"):
"""Agent that sends observations to an RL model and receives actions from that model.""" """Agent that sends observations to an RL model and receives actions from that model."""
config: "ProxyAgent.ConfigSchema" = Field(default_factory=lambda: ProxyAgent.ConfigSchema()) config: "ProxyAgent.ConfigSchema" = Field(default_factory=lambda: ProxyAgent.ConfigSchema())

2
src/primaite/game/agent/observations/acl_observation.py
View File

@@ -16,7 +16,7 @@ from primaite.utils.validation.port import Port
_LOGGER = getLogger(__name__) _LOGGER = getLogger(__name__)
class ACLObservation(AbstractObservation, identifier="ACL"): class ACLObservation(AbstractObservation, discriminator="ACL"):
"""ACL observation, provides information about access control lists within the simulation environment.""" """ACL observation, provides information about access control lists within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema): class ConfigSchema(AbstractObservation.ConfigSchema):

4
src/primaite/game/agent/observations/file_system_observations.py
View File

@@ -13,7 +13,7 @@ from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_ST
_LOGGER = getLogger(__name__) _LOGGER = getLogger(__name__)
class FileObservation(AbstractObservation, identifier="FILE"): class FileObservation(AbstractObservation, discriminator="FILE"):
"""File observation, provides status information about a file within the simulation environment.""" """File observation, provides status information about a file within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema): class ConfigSchema(AbstractObservation.ConfigSchema):
@@ -125,7 +125,7 @@ class FileObservation(AbstractObservation, identifier="FILE"):
) )
class FolderObservation(AbstractObservation, identifier="FOLDER"): class FolderObservation(AbstractObservation, discriminator="FOLDER"):
"""Folder observation, provides status information about a folder within the simulation environment.""" """Folder observation, provides status information about a folder within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema): class ConfigSchema(AbstractObservation.ConfigSchema):

2
src/primaite/game/agent/observations/firewall_observation.py
View File

@@ -18,7 +18,7 @@ from primaite.utils.validation.port import Port
_LOGGER = getLogger(__name__) _LOGGER = getLogger(__name__)
class FirewallObservation(AbstractObservation, identifier="FIREWALL"): class FirewallObservation(AbstractObservation, discriminator="FIREWALL"):
"""Firewall observation, provides status information about a firewall within the simulation environment.""" """Firewall observation, provides status information about a firewall within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema): class ConfigSchema(AbstractObservation.ConfigSchema):

2
src/primaite/game/agent/observations/host_observations.py
View File

@@ -18,7 +18,7 @@ from primaite.utils.validation.port import Port
_LOGGER = getLogger(__name__) _LOGGER = getLogger(__name__)
class HostObservation(AbstractObservation, identifier="HOST"): class HostObservation(AbstractObservation, discriminator="HOST"):
"""Host observation, provides status information about a host within the simulation environment.""" """Host observation, provides status information about a host within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema): class ConfigSchema(AbstractObservation.ConfigSchema):

4
src/primaite/game/agent/observations/link_observation.py
View File

@@ -13,7 +13,7 @@ from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_ST
_LOGGER = getLogger(__name__) _LOGGER = getLogger(__name__)
class LinkObservation(AbstractObservation, identifier="LINK"): class LinkObservation(AbstractObservation, discriminator="LINK"):
"""Link observation, providing information about a specific link within the simulation environment.""" """Link observation, providing information about a specific link within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema): class ConfigSchema(AbstractObservation.ConfigSchema):
@@ -90,7 +90,7 @@ class LinkObservation(AbstractObservation, identifier="LINK"):
return cls(where=where) return cls(where=where)
class LinksObservation(AbstractObservation, identifier="LINKS"): class LinksObservation(AbstractObservation, discriminator="LINKS"):
"""Collection of link observations representing multiple links within the simulation environment.""" """Collection of link observations representing multiple links within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema): class ConfigSchema(AbstractObservation.ConfigSchema):

4
src/primaite/game/agent/observations/nic_observations.py
View File

@@ -12,7 +12,7 @@ from primaite.utils.validation.ip_protocol import IPProtocol
from primaite.utils.validation.port import Port from primaite.utils.validation.port import Port
class NICObservation(AbstractObservation, identifier="NETWORK_INTERFACE"): class NICObservation(AbstractObservation, discriminator="NETWORK_INTERFACE"):
"""Status information about a network interface within the simulation environment.""" """Status information about a network interface within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema): class ConfigSchema(AbstractObservation.ConfigSchema):
@@ -227,7 +227,7 @@ class NICObservation(AbstractObservation, identifier="NETWORK_INTERFACE"):
) )
class PortObservation(AbstractObservation, identifier="PORT"): class PortObservation(AbstractObservation, discriminator="PORT"):
"""Port observation, provides status information about a network port within the simulation environment.""" """Port observation, provides status information about a network port within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema): class ConfigSchema(AbstractObservation.ConfigSchema):

2
src/primaite/game/agent/observations/node_observations.py
View File

@@ -19,7 +19,7 @@ from primaite.utils.validation.port import Port
_LOGGER = getLogger(__name__) _LOGGER = getLogger(__name__)
class NodesObservation(AbstractObservation, identifier="NODES"): class NodesObservation(AbstractObservation, discriminator="NODES"):
"""Nodes observation, provides status information about nodes within the simulation environment.""" """Nodes observation, provides status information about nodes within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema): class ConfigSchema(AbstractObservation.ConfigSchema):

10
src/primaite/game/agent/observations/observation_manager.py
View File

@@ -11,7 +11,7 @@ from pydantic import BaseModel, computed_field, ConfigDict, Field, model_validat
from primaite.game.agent.observations.observations import AbstractObservation, WhereType from primaite.game.agent.observations.observations import AbstractObservation, WhereType
class NestedObservation(AbstractObservation, identifier="CUSTOM"): class NestedObservation(AbstractObservation, discriminator="CUSTOM"):
"""Observation type that allows combining other observations into a gymnasium.spaces.Dict space.""" """Observation type that allows combining other observations into a gymnasium.spaces.Dict space."""
class NestedObservationItem(BaseModel): class NestedObservationItem(BaseModel):
@@ -19,7 +19,7 @@ class NestedObservation(AbstractObservation, identifier="CUSTOM"):
model_config = ConfigDict(extra="forbid") model_config = ConfigDict(extra="forbid")
type: str type: str
"""Select observation class. It maps to the identifier of the obs class by checking the registry.""" """Select observation class. It maps to the discriminator of the obs class by checking the registry."""
label: str label: str
"""Dict key in the final observation space.""" """Dict key in the final observation space."""
options: Dict options: Dict
@@ -119,7 +119,7 @@ class NestedObservation(AbstractObservation, identifier="CUSTOM"):
return cls(components=instances) return cls(components=instances)
class NullObservation(AbstractObservation, identifier="NONE"): class NullObservation(AbstractObservation, discriminator="NONE"):
"""Empty observation that acts as a placeholder.""" """Empty observation that acts as a placeholder."""
def __init__(self) -> None: def __init__(self) -> None:
@@ -158,7 +158,7 @@ class ObservationManager(BaseModel):
model_config = ConfigDict(extra="forbid") model_config = ConfigDict(extra="forbid")
type: str = "NONE" type: str = "NONE"
"""Identifier name for the top-level observation.""" """discriminator name for the top-level observation."""
options: AbstractObservation.ConfigSchema = Field( options: AbstractObservation.ConfigSchema = Field(
default_factory=lambda: NullObservation.ConfigSchema(), validate_default=True default_factory=lambda: NullObservation.ConfigSchema(), validate_default=True
) )
@@ -235,7 +235,7 @@ class ObservationManager(BaseModel):
:param config: Dictionary containing the configuration for this observation space. :param config: Dictionary containing the configuration for this observation space.
If None, a blank observation space is created. If None, a blank observation space is created.
Otherwise, this must be a Dict with a type field and options field. Otherwise, this must be a Dict with a type field and options field.
type: string that corresponds to one of the observation identifiers that are provided when subclassing type: string that corresponds to one of the observation discriminators that are provided when subclassing
AbstractObservation AbstractObservation
options: this must adhere to the chosen observation type's ConfigSchema nested class. options: this must adhere to the chosen observation type's ConfigSchema nested class.
:type config: Dict :type config: Dict

14
src/primaite/game/agent/observations/observations.py
View File

@@ -31,20 +31,20 @@ class AbstractObservation(ABC):
"""Initialise an observation. This method must be overwritten.""" """Initialise an observation. This method must be overwritten."""
self.default_observation: ObsType self.default_observation: ObsType
def __init_subclass__(cls, identifier: Optional[str] = None, **kwargs: Any) -> None: def __init_subclass__(cls, discriminator: Optional[str] = None, **kwargs: Any) -> None:
""" """
Register an observation type. Register an observation type.
:param identifier: Identifier used to uniquely specify observation component types. :param discriminator: discriminator used to uniquely specify observation component types.
:type identifier: str :type discriminator: str
:raises ValueError: When attempting to create a component with a name that is already in use. :raises ValueError: When attempting to create a component with a name that is already in use.
""" """
super().__init_subclass__(**kwargs) super().__init_subclass__(**kwargs)
if identifier is None: if discriminator is None:
return return
if identifier in cls._registry: if discriminator in cls._registry:
raise ValueError(f"Duplicate observation component type {identifier}") raise ValueError(f"Duplicate observation component type {discriminator}")
cls._registry[identifier] = cls cls._registry[discriminator] = cls
@abstractmethod @abstractmethod
def observe(self, state: Dict) -> Any: def observe(self, state: Dict) -> Any:

2
src/primaite/game/agent/observations/router_observation.py
View File

@@ -18,7 +18,7 @@ from primaite.utils.validation.port import Port
_LOGGER = getLogger(__name__) _LOGGER = getLogger(__name__)
class RouterObservation(AbstractObservation, identifier="ROUTER"): class RouterObservation(AbstractObservation, discriminator="ROUTER"):
"""Router observation, provides status information about a router within the simulation environment.""" """Router observation, provides status information about a router within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema): class ConfigSchema(AbstractObservation.ConfigSchema):

4
src/primaite/game/agent/observations/software_observation.py
View File

@@ -10,7 +10,7 @@ from primaite.game.agent.observations.observations import AbstractObservation, W
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
class ServiceObservation(AbstractObservation, identifier="SERVICE"): class ServiceObservation(AbstractObservation, discriminator="SERVICE"):
"""Service observation, shows status of a service in the simulation environment.""" """Service observation, shows status of a service in the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema): class ConfigSchema(AbstractObservation.ConfigSchema):
@@ -73,7 +73,7 @@ class ServiceObservation(AbstractObservation, identifier="SERVICE"):
return cls(where=parent_where + ["services", config.service_name]) return cls(where=parent_where + ["services", config.service_name])
class ApplicationObservation(AbstractObservation, identifier="APPLICATION"): class ApplicationObservation(AbstractObservation, discriminator="APPLICATION"):
"""Application observation, shows the status of an application within the simulation environment.""" """Application observation, shows the status of an application within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema): class ConfigSchema(AbstractObservation.ConfigSchema):

24
src/primaite/game/agent/rewards.py
View File

@@ -55,13 +55,13 @@ class AbstractReward(BaseModel):
_registry: ClassVar[Dict[str, Type["AbstractReward"]]] = {} _registry: ClassVar[Dict[str, Type["AbstractReward"]]] = {}
def __init_subclass__(cls, identifier: Optional[str] = None, **kwargs: Any) -> None: def __init_subclass__(cls, discriminator: Optional[str] = None, **kwargs: Any) -> None:
super().__init_subclass__(**kwargs) super().__init_subclass__(**kwargs)
if identifier is None: if discriminator is None:
return return
if identifier in cls._registry: if discriminator in cls._registry:
raise ValueError(f"Duplicate reward {identifier}") raise ValueError(f"Duplicate reward {discriminator}")
cls._registry[identifier] = cls cls._registry[discriminator] = cls
@classmethod @classmethod
def from_config(cls, config: Dict) -> "AbstractReward": def from_config(cls, config: Dict) -> "AbstractReward":
@@ -92,7 +92,7 @@ class AbstractReward(BaseModel):
return 0.0 return 0.0
class DummyReward(AbstractReward, identifier="DUMMY"): class DummyReward(AbstractReward, discriminator="DUMMY"):
"""Dummy reward function component which always returns 0.0.""" """Dummy reward function component which always returns 0.0."""
def calculate(self, state: Dict, last_action_response: "AgentHistoryItem") -> float: def calculate(self, state: Dict, last_action_response: "AgentHistoryItem") -> float:
@@ -108,7 +108,7 @@ class DummyReward(AbstractReward, identifier="DUMMY"):
return 0.0 return 0.0
class DatabaseFileIntegrity(AbstractReward, identifier="DATABASE_FILE_INTEGRITY"): class DatabaseFileIntegrity(AbstractReward, discriminator="DATABASE_FILE_INTEGRITY"):
"""Reward function component which rewards the agent for maintaining the integrity of a database file.""" """Reward function component which rewards the agent for maintaining the integrity of a database file."""
config: "DatabaseFileIntegrity.ConfigSchema" config: "DatabaseFileIntegrity.ConfigSchema"
@@ -161,7 +161,7 @@ class DatabaseFileIntegrity(AbstractReward, identifier="DATABASE_FILE_INTEGRITY"
return 0 return 0
class WebServer404Penalty(AbstractReward, identifier="WEB_SERVER_404_PENALTY"): class WebServer404Penalty(AbstractReward, discriminator="WEB_SERVER_404_PENALTY"):
"""Reward function component which penalises the agent when the web server returns a 404 error.""" """Reward function component which penalises the agent when the web server returns a 404 error."""
config: "WebServer404Penalty.ConfigSchema" config: "WebServer404Penalty.ConfigSchema"
@@ -215,7 +215,7 @@ class WebServer404Penalty(AbstractReward, identifier="WEB_SERVER_404_PENALTY"):
return self.reward return self.reward
class WebpageUnavailablePenalty(AbstractReward, identifier="WEBPAGE_UNAVAILABLE_PENALTY"): class WebpageUnavailablePenalty(AbstractReward, discriminator="WEBPAGE_UNAVAILABLE_PENALTY"):
"""Penalises the agent when the web browser fails to fetch a webpage.""" """Penalises the agent when the web browser fails to fetch a webpage."""
config: "WebpageUnavailablePenalty.ConfigSchema" config: "WebpageUnavailablePenalty.ConfigSchema"
@@ -289,7 +289,7 @@ class WebpageUnavailablePenalty(AbstractReward, identifier="WEBPAGE_UNAVAILABLE_
return self.reward return self.reward
class GreenAdminDatabaseUnreachablePenalty(AbstractReward, identifier="GREEN_ADMIN_DATABASE_UNREACHABLE_PENALTY"): class GreenAdminDatabaseUnreachablePenalty(AbstractReward, discriminator="GREEN_ADMIN_DATABASE_UNREACHABLE_PENALTY"):
"""Penalises the agent when the green db clients fail to connect to the database.""" """Penalises the agent when the green db clients fail to connect to the database."""
config: "GreenAdminDatabaseUnreachablePenalty.ConfigSchema" config: "GreenAdminDatabaseUnreachablePenalty.ConfigSchema"
@@ -339,7 +339,7 @@ class GreenAdminDatabaseUnreachablePenalty(AbstractReward, identifier="GREEN_ADM
return self.reward return self.reward
class SharedReward(AbstractReward, identifier="SHARED_REWARD"): class SharedReward(AbstractReward, discriminator="SHARED_REWARD"):
"""Adds another agent's reward to the overall reward.""" """Adds another agent's reward to the overall reward."""
config: "SharedReward.ConfigSchema" config: "SharedReward.ConfigSchema"
@@ -376,7 +376,7 @@ class SharedReward(AbstractReward, identifier="SHARED_REWARD"):
return self.callback(self.config.agent_name) return self.callback(self.config.agent_name)
class ActionPenalty(AbstractReward, identifier="ACTION_PENALTY"): class ActionPenalty(AbstractReward, discriminator="ACTION_PENALTY"):
"""Apply a negative reward when taking any action except do_nothing.""" """Apply a negative reward when taking any action except do_nothing."""
config: "ActionPenalty.ConfigSchema" config: "ActionPenalty.ConfigSchema"

2
src/primaite/game/agent/scripted_agents/abstract_tap.py
View File

@@ -13,7 +13,7 @@ from primaite.game.agent.scripted_agents.random_agent import PeriodicAgent
__all__ = "AbstractTAPAgent" __all__ = "AbstractTAPAgent"
class AbstractTAPAgent(PeriodicAgent, identifier="AbstractTAP"): class AbstractTAPAgent(PeriodicAgent, discriminator="AbstractTAP"):
"""Base class for TAP agents to inherit from.""" """Base class for TAP agents to inherit from."""
config: "AbstractTAPAgent.ConfigSchema" = Field(default_factory=lambda: AbstractTAPAgent.ConfigSchema()) config: "AbstractTAPAgent.ConfigSchema" = Field(default_factory=lambda: AbstractTAPAgent.ConfigSchema())

2
src/primaite/game/agent/scripted_agents/data_manipulation_bot.py
View File

@@ -9,7 +9,7 @@ from primaite.game.agent.scripted_agents.random_agent import PeriodicAgent
__all__ = "DataManipulationAgent" __all__ = "DataManipulationAgent"
class DataManipulationAgent(PeriodicAgent, identifier="RedDatabaseCorruptingAgent"): class DataManipulationAgent(PeriodicAgent, discriminator="RedDatabaseCorruptingAgent"):
"""Agent that uses a DataManipulationBot to perform an SQL injection attack.""" """Agent that uses a DataManipulationBot to perform an SQL injection attack."""
class AgentSettingsSchema(PeriodicAgent.AgentSettingsSchema): class AgentSettingsSchema(PeriodicAgent.AgentSettingsSchema):

2
src/primaite/game/agent/scripted_agents/probabilistic_agent.py
View File

@@ -13,7 +13,7 @@ from primaite.game.agent.interface import AbstractScriptedAgent
__all__ = "ProbabilisticAgent" __all__ = "ProbabilisticAgent"
class ProbabilisticAgent(AbstractScriptedAgent, identifier="ProbabilisticAgent"): class ProbabilisticAgent(AbstractScriptedAgent, discriminator="ProbabilisticAgent"):
"""Scripted agent which randomly samples its action space with prescribed probabilities for each action.""" """Scripted agent which randomly samples its action space with prescribed probabilities for each action."""
rng: Generator = Field(default_factory=lambda: np.random.default_rng(np.random.randint(0, 65535))) rng: Generator = Field(default_factory=lambda: np.random.default_rng(np.random.randint(0, 65535)))

4
src/primaite/game/agent/scripted_agents/random_agent.py
View File

@@ -11,7 +11,7 @@ from primaite.game.agent.interface import AbstractScriptedAgent
__all__ = ("RandomAgent", "PeriodicAgent") __all__ = ("RandomAgent", "PeriodicAgent")
class RandomAgent(AbstractScriptedAgent, identifier="RandomAgent"): class RandomAgent(AbstractScriptedAgent, discriminator="RandomAgent"):
"""Agent that ignores its observation and acts completely at random.""" """Agent that ignores its observation and acts completely at random."""
config: "RandomAgent.ConfigSchema" = Field(default_factory=lambda: RandomAgent.ConfigSchema()) config: "RandomAgent.ConfigSchema" = Field(default_factory=lambda: RandomAgent.ConfigSchema())
@@ -34,7 +34,7 @@ class RandomAgent(AbstractScriptedAgent, identifier="RandomAgent"):
return self.action_manager.get_action(self.action_manager.space.sample()) return self.action_manager.get_action(self.action_manager.space.sample())
class PeriodicAgent(AbstractScriptedAgent, identifier="PeriodicAgent"): class PeriodicAgent(AbstractScriptedAgent, discriminator="PeriodicAgent"):
"""Agent that does nothing most of the time, but executes application at regular intervals (with variance).""" """Agent that does nothing most of the time, but executes application at regular intervals (with variance)."""
config: "PeriodicAgent.ConfigSchema" = Field(default_factory=lambda: PeriodicAgent.ConfigSchema()) config: "PeriodicAgent.ConfigSchema" = Field(default_factory=lambda: PeriodicAgent.ConfigSchema())

2
src/primaite/simulator/_package_data/create-simulation_demo.ipynb
View File

@@ -172,7 +172,7 @@
"\n", "\n",
"\n", "\n",
"# no applications exist yet so we will create our own.\n", "# no applications exist yet so we will create our own.\n",
"class MSPaint(Application, identifier=\"MSPaint\"):\n", "class MSPaint(Application, discriminator=\"MSPaint\"):\n",
" def describe_state(self):\n", " def describe_state(self):\n",
" return super().describe_state()" " return super().describe_state()"
] ]

2
src/primaite/simulator/core.py
View File

@@ -244,7 +244,7 @@ class SimComponent(BaseModel):
..code::python ..code::python
class WebBrowser(Application, identifier="WebBrowser"): class WebBrowser(Application, discriminator="WebBrowser"):
def _init_request_manager(self) -> RequestManager: def _init_request_manager(self) -> RequestManager:
rm = super()._init_request_manager() # all requests generic to any Application get initialised rm = super()._init_request_manager() # all requests generic to any Application get initialised
rm.add_request(...) # initialise any requests specific to the web browser rm.add_request(...) # initialise any requests specific to the web browser

2
src/primaite/simulator/network/container.py
View File

@@ -180,7 +180,7 @@ class Network(SimComponent):
table.align = "l" table.align = "l"
table.title = "Nodes" table.title = "Nodes"
for node in self.nodes.values(): for node in self.nodes.values():
table.add_row((node.hostname, type(node)._identifier, node.operating_state.name)) table.add_row((node.hostname, type(node)._discriminator, node.operating_state.name))
print(table) print(table)
if ip_addresses: if ip_addresses:

22
src/primaite/simulator/network/creation.py
View File

@@ -22,7 +22,7 @@ class NetworkNodeAdder(BaseModel):
Here is a template that users can use to define custom node adders: Here is a template that users can use to define custom node adders:
``` ```
class YourNodeAdder(NetworkNodeAdder, identifier="your_name"): class YourNodeAdder(NetworkNodeAdder, discriminator="your_name"):
class ConfigSchema(NetworkNodeAdder.ConfigSchema): class ConfigSchema(NetworkNodeAdder.ConfigSchema):
property_1 : str property_1 : str
property_2 : int property_2 : int
@@ -40,8 +40,8 @@ class NetworkNodeAdder(BaseModel):
""" """
Base schema for node adders. Base schema for node adders.
Child classes of NetworkNodeAdder must define a schema which inherits from this schema. The identifier is used Child classes of NetworkNodeAdder must define a schema which inherits from this schema. The discriminator is
by the from_config method to select the correct node adder at runtime. used by the from_config method to select the correct node adder at runtime.
""" """
model_config = ConfigDict(extra="forbid") model_config = ConfigDict(extra="forbid")
@@ -50,20 +50,20 @@ class NetworkNodeAdder(BaseModel):
_registry: ClassVar[Dict[str, Type["NetworkNodeAdder"]]] = {} _registry: ClassVar[Dict[str, Type["NetworkNodeAdder"]]] = {}
def __init_subclass__(cls, identifier: Optional[str], **kwargs: Any) -> None: def __init_subclass__(cls, discriminator: Optional[str], **kwargs: Any) -> None:
""" """
Register a network node adder class. Register a network node adder class.
:param identifier: Unique name for the node adder to use for matching against primaite config entries. :param discriminator: Unique name for the node adder to use for matching against primaite config entries.
:type identifier: str :type discriminator: str
:raises ValueError: When attempting to register a name that is already reserved. :raises ValueError: When attempting to register a name that is already reserved.
""" """
super().__init_subclass__(**kwargs) super().__init_subclass__(**kwargs)
if identifier is None: if discriminator is None:
return return
if identifier in cls._registry: if discriminator in cls._registry:
raise ValueError(f"Duplicate node adder {identifier}") raise ValueError(f"Duplicate node adder {discriminator}")
cls._registry[identifier] = cls cls._registry[discriminator] = cls
@classmethod @classmethod
@abstractmethod @abstractmethod
@@ -99,7 +99,7 @@ class NetworkNodeAdder(BaseModel):
adder_class.add_nodes_to_net(config=adder_class.ConfigSchema(**config), network=network) adder_class.add_nodes_to_net(config=adder_class.ConfigSchema(**config), network=network)
class OfficeLANAdder(NetworkNodeAdder, identifier="office_lan"): class OfficeLANAdder(NetworkNodeAdder, discriminator="office_lan"):
"""Creates an office LAN.""" """Creates an office LAN."""
class ConfigSchema(NetworkNodeAdder.ConfigSchema): class ConfigSchema(NetworkNodeAdder.ConfigSchema):

29
src/primaite/simulator/network/hardware/base.py
View File

@@ -824,7 +824,7 @@ class User(SimComponent):
return self.model_dump() return self.model_dump()
class UserManager(Service, identifier="UserManager"): class UserManager(Service, discriminator="UserManager"):
""" """
Manages users within the PrimAITE system, handling creation, authentication, and administration. Manages users within the PrimAITE system, handling creation, authentication, and administration.
@@ -1137,7 +1137,7 @@ class RemoteUserSession(UserSession):
return state return state
class UserSessionManager(Service, identifier="UserSessionManager"): class UserSessionManager(Service, discriminator="UserSessionManager"):
""" """
Manages user sessions on a Node, including local and remote sessions. Manages user sessions on a Node, including local and remote sessions.
@@ -1483,7 +1483,7 @@ class UserSessionManager(Service, identifier="UserSessionManager"):
return self.local_session is not None return self.local_session is not None
class Node(SimComponent): class Node(SimComponent, ABC):
""" """
A basic Node class that represents a node on the network. A basic Node class that represents a node on the network.
@@ -1556,25 +1556,26 @@ class Node(SimComponent):
_registry: ClassVar[Dict[str, Type["Node"]]] = {} _registry: ClassVar[Dict[str, Type["Node"]]] = {}
"""Registry of application types. Automatically populated when subclasses are defined.""" """Registry of application types. Automatically populated when subclasses are defined."""
_identifier: ClassVar[str] = "unknown" # TODO: this should not be set for abstract classes.
"""Identifier for this particular class, used for printing and logging. Each subclass redefines this.""" _discriminator: ClassVar[str]
"""discriminator for this particular class, used for printing and logging. Each subclass redefines this."""
def __init_subclass__(cls, identifier: Optional[str] = None, **kwargs: Any) -> None: def __init_subclass__(cls, discriminator: Optional[str] = None, **kwargs: Any) -> None:
""" """
Register a node type. Register a node type.
:param identifier: Uniquely specifies an node class by name. Used for finding items by config. :param discriminator: Uniquely specifies an node class by name. Used for finding items by config.
:type identifier: str :type discriminator: str
:raises ValueError: When attempting to register an node with a name that is already allocated. :raises ValueError: When attempting to register an node with a name that is already allocated.
""" """
super().__init_subclass__(**kwargs) super().__init_subclass__(**kwargs)
if identifier is None: if discriminator is None:
return return
identifier = identifier.lower() discriminator = discriminator.lower()
if identifier in cls._registry: if discriminator in cls._registry:
raise ValueError(f"Tried to define new node {identifier}, but this name is already reserved.") raise ValueError(f"Tried to define new node {discriminator}, but this name is already reserved.")
cls._registry[identifier] = cls cls._registry[discriminator] = cls
cls._identifier = identifier cls._discriminator = discriminator
def __init__(self, **kwargs): def __init__(self, **kwargs):
""" """

2
src/primaite/simulator/network/hardware/nodes/host/computer.py
View File

@@ -5,7 +5,7 @@ from primaite.simulator.network.hardware.nodes.host.host_node import HostNode
from primaite.simulator.system.services.ftp.ftp_client import FTPClient from primaite.simulator.system.services.ftp.ftp_client import FTPClient
class Computer(HostNode, identifier="computer"): class Computer(HostNode, discriminator="computer"):
""" """
A basic Computer class. A basic Computer class.

2
src/primaite/simulator/network/hardware/nodes/host/host_node.py
View File

@@ -262,7 +262,7 @@ class NIC(IPWiredNetworkInterface):
return f"Port {self.port_name if self.port_name else self.port_num}: {self.mac_address}/{self.ip_address}" return f"Port {self.port_name if self.port_name else self.port_num}: {self.mac_address}/{self.ip_address}"
class HostNode(Node, identifier="HostNode"): class HostNode(Node, discriminator="HostNode"):
""" """
Represents a host node in the network. Represents a host node in the network.

4
src/primaite/simulator/network/hardware/nodes/host/server.py
View File

@@ -2,7 +2,7 @@
from primaite.simulator.network.hardware.nodes.host.host_node import HostNode from primaite.simulator.network.hardware.nodes.host.host_node import HostNode
class Server(HostNode, identifier="server"): class Server(HostNode, discriminator="server"):
""" """
A basic Server class. A basic Server class.
@@ -31,7 +31,7 @@ class Server(HostNode, identifier="server"):
""" """
class Printer(HostNode, identifier="printer"): class Printer(HostNode, discriminator="printer"):
"""Printer? I don't even know her!.""" """Printer? I don't even know her!."""
# TODO: Implement printer-specific behaviour # TODO: Implement printer-specific behaviour

2
src/primaite/simulator/network/hardware/nodes/network/firewall.py
View File

@@ -27,7 +27,7 @@ DMZ_PORT_ID: Final[int] = 3
"""The Firewall port ID of the DMZ port.""" """The Firewall port ID of the DMZ port."""
class Firewall(Router, identifier="firewall"): class Firewall(Router, discriminator="firewall"):
""" """
A Firewall class that extends the functionality of a Router. A Firewall class that extends the functionality of a Router.

2
src/primaite/simulator/network/hardware/nodes/network/network_node.py
View File

@@ -7,7 +7,7 @@ from primaite.simulator.network.transmission.data_link_layer import Frame
from primaite.simulator.system.services.arp.arp import ARP from primaite.simulator.system.services.arp.arp import ARP
class NetworkNode(Node, identifier="NetworkNode"): class NetworkNode(Node, discriminator="NetworkNode"):
""" """
Represents an abstract base class for a network node that can receive and process network frames. Represents an abstract base class for a network node that can receive and process network frames.

2
src/primaite/simulator/network/hardware/nodes/network/router.py
View File

@@ -1184,7 +1184,7 @@ class RouterSessionManager(SessionManager):
return outbound_network_interface, dst_mac_address, dst_ip_address, src_port, dst_port, protocol, is_broadcast return outbound_network_interface, dst_mac_address, dst_ip_address, src_port, dst_port, protocol, is_broadcast
class Router(NetworkNode, identifier="router"): class Router(NetworkNode, discriminator="router"):
""" """
Represents a network router, managing routing and forwarding of IP packets across network interfaces. Represents a network router, managing routing and forwarding of IP packets across network interfaces.

2
src/primaite/simulator/network/hardware/nodes/network/switch.py
View File

@@ -87,7 +87,7 @@ class SwitchPort(WiredNetworkInterface):
return False return False
class Switch(NetworkNode, identifier="switch"): class Switch(NetworkNode, discriminator="switch"):
""" """
A class representing a Layer 2 network switch. A class representing a Layer 2 network switch.

2
src/primaite/simulator/network/hardware/nodes/network/wireless_router.py
View File

@@ -91,7 +91,7 @@ class WirelessAccessPoint(IPWirelessNetworkInterface):
) )
class WirelessRouter(Router, identifier="wireless_router"): class WirelessRouter(Router, discriminator="wireless_router"):
""" """
A WirelessRouter class that extends the functionality of a standard Router to include wireless capabilities. A WirelessRouter class that extends the functionality of a standard Router to include wireless capabilities.

14
src/primaite/simulator/system/applications/application.py
View File

@@ -53,20 +53,20 @@ class Application(IOSoftware, ABC):
_registry: ClassVar[Dict[str, Type["Application"]]] = {} _registry: ClassVar[Dict[str, Type["Application"]]] = {}
"""Registry of application types. Automatically populated when subclasses are defined.""" """Registry of application types. Automatically populated when subclasses are defined."""
def __init_subclass__(cls, identifier: Optional[str] = None, **kwargs: Any) -> None: def __init_subclass__(cls, discriminator: Optional[str] = None, **kwargs: Any) -> None:
""" """
Register an application type. Register an application type.
:param identifier: Uniquely specifies an application class by name. Used for finding items by config. :param discriminator: Uniquely specifies an application class by name. Used for finding items by config.
:type identifier: Optional[str] :type discriminator: Optional[str]
:raises ValueError: When attempting to register an application with a name that is already allocated. :raises ValueError: When attempting to register an application with a name that is already allocated.
""" """
super().__init_subclass__(**kwargs) super().__init_subclass__(**kwargs)
if identifier is None: if discriminator is None:
return return
if identifier in cls._registry: if discriminator in cls._registry:
raise ValueError(f"Tried to define new application {identifier}, but this name is already reserved.") raise ValueError(f"Tried to define new application {discriminator}, but this name is already reserved.")
cls._registry[identifier] = cls cls._registry[discriminator] = cls
@classmethod @classmethod
def from_config(cls, config: Dict) -> "Application": def from_config(cls, config: Dict) -> "Application":

2
src/primaite/simulator/system/applications/database_client.py
View File

@@ -61,7 +61,7 @@ class DatabaseClientConnection(BaseModel):
return str(self) return str(self)
class DatabaseClient(Application, identifier="DatabaseClient"): class DatabaseClient(Application, discriminator="DatabaseClient"):
""" """
A DatabaseClient application. A DatabaseClient application.

2
src/primaite/simulator/system/applications/nmap.py
View File

@@ -44,7 +44,7 @@ class PortScanPayload(SimComponent):
return state return state
class NMAP(Application, identifier="NMAP"): class NMAP(Application, discriminator="NMAP"):
""" """
A class representing the NMAP application for network scanning. A class representing the NMAP application for network scanning.

2
src/primaite/simulator/system/applications/red_applications/c2/c2_beacon.py
View File

@@ -17,7 +17,7 @@ from primaite.utils.validation.ipv4_address import IPV4Address
from primaite.utils.validation.port import Port, PORT_LOOKUP from primaite.utils.validation.port import Port, PORT_LOOKUP
class C2Beacon(AbstractC2, identifier="C2Beacon"): class C2Beacon(AbstractC2, discriminator="C2Beacon"):
""" """
C2 Beacon Application. C2 Beacon Application.

2
src/primaite/simulator/system/applications/red_applications/c2/c2_server.py
View File

@@ -16,7 +16,7 @@ from primaite.simulator.system.applications.red_applications.c2 import (
from primaite.simulator.system.applications.red_applications.c2.abstract_c2 import AbstractC2, C2Command, C2Payload from primaite.simulator.system.applications.red_applications.c2.abstract_c2 import AbstractC2, C2Command, C2Payload
class C2Server(AbstractC2, identifier="C2Server"): class C2Server(AbstractC2, discriminator="C2Server"):
""" """
C2 Server Application. C2 Server Application.

2
src/primaite/simulator/system/applications/red_applications/data_manipulation_bot.py
View File

@@ -40,7 +40,7 @@ class DataManipulationAttackStage(IntEnum):
"Signifies that the attack has failed." "Signifies that the attack has failed."
class DataManipulationBot(Application, identifier="DataManipulationBot"): class DataManipulationBot(Application, discriminator="DataManipulationBot"):
"""A bot that simulates a script which performs a SQL injection attack.""" """A bot that simulates a script which performs a SQL injection attack."""
class ConfigSchema(Application.ConfigSchema): class ConfigSchema(Application.ConfigSchema):

2
src/primaite/simulator/system/applications/red_applications/dos_bot.py
View File

@@ -32,7 +32,7 @@ class DoSAttackStage(IntEnum):
"Attack is completed." "Attack is completed."
class DoSBot(DatabaseClient, identifier="DoSBot"): class DoSBot(DatabaseClient, discriminator="DoSBot"):
"""A bot that simulates a Denial of Service attack.""" """A bot that simulates a Denial of Service attack."""
class ConfigSchema(DatabaseClient.ConfigSchema): class ConfigSchema(DatabaseClient.ConfigSchema):

2
src/primaite/simulator/system/applications/red_applications/ransomware_script.py
View File

@@ -14,7 +14,7 @@ from primaite.utils.validation.ipv4_address import IPV4Address
from primaite.utils.validation.port import PORT_LOOKUP from primaite.utils.validation.port import PORT_LOOKUP
class RansomwareScript(Application, identifier="RansomwareScript"): class RansomwareScript(Application, discriminator="RansomwareScript"):
"""Ransomware Kill Chain - Designed to be used by the TAP001 Agent on the example layout Network. """Ransomware Kill Chain - Designed to be used by the TAP001 Agent on the example layout Network.
:ivar payload: The attack stage query payload. (Default ENCRYPT) :ivar payload: The attack stage query payload. (Default ENCRYPT)

2
src/primaite/simulator/system/applications/web_browser.py
View File

@@ -23,7 +23,7 @@ from primaite.utils.validation.port import Port, PORT_LOOKUP
_LOGGER = getLogger(__name__) _LOGGER = getLogger(__name__)
class WebBrowser(Application, identifier="WebBrowser"): class WebBrowser(Application, discriminator="WebBrowser"):
""" """
Represents a web browser in the simulation environment. Represents a web browser in the simulation environment.

2
src/primaite/simulator/system/services/arp/arp.py
View File

@@ -15,7 +15,7 @@ from primaite.utils.validation.ipv4_address import IPV4Address
from primaite.utils.validation.port import PORT_LOOKUP from primaite.utils.validation.port import PORT_LOOKUP
class ARP(Service, identifier="ARP"): class ARP(Service, discriminator="ARP"):
""" """
The ARP (Address Resolution Protocol) Service. The ARP (Address Resolution Protocol) Service.

2
src/primaite/simulator/system/services/database/database_service.py
View File

@@ -19,7 +19,7 @@ from primaite.utils.validation.port import PORT_LOOKUP
_LOGGER = getLogger(__name__) _LOGGER = getLogger(__name__)
class DatabaseService(Service, identifier="DatabaseService"): class DatabaseService(Service, discriminator="DatabaseService"):
""" """
A class for simulating a generic SQL Server service. A class for simulating a generic SQL Server service.

2
src/primaite/simulator/system/services/dns/dns_client.py
View File

@@ -18,7 +18,7 @@ if TYPE_CHECKING:
_LOGGER = getLogger(__name__) _LOGGER = getLogger(__name__)
class DNSClient(Service, identifier="DNSClient"): class DNSClient(Service, discriminator="DNSClient"):
"""Represents a DNS Client as a Service.""" """Represents a DNS Client as a Service."""
class ConfigSchema(Service.ConfigSchema): class ConfigSchema(Service.ConfigSchema):

2
src/primaite/simulator/system/services/dns/dns_server.py
View File

@@ -14,7 +14,7 @@ from primaite.utils.validation.port import PORT_LOOKUP
_LOGGER = getLogger(__name__) _LOGGER = getLogger(__name__)
class DNSServer(Service, identifier="DNSServer"): class DNSServer(Service, discriminator="DNSServer"):
"""Represents a DNS Server as a Service.""" """Represents a DNS Server as a Service."""
class ConfigSchema(Service.ConfigSchema): class ConfigSchema(Service.ConfigSchema):

2
src/primaite/simulator/system/services/ftp/ftp_client.py
View File

@@ -18,7 +18,7 @@ from primaite.utils.validation.port import Port, PORT_LOOKUP
_LOGGER = getLogger(__name__) _LOGGER = getLogger(__name__)
class FTPClient(FTPServiceABC, identifier="FTPClient"): class FTPClient(FTPServiceABC, discriminator="FTPClient"):
""" """
A class for simulating an FTP client service. A class for simulating an FTP client service.

2
src/primaite/simulator/system/services/ftp/ftp_server.py
View File

@@ -12,7 +12,7 @@ from primaite.utils.validation.port import is_valid_port, PORT_LOOKUP
_LOGGER = getLogger(__name__) _LOGGER = getLogger(__name__)
class FTPServer(FTPServiceABC, identifier="FTPServer"): class FTPServer(FTPServiceABC, discriminator="FTPServer"):
""" """
A class for simulating an FTP server service. A class for simulating an FTP server service.

2
src/primaite/simulator/system/services/icmp/icmp.py
View File

@@ -16,7 +16,7 @@ from primaite.utils.validation.port import PORT_LOOKUP
_LOGGER = getLogger(__name__) _LOGGER = getLogger(__name__)
class ICMP(Service, identifier="ICMP"): class ICMP(Service, discriminator="ICMP"):
""" """
The Internet Control Message Protocol (ICMP) service. The Internet Control Message Protocol (ICMP) service.

2
src/primaite/simulator/system/services/ntp/ntp_client.py
View File

@@ -15,7 +15,7 @@ from primaite.utils.validation.port import Port, PORT_LOOKUP
_LOGGER = getLogger(__name__) _LOGGER = getLogger(__name__)
class NTPClient(Service, identifier="NTPClient"): class NTPClient(Service, discriminator="NTPClient"):
"""Represents a NTP client as a service.""" """Represents a NTP client as a service."""
class ConfigSchema(Service.ConfigSchema): class ConfigSchema(Service.ConfigSchema):

2
src/primaite/simulator/system/services/ntp/ntp_server.py
View File

@@ -13,7 +13,7 @@ from primaite.utils.validation.port import PORT_LOOKUP
_LOGGER = getLogger(__name__) _LOGGER = getLogger(__name__)
class NTPServer(Service, identifier="NTPServer"): class NTPServer(Service, discriminator="NTPServer"):
"""Represents a NTP server as a service.""" """Represents a NTP server as a service."""
class ConfigSchema(Service.ConfigSchema): class ConfigSchema(Service.ConfigSchema):

16
src/primaite/simulator/system/services/service.py
View File

@@ -61,22 +61,22 @@ class Service(IOSoftware):
def __init__(self, **kwargs): def __init__(self, **kwargs):
super().__init__(**kwargs) super().__init__(**kwargs)
def __init_subclass__(cls, identifier: Optional[str] = None, **kwargs: Any) -> None: def __init_subclass__(cls, discriminator: Optional[str] = None, **kwargs: Any) -> None:
""" """
Register a hostnode type. Register a hostnode type.
:param identifier: Uniquely specifies an hostnode class by name. Used for finding items by config. :param discriminator: Uniquely specifies an hostnode class by name. Used for finding items by config.
:type identifier: str :type discriminator: str
:raises ValueError: When attempting to register an hostnode with a name that is already allocated. :raises ValueError: When attempting to register an hostnode with a name that is already allocated.
""" """
super().__init_subclass__(**kwargs) super().__init_subclass__(**kwargs)
if identifier is None: if discriminator is None:
return return
# Enforce lowercase registry entries because it makes comparisons everywhere else much easier. # Enforce lowercase registry entries because it makes comparisons everywhere else much easier.
identifier = identifier.lower() discriminator = discriminator.lower()
if identifier in cls._registry: if discriminator in cls._registry:
raise ValueError(f"Tried to define new hostnode {identifier}, but this name is already reserved.") raise ValueError(f"Tried to define new hostnode {discriminator}, but this name is already reserved.")
cls._registry[identifier] = cls cls._registry[discriminator] = cls
@classmethod @classmethod
def from_config(cls, config: Dict) -> "Service": def from_config(cls, config: Dict) -> "Service":

2
src/primaite/simulator/system/services/terminal/terminal.py
View File

@@ -129,7 +129,7 @@ class RemoteTerminalConnection(TerminalClientConnection):
return self.parent_terminal.send(payload=payload, session_id=self.ssh_session_id) return self.parent_terminal.send(payload=payload, session_id=self.ssh_session_id)
class Terminal(Service, identifier="Terminal"): class Terminal(Service, discriminator="Terminal"):
"""Class used to simulate a generic terminal service. Can be interacted with by other terminals via SSH.""" """Class used to simulate a generic terminal service. Can be interacted with by other terminals via SSH."""
class ConfigSchema(Service.ConfigSchema): class ConfigSchema(Service.ConfigSchema):

2
src/primaite/simulator/system/services/web_server/web_server.py
View File

@@ -21,7 +21,7 @@ from primaite.utils.validation.port import Port, PORT_LOOKUP
_LOGGER = getLogger(__name__) _LOGGER = getLogger(__name__)
class WebServer(Service, identifier="WebServer"): class WebServer(Service, discriminator="WebServer"):
"""Class used to represent a Web Server Service in simulation.""" """Class used to represent a Web Server Service in simulation."""
class ConfigSchema(Service.ConfigSchema): class ConfigSchema(Service.ConfigSchema):

6
tests/conftest.py
View File

@@ -39,7 +39,7 @@ ACTION_SPACE_NODE_ACTION_VALUES = 1
_LOGGER = getLogger(__name__) _LOGGER = getLogger(__name__)
class DummyService(Service, identifier="DummyService"): class DummyService(Service, discriminator="DummyService"):
"""Test Service class""" """Test Service class"""
class ConfigSchema(Service.ConfigSchema): class ConfigSchema(Service.ConfigSchema):
@@ -62,7 +62,7 @@ class DummyService(Service, identifier="DummyService"):
pass pass
class DummyApplication(Application, identifier="DummyApplication"): class DummyApplication(Application, discriminator="DummyApplication"):
"""Test Application class""" """Test Application class"""
class ConfigSchema(Application.ConfigSchema): class ConfigSchema(Application.ConfigSchema):
@@ -280,7 +280,7 @@ def example_network() -> Network:
return network return network
class ControlledAgent(AbstractAgent, identifier="ControlledAgent"): class ControlledAgent(AbstractAgent, discriminator="ControlledAgent"):
"""Agent that can be controlled by the tests.""" """Agent that can be controlled by the tests."""
config: "ControlledAgent.ConfigSchema" = Field(default_factory=lambda: ControlledAgent.ConfigSchema()) config: "ControlledAgent.ConfigSchema" = Field(default_factory=lambda: ControlledAgent.ConfigSchema())

2
tests/integration_tests/extensions/applications/extended_application.py
View File

@@ -24,7 +24,7 @@ from primaite.utils.validation.port import PORT_LOOKUP
_LOGGER = getLogger(__name__) _LOGGER = getLogger(__name__)
class ExtendedApplication(Application, identifier="ExtendedApplication"): class ExtendedApplication(Application, discriminator="ExtendedApplication"):
""" """
Clone of web browser that uses the extension framework instead of being part of PrimAITE directly. Clone of web browser that uses the extension framework instead of being part of PrimAITE directly.

2
tests/integration_tests/extensions/nodes/giga_switch.py
View File

@@ -11,7 +11,7 @@ from primaite.simulator.network.hardware.nodes.network.switch import SwitchPort
from primaite.simulator.network.transmission.data_link_layer import Frame from primaite.simulator.network.transmission.data_link_layer import Frame
class GigaSwitch(NetworkNode, identifier="gigaswitch"): class GigaSwitch(NetworkNode, discriminator="gigaswitch"):
""" """
A class representing a Layer 2 network switch. A class representing a Layer 2 network switch.

2
tests/integration_tests/extensions/nodes/super_computer.py
View File

@@ -6,7 +6,7 @@ from primaite.simulator.system.services.ftp.ftp_client import FTPClient
from primaite.utils.validation.ipv4_address import IPV4Address from primaite.utils.validation.ipv4_address import IPV4Address
class SuperComputer(HostNode, identifier="supercomputer"): class SuperComputer(HostNode, discriminator="supercomputer"):
""" """
A basic Computer class. A basic Computer class.

2
tests/integration_tests/extensions/services/extended_service.py
View File

@@ -19,7 +19,7 @@ from primaite.utils.validation.port import PORT_LOOKUP
_LOGGER = getLogger(__name__) _LOGGER = getLogger(__name__)
class ExtendedService(Service, identifier="ExtendedService"): class ExtendedService(Service, discriminator="ExtendedService"):
""" """
A copy of DatabaseService that uses the extension framework instead of being part of PrimAITE. A copy of DatabaseService that uses the extension framework instead of being part of PrimAITE.

4
tests/integration_tests/network/test_broadcast.py
View File

@@ -15,7 +15,7 @@ from primaite.utils.validation.ip_protocol import PROTOCOL_LOOKUP
from primaite.utils.validation.port import PORT_LOOKUP from primaite.utils.validation.port import PORT_LOOKUP
class BroadcastTestService(Service, identifier="BroadcastTestService"): class BroadcastTestService(Service, discriminator="BroadcastTestService"):
"""A service for sending broadcast and unicast messages over a network.""" """A service for sending broadcast and unicast messages over a network."""
class ConfigSchema(Service.ConfigSchema): class ConfigSchema(Service.ConfigSchema):
@@ -51,7 +51,7 @@ class BroadcastTestService(Service, identifier="BroadcastTestService"):
) )
class BroadcastTestClient(Application, identifier="BroadcastTestClient"): class BroadcastTestClient(Application, discriminator="BroadcastTestClient"):
"""A client application to receive broadcast and unicast messages.""" """A client application to receive broadcast and unicast messages."""
class ConfigSchema(Service.ConfigSchema): class ConfigSchema(Service.ConfigSchema):

2
tests/integration_tests/system/test_service_listening_on_ports.py
View File

@@ -14,7 +14,7 @@ from primaite.utils.validation.port import PORT_LOOKUP
from tests import TEST_ASSETS_ROOT from tests import TEST_ASSETS_ROOT
class _DatabaseListener(Service, identifier="_DatabaseListener"): class _DatabaseListener(Service, discriminator="_DatabaseListener"):
class ConfigSchema(Service.ConfigSchema): class ConfigSchema(Service.ConfigSchema):
"""ConfigSchema for _DatabaseListener.""" """ConfigSchema for _DatabaseListener."""

4
tests/unit_tests/_primaite/_simulator/_system/_applications/test_application_registry.py
View File

@@ -5,14 +5,14 @@ from primaite.simulator.system.applications.application import Application
def test_adding_to_app_registry(): def test_adding_to_app_registry():
class temp_application(Application, identifier="temp_app"): class temp_application(Application, discriminator="temp_app"):
pass pass
assert Application._registry["temp_app"] is temp_application assert Application._registry["temp_app"] is temp_application
with pytest.raises(ValueError): with pytest.raises(ValueError):
class another_application(Application, identifier="temp_app"): class another_application(Application, discriminator="temp_app"):
pass pass
# This is kinda evil... # This is kinda evil...

2
tests/unit_tests/_primaite/_simulator/_system/test_software.py
View File

@@ -11,7 +11,7 @@ from primaite.utils.validation.ip_protocol import PROTOCOL_LOOKUP
from primaite.utils.validation.port import PORT_LOOKUP from primaite.utils.validation.port import PORT_LOOKUP
class TestSoftware(Service, identifier="TestSoftware"): class TestSoftware(Service, discriminator="TestSoftware"):
class ConfigSchema(Service.ConfigSchema): class ConfigSchema(Service.ConfigSchema):
"""ConfigSChema for TestSoftware.""" """ConfigSChema for TestSoftware."""