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

#2658: added monitored traffic into config + default obs space

Browse Source
This commit is contained in:
Czar Echavez
2024-06-13 11:48:13 +01:00
parent 969daeb490
commit 539577ddc3
8 changed files with 216 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -43,6 +43,8 @@ class HostObservation(AbstractObservation, identifier="HOST"):
"""Number of spaces for network interface observations on this host."""
include_nmne: Optional[bool] = None
"""Whether network interface observations should include number of malicious network events."""
monitored_traffic: Optional[Dict] = None
"""A dict containing which traffic types are to be included in the observation."""
include_num_access: Optional[bool] = None
"""Whether to include the number of accesses to files observations on this host."""
@@ -59,6 +61,7 @@ class HostObservation(AbstractObservation, identifier="HOST"):
num_files: int,
num_nics: int,
include_nmne: bool,
monitored_traffic: Optional[Dict],
include_num_access: bool,
) -> None:
"""
@@ -87,6 +90,8 @@ class HostObservation(AbstractObservation, identifier="HOST"):
:type num_nics: int
:param include_nmne: Flag to include network metrics and errors.
:type include_nmne: bool
:param monitored_traffic: Dict which contains the protocol and ports to observe
:type monitored_traffic: Dict
:param include_num_access: Flag to include the number of accesses to files.
:type include_num_access: bool
"""
@@ -123,7 +128,7 @@ class HostObservation(AbstractObservation, identifier="HOST"):
self.nics: List[NICObservation] = network_interfaces
while len(self.nics) < num_nics:
self.nics.append(NICObservation(where=None, include_nmne=include_nmne))
self.nics.append(NICObservation(where=None, include_nmne=include_nmne, monitored_traffic=monitored_traffic))
while len(self.nics) > num_nics:
truncated_nic = self.nics.pop()
msg = f"Too many network_interfaces in Node observation space for node. Truncating {truncated_nic.where}"
@@ -231,7 +236,9 @@ class HostObservation(AbstractObservation, identifier="HOST"):
# monitor the first N interfaces. Network interface numbering starts at 1.
count = 1
while len(nics) < config.num_nics:
nic_config = NICObservation.ConfigSchema(nic_num=count, include_nmne=config.include_nmne)
nic_config = NICObservation.ConfigSchema(
nic_num=count, include_nmne=config.include_nmne, monitored_traffic=config.monitored_traffic
)
nics.append(NICObservation.from_config(config=nic_config, parent_where=where))
count += 1
@@ -247,5 +254,6 @@ class HostObservation(AbstractObservation, identifier="HOST"):
num_files=config.num_files,
num_nics=config.num_nics,
include_nmne=config.include_nmne,
monitored_traffic=config.monitored_traffic,
include_num_access=config.include_num_access,
)

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

@@ -7,6 +7,7 @@ from gymnasium.core import ObsType
from primaite.game.agent.observations.observations import AbstractObservation, WhereType
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
from primaite.simulator.network.transmission.transport_layer import Port
class NICObservation(AbstractObservation, identifier="NETWORK_INTERFACE"):
@@ -19,12 +20,10 @@ class NICObservation(AbstractObservation, identifier="NETWORK_INTERFACE"):
"""Number of the network interface."""
include_nmne: Optional[bool] = None
"""Whether to include number of malicious network events (NMNE) in the observation."""
monitored_traffic: Optional[Dict] = None
"""A dict containing which traffic types are to be included in the observation."""
def __init__(
self,
where: WhereType,
include_nmne: bool,
) -> None:
def __init__(self, where: WhereType, include_nmne: bool, monitored_traffic: Optional[Dict] = None) -> None:
"""
Initialise a network interface observation instance.
@@ -49,6 +48,26 @@ class NICObservation(AbstractObservation, identifier="NETWORK_INTERFACE"):
self.med_nmne_threshold = 5
self.low_nmne_threshold = 0
self.monitored_traffic = monitored_traffic
if self.monitored_traffic:
self.default_observation.update(
self._default_monitored_traffic_observation(monitored_traffic_config=monitored_traffic)
)
def _default_monitored_traffic_observation(self, monitored_traffic_config: Dict) -> Dict:
default_traffic_obs = {"monitored_traffic": {}}
for protocol in monitored_traffic_config:
default_traffic_obs["monitored_traffic"][str(protocol).lower()] = {}
if protocol == "icmp":
default_traffic_obs["monitored_traffic"]["icmp"] = {"inbound": 0, "outbound": 0}
else:
for port in monitored_traffic_config[protocol]:
default_traffic_obs["monitored_traffic"][protocol][Port[port].value] = {"inbound": 0, "outbound": 0}
return default_traffic_obs
def _categorise_mne_count(self, nmne_count: int) -> int:
"""
Categorise the number of Malicious Network Events (NMNEs) into discrete bins.
@@ -72,6 +91,16 @@ class NICObservation(AbstractObservation, identifier="NETWORK_INTERFACE"):
return 1
return 0
def _categorise_traffic(self, traffic_value: float, nic_state: Dict) -> int:
"""Categorise the traffic into discrete categories."""
if traffic_value == 0:
return 0
nic_max_bandwidth = nic_state.get("speed")
bandwidth_utilisation = traffic_value / nic_max_bandwidth
return int(bandwidth_utilisation * 9) + 1
def observe(self, state: Dict) -> ObsType:
"""
Generate observation based on the current state of the simulation.
@@ -87,6 +116,50 @@ class NICObservation(AbstractObservation, identifier="NETWORK_INTERFACE"):
return self.default_observation
obs = {"nic_status": 1 if nic_state["enabled"] else 2}
# if the observation was configured to monitor traffic from ports/protocols
if self.monitored_traffic:
obs["monitored_traffic"] = {}
# iterate through the protocols
for protocol in self.monitored_traffic:
obs["monitored_traffic"][str(protocol).lower()] = {}
# check if the nic has seen traffic with this protocol
if nic_state["traffic"].get(protocol):
# deal with icmp
if str(protocol).lower() == "icmp":
obs["monitored_traffic"][protocol] = {
"inbound": self._categorise_traffic(
traffic_value=nic_state["traffic"]["icmp"]["inbound"], nic_state=nic_state
),
"outbound": self._categorise_traffic(
traffic_value=nic_state["traffic"]["icmp"]["outbound"], nic_state=nic_state
),
}
else:
for port in self.monitored_traffic[protocol]:
port_enum = Port[port]
obs["monitored_traffic"][str(protocol).lower()][port_enum.value] = {}
traffic = {"inbound": 0, "outbound": 0}
if nic_state["traffic"][protocol].get(port_enum.value) is not None:
traffic = nic_state["traffic"][protocol][port_enum.value]
obs["monitored_traffic"][protocol][port_enum.value]["inbound"] = self._categorise_traffic(
traffic_value=traffic["inbound"], nic_state=nic_state
)
obs["monitored_traffic"][protocol][port_enum.value]["outbound"] = self._categorise_traffic(
traffic_value=traffic["outbound"], nic_state=nic_state
)
# set all the ports under the protocol to 0
else:
if str(protocol).lower() == "icmp":
obs["monitored_traffic"]["icmp"] = {"inbound": 0, "outbound": 0}
else:
for port in self.monitored_traffic[protocol]:
obs["monitored_traffic"][protocol][Port[port].value] = {"inbound": 0, "outbound": 0}
if self.include_nmne:
obs.update({"NMNE": {}})
direction_dict = nic_state["nmne"].get("direction", {})
@@ -128,7 +201,11 @@ class NICObservation(AbstractObservation, identifier="NETWORK_INTERFACE"):
:return: Constructed network interface observation instance.
:rtype: NICObservation
"""
return cls(where=parent_where + ["NICs", config.nic_num], include_nmne=config.include_nmne)
return cls(
where=parent_where + ["NICs", config.nic_num],
include_nmne=config.include_nmne,
monitored_traffic=config.monitored_traffic,
)
class PortObservation(AbstractObservation, identifier="PORT"):

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

@@ -39,6 +39,8 @@ class NodesObservation(AbstractObservation, identifier="NODES"):
"""Number of network interface cards (NICs)."""
include_nmne: Optional[bool] = None
"""Flag to include nmne."""
monitored_traffic: Optional[Dict] = None
"""A dict containing which traffic types are to be included in the observation."""
include_num_access: Optional[bool] = None
"""Flag to include the number of accesses."""
num_ports: Optional[int] = None
@@ -180,6 +182,8 @@ class NodesObservation(AbstractObservation, identifier="NODES"):
host_config.num_nics = config.num_nics
if host_config.include_nmne is None:
host_config.include_nmne = config.include_nmne
if host_config.monitored_traffic is None:
host_config.monitored_traffic = config.monitored_traffic
if host_config.include_num_access is None:
host_config.include_num_access = config.include_num_access