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

#2445: apply PR suggestions

Browse Source
This commit is contained in:
Czar Echavez
2024-09-25 10:50:26 +01:00
parent 171dd83f2f
commit b9df2bd6a8
9 changed files with 103 additions and 81 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -55,21 +55,35 @@ class FileObservation(AbstractObservation, identifier="FILE"):
self.default_observation["num_access"] = 0
if thresholds.get("file_access") is None:
self.low_threshold = 0
self.med_threshold = 5
self.high_threshold = 10
self.low_file_access_threshold = 0
self.med_file_access_threshold = 5
self.high_file_access_threshold = 10
else:
if self._validate_thresholds(
self._set_file_access_threshold(
thresholds=[
thresholds.get("file_access")["low"],
thresholds.get("file_access")["medium"],
thresholds.get("file_access")["high"],
],
threshold_identifier="file_access",
):
self.low_threshold = thresholds.get("file_access")["low"]
self.med_threshold = thresholds.get("file_access")["medium"]
self.high_threshold = thresholds.get("file_access")["high"]
]
)
def _set_file_access_threshold(self, thresholds: List[int]):
"""
Method that validates and then sets the file access threshold.
:param: thresholds: The file access threshold to validate and set.
"""
if self._validate_thresholds(
thresholds=[
thresholds[0],
thresholds[1],
thresholds[2],
],
threshold_identifier="file_access",
):
self.low_file_access_threshold = thresholds[0]
self.med_file_access_threshold = thresholds[1]
self.high_file_access_threshold = thresholds[2]
def _categorise_num_access(self, num_access: int) -> int:
"""
@@ -78,11 +92,11 @@ class FileObservation(AbstractObservation, identifier="FILE"):
:param num_access: Number of file accesses.
:return: Bin number corresponding to the number of accesses.
"""
if num_access > self.high_threshold:
if num_access > self.high_file_access_threshold:
return 3
elif num_access > self.med_threshold:
elif num_access > self.med_file_access_threshold:
return 2
elif num_access > self.low_threshold:
elif num_access > self.low_file_access_threshold:
return 1
return 0
@@ -190,23 +204,6 @@ class FolderObservation(AbstractObservation, identifier="FOLDER"):
self.file_system_requires_scan: bool = file_system_requires_scan
if thresholds.get("file_access") is None:
self.low_threshold = 0
self.med_threshold = 5
self.high_threshold = 10
else:
if self._validate_thresholds(
thresholds=[
thresholds.get("file_access")["low"],
thresholds.get("file_access")["medium"],
thresholds.get("file_access")["high"],
],
threshold_identifier="file_access",
):
self.low_threshold = thresholds.get("file_access")["low"]
self.med_threshold = thresholds.get("file_access")["medium"]
self.high_threshold = thresholds.get("file_access")["high"]
self.files: List[FileObservation] = files
while len(self.files) < num_files:
self.files.append(

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

@@ -1,7 +1,7 @@
# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
from __future__ import annotations
from typing import ClassVar, Dict, Optional
from typing import ClassVar, Dict, List, Optional
from gymnasium import spaces
from gymnasium.core import ObsType
@@ -55,21 +55,17 @@ class NICObservation(AbstractObservation, identifier="NETWORK_INTERFACE"):
self.nmne_outbound_last_step: int = 0
if thresholds.get("nmne") is None:
self.low_threshold = 0
self.med_threshold = 5
self.high_threshold = 10
self.low_nmne_threshold = 0
self.med_nmne_threshold = 5
self.high_nmne_threshold = 10
else:
if self._validate_thresholds(
self._set_nmne_threshold(
thresholds=[
thresholds.get("nmne")["low"],
thresholds.get("nmne")["medium"],
thresholds.get("nmne")["high"],
],
threshold_identifier="nmne",
):
self.low_threshold = thresholds.get("nmne")["low"]
self.med_threshold = thresholds.get("nmne")["medium"]
self.high_threshold = thresholds.get("nmne")["high"]
]
)
self.monitored_traffic = monitored_traffic
if self.monitored_traffic:
@@ -108,11 +104,11 @@ class NICObservation(AbstractObservation, identifier="NETWORK_INTERFACE"):
:param nmne_count: Number of MNEs detected.
:return: Bin number corresponding to the number of MNEs. Returns 0, 1, 2, or 3 based on the detected MNE count.
"""
if nmne_count > self.high_threshold:
if nmne_count > self.high_nmne_threshold:
return 3
elif nmne_count > self.med_threshold:
elif nmne_count > self.med_nmne_threshold:
return 2
elif nmne_count > self.low_threshold:
elif nmne_count > self.low_nmne_threshold:
return 1
return 0
@@ -126,6 +122,20 @@ class NICObservation(AbstractObservation, identifier="NETWORK_INTERFACE"):
bandwidth_utilisation = traffic_value / nic_max_bandwidth
return int(bandwidth_utilisation * 9) + 1
def _set_nmne_threshold(self, thresholds: List[int]):
"""
Method that validates and then sets the NMNE threshold.
:param: thresholds: The NMNE threshold to validate and set.
"""
if self._validate_thresholds(
thresholds=thresholds,
threshold_identifier="nmne",
):
self.low_nmne_threshold = thresholds[0]
self.med_nmne_threshold = thresholds[1]
self.high_nmne_threshold = thresholds[2]
def observe(self, state: Dict) -> ObsType:
"""
Generate observation based on the current state of the simulation.

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

@@ -97,7 +97,7 @@ class AbstractObservation(ABC):
if thresholds[idx] <= thresholds[idx - 1]:
raise Exception(
f"{threshold_identifier} threshold ({thresholds[idx]}) "
f"is greater than or equal to ({thresholds[idx - 1]}.)"
f"{threshold_identifier} threshold ({thresholds[idx - 1]}) "
f"is greater than or equal to ({thresholds[idx]}.)"
)
return True

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

@@ -1,7 +1,7 @@
# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
from __future__ import annotations
from typing import Dict, Optional
from typing import Dict, List, Optional
from gymnasium import spaces
from gymnasium.core import ObsType
@@ -109,21 +109,35 @@ class ApplicationObservation(AbstractObservation, identifier="APPLICATION"):
self.default_observation = {"operating_status": 0, "health_status": 0, "num_executions": 0}
if thresholds.get("app_executions") is None:
self.low_threshold = 0
self.med_threshold = 5
self.high_threshold = 10
self.low_app_execution_threshold = 0
self.med_app_execution_threshold = 5
self.high_app_execution_threshold = 10
else:
if self._validate_thresholds(
self._set_application_execution_thresholds(
thresholds=[
thresholds.get("app_executions")["low"],
thresholds.get("app_executions")["medium"],
thresholds.get("app_executions")["high"],
],
threshold_identifier="app_executions",
):
self.low_threshold = thresholds.get("app_executions")["low"]
self.med_threshold = thresholds.get("app_executions")["medium"]
self.high_threshold = thresholds.get("app_executions")["high"]
]
)
def _set_application_execution_thresholds(self, thresholds: List[int]):
"""
Method that validates and then sets the application execution threshold.
:param: thresholds: The application execution threshold to validate and set.
"""
if self._validate_thresholds(
thresholds=[
thresholds[0],
thresholds[1],
thresholds[2],
],
threshold_identifier="app_executions",
):
self.low_app_execution_threshold = thresholds[0]
self.med_app_execution_threshold = thresholds[1]
self.high_app_execution_threshold = thresholds[2]
def _categorise_num_executions(self, num_executions: int) -> int:
"""
@@ -132,11 +146,11 @@ class ApplicationObservation(AbstractObservation, identifier="APPLICATION"):
:param num_access: Number of application executions.
:return: Bin number corresponding to the number of accesses.
"""
if num_executions > self.high_threshold:
if num_executions > self.high_app_execution_threshold:
return 3
elif num_executions > self.med_threshold:
elif num_executions > self.med_app_execution_threshold:
return 2
elif num_executions > self.low_threshold:
elif num_executions > self.low_app_execution_threshold:
return 1
return 0