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

port refactor - eod commit, airspace still broken

Browse Source
This commit is contained in:
Marek Wolan
2024-09-18 16:02:25 +01:00
parent 310876cd3b
commit dd931d900b
92 changed files with 957 additions and 682 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -10,6 +10,8 @@ from gymnasium.core import ObsType
from primaite import getLogger
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.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port
_LOGGER = getLogger(__name__)
@@ -61,7 +63,7 @@ class ACLObservation(AbstractObservation, identifier="ACL"):
self.ip_to_id: Dict[str, int] = {p: i + 2 for i, p in enumerate(ip_list)}
self.wildcard_to_id: Dict[str, int] = {p: i + 2 for i, p in enumerate(wildcard_list)}
self.port_to_id: Dict[int, int] = {p: i + 2 for i, p in enumerate(port_list)}
self.protocol_to_id: Dict[str, int] = {p: i + 2 for i, p in enumerate(protocol_list)}
self.protocol_to_id: Dict[str, int] = {IPProtocol[p]: i + 2 for i, p in enumerate(protocol_list)}
self.default_observation: Dict = {
i
+ 1: {

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

@@ -5,6 +5,7 @@ from typing import Dict, List, Optional
from gymnasium import spaces
from gymnasium.core import ObsType
from pydantic import field_validator
from primaite import getLogger
from primaite.game.agent.observations.file_system_observations import FolderObservation
@@ -12,6 +13,8 @@ from primaite.game.agent.observations.nic_observations import NICObservation
from primaite.game.agent.observations.observations import AbstractObservation, WhereType
from primaite.game.agent.observations.software_observation import ApplicationObservation, ServiceObservation
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port
_LOGGER = getLogger(__name__)
@@ -55,6 +58,21 @@ class HostObservation(AbstractObservation, identifier="HOST"):
include_users: Optional[bool] = True
"""If True, report user session information."""
@field_validator('monitored_traffic', mode='before')
def traffic_lookup(cls, val:Optional[Dict]) -> Optional[Dict]:
if val is None:
return val
new_val = {}
for proto, port_list in val.items():
# convert protocol, for instance ICMP becomes "icmp"
proto = IPProtocol[proto] if proto in IPProtocol else proto
new_val[proto] = []
for port in port_list:
# convert ports, for instance "HTTP" becomes 80
port = Port[port] if port in Port else port
new_val[proto].append(port)
return new_val
def __init__(
self,
where: WhereType,

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

@@ -5,9 +5,11 @@ from typing import Dict, Optional
from gymnasium import spaces
from gymnasium.core import ObsType
from pydantic import field_validator
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.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port
@@ -24,6 +26,22 @@ class NICObservation(AbstractObservation, identifier="NETWORK_INTERFACE"):
monitored_traffic: Optional[Dict] = None
"""A dict containing which traffic types are to be included in the observation."""
@field_validator('monitored_traffic', mode='before')
def traffic_lookup(cls, val:Optional[Dict]) -> Optional[Dict]:
if val is None:
return val
new_val = {}
for proto, port_list in val.items():
# convert protocol, for instance ICMP becomes "icmp"
proto = IPProtocol[proto] if proto in IPProtocol else proto
new_val[proto] = []
for port in port_list:
# convert ports, for instance "HTTP" becomes 80
port = Port[port] if port in Port else port
new_val[proto].append(port)
return new_val
def __init__(self, where: WhereType, include_nmne: bool, monitored_traffic: Optional[Dict] = None) -> None:
"""
Initialise a network interface observation instance.
@@ -67,7 +85,7 @@ class NICObservation(AbstractObservation, identifier="NETWORK_INTERFACE"):
else:
default_traffic_obs["TRAFFIC"][protocol] = {}
for port in monitored_traffic_config[protocol]:
default_traffic_obs["TRAFFIC"][protocol][Port[port].value] = {"inbound": 0, "outbound": 0}
default_traffic_obs["TRAFFIC"][protocol] = {"inbound": 0, "outbound": 0}
return default_traffic_obs
@@ -142,17 +160,16 @@ class NICObservation(AbstractObservation, identifier="NETWORK_INTERFACE"):
}
else:
for port in self.monitored_traffic[protocol]:
port_enum = Port[port]
obs["TRAFFIC"][protocol][port_enum.value] = {}
obs["TRAFFIC"][protocol][port] = {}
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]
if nic_state["traffic"][protocol].get(port) is not None:
traffic = nic_state["traffic"][protocol][port]
obs["TRAFFIC"][protocol][port_enum.value]["inbound"] = self._categorise_traffic(
obs["TRAFFIC"][protocol][port]["inbound"] = self._categorise_traffic(
traffic_value=traffic["inbound"], nic_state=nic_state
)
obs["TRAFFIC"][protocol][port_enum.value]["outbound"] = self._categorise_traffic(
obs["TRAFFIC"][protocol][port]["outbound"] = self._categorise_traffic(
traffic_value=traffic["outbound"], nic_state=nic_state
)
@@ -162,7 +179,7 @@ class NICObservation(AbstractObservation, identifier="NETWORK_INTERFACE"):
obs["TRAFFIC"]["icmp"] = {"inbound": 0, "outbound": 0}
else:
for port in self.monitored_traffic[protocol]:
obs["TRAFFIC"][protocol][Port[port].value] = {"inbound": 0, "outbound": 0}
obs["TRAFFIC"][protocol][port] = {"inbound": 0, "outbound": 0}
if self.include_nmne:
obs.update({"NMNE": {}})
@@ -201,7 +218,7 @@ class NICObservation(AbstractObservation, identifier="NETWORK_INTERFACE"):
else:
space["TRAFFIC"][protocol] = spaces.Dict({})
for port in self.monitored_traffic[protocol]:
space["TRAFFIC"][protocol][Port[port].value] = spaces.Dict(
space["TRAFFIC"][protocol][port] = spaces.Dict(
{"inbound": spaces.Discrete(11), "outbound": spaces.Discrete(11)}
)

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

@@ -5,13 +5,15 @@ from typing import Dict, List, Optional
from gymnasium import spaces
from gymnasium.core import ObsType
from pydantic import model_validator
from pydantic import field_validator, model_validator
from primaite import getLogger
from primaite.game.agent.observations.firewall_observation import FirewallObservation
from primaite.game.agent.observations.host_observations import HostObservation
from primaite.game.agent.observations.observations import AbstractObservation, WhereType
from primaite.game.agent.observations.router_observation import RouterObservation
from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port
_LOGGER = getLogger(__name__)
@@ -61,6 +63,21 @@ class NodesObservation(AbstractObservation, identifier="NODES"):
num_rules: Optional[int] = None
"""Number of rules ACL rules to show."""
@field_validator('monitored_traffic', mode='before')
def traffic_lookup(cls, val:Optional[Dict]) -> Optional[Dict]:
if val is None:
return val
new_val = {}
for proto, port_list in val.items():
# convert protocol, for instance ICMP becomes "icmp"
proto = IPProtocol[proto] if proto in IPProtocol else proto
new_val[proto] = []
for port in port_list:
# convert ports, for instance "HTTP" becomes 80
port = Port[port] if port in Port else port
new_val[proto].append(port)
return new_val
@model_validator(mode="after")
def force_optional_fields(self) -> NodesObservation.ConfigSchema:
"""Check that options are specified only if they are needed for the nodes that are part of the config."""