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

Fix airspace and remaining port problems from refactor

Browse Source
This commit is contained in:
Marek Wolan
2024-09-19 15:06:29 +01:00
parent dd931d900b
commit 08f1cf1fbd
34 changed files with 227 additions and 177 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,8 +10,6 @@ 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__)
@@ -63,7 +61,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] = {IPProtocol[p]: i + 2 for i, p in enumerate(protocol_list)}
self.protocol_to_id: Dict[str, int] = {p: i + 2 for i, p in enumerate(protocol_list)}
self.default_observation: Dict = {
i
+ 1: {

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

@@ -58,8 +58,14 @@ 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]:
@field_validator("monitored_traffic", mode="before")
def traffic_lookup(cls, val: Optional[Dict]) -> Optional[Dict]:
"""
Convert monitored_traffic by lookup against Port and Protocol dicts.
This is necessary for retaining compatiblility with configs written for PrimAITE <=3.3.
This method will be removed in PrimAITE >= 4.0
"""
if val is None:
return val
new_val = {}

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

@@ -26,8 +26,14 @@ 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]:
@field_validator("monitored_traffic", mode="before")
def traffic_lookup(cls, val: Optional[Dict]) -> Optional[Dict]:
"""
Convert monitored_traffic by lookup against Port and Protocol dicts.
This is necessary for retaining compatiblility with configs written for PrimAITE <=3.3.
This method will be removed in PrimAITE >= 4.0
"""
if val is None:
return val
new_val = {}
@@ -41,7 +47,6 @@ class NICObservation(AbstractObservation, identifier="NETWORK_INTERFACE"):
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.
@@ -76,7 +81,7 @@ class NICObservation(AbstractObservation, identifier="NETWORK_INTERFACE"):
def _default_monitored_traffic_observation(self, monitored_traffic_config: Dict) -> Dict:
default_traffic_obs = {"TRAFFIC": {}}
for protocol in monitored_traffic_config:
for protocol in self.monitored_traffic:
protocol = str(protocol).lower()
default_traffic_obs["TRAFFIC"][protocol] = {}
@@ -84,8 +89,8 @@ class NICObservation(AbstractObservation, identifier="NETWORK_INTERFACE"):
default_traffic_obs["TRAFFIC"]["icmp"] = {"inbound": 0, "outbound": 0}
else:
default_traffic_obs["TRAFFIC"][protocol] = {}
for port in monitored_traffic_config[protocol]:
default_traffic_obs["TRAFFIC"][protocol] = {"inbound": 0, "outbound": 0}
for port in self.monitored_traffic[protocol]:
default_traffic_obs["TRAFFIC"][protocol][port] = {"inbound": 0, "outbound": 0}
return default_traffic_obs

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

@@ -63,8 +63,14 @@ 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]:
@field_validator("monitored_traffic", mode="before")
def traffic_lookup(cls, val: Optional[Dict]) -> Optional[Dict]:
"""
Convert monitored_traffic by lookup against Port and Protocol dicts.
This is necessary for retaining compatiblility with configs written for PrimAITE <=3.3.
This method will be removed in PrimAITE >= 4.0
"""
if val is None:
return val
new_val = {}

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

@@ -17,10 +17,9 @@ from primaite.game.agent.scripted_agents.random_agent import PeriodicAgent
from primaite.game.agent.scripted_agents.tap001 import TAP001
from primaite.game.science import graph_has_cycle, topological_sort
from primaite.simulator import SIM_OUTPUT
from primaite.simulator.network.airspace import AirSpaceFrequency
from primaite.simulator.network.hardware.base import NetworkInterface, NodeOperatingState, UserManager
from primaite.simulator.network.hardware.nodes.host.computer import Computer
from primaite.simulator.network.hardware.nodes.host.host_node import NIC, HostNode
from primaite.simulator.network.hardware.nodes.host.host_node import HostNode, NIC
from primaite.simulator.network.hardware.nodes.host.server import Printer, Server
from primaite.simulator.network.hardware.nodes.network.firewall import Firewall
from primaite.simulator.network.hardware.nodes.network.network_node import NetworkNode
@@ -89,8 +88,8 @@ class PrimaiteGameOptions(BaseModel):
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]:
@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.
@@ -102,8 +101,8 @@ class PrimaiteGameOptions(BaseModel):
vals[i] = Port[port_val]
return vals
@field_validator('protocols', mode='before')
def protocols_str2int(cls, vals:List[str]) -> List[str]:
@field_validator("protocols", mode="before")
def protocols_str2int(cls, vals: List[str]) -> List[str]:
"""
Convert old-style named protocols to their proper values.
@@ -116,7 +115,6 @@ class PrimaiteGameOptions(BaseModel):
return vals
class PrimaiteGame:
"""
Primaite game encapsulates the simulation and agents which interact with it.
@@ -294,10 +292,7 @@ class PrimaiteGame:
network_config = simulation_config.get("network", {})
airspace_cfg = network_config.get("airspace", {})
frequency_max_capacity_mbps_cfg = airspace_cfg.get("frequency_max_capacity_mbps", {})
frequency_max_capacity_mbps_cfg = {AirSpaceFrequency[k]: v for k, v in frequency_max_capacity_mbps_cfg.items()}
net.airspace.frequency_max_capacity_mbps_ = frequency_max_capacity_mbps_cfg
net.airspace.set_frequency_max_capacity_mbps(frequency_max_capacity_mbps_cfg)
nodes_cfg = network_config.get("nodes", [])
links_cfg = network_config.get("links", [])
@@ -318,11 +313,10 @@ class PrimaiteGame:
dns_server=node_cfg.get("dns_server", None),
operating_state=NodeOperatingState.ON
if not (p := node_cfg.get("operating_state"))
else NodeOperatingState[p.upper()])
elif n_type in NetworkNode._registry:
new_node = NetworkNode._registry[n_type](
**node_cfg
else NodeOperatingState[p.upper()],
)
elif n_type in NetworkNode._registry:
new_node = NetworkNode._registry[n_type](**node_cfg)
# Default PrimAITE nodes
elif n_type == "computer":
new_node = Computer(
@@ -502,7 +496,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", "POSTGRES_SERVER")],
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")),