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

Add port and protocol custom validators

Browse Source
This commit is contained in:
Marek Wolan
2024-09-20 11:21:28 +01:00
parent 08f1cf1fbd
commit 695891f55c
77 changed files with 616 additions and 613 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -5,7 +5,6 @@ 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
@@ -13,8 +12,7 @@ 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
from primaite.utils.validators import IPProtocol, Port
_LOGGER = getLogger(__name__)
@@ -47,7 +45,7 @@ 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
monitored_traffic: Optional[Dict[IPProtocol, List[Port]]] = 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."""
@@ -58,26 +56,26 @@ 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]:
"""
Convert monitored_traffic by lookup against Port and Protocol dicts.
# @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 = {}
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
# 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 = {}
# for proto, port_list in val.items():
# # convert protocol, for instance ICMP becomes "icmp"
# proto = PROTOCOL_LOOKUP[proto] if proto in PROTOCOL_LOOKUP else proto
# new_val[proto] = []
# for port in port_list:
# # convert ports, for instance "HTTP" becomes 80
# port = PORT_LOOKUP[port] if port in PORT_LOOKUP else port
# new_val[proto].append(port)
# return new_val
def __init__(
self,

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

@@ -9,8 +9,8 @@ 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
from primaite.simulator.network.transmission.transport_layer import PORT_LOOKUP
from primaite.utils.validators import PROTOCOL_LOOKUP
class NICObservation(AbstractObservation, identifier="NETWORK_INTERFACE"):
@@ -39,11 +39,11 @@ class NICObservation(AbstractObservation, identifier="NETWORK_INTERFACE"):
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
proto = PROTOCOL_LOOKUP[proto] if proto in PROTOCOL_LOOKUP 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
port = PORT_LOOKUP[port] if port in PORT_LOOKUP else port
new_val[proto].append(port)
return new_val

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

@@ -12,8 +12,8 @@ from primaite.game.agent.observations.firewall_observation import FirewallObserv
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
from primaite.simulator.network.transmission.transport_layer import PORT_LOOKUP
from primaite.utils.validators import PROTOCOL_LOOKUP
_LOGGER = getLogger(__name__)
@@ -76,11 +76,11 @@ class NodesObservation(AbstractObservation, identifier="NODES"):
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
proto = PROTOCOL_LOOKUP[proto] if proto in PROTOCOL_LOOKUP 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
port = PORT_LOOKUP[port] if port in PORT_LOOKUP else port
new_val[proto].append(port)
return new_val

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

@@ -27,8 +27,7 @@ from primaite.simulator.network.hardware.nodes.network.router import Router
from primaite.simulator.network.hardware.nodes.network.switch import Switch
from primaite.simulator.network.hardware.nodes.network.wireless_router import WirelessRouter
from primaite.simulator.network.nmne import NMNEConfig
from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.network.transmission.transport_layer import PORT_LOOKUP
from primaite.simulator.sim_container import Simulation
from primaite.simulator.system.applications.application import Application
from primaite.simulator.system.applications.database_client import DatabaseClient # noqa: F401
@@ -51,6 +50,7 @@ from primaite.simulator.system.services.service import Service
from primaite.simulator.system.services.terminal.terminal import Terminal
from primaite.simulator.system.services.web_server.web_server import WebServer
from primaite.simulator.system.software import Software
from primaite.utils.validators import PROTOCOL_LOOKUP
_LOGGER = getLogger(__name__)
@@ -97,8 +97,8 @@ class PrimaiteGameOptions(BaseModel):
: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]
if port_val in PORT_LOOKUP:
vals[i] = PORT_LOOKUP[port_val]
return vals
@field_validator("protocols", mode="before")
@@ -110,8 +110,8 @@ class PrimaiteGameOptions(BaseModel):
: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]
if proto_val in PROTOCOL_LOOKUP:
vals[i] = PROTOCOL_LOOKUP[proto_val]
return vals
@@ -381,7 +381,7 @@ class PrimaiteGame:
if isinstance(port_id, int):
port = port_id
elif isinstance(port_id, str):
port = Port[port_id]
port = PORT_LOOKUP[port_id]
if port:
listen_on_ports.append(port)
software.listen_on_ports = set(listen_on_ports)
@@ -496,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_LOOKUP[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")),
@@ -509,8 +509,10 @@ 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=PROTOCOL_LOOKUP[
(opt.get("masquerade_protocol", PROTOCOL_LOOKUP["TCP"]))
],
masquerade_port=PORT_LOOKUP[(opt.get("masquerade_port", PORT_LOOKUP["HTTP"]))],
)
if "network_interfaces" in node_cfg:
for nic_num, nic_cfg in node_cfg["network_interfaces"].items():