#2676: Code review changes

src/primaite/game/game.py

@@ -26,7 +26,7 @@ from primaite.simulator.network.hardware.nodes.network.firewall import Firewall
 from primaite.simulator.network.hardware.nodes.network.router import Router
 from primaite.simulator.network.hardware.nodes.network.switch import Switch
 from primaite.simulator.network.hardware.nodes.network.wireless_router import WirelessRouter
-from primaite.simulator.network.nmne import NmneData, store_nmne_config
+from primaite.simulator.network.nmne import store_nmne_config
 from primaite.simulator.network.transmission.transport_layer import Port
 from primaite.simulator.sim_container import Simulation
 from primaite.simulator.system.applications.application import Application
@@ -110,9 +110,6 @@ class PrimaiteGame:
         self._reward_calculation_order: List[str] = [name for name in self.agents]
         """Agent order for reward evaluation, as some rewards can be dependent on other agents' rewards."""
 
-        self.nmne_config: NmneData = None
-        """ Config data from Number of Malicious Network Events."""
-
     def step(self):
         """
         Perform one step of the simulation/agent loop.

src/primaite/simulator/network/hardware/base.py

@@ -19,7 +19,7 @@ from primaite.simulator.core import RequestFormat, RequestManager, RequestPermis
 from primaite.simulator.domain.account import Account
 from primaite.simulator.file_system.file_system import FileSystem
 from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
-from primaite.simulator.network.nmne import NmneData
+from primaite.simulator.network.nmne import NMNEConfig
 from primaite.simulator.network.transmission.data_link_layer import Frame
 from primaite.simulator.network.transmission.network_layer import IPProtocol
 from primaite.simulator.system.applications.application import Application
@@ -99,7 +99,7 @@ class NetworkInterface(SimComponent, ABC):
     pcap: Optional[PacketCapture] = None
     "A PacketCapture instance for capturing and analysing packets passing through this interface."
 
-    nmne_config: ClassVar[NmneData] = None
+    nmne_config: ClassVar[NMNEConfig] = None
     "A dataclass defining malicious network events to be captured."
 
     nmne: Dict = Field(default_factory=lambda: {})

src/primaite/simulator/network/nmne.py

@@ -1,15 +1,14 @@
 # © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
-from dataclasses import dataclass, field
+from pydantic import BaseModel
 from typing import Dict, List
 
 
-@dataclass
-class NmneData:
+class NMNEConfig(BaseModel):
     """Store all the information to perform NMNE operations."""
 
     capture_nmne: bool = True
     """Indicates whether Malicious Network Events (MNEs) should be captured."""
-    nmne_capture_keywords: List[str] = field(default_factory=list)
+    nmne_capture_keywords: List[str] = []
     """List of keywords to identify malicious network events."""
     capture_by_direction: bool = True
     """Captures should be organized by traffic direction (inbound/outbound)."""
@@ -23,7 +22,7 @@ class NmneData:
     """Captures should be filtered and categorised based on specific keywords."""
 
 
-def store_nmne_config(nmne_config: Dict) -> NmneData:
+def store_nmne_config(nmne_config: Dict) -> NMNEConfig:
     """
     Store configuration for capturing Malicious Network Events (MNEs).
 
@@ -51,4 +50,4 @@ def store_nmne_config(nmne_config: Dict) -> NmneData:
     if not isinstance(nmne_capture_keywords, list):
         nmne_capture_keywords = []  # Reset to empty list if the provided value is not a list
 
-    return NmneData(capture_nmne=capture_nmne, nmne_capture_keywords=nmne_capture_keywords)
+    return NMNEConfig(capture_nmne=capture_nmne, nmne_capture_keywords=nmne_capture_keywords)