From 43617340148051973518f72b45dce01cbb6f40cf Mon Sep 17 00:00:00 2001 From: Nick Todd Date: Wed, 17 Jul 2024 17:50:55 +0100 Subject: [PATCH] #2676: Code review changes --- src/primaite/game/game.py | 5 +---- src/primaite/simulator/network/hardware/base.py | 4 ++-- src/primaite/simulator/network/nmne.py | 11 +++++------ 3 files changed, 8 insertions(+), 12 deletions(-) diff --git a/src/primaite/game/game.py b/src/primaite/game/game.py index aca75b63..0c1b3192 100644 --- a/src/primaite/game/game.py +++ b/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. diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py index f85d3f2e..50549389 100644 --- a/src/primaite/simulator/network/hardware/base.py +++ b/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: {}) diff --git a/src/primaite/simulator/network/nmne.py b/src/primaite/simulator/network/nmne.py index 947f27ac..431ec07d 100644 --- a/src/primaite/simulator/network/nmne.py +++ b/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)