Merge remote-tracking branch 'origin/feature/2438-application-install-action-refactor' into feature/2700-application-configure-actions

src/primaite/game/game.py

@@ -26,11 +26,14 @@ from primaite.simulator.network.hardware.nodes.network.wireless_router import Wi
 from primaite.simulator.network.nmne import set_nmne_config
 from primaite.simulator.network.transmission.transport_layer import Port
 from primaite.simulator.sim_container import Simulation
-from primaite.simulator.system.applications.database_client import DatabaseClient
-from primaite.simulator.system.applications.red_applications.data_manipulation_bot import DataManipulationBot
-from primaite.simulator.system.applications.red_applications.dos_bot import DoSBot
-from primaite.simulator.system.applications.red_applications.ransomware_script import RansomwareScript
-from primaite.simulator.system.applications.web_browser import WebBrowser
+from primaite.simulator.system.applications.application import Application
+from primaite.simulator.system.applications.database_client import DatabaseClient  # noqa: F401
+from primaite.simulator.system.applications.red_applications.data_manipulation_bot import (  # noqa: F401
+    DataManipulationBot,
+)
+from primaite.simulator.system.applications.red_applications.dos_bot import DoSBot  # noqa: F401
+from primaite.simulator.system.applications.red_applications.ransomware_script import RansomwareScript  # noqa: F401
+from primaite.simulator.system.applications.web_browser import WebBrowser  # noqa: F401
 from primaite.simulator.system.services.database.database_service import DatabaseService
 from primaite.simulator.system.services.dns.dns_client import DNSClient
 from primaite.simulator.system.services.dns.dns_server import DNSServer
@@ -42,15 +45,6 @@ from primaite.simulator.system.services.web_server.web_server import WebServer
 
 _LOGGER = getLogger(__name__)
 
-APPLICATION_TYPES_MAPPING = {
-    "WebBrowser": WebBrowser,
-    "DatabaseClient": DatabaseClient,
-    "DataManipulationBot": DataManipulationBot,
-    "DoSBot": DoSBot,
-    "RansomwareScript": RansomwareScript,
-}
-"""List of available applications that can be installed on nodes in the PrimAITE Simulation."""
-
 SERVICE_TYPES_MAPPING = {
     "DNSClient": DNSClient,
     "DNSServer": DNSServer,
@@ -333,9 +327,9 @@ class PrimaiteGame:
                     new_application = None
                     application_type = application_cfg["type"]
 
-                    if application_type in APPLICATION_TYPES_MAPPING:
-                        new_node.software_manager.install(APPLICATION_TYPES_MAPPING[application_type])
-                        new_application = new_node.software_manager.software[application_type]
+                    if application_type in Application._application_registry:
+                        new_node.software_manager.install(Application._application_registry[application_type])
+                        new_application = new_node.software_manager.software[application_type]  # grab the instance
                     else:
                         msg = f"Configuration contains an invalid application type: {application_type}"
                         _LOGGER.error(msg)

src/primaite/simulator/_package_data/create-simulation_demo.ipynb

@@ -171,7 +171,7 @@
     "from primaite.simulator.file_system.file_system import FileSystem\n",
     "\n",
     "# no applications exist yet so we will create our own.\n",
-    "class MSPaint(Application):\n",
+    "class MSPaint(Application, identifier=\"MSPaint\"):\n",
     "    def describe_state(self):\n",
     "        return super().describe_state()"
    ]

src/primaite/simulator/core.py

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

src/primaite/simulator/system/applications/application.py

@@ -1,7 +1,7 @@
 # © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
 from abc import abstractmethod
 from enum import Enum
-from typing import Any, Dict, Optional, Set
+from typing import Any, ClassVar, Dict, Optional, Set, Type
 
 from primaite.interface.request import RequestResponse
 from primaite.simulator.core import RequestManager, RequestType
@@ -39,6 +39,22 @@ class Application(IOSoftware):
     install_countdown: Optional[int] = None
     "The countdown to the end of the installation process. None if not currently installing"
 
+    _application_registry: ClassVar[Dict[str, Type["Application"]]] = {}
+    """Registry of application types. Automatically populated when subclasses are defined."""
+
+    def __init_subclass__(cls, identifier: str, **kwargs: Any) -> None:
+        """
+        Register an application type.
+
+        :param identifier: Uniquely specifies an application class by name. Used for finding items by config.
+        :type identifier: str
+        :raises ValueError: When attempting to register an application with a name that is already allocated.
+        """
+        super().__init_subclass__(**kwargs)
+        if identifier in cls._application_registry:
+            raise ValueError(f"Tried to define new application {identifier}, but this name is already reserved.")
+        cls._application_registry[identifier] = cls
+
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
 

src/primaite/simulator/system/applications/database_client.py

@@ -55,7 +55,7 @@ class DatabaseClientConnection(BaseModel):
             self.client._disconnect(self.connection_id)  # noqa
 
 
-class DatabaseClient(Application):
+class DatabaseClient(Application, identifier="DatabaseClient"):
     """
     A DatabaseClient application.
 

src/primaite/simulator/system/applications/nmap.py

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

src/primaite/simulator/system/applications/red_applications/data_manipulation_bot.py

@@ -37,7 +37,7 @@ class DataManipulationAttackStage(IntEnum):
     "Signifies that the attack has failed."
 
 
-class DataManipulationBot(Application):
+class DataManipulationBot(Application, identifier="DataManipulationBot"):
     """A bot that simulates a script which performs a SQL injection attack."""
 
     payload: Optional[str] = None

src/primaite/simulator/system/applications/red_applications/dos_bot.py

@@ -29,7 +29,7 @@ class DoSAttackStage(IntEnum):
     "Attack is completed."
 
 
-class DoSBot(DatabaseClient):
+class DoSBot(DatabaseClient, identifier="DoSBot"):
     """A bot that simulates a Denial of Service attack."""
 
     target_ip_address: Optional[IPv4Address] = None

src/primaite/simulator/system/applications/red_applications/ransomware_script.py

@@ -10,7 +10,7 @@ from primaite.simulator.system.applications.application import Application
 from primaite.simulator.system.applications.database_client import DatabaseClient, DatabaseClientConnection
 
 
-class RansomwareScript(Application):
+class RansomwareScript(Application, identifier="RansomwareScript"):
     """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)

src/primaite/simulator/system/applications/web_browser.py

@@ -23,7 +23,7 @@ from primaite.simulator.system.services.dns.dns_client import DNSClient
 _LOGGER = getLogger(__name__)
 
 
-class WebBrowser(Application):
+class WebBrowser(Application, identifier="WebBrowser"):
     """
     Represents a web browser in the simulation environment.