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."""

39
src/primaite/game/game.py
View File

@@ -4,7 +4,7 @@ from ipaddress import IPv4Address
from typing import Dict, List, Optional, Union
import numpy as np
from pydantic import BaseModel, ConfigDict
from pydantic import BaseModel, ConfigDict, field_validator
from primaite import DEFAULT_BANDWIDTH, getLogger
from primaite.game.agent.actions import ActionManager
@@ -82,13 +82,40 @@ class PrimaiteGameOptions(BaseModel):
"""Random number seed for RNGs."""
max_episode_length: int = 256
"""Maximum number of episodes for the PrimAITE game."""
ports: List[str]
ports: List[int]
"""A whitelist of available ports in the simulation."""
protocols: List[str]
"""A whitelist of available protocols in the simulation."""
thresholds: Optional[Dict] = {}
"""A dict containing the thresholds used for determining what is acceptable during observations."""
@field_validator('ports', mode='before')
def ports_str2int(cls, vals:Union[List[str],List[int]]) -> List[int]:
"""
Convert named port strings to port integer values. Integer ports remain unaffected.
This is necessary to retain backwards compatibility with configs written for PrimAITE<=3.3.
:warning: This will be deprecated in PrimAITE 4.0 and configs will need to be converted.
"""
for i, port_val in enumerate(vals):
if port_val in Port:
vals[i] = Port[port_val]
return vals
@field_validator('protocols', mode='before')
def protocols_str2int(cls, vals:List[str]) -> List[str]:
"""
Convert old-style named protocols to their proper values.
This is necessary to retain backwards compatibility with configs written for PrimAITE<=3.3.
:warning: This will be deprecated in PrimAITE 4.0 and configs will need to be converted.
"""
for i, proto_val in enumerate(vals):
if proto_val in IPProtocol:
vals[i] = IPProtocol[proto_val]
return vals
class PrimaiteGame:
"""
@@ -358,7 +385,7 @@ class PrimaiteGame:
for port_id in set(software_cfg.get("options", {}).get("listen_on_ports", [])):
port = None
if isinstance(port_id, int):
port = Port(port_id)
port = port_id
elif isinstance(port_id, str):
port = Port[port_id]
if port:
@@ -475,7 +502,7 @@ class PrimaiteGame:
opt = application_cfg["options"]
new_application.configure(
target_ip_address=IPv4Address(opt.get("target_ip_address")),
target_port=Port(opt.get("target_port", Port.POSTGRES_SERVER.value)),
target_port = Port[opt.get("target_port", "POSTGRES_SERVER")],
payload=opt.get("payload"),
repeat=bool(opt.get("repeat")),
port_scan_p_of_success=float(opt.get("port_scan_p_of_success", "0.1")),
@@ -488,8 +515,8 @@ class PrimaiteGame:
new_application.configure(
c2_server_ip_address=IPv4Address(opt.get("c2_server_ip_address")),
keep_alive_frequency=(opt.get("keep_alive_frequency", 5)),
masquerade_protocol=IPProtocol[(opt.get("masquerade_protocol", IPProtocol.TCP))],
masquerade_port=Port[(opt.get("masquerade_port", Port.HTTP))],
masquerade_protocol=IPProtocol[(opt.get("masquerade_protocol", IPProtocol["TCP"]))],
masquerade_port=Port[(opt.get("masquerade_port", Port["HTTP"]))],
)
if "network_interfaces" in node_cfg:
for nic_num, nic_cfg in node_cfg["network_interfaces"].items():