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():

10
src/primaite/simulator/network/creation.py
View File

@@ -6,8 +6,8 @@ from primaite.simulator.network.container import Network
from primaite.simulator.network.hardware.nodes.host.computer import Computer
from primaite.simulator.network.hardware.nodes.network.router import ACLAction, Router
from primaite.simulator.network.hardware.nodes.network.switch import Switch
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
def num_of_switches_required(num_nodes: int, max_network_interface: int = 24) -> int:
@@ -98,8 +98,10 @@ def create_office_lan(
default_gateway = IPv4Address(f"192.168.{subnet_base}.1")
router = Router(hostname=f"router_{lan_name}", start_up_duration=0)
router.power_on()
router.acl.add_rule(action=ACLAction.PERMIT, src_port=Port["ARP"], dst_port=Port["ARP"], position=22)
router.acl.add_rule(action=ACLAction.PERMIT, protocol=IPProtocol["ICMP"], position=23)
router.acl.add_rule(
action=ACLAction.PERMIT, src_port=PORT_LOOKUP["ARP"], dst_port=PORT_LOOKUP["ARP"], position=22
)
router.acl.add_rule(action=ACLAction.PERMIT, protocol=PROTOCOL_LOOKUP["ICMP"], position=23)
network.add_node(router)
router.configure_port(port=1, ip_address=default_gateway, subnet_mask="255.255.255.0")
router.enable_port(1)

23
src/primaite/simulator/network/hardware/base.py
View File

@@ -21,8 +21,7 @@ from primaite.simulator.file_system.file_system import FileSystem
from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
from primaite.simulator.network.nmne import NMNEConfig
from primaite.simulator.network.transmission.data_link_layer import Frame
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.system.applications.application import Application
from primaite.simulator.system.core.packet_capture import PacketCapture
from primaite.simulator.system.core.session_manager import SessionManager
@@ -33,7 +32,7 @@ from primaite.simulator.system.services.service import Service
from primaite.simulator.system.services.terminal.terminal import Terminal
from primaite.simulator.system.software import IOSoftware, Software
from primaite.utils.converters import convert_dict_enum_keys_to_enum_values
from primaite.utils.validators import IPV4Address
from primaite.utils.validators import IPV4Address, PROTOCOL_LOOKUP
IOSoftwareClass = TypeVar("IOSoftwareClass", bound=IOSoftware)
@@ -274,20 +273,20 @@ class NetworkInterface(SimComponent, ABC):
# Identify the protocol and port from the frame
if frame.tcp:
protocol = IPProtocol["TCP"]
protocol = PROTOCOL_LOOKUP["TCP"]
port = frame.tcp.dst_port
elif frame.udp:
protocol = IPProtocol["UDP"]
protocol = PROTOCOL_LOOKUP["UDP"]
port = frame.udp.dst_port
elif frame.icmp:
protocol = IPProtocol["ICMP"]
protocol = PROTOCOL_LOOKUP["ICMP"]
# Ensure the protocol is in the capture dict
if protocol not in self.traffic:
self.traffic[protocol] = {}
# Handle non-ICMP protocols that use ports
if protocol != IPProtocol["ICMP"]:
if protocol != PROTOCOL_LOOKUP["ICMP"]:
if port not in self.traffic[protocol]:
self.traffic[protocol][port] = {"inbound": 0, "outbound": 0}
self.traffic[protocol][port][direction] += frame.size_Mbits
@@ -843,8 +842,8 @@ class UserManager(Service):
:param password: The password for the default admin user
"""
kwargs["name"] = "UserManager"
kwargs["port"] = Port["NONE"]
kwargs["protocol"] = IPProtocol["NONE"]
kwargs["port"] = PORT_LOOKUP["NONE"]
kwargs["protocol"] = PROTOCOL_LOOKUP["NONE"]
super().__init__(**kwargs)
self.start()
@@ -1166,8 +1165,8 @@ class UserSessionManager(Service):
:param password: The password for the default admin user
"""
kwargs["name"] = "UserSessionManager"
kwargs["port"] = Port["NONE"]
kwargs["protocol"] = IPProtocol["NONE"]
kwargs["port"] = PORT_LOOKUP["NONE"]
kwargs["protocol"] = PROTOCOL_LOOKUP["NONE"]
super().__init__(**kwargs)
self.start()
@@ -1312,7 +1311,7 @@ class UserSessionManager(Service):
software_manager: SoftwareManager = self.software_manager
software_manager.send_payload_to_session_manager(
payload={"type": "user_timeout", "connection_id": session.uuid},
dest_port=Port["SSH"],
dest_port=PORT_LOOKUP["SSH"],
dest_ip_address=session.remote_ip_address,
)

41
src/primaite/simulator/network/hardware/nodes/network/firewall.py
View File

@@ -14,10 +14,9 @@ from primaite.simulator.network.hardware.nodes.network.router import (
RouterInterface,
)
from primaite.simulator.network.transmission.data_link_layer import Frame
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.system.core.sys_log import SysLog
from primaite.utils.validators import IPV4Address
from primaite.utils.validators import IPV4Address, PROTOCOL_LOOKUP
EXTERNAL_PORT_ID: Final[int] = 1
"""The Firewall port ID of the external port."""
@@ -596,9 +595,9 @@ class Firewall(Router):
for r_num, r_cfg in cfg["acl"]["internal_inbound_acl"].items():
firewall.internal_inbound_acl.add_rule(
action=ACLAction[r_cfg["action"]],
src_port=None if not (p := r_cfg.get("src_port")) else Port[p],
dst_port=None if not (p := r_cfg.get("dst_port")) else Port[p],
protocol=None if not (p := r_cfg.get("protocol")) else IPProtocol[p],
src_port=None if not (p := r_cfg.get("src_port")) else PORT_LOOKUP[p],
dst_port=None if not (p := r_cfg.get("dst_port")) else PORT_LOOKUP[p],
protocol=None if not (p := r_cfg.get("protocol")) else PROTOCOL_LOOKUP[p],
src_ip_address=r_cfg.get("src_ip"),
src_wildcard_mask=r_cfg.get("src_wildcard_mask"),
dst_ip_address=r_cfg.get("dst_ip"),
@@ -611,9 +610,9 @@ class Firewall(Router):
for r_num, r_cfg in cfg["acl"]["internal_outbound_acl"].items():
firewall.internal_outbound_acl.add_rule(
action=ACLAction[r_cfg["action"]],
src_port=None if not (p := r_cfg.get("src_port")) else Port[p],
dst_port=None if not (p := r_cfg.get("dst_port")) else Port[p],
protocol=None if not (p := r_cfg.get("protocol")) else IPProtocol[p],
src_port=None if not (p := r_cfg.get("src_port")) else PORT_LOOKUP[p],
dst_port=None if not (p := r_cfg.get("dst_port")) else PORT_LOOKUP[p],
protocol=None if not (p := r_cfg.get("protocol")) else PROTOCOL_LOOKUP[p],
src_ip_address=r_cfg.get("src_ip"),
src_wildcard_mask=r_cfg.get("src_wildcard_mask"),
dst_ip_address=r_cfg.get("dst_ip"),
@@ -626,9 +625,9 @@ class Firewall(Router):
for r_num, r_cfg in cfg["acl"]["dmz_inbound_acl"].items():
firewall.dmz_inbound_acl.add_rule(
action=ACLAction[r_cfg["action"]],
src_port=None if not (p := r_cfg.get("src_port")) else Port[p],
dst_port=None if not (p := r_cfg.get("dst_port")) else Port[p],
protocol=None if not (p := r_cfg.get("protocol")) else IPProtocol[p],
src_port=None if not (p := r_cfg.get("src_port")) else PORT_LOOKUP[p],
dst_port=None if not (p := r_cfg.get("dst_port")) else PORT_LOOKUP[p],
protocol=None if not (p := r_cfg.get("protocol")) else PROTOCOL_LOOKUP[p],
src_ip_address=r_cfg.get("src_ip"),
src_wildcard_mask=r_cfg.get("src_wildcard_mask"),
dst_ip_address=r_cfg.get("dst_ip"),
@@ -641,9 +640,9 @@ class Firewall(Router):
for r_num, r_cfg in cfg["acl"]["dmz_outbound_acl"].items():
firewall.dmz_outbound_acl.add_rule(
action=ACLAction[r_cfg["action"]],
src_port=None if not (p := r_cfg.get("src_port")) else Port[p],
dst_port=None if not (p := r_cfg.get("dst_port")) else Port[p],
protocol=None if not (p := r_cfg.get("protocol")) else IPProtocol[p],
src_port=None if not (p := r_cfg.get("src_port")) else PORT_LOOKUP[p],
dst_port=None if not (p := r_cfg.get("dst_port")) else PORT_LOOKUP[p],
protocol=None if not (p := r_cfg.get("protocol")) else PROTOCOL_LOOKUP[p],
src_ip_address=r_cfg.get("src_ip"),
src_wildcard_mask=r_cfg.get("src_wildcard_mask"),
dst_ip_address=r_cfg.get("dst_ip"),
@@ -656,9 +655,9 @@ class Firewall(Router):
for r_num, r_cfg in cfg["acl"]["external_inbound_acl"].items():
firewall.external_inbound_acl.add_rule(
action=ACLAction[r_cfg["action"]],
src_port=None if not (p := r_cfg.get("src_port")) else Port[p],
dst_port=None if not (p := r_cfg.get("dst_port")) else Port[p],
protocol=None if not (p := r_cfg.get("protocol")) else IPProtocol[p],
src_port=None if not (p := r_cfg.get("src_port")) else PORT_LOOKUP[p],
dst_port=None if not (p := r_cfg.get("dst_port")) else PORT_LOOKUP[p],
protocol=None if not (p := r_cfg.get("protocol")) else PROTOCOL_LOOKUP[p],
src_ip_address=r_cfg.get("src_ip"),
src_wildcard_mask=r_cfg.get("src_wildcard_mask"),
dst_ip_address=r_cfg.get("dst_ip"),
@@ -671,9 +670,9 @@ class Firewall(Router):
for r_num, r_cfg in cfg["acl"]["external_outbound_acl"].items():
firewall.external_outbound_acl.add_rule(
action=ACLAction[r_cfg["action"]],
src_port=None if not (p := r_cfg.get("src_port")) else Port[p],
dst_port=None if not (p := r_cfg.get("dst_port")) else Port[p],
protocol=None if not (p := r_cfg.get("protocol")) else IPProtocol[p],
src_port=None if not (p := r_cfg.get("src_port")) else PORT_LOOKUP[p],
dst_port=None if not (p := r_cfg.get("dst_port")) else PORT_LOOKUP[p],
protocol=None if not (p := r_cfg.get("protocol")) else PROTOCOL_LOOKUP[p],
src_ip_address=r_cfg.get("src_ip"),
src_wildcard_mask=r_cfg.get("src_wildcard_mask"),
dst_ip_address=r_cfg.get("dst_ip"),

25
src/primaite/simulator/network/hardware/nodes/network/router.py
View File

@@ -17,15 +17,14 @@ from primaite.simulator.network.hardware.nodes.network.network_node import Netwo
from primaite.simulator.network.protocols.arp import ARPPacket
from primaite.simulator.network.protocols.icmp import ICMPPacket, ICMPType
from primaite.simulator.network.transmission.data_link_layer import Frame
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.system.applications.nmap import NMAP
from primaite.simulator.system.core.session_manager import SessionManager
from primaite.simulator.system.core.sys_log import SysLog
from primaite.simulator.system.services.arp.arp import ARP
from primaite.simulator.system.services.icmp.icmp import ICMP
from primaite.simulator.system.services.terminal.terminal import Terminal
from primaite.utils.validators import IPV4Address
from primaite.utils.validators import IPV4Address, PROTOCOL_LOOKUP
@validate_call()
@@ -134,14 +133,14 @@ class ACLRule(SimComponent):
def protocol_valid(cls, val: Optional[str]) -> Optional[str]:
"""Assert that the protocol for the rule is predefined in the IPProtocol lookup."""
if val is not None:
assert val in IPProtocol.values(), f"Cannot create ACL rule with invalid protocol {val}"
assert val in PROTOCOL_LOOKUP.values(), f"Cannot create ACL rule with invalid protocol {val}"
return val
@field_validator("src_port", "dst_port", mode="before")
def ports_valid(cls, val: Optional[int]) -> Optional[int]:
"""Assert that the port for the rule is predefined in the Port lookup."""
if val is not None:
assert val in Port.values(), f"Cannot create ACL rule with invalid port {val}"
assert val in PORT_LOOKUP.values(), f"Cannot create ACL rule with invalid port {val}"
return val
def __str__(self) -> str:
@@ -1271,8 +1270,10 @@ class Router(NetworkNode):
Initializes the router's ACL (Access Control List) with default rules, permitting essential protocols like ARP
and ICMP, which are necessary for basic network operations and diagnostics.
"""
self.acl.add_rule(action=ACLAction.PERMIT, src_port=Port["ARP"], dst_port=Port["ARP"], position=22)
self.acl.add_rule(action=ACLAction.PERMIT, protocol=IPProtocol["ICMP"], position=23)
self.acl.add_rule(
action=ACLAction.PERMIT, src_port=PORT_LOOKUP["ARP"], dst_port=PORT_LOOKUP["ARP"], position=22
)
self.acl.add_rule(action=ACLAction.PERMIT, protocol=PROTOCOL_LOOKUP["ICMP"], position=23)
def setup_for_episode(self, episode: int):
"""
@@ -1371,9 +1372,9 @@ class Router(NetworkNode):
"""
dst_ip_address = frame.ip.dst_ip_address
dst_port = None
if frame.ip.protocol == IPProtocol["TCP"]:
if frame.ip.protocol == PROTOCOL_LOOKUP["TCP"]:
dst_port = frame.tcp.dst_port
elif frame.ip.protocol == IPProtocol["UDP"]:
elif frame.ip.protocol == PROTOCOL_LOOKUP["UDP"]:
dst_port = frame.udp.dst_port
if self.ip_is_router_interface(dst_ip_address) and (
@@ -1646,9 +1647,9 @@ class Router(NetworkNode):
for r_num, r_cfg in cfg["acl"].items():
router.acl.add_rule(
action=ACLAction[r_cfg["action"]],
src_port=None if not (p := r_cfg.get("src_port")) else Port[p],
dst_port=None if not (p := r_cfg.get("dst_port")) else Port[p],
protocol=None if not (p := r_cfg.get("protocol")) else IPProtocol[p],
src_port=None if not (p := r_cfg.get("src_port")) else PORT_LOOKUP[p],
dst_port=None if not (p := r_cfg.get("dst_port")) else PORT_LOOKUP[p],
protocol=None if not (p := r_cfg.get("protocol")) else PROTOCOL_LOOKUP[p],
src_ip_address=r_cfg.get("src_ip"),
src_wildcard_mask=r_cfg.get("src_wildcard_mask"),
dst_ip_address=r_cfg.get("dst_ip"),

11
src/primaite/simulator/network/hardware/nodes/network/wireless_router.py
View File

@@ -8,9 +8,8 @@ from primaite.simulator.network.airspace import AirSpace, IPWirelessNetworkInter
from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
from primaite.simulator.network.hardware.nodes.network.router import ACLAction, Router, RouterInterface
from primaite.simulator.network.transmission.data_link_layer import Frame
from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.utils.validators import IPV4Address
from primaite.simulator.network.transmission.transport_layer import PORT_LOOKUP
from primaite.utils.validators import IPV4Address, PROTOCOL_LOOKUP
class WirelessAccessPoint(IPWirelessNetworkInterface):
@@ -271,9 +270,9 @@ class WirelessRouter(Router):
for r_num, r_cfg in cfg["acl"].items():
router.acl.add_rule(
action=ACLAction[r_cfg["action"]],
src_port=None if not (p := r_cfg.get("src_port")) else Port[p],
dst_port=None if not (p := r_cfg.get("dst_port")) else Port[p],
protocol=None if not (p := r_cfg.get("protocol")) else IPProtocol[p],
src_port=None if not (p := r_cfg.get("src_port")) else PORT_LOOKUP[p],
dst_port=None if not (p := r_cfg.get("dst_port")) else PORT_LOOKUP[p],
protocol=None if not (p := r_cfg.get("protocol")) else PROTOCOL_LOOKUP[p],
src_ip_address=r_cfg.get("src_ip"),
dst_ip_address=r_cfg.get("dst_ip"),
src_wildcard_mask=r_cfg.get("src_wildcard_mask"),

29
src/primaite/simulator/network/networks.py
View File

@@ -12,14 +12,14 @@ from primaite.simulator.network.hardware.nodes.host.host_node import NIC
from primaite.simulator.network.hardware.nodes.host.server import Server
from primaite.simulator.network.hardware.nodes.network.router import ACLAction, Router
from primaite.simulator.network.hardware.nodes.network.switch import Switch
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.system.applications.database_client import DatabaseClient
from primaite.simulator.system.applications.red_applications.data_manipulation_bot import DataManipulationBot
from primaite.simulator.system.services.database.database_service import DatabaseService
from primaite.simulator.system.services.dns.dns_server import DNSServer
from primaite.simulator.system.services.ftp.ftp_server import FTPServer
from primaite.simulator.system.services.web_server.web_server import WebServer
from primaite.utils.validators import PROTOCOL_LOOKUP
_LOGGER = getLogger(__name__)
@@ -79,9 +79,11 @@ def client_server_routed() -> Network:
server_1.power_on()
network.connect(endpoint_b=server_1.network_interface[1], endpoint_a=switch_1.network_interface[1])
router_1.acl.add_rule(action=ACLAction.PERMIT, src_port=Port["ARP"], dst_port=Port["ARP"], position=22)
router_1.acl.add_rule(
action=ACLAction.PERMIT, src_port=PORT_LOOKUP["ARP"], dst_port=PORT_LOOKUP["ARP"], position=22
)
router_1.acl.add_rule(action=ACLAction.PERMIT, protocol=IPProtocol["ICMP"], position=23)
router_1.acl.add_rule(action=ACLAction.PERMIT, protocol=PROTOCOL_LOOKUP["ICMP"], position=23)
return network
@@ -271,23 +273,30 @@ def arcd_uc2_network() -> Network:
security_suite.connect_nic(NIC(ip_address="192.168.10.110", subnet_mask="255.255.255.0"))
network.connect(endpoint_b=security_suite.network_interface[2], endpoint_a=switch_2.network_interface[7])
router_1.acl.add_rule(action=ACLAction.PERMIT, src_port=Port["ARP"], dst_port=Port["ARP"], position=22)
router_1.acl.add_rule(
action=ACLAction.PERMIT, src_port=PORT_LOOKUP["ARP"], dst_port=PORT_LOOKUP["ARP"], position=22
)
router_1.acl.add_rule(action=ACLAction.PERMIT, protocol=IPProtocol["ICMP"], position=23)
router_1.acl.add_rule(action=ACLAction.PERMIT, protocol=PROTOCOL_LOOKUP["ICMP"], position=23)
# Allow PostgreSQL requests
router_1.acl.add_rule(
action=ACLAction.PERMIT, src_port=Port["POSTGRES_SERVER"], dst_port=Port["POSTGRES_SERVER"], position=0
action=ACLAction.PERMIT,
src_port=PORT_LOOKUP["POSTGRES_SERVER"],
dst_port=PORT_LOOKUP["POSTGRES_SERVER"],
position=0,
)
# Allow DNS requests
router_1.acl.add_rule(action=ACLAction.PERMIT, src_port=Port["DNS"], dst_port=Port["DNS"], position=1)
router_1.acl.add_rule(action=ACLAction.PERMIT, src_port=PORT_LOOKUP["DNS"], dst_port=PORT_LOOKUP["DNS"], position=1)
# Allow FTP requests
router_1.acl.add_rule(action=ACLAction.PERMIT, src_port=Port["FTP"], dst_port=Port["FTP"], position=2)
router_1.acl.add_rule(action=ACLAction.PERMIT, src_port=PORT_LOOKUP["FTP"], dst_port=PORT_LOOKUP["FTP"], position=2)
# Open port 80 for web server
router_1.acl.add_rule(action=ACLAction.PERMIT, src_port=Port["HTTP"], dst_port=Port["HTTP"], position=3)
router_1.acl.add_rule(
action=ACLAction.PERMIT, src_port=PORT_LOOKUP["HTTP"], dst_port=PORT_LOOKUP["HTTP"], position=3
)
return network

13
src/primaite/simulator/network/transmission/data_link_layer.py
View File

@@ -7,10 +7,11 @@ from pydantic import BaseModel
from primaite import getLogger
from primaite.simulator.network.protocols.icmp import ICMPPacket
from primaite.simulator.network.protocols.packet import DataPacket
from primaite.simulator.network.transmission.network_layer import IPPacket, IPProtocol
from primaite.simulator.network.transmission.network_layer import IPPacket
from primaite.simulator.network.transmission.primaite_layer import PrimaiteHeader
from primaite.simulator.network.transmission.transport_layer import Port, TCPHeader, UDPHeader
from primaite.simulator.network.transmission.transport_layer import PORT_LOOKUP, TCPHeader, UDPHeader
from primaite.simulator.network.utils import convert_bytes_to_megabits
from primaite.utils.validators import PROTOCOL_LOOKUP
_LOGGER = getLogger(__name__)
@@ -70,15 +71,15 @@ class Frame(BaseModel):
msg = "Network Frame cannot have both a TCP header and a UDP header"
_LOGGER.error(msg)
raise ValueError(msg)
if kwargs["ip"].protocol == IPProtocol["TCP"] and not kwargs.get("tcp"):
if kwargs["ip"].protocol == PROTOCOL_LOOKUP["TCP"] and not kwargs.get("tcp"):
msg = "Cannot build a Frame using the TCP IP Protocol without a TCPHeader"
_LOGGER.error(msg)
raise ValueError(msg)
if kwargs["ip"].protocol == IPProtocol["UDP"] and not kwargs.get("udp"):
if kwargs["ip"].protocol == PROTOCOL_LOOKUP["UDP"] and not kwargs.get("udp"):
msg = "Cannot build a Frame using the UDP IP Protocol without a UDPHeader"
_LOGGER.error(msg)
raise ValueError(msg)
if kwargs["ip"].protocol == IPProtocol["ICMP"] and not kwargs.get("icmp"):
if kwargs["ip"].protocol == PROTOCOL_LOOKUP["ICMP"] and not kwargs.get("icmp"):
msg = "Cannot build a Frame using the ICMP IP Protocol without a ICMPPacket"
_LOGGER.error(msg)
raise ValueError(msg)
@@ -165,7 +166,7 @@ class Frame(BaseModel):
:return: True if the Frame is an ARP packet, otherwise False.
"""
return self.udp.dst_port == Port["ARP"]
return self.udp.dst_port == PORT_LOOKUP["ARP"]
@property
def is_icmp(self) -> bool:

32
src/primaite/simulator/network/transmission/network_layer.py
View File

@@ -4,39 +4,11 @@ from enum import Enum
from pydantic import BaseModel
from primaite import getLogger
from primaite.utils.validators import IPV4Address
from primaite.utils.validators import IPProtocol, IPV4Address, PROTOCOL_LOOKUP
_LOGGER = getLogger(__name__)
IPProtocol: dict[str, str] = dict(
NONE="none",
TCP="tcp",
UDP="udp",
ICMP="icmp",
)
# class IPProtocol(Enum):
# """
# Enum representing transport layer protocols in IP header.
# .. _List of IPProtocols:
# """
# NONE = "none"
# """Placeholder for a non-protocol."""
# TCP = "tcp"
# """Transmission Control Protocol."""
# UDP = "udp"
# """User Datagram Protocol."""
# ICMP = "icmp"
# """Internet Control Message Protocol."""
# def model_dump(self) -> str:
# """Return as JSON-serialisable string."""
# return self.name
class Precedence(Enum):
"""
Enum representing the Precedence levels in Quality of Service (QoS) for IP packets.
@@ -98,7 +70,7 @@ class IPPacket(BaseModel):
"Source IP address."
dst_ip_address: IPV4Address
"Destination IP address."
protocol: str = IPProtocol["TCP"]
protocol: IPProtocol = PROTOCOL_LOOKUP["TCP"]
"IPProtocol."
ttl: int = 64
"Time to Live (TTL) for the packet."

77
src/primaite/simulator/network/transmission/transport_layer.py
View File

@@ -4,7 +4,7 @@ from typing import List
from pydantic import BaseModel
Port: dict[str, int] = dict(
PORT_LOOKUP: dict[str, int] = dict(
UNUSED=-1,
NONE=0,
WOL=9,
@@ -36,81 +36,6 @@ Port: dict[str, int] = dict(
POSTGRES_SERVER=5432,
)
# class Port():
# def __getattr__()
# class Port(Enum):
# """
# Enumeration of common known TCP/UDP ports used by protocols for operation of network applications.
# .. _List of Ports:
# """
# UNUSED = -1
# "An unused port stub."
# NONE = 0
# "Place holder for a non-port."
# WOL = 9
# "Wake-on-Lan (WOL) - Used to turn or awaken a computer from sleep mode by a network message."
# FTP_DATA = 20
# "File Transfer [Default Data]"
# FTP = 21
# "File Transfer Protocol (FTP) - FTP control (command)"
# SSH = 22
# "Secure Shell (SSH) - Used for secure remote access and command execution."
# SMTP = 25
# "Simple Mail Transfer Protocol (SMTP) - Used for email delivery between servers."
# DNS = 53
# "Domain Name System (DNS) - Used for translating domain names to IP addresses."
# HTTP = 80
# "HyperText Transfer Protocol (HTTP) - Used for web traffic."
# POP3 = 110
# "Post Office Protocol version 3 (POP3) - Used for retrieving emails from a mail server."
# SFTP = 115
# "Secure File Transfer Protocol (SFTP) - Used for secure file transfer over SSH."
# NTP = 123
# "Network Time Protocol (NTP) - Used for clock synchronization between computer systems."
# IMAP = 143
# "Internet Message Access Protocol (IMAP) - Used for retrieving emails from a mail server."
# SNMP = 161
# "Simple Network Management Protocol (SNMP) - Used for network device management."
# SNMP_TRAP = 162
# "SNMP Trap - Used for sending SNMP notifications (traps) to a network management system."
# ARP = 219
# "Address resolution Protocol - Used to connect a MAC address to an IP address."
# LDAP = 389
# "Lightweight Directory Access Protocol (LDAP) - Used for accessing and modifying directory information."
# HTTPS = 443
# "HyperText Transfer Protocol Secure (HTTPS) - Used for secure web traffic."
# SMB = 445
# "Server Message Block (SMB) - Used for file sharing and printer sharing in Windows environments."
# IPP = 631
# "Internet Printing Protocol (IPP) - Used for printing over the internet or an intranet."
# SQL_SERVER = 1433
# "Microsoft SQL Server Database Engine - Used for communication with the SQL Server."
# MYSQL = 3306
# "MySQL Database Server - Used for MySQL database communication."
# RDP = 3389
# "Remote Desktop Protocol (RDP) - Used for remote desktop access to Windows machines."
# RTP = 5004
# "Real-time Transport Protocol (RTP) - Used for transmitting real-time media, e.g., audio and video."
# RTP_ALT = 5005
# "Alternative port for RTP (RTP_ALT) - Used in some configurations for transmitting real-time media."
# DNS_ALT = 5353
# "Alternative port for DNS (DNS_ALT) - Used in some configurations for DNS service."
# HTTP_ALT = 8080
# "Alternative port for HTTP (HTTP_ALT) - Often used as an alternative HTTP port for web applications."
# HTTPS_ALT = 8443
# "Alternative port for HTTPS (HTTPS_ALT) - Used in some configurations for secure web traffic."
# POSTGRES_SERVER = 5432
# "Postgres SQL Server."
# def model_dump(self) -> str:
# """Return a json-serialisable string."""
# return self.name
class UDPHeader(BaseModel):
"""

9
src/primaite/simulator/system/applications/database_client.py
View File

@@ -11,11 +11,10 @@ from pydantic import BaseModel
from primaite.interface.request import RequestFormat, RequestResponse
from primaite.simulator.core import RequestManager, RequestType
from primaite.simulator.network.hardware.nodes.host.host_node import HostNode
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.system.applications.application import Application
from primaite.simulator.system.core.software_manager import SoftwareManager
from primaite.utils.validators import IPV4Address
from primaite.utils.validators import IPV4Address, PROTOCOL_LOOKUP
class DatabaseClientConnection(BaseModel):
@@ -90,8 +89,8 @@ class DatabaseClient(Application, identifier="DatabaseClient"):
def __init__(self, **kwargs):
kwargs["name"] = "DatabaseClient"
kwargs["port"] = Port["POSTGRES_SERVER"]
kwargs["protocol"] = IPProtocol["TCP"]
kwargs["port"] = PORT_LOOKUP["POSTGRES_SERVER"]
kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"]
super().__init__(**kwargs)
def _init_request_manager(self) -> RequestManager:

13
src/primaite/simulator/system/applications/nmap.py
View File

@@ -7,10 +7,9 @@ from pydantic import validate_call
from primaite.interface.request import RequestResponse
from primaite.simulator.core import RequestManager, RequestType, SimComponent
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.system.applications.application import Application
from primaite.utils.validators import IPV4Address
from primaite.utils.validators import IPV4Address, PROTOCOL_LOOKUP
class PortScanPayload(SimComponent):
@@ -64,8 +63,8 @@ class NMAP(Application, identifier="NMAP"):
def __init__(self, **kwargs):
kwargs["name"] = "NMAP"
kwargs["port"] = Port["NONE"]
kwargs["protocol"] = IPProtocol["NONE"]
kwargs["port"] = PORT_LOOKUP["NONE"]
kwargs["protocol"] = PROTOCOL_LOOKUP["NONE"]
super().__init__(**kwargs)
def _can_perform_network_action(self) -> bool:
@@ -348,12 +347,12 @@ class NMAP(Application, identifier="NMAP"):
if isinstance(target_port, int):
target_port = [target_port]
elif target_port is None:
target_port = [port for port in Port if port not in {Port["NONE"], Port["UNUSED"]}]
target_port = [port for port in PORT_LOOKUP if port not in {PORT_LOOKUP["NONE"], PORT_LOOKUP["UNUSED"]}]
if isinstance(target_protocol, str):
target_protocol = [target_protocol]
elif target_protocol is None:
target_protocol = [IPProtocol["TCP"], IPProtocol["UDP"]]
target_protocol = [PROTOCOL_LOOKUP["TCP"], PROTOCOL_LOOKUP["UDP"]]
scan_type = self._determine_port_scan_type(list(ip_addresses), target_port)
active_ports = {}

18
src/primaite/simulator/system/applications/red_applications/c2/abstract_c2.py
View File

@@ -9,14 +9,14 @@ from pydantic import BaseModel, Field, validate_call
from primaite.interface.request import RequestResponse
from primaite.simulator.file_system.file_system import FileSystem, Folder
from primaite.simulator.network.protocols.masquerade import C2Packet
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.system.applications.application import Application, ApplicationOperatingState
from primaite.simulator.system.core.session_manager import Session
from primaite.simulator.system.services.ftp.ftp_client import FTPClient
from primaite.simulator.system.services.ftp.ftp_server import FTPServer
from primaite.simulator.system.services.service import ServiceOperatingState
from primaite.simulator.system.software import SoftwareHealthState
from primaite.utils.validators import PROTOCOL_LOOKUP
class C2Command(Enum):
@@ -81,10 +81,10 @@ class AbstractC2(Application, identifier="AbstractC2"):
keep_alive_frequency: int = Field(default=5, ge=1)
"""The frequency at which ``Keep Alive`` packets are sent to the C2 Server from the C2 Beacon."""
masquerade_protocol: str = Field(default=IPProtocol["TCP"])
masquerade_protocol: str = Field(default=PROTOCOL_LOOKUP["TCP"])
"""The currently chosen protocol that the C2 traffic is masquerading as. Defaults as TCP."""
masquerade_port: int = Field(default=Port["HTTP"])
masquerade_port: int = Field(default=PORT_LOOKUP["HTTP"])
"""The currently chosen port that the C2 traffic is masquerading as. Defaults at HTTP."""
c2_config: _C2Opts = _C2Opts()
@@ -142,9 +142,9 @@ class AbstractC2(Application, identifier="AbstractC2"):
def __init__(self, **kwargs):
"""Initialise the C2 applications to by default listen for HTTP traffic."""
kwargs["listen_on_ports"] = {Port["HTTP"], Port["FTP"], Port["DNS"]}
kwargs["port"] = Port["NONE"]
kwargs["protocol"] = IPProtocol["TCP"]
kwargs["listen_on_ports"] = {PORT_LOOKUP["HTTP"], PORT_LOOKUP["FTP"], PORT_LOOKUP["DNS"]}
kwargs["port"] = PORT_LOOKUP["NONE"]
kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"]
super().__init__(**kwargs)
@property
@@ -410,8 +410,8 @@ class AbstractC2(Application, identifier="AbstractC2"):
self.keep_alive_inactivity = 0
self.keep_alive_frequency = 5
self.c2_remote_connection = None
self.c2_config.masquerade_port = Port["HTTP"]
self.c2_config.masquerade_protocol = IPProtocol["TCP"]
self.c2_config.masquerade_port = PORT_LOOKUP["HTTP"]
self.c2_config.masquerade_protocol = PROTOCOL_LOOKUP["TCP"]
@abstractmethod
def _confirm_remote_connection(self, timestep: int) -> bool:

12
src/primaite/simulator/system/applications/red_applications/c2/c2_beacon.py
View File

@@ -8,12 +8,12 @@ from pydantic import validate_call
from primaite.interface.request import RequestFormat, RequestResponse
from primaite.simulator.core import RequestManager, RequestType
from primaite.simulator.network.protocols.masquerade import C2Packet
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.system.applications.red_applications.c2 import ExfilOpts, RansomwareOpts, TerminalOpts
from primaite.simulator.system.applications.red_applications.c2.abstract_c2 import AbstractC2, C2Command, C2Payload
from primaite.simulator.system.applications.red_applications.ransomware_script import RansomwareScript
from primaite.simulator.system.services.terminal.terminal import Terminal, TerminalClientConnection
from primaite.utils.validators import PROTOCOL_LOOKUP
class C2Beacon(AbstractC2, identifier="C2Beacon"):
@@ -111,8 +111,8 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"):
self.configure(
c2_server_ip_address=c2_remote_ip,
keep_alive_frequency=frequency,
masquerade_protocol=IPProtocol[protocol],
masquerade_port=Port[port],
masquerade_protocol=PROTOCOL_LOOKUP[protocol],
masquerade_port=PORT_LOOKUP[port],
)
)
@@ -129,8 +129,8 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"):
self,
c2_server_ip_address: IPv4Address = None,
keep_alive_frequency: int = 5,
masquerade_protocol: str = IPProtocol["TCP"],
masquerade_port: int = Port["HTTP"],
masquerade_protocol: str = PROTOCOL_LOOKUP["TCP"],
masquerade_port: int = PORT_LOOKUP["HTTP"],
) -> bool:
"""
Configures the C2 beacon to communicate with the C2 server.

8
src/primaite/simulator/system/applications/red_applications/data_manipulation_bot.py
View File

@@ -7,10 +7,10 @@ from primaite import getLogger
from primaite.game.science import simulate_trial
from primaite.interface.request import RequestResponse
from primaite.simulator.core import RequestManager, RequestType
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.system.applications.application import Application
from primaite.simulator.system.applications.database_client import DatabaseClient, DatabaseClientConnection
from primaite.utils.validators import PROTOCOL_LOOKUP
_LOGGER = getLogger(__name__)
@@ -50,8 +50,8 @@ class DataManipulationBot(Application, identifier="DataManipulationBot"):
def __init__(self, **kwargs):
kwargs["name"] = "DataManipulationBot"
kwargs["port"] = Port["NONE"]
kwargs["protocol"] = IPProtocol["NONE"]
kwargs["port"] = PORT_LOOKUP["NONE"]
kwargs["protocol"] = PROTOCOL_LOOKUP["NONE"]
super().__init__(**kwargs)
self._db_connection: Optional[DatabaseClientConnection] = None

6
src/primaite/simulator/system/applications/red_applications/dos_bot.py
View File

@@ -7,7 +7,7 @@ from primaite import getLogger
from primaite.game.science import simulate_trial
from primaite.interface.request import RequestFormat, RequestResponse
from primaite.simulator.core import RequestManager, RequestType
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.network.transmission.transport_layer import PORT_LOOKUP
from primaite.simulator.system.applications.database_client import DatabaseClient
_LOGGER = getLogger(__name__)
@@ -85,7 +85,7 @@ class DoSBot(DatabaseClient, identifier="DoSBot"):
if "target_ip_address" in request[-1]:
request[-1]["target_ip_address"] = IPv4Address(request[-1]["target_ip_address"])
if "target_port" in request[-1]:
request[-1]["target_port"] = Port[request[-1]["target_port"]]
request[-1]["target_port"] = PORT_LOOKUP[request[-1]["target_port"]]
return RequestResponse.from_bool(self.configure(**request[-1]))
rm.add_request("configure", request_type=RequestType(func=_configure))
@@ -94,7 +94,7 @@ class DoSBot(DatabaseClient, identifier="DoSBot"):
def configure(
self,
target_ip_address: IPv4Address,
target_port: Optional[int] = Port["POSTGRES_SERVER"],
target_port: Optional[int] = PORT_LOOKUP["POSTGRES_SERVER"],
payload: Optional[str] = None,
repeat: bool = False,
port_scan_p_of_success: float = 0.1,

8
src/primaite/simulator/system/applications/red_applications/ransomware_script.py
View File

@@ -6,10 +6,10 @@ from prettytable import MARKDOWN, PrettyTable
from primaite.interface.request import RequestFormat, RequestResponse
from primaite.simulator.core import RequestManager, RequestType
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.system.applications.application import Application
from primaite.simulator.system.applications.database_client import DatabaseClient, DatabaseClientConnection
from primaite.utils.validators import PROTOCOL_LOOKUP
class RansomwareScript(Application, identifier="RansomwareScript"):
@@ -27,8 +27,8 @@ class RansomwareScript(Application, identifier="RansomwareScript"):
def __init__(self, **kwargs):
kwargs["name"] = "RansomwareScript"
kwargs["port"] = Port["NONE"]
kwargs["protocol"] = IPProtocol["NONE"]
kwargs["port"] = PORT_LOOKUP["NONE"]
kwargs["protocol"] = PROTOCOL_LOOKUP["NONE"]
super().__init__(**kwargs)
self._db_connection: Optional[DatabaseClientConnection] = None

12
src/primaite/simulator/system/applications/web_browser.py
View File

@@ -15,10 +15,10 @@ from primaite.simulator.network.protocols.http import (
HttpResponsePacket,
HttpStatusCode,
)
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.system.applications.application import Application
from primaite.simulator.system.services.dns.dns_client import DNSClient
from primaite.utils.validators import PROTOCOL_LOOKUP
_LOGGER = getLogger(__name__)
@@ -43,10 +43,10 @@ class WebBrowser(Application, identifier="WebBrowser"):
def __init__(self, **kwargs):
kwargs["name"] = "WebBrowser"
kwargs["protocol"] = IPProtocol["TCP"]
kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"]
# default for web is port 80
if kwargs.get("port") is None:
kwargs["port"] = Port["HTTP"]
kwargs["port"] = PORT_LOOKUP["HTTP"]
super().__init__(**kwargs)
self.run()
@@ -126,7 +126,7 @@ class WebBrowser(Application, identifier="WebBrowser"):
if self.send(
payload=payload,
dest_ip_address=self.domain_name_ip_address,
dest_port=parsed_url.port if parsed_url.port else Port["HTTP"],
dest_port=parsed_url.port if parsed_url.port else PORT_LOOKUP["HTTP"],
):
self.sys_log.info(
f"{self.name}: Received HTTP {payload.request_method.name} "
@@ -154,7 +154,7 @@ class WebBrowser(Application, identifier="WebBrowser"):
self,
payload: HttpRequestPacket,
dest_ip_address: Optional[IPv4Address] = None,
dest_port: Optional[int] = Port["HTTP"],
dest_port: Optional[int] = PORT_LOOKUP["HTTP"],
session_id: Optional[str] = None,
**kwargs,
) -> bool:

19
src/primaite/simulator/system/core/session_manager.py
View File

@@ -10,8 +10,9 @@ from primaite.simulator.core import SimComponent
from primaite.simulator.network.protocols.arp import ARPPacket
from primaite.simulator.network.protocols.icmp import ICMPPacket
from primaite.simulator.network.transmission.data_link_layer import EthernetHeader, Frame
from primaite.simulator.network.transmission.network_layer import IPPacket, IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port, TCPHeader, UDPHeader
from primaite.simulator.network.transmission.network_layer import IPPacket
from primaite.simulator.network.transmission.transport_layer import PORT_LOOKUP, TCPHeader, UDPHeader
from primaite.utils.validators import PROTOCOL_LOOKUP
if TYPE_CHECKING:
from primaite.simulator.network.hardware.base import NetworkInterface
@@ -117,7 +118,7 @@ class SessionManager:
"""
protocol = frame.ip.protocol
with_ip_address = frame.ip.src_ip_address
if protocol == IPProtocol["TCP"]:
if protocol == PROTOCOL_LOOKUP["TCP"]:
if inbound_frame:
src_port = frame.tcp.src_port
dst_port = frame.tcp.dst_port
@@ -125,7 +126,7 @@ class SessionManager:
dst_port = frame.tcp.src_port
src_port = frame.tcp.dst_port
with_ip_address = frame.ip.dst_ip_address
elif protocol == IPProtocol["UDP"]:
elif protocol == PROTOCOL_LOOKUP["UDP"]:
if inbound_frame:
src_port = frame.udp.src_port
dst_port = frame.udp.dst_port
@@ -260,7 +261,7 @@ class SessionManager:
src_port: Optional[int] = None,
dst_port: Optional[int] = None,
session_id: Optional[str] = None,
ip_protocol: str = IPProtocol["TCP"],
ip_protocol: str = PROTOCOL_LOOKUP["TCP"],
icmp_packet: Optional[ICMPPacket] = None,
) -> Union[Any, None]:
"""
@@ -284,7 +285,7 @@ class SessionManager:
dst_mac_address = payload.target_mac_addr
outbound_network_interface = self.resolve_outbound_network_interface(payload.target_ip_address)
is_broadcast = payload.request
ip_protocol = IPProtocol["UDP"]
ip_protocol = PROTOCOL_LOOKUP["UDP"]
else:
vals = self.resolve_outbound_transmission_details(
dst_ip_address=dst_ip_address,
@@ -316,12 +317,12 @@ class SessionManager:
tcp_header = None
udp_header = None
if ip_protocol == IPProtocol["TCP"]:
if ip_protocol == PROTOCOL_LOOKUP["TCP"]:
tcp_header = TCPHeader(
src_port=dst_port,
dst_port=dst_port,
)
elif ip_protocol == IPProtocol["UDP"]:
elif ip_protocol == PROTOCOL_LOOKUP["UDP"]:
udp_header = UDPHeader(
src_port=dst_port,
dst_port=dst_port,
@@ -385,7 +386,7 @@ class SessionManager:
elif frame.udp:
dst_port = frame.udp.dst_port
elif frame.icmp:
dst_port = Port["NONE"]
dst_port = PORT_LOOKUP["NONE"]
self.software_manager.receive_payload_from_session_manager(
payload=frame.payload,
port=dst_port,

8
src/primaite/simulator/system/core/software_manager.py
View File

@@ -8,12 +8,12 @@ from prettytable import MARKDOWN, PrettyTable
from primaite.simulator.core import RequestType
from primaite.simulator.file_system.file_system import FileSystem
from primaite.simulator.network.transmission.data_link_layer import Frame
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.system.applications.application import Application, ApplicationOperatingState
from primaite.simulator.system.core.sys_log import SysLog
from primaite.simulator.system.services.service import Service, ServiceOperatingState
from primaite.simulator.system.software import IOSoftware
from primaite.utils.validators import PROTOCOL_LOOKUP
if TYPE_CHECKING:
from primaite.simulator.system.core.session_manager import SessionManager
@@ -191,7 +191,7 @@ class SoftwareManager:
dest_ip_address: Optional[Union[IPv4Address, IPv4Network]] = None,
src_port: Optional[int] = None,
dest_port: Optional[int] = None,
ip_protocol: str = IPProtocol["TCP"],
ip_protocol: str = PROTOCOL_LOOKUP["TCP"],
session_id: Optional[str] = None,
) -> bool:
"""
@@ -275,7 +275,7 @@ class SoftwareManager:
software_type,
software.operating_state.name,
software.health_state_actual.name,
software.port if software.port != Port["NONE"] else None,
software.port if software.port != PORT_LOOKUP["NONE"] else None,
software.protocol,
]
)

9
src/primaite/simulator/system/services/arp/arp.py
View File

@@ -8,10 +8,9 @@ from prettytable import MARKDOWN, PrettyTable
from primaite.simulator.network.hardware.base import NetworkInterface
from primaite.simulator.network.protocols.arp import ARPEntry, ARPPacket
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.system.services.service import Service
from primaite.utils.validators import IPV4Address
from primaite.utils.validators import IPV4Address, PROTOCOL_LOOKUP
class ARP(Service):
@@ -26,8 +25,8 @@ class ARP(Service):
def __init__(self, **kwargs):
kwargs["name"] = "ARP"
kwargs["port"] = Port["ARP"]
kwargs["protocol"] = IPProtocol["UDP"]
kwargs["port"] = PORT_LOOKUP["ARP"]
kwargs["protocol"] = PROTOCOL_LOOKUP["UDP"]
super().__init__(**kwargs)
def describe_state(self) -> Dict:

8
src/primaite/simulator/system/services/database/database_service.py
View File

@@ -7,12 +7,12 @@ from primaite import getLogger
from primaite.simulator.file_system.file_system import File
from primaite.simulator.file_system.file_system_item_abc import FileSystemItemHealthStatus
from primaite.simulator.file_system.folder import Folder
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.system.core.software_manager import SoftwareManager
from primaite.simulator.system.services.ftp.ftp_client import FTPClient
from primaite.simulator.system.services.service import Service, ServiceOperatingState
from primaite.simulator.system.software import SoftwareHealthState
from primaite.utils.validators import PROTOCOL_LOOKUP
_LOGGER = getLogger(__name__)
@@ -38,8 +38,8 @@ class DatabaseService(Service):
def __init__(self, **kwargs):
kwargs["name"] = "DatabaseService"
kwargs["port"] = Port["POSTGRES_SERVER"]
kwargs["protocol"] = IPProtocol["TCP"]
kwargs["port"] = PORT_LOOKUP["POSTGRES_SERVER"]
kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"]
super().__init__(**kwargs)
self._create_db_file()

10
src/primaite/simulator/system/services/dns/dns_client.py
View File

@@ -4,10 +4,10 @@ from typing import Dict, Optional
from primaite import getLogger
from primaite.simulator.network.protocols.dns import DNSPacket, DNSRequest
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.system.core.software_manager import SoftwareManager
from primaite.simulator.system.services.service import Service
from primaite.utils.validators import PROTOCOL_LOOKUP
_LOGGER = getLogger(__name__)
@@ -22,11 +22,11 @@ class DNSClient(Service):
def __init__(self, **kwargs):
kwargs["name"] = "DNSClient"
kwargs["port"] = Port["DNS"]
kwargs["port"] = PORT_LOOKUP["DNS"]
# DNS uses UDP by default
# it switches to TCP when the bytes exceed 512 (or 4096) bytes
# TCP for now
kwargs["protocol"] = IPProtocol["TCP"]
kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"]
super().__init__(**kwargs)
self.start()
@@ -95,7 +95,7 @@ class DNSClient(Service):
# send a request to check if domain name exists in the DNS Server
software_manager: SoftwareManager = self.software_manager
software_manager.send_payload_to_session_manager(
payload=payload, dest_ip_address=self.dns_server, dest_port=Port["DNS"]
payload=payload, dest_ip_address=self.dns_server, dest_port=PORT_LOOKUP["DNS"]
)
# recursively re-call the function passing is_reattempt=True

8
src/primaite/simulator/system/services/dns/dns_server.py
View File

@@ -6,9 +6,9 @@ from prettytable import MARKDOWN, PrettyTable
from primaite import getLogger
from primaite.simulator.network.protocols.dns import DNSPacket
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.system.services.service import Service
from primaite.utils.validators import PROTOCOL_LOOKUP
_LOGGER = getLogger(__name__)
@@ -21,11 +21,11 @@ class DNSServer(Service):
def __init__(self, **kwargs):
kwargs["name"] = "DNSServer"
kwargs["port"] = Port["DNS"]
kwargs["port"] = PORT_LOOKUP["DNS"]
# DNS uses UDP by default
# it switches to TCP when the bytes exceed 512 (or 4096) bytes
# TCP for now
kwargs["protocol"] = IPProtocol["TCP"]
kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"]
super().__init__(**kwargs)
self.start()

18
src/primaite/simulator/system/services/ftp/ftp_client.py
View File

@@ -7,10 +7,10 @@ from primaite.interface.request import RequestFormat, RequestResponse
from primaite.simulator.core import RequestManager, RequestType
from primaite.simulator.file_system.file_system import File
from primaite.simulator.network.protocols.ftp import FTPCommand, FTPPacket, FTPStatusCode
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.system.core.software_manager import SoftwareManager
from primaite.simulator.system.services.ftp.ftp_service import FTPServiceABC
from primaite.utils.validators import PROTOCOL_LOOKUP
_LOGGER = getLogger(__name__)
@@ -25,8 +25,8 @@ class FTPClient(FTPServiceABC):
def __init__(self, **kwargs):
kwargs["name"] = "FTPClient"
kwargs["port"] = Port["FTP"]
kwargs["protocol"] = IPProtocol["TCP"]
kwargs["port"] = PORT_LOOKUP["FTP"]
kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"]
super().__init__(**kwargs)
self.start()
@@ -104,7 +104,7 @@ class FTPClient(FTPServiceABC):
def _connect_to_server(
self,
dest_ip_address: Optional[IPv4Address] = None,
dest_port: Optional[int] = Port["FTP"],
dest_port: Optional[int] = PORT_LOOKUP["FTP"],
session_id: Optional[str] = None,
is_reattempt: Optional[bool] = False,
) -> bool:
@@ -124,7 +124,7 @@ class FTPClient(FTPServiceABC):
# normally FTP will choose a random port for the transfer, but using the FTP command port will do for now
# create FTP packet
payload: FTPPacket = FTPPacket(ftp_command=FTPCommand.PORT, ftp_command_args=Port["FTP"])
payload: FTPPacket = FTPPacket(ftp_command=FTPCommand.PORT, ftp_command_args=PORT_LOOKUP["FTP"])
if self.send(payload=payload, dest_ip_address=dest_ip_address, dest_port=dest_port, session_id=session_id):
if payload.status_code == FTPStatusCode.OK:
@@ -152,7 +152,7 @@ class FTPClient(FTPServiceABC):
return False
def _disconnect_from_server(
self, dest_ip_address: Optional[IPv4Address] = None, dest_port: Optional[int] = Port["FTP"]
self, dest_ip_address: Optional[IPv4Address] = None, dest_port: Optional[int] = PORT_LOOKUP["FTP"]
) -> bool:
"""
Connects the client from a given FTP server.
@@ -179,7 +179,7 @@ class FTPClient(FTPServiceABC):
src_file_name: str,
dest_folder_name: str,
dest_file_name: str,
dest_port: Optional[int] = Port["FTP"],
dest_port: Optional[int] = PORT_LOOKUP["FTP"],
session_id: Optional[str] = None,
) -> bool:
"""
@@ -241,7 +241,7 @@ class FTPClient(FTPServiceABC):
src_file_name: str,
dest_folder_name: str,
dest_file_name: str,
dest_port: Optional[int] = Port["FTP"],
dest_port: Optional[int] = PORT_LOOKUP["FTP"],
) -> bool:
"""
Request a file from a target IP address.

8
src/primaite/simulator/system/services/ftp/ftp_server.py
View File

@@ -3,9 +3,9 @@ from typing import Any, Optional
from primaite import getLogger
from primaite.simulator.network.protocols.ftp import FTPCommand, FTPPacket, FTPStatusCode
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.system.services.ftp.ftp_service import FTPServiceABC
from primaite.utils.validators import PROTOCOL_LOOKUP
_LOGGER = getLogger(__name__)
@@ -23,8 +23,8 @@ class FTPServer(FTPServiceABC):
def __init__(self, **kwargs):
kwargs["name"] = "FTPServer"
kwargs["port"] = Port["FTP"]
kwargs["protocol"] = IPProtocol["TCP"]
kwargs["port"] = PORT_LOOKUP["FTP"]
kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"]
super().__init__(**kwargs)
self.start()

8
src/primaite/simulator/system/services/icmp/icmp.py
View File

@@ -7,9 +7,9 @@ from primaite import getLogger
from primaite.simulator.network.hardware.base import NetworkInterface
from primaite.simulator.network.protocols.icmp import ICMPPacket, ICMPType
from primaite.simulator.network.transmission.data_link_layer import Frame
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.system.services.service import Service
from primaite.utils.validators import PROTOCOL_LOOKUP
_LOGGER = getLogger(__name__)
@@ -26,8 +26,8 @@ class ICMP(Service):
def __init__(self, **kwargs):
kwargs["name"] = "ICMP"
kwargs["port"] = Port["NONE"]
kwargs["protocol"] = IPProtocol["ICMP"]
kwargs["port"] = PORT_LOOKUP["NONE"]
kwargs["protocol"] = PROTOCOL_LOOKUP["ICMP"]
super().__init__(**kwargs)
def describe_state(self) -> Dict:

10
src/primaite/simulator/system/services/ntp/ntp_client.py
View File

@@ -5,9 +5,9 @@ from typing import Dict, Optional
from primaite import getLogger
from primaite.simulator.network.protocols.ntp import NTPPacket
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.system.services.service import Service, ServiceOperatingState
from primaite.utils.validators import PROTOCOL_LOOKUP
_LOGGER = getLogger(__name__)
@@ -21,8 +21,8 @@ class NTPClient(Service):
def __init__(self, **kwargs):
kwargs["name"] = "NTPClient"
kwargs["port"] = Port["NTP"]
kwargs["protocol"] = IPProtocol["UDP"]
kwargs["port"] = PORT_LOOKUP["NTP"]
kwargs["protocol"] = PROTOCOL_LOOKUP["UDP"]
super().__init__(**kwargs)
self.start()
@@ -55,7 +55,7 @@ class NTPClient(Service):
payload: NTPPacket,
session_id: Optional[str] = None,
dest_ip_address: IPv4Address = None,
dest_port: int = Port["NTP"],
dest_port: int = PORT_LOOKUP["NTP"],
**kwargs,
) -> bool:
"""Requests NTP data from NTP server.

8
src/primaite/simulator/system/services/ntp/ntp_server.py
View File

@@ -4,9 +4,9 @@ from typing import Dict, Optional
from primaite import getLogger
from primaite.simulator.network.protocols.ntp import NTPPacket
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.system.services.service import Service
from primaite.utils.validators import PROTOCOL_LOOKUP
_LOGGER = getLogger(__name__)
@@ -16,8 +16,8 @@ class NTPServer(Service):
def __init__(self, **kwargs):
kwargs["name"] = "NTPServer"
kwargs["port"] = Port["NTP"]
kwargs["protocol"] = IPProtocol["UDP"]
kwargs["port"] = PORT_LOOKUP["NTP"]
kwargs["protocol"] = PROTOCOL_LOOKUP["UDP"]
super().__init__(**kwargs)
self.start()

8
src/primaite/simulator/system/services/terminal/terminal.py
View File

@@ -17,10 +17,10 @@ from primaite.simulator.network.protocols.ssh import (
SSHTransportMessage,
SSHUserCredentials,
)
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.system.core.software_manager import SoftwareManager
from primaite.simulator.system.services.service import Service, ServiceOperatingState
from primaite.utils.validators import PROTOCOL_LOOKUP
# TODO 2824: Since remote terminal connections and remote user sessions are the same thing, we could refactor
@@ -137,8 +137,8 @@ class Terminal(Service):
def __init__(self, **kwargs):
kwargs["name"] = "Terminal"
kwargs["port"] = Port["SSH"]
kwargs["protocol"] = IPProtocol["TCP"]
kwargs["port"] = PORT_LOOKUP["SSH"]
kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"]
super().__init__(**kwargs)
def describe_state(self) -> Dict:

8
src/primaite/simulator/system/services/web_server/web_server.py
View File

@@ -10,11 +10,11 @@ from primaite.simulator.network.protocols.http import (
HttpResponsePacket,
HttpStatusCode,
)
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.system.applications.database_client import DatabaseClientConnection
from primaite.simulator.system.services.service import Service
from primaite.simulator.system.software import SoftwareHealthState
from primaite.utils.validators import PROTOCOL_LOOKUP
_LOGGER = getLogger(__name__)
@@ -49,10 +49,10 @@ class WebServer(Service):
def __init__(self, **kwargs):
kwargs["name"] = "WebServer"
kwargs["protocol"] = IPProtocol["TCP"]
kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"]
# default for web is port 80
if kwargs.get("port") is None:
kwargs["port"] = Port["HTTP"]
kwargs["port"] = PORT_LOOKUP["HTTP"]
super().__init__(**kwargs)
self._install_web_files()

4
src/primaite/simulator/system/software.py
View File

@@ -13,9 +13,9 @@ from primaite.interface.request import RequestResponse
from primaite.simulator.core import RequestManager, RequestType, SimComponent
from primaite.simulator.file_system.file_system import FileSystem, Folder
from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.system.core.session_manager import Session
from primaite.simulator.system.core.sys_log import SysLog
from primaite.utils.validators import PROTOCOL_LOOKUP
if TYPE_CHECKING:
from primaite.simulator.system.core.software_manager import SoftwareManager
@@ -386,7 +386,7 @@ class IOSoftware(Software):
session_id: Optional[str] = None,
dest_ip_address: Optional[Union[IPv4Address, IPv4Network]] = None,
dest_port: Optional[int] = None,
ip_protocol: str = IPProtocol["TCP"],
ip_protocol: str = PROTOCOL_LOOKUP["TCP"],
**kwargs,
) -> bool:
"""

42
src/primaite/utils/validators.py
View File

@@ -6,6 +6,9 @@ from pydantic import BeforeValidator
from typing_extensions import Annotated
# Define a custom type IPV4Address using the typing_extensions.Annotated.
# Annotated is used to attach metadata to type hints. In this case, it's used to associate the ipv4_validator
# with the IPv4Address type, ensuring that any usage of IPV4Address undergoes validation before assignment.
def ipv4_validator(v: Any) -> IPv4Address:
"""
Validate the input and ensure it can be converted to an IPv4Address instance.
@@ -24,9 +27,6 @@ def ipv4_validator(v: Any) -> IPv4Address:
return IPv4Address(v)
# Define a custom type IPV4Address using the typing_extensions.Annotated.
# Annotated is used to attach metadata to type hints. In this case, it's used to associate the ipv4_validator
# with the IPv4Address type, ensuring that any usage of IPV4Address undergoes validation before assignment.
IPV4Address: Final[Annotated] = Annotated[IPv4Address, BeforeValidator(ipv4_validator)]
"""
IPv4Address with with IPv4Address with with pre-validation and auto-conversion from str using ipv4_validator..
@@ -37,3 +37,39 @@ will automatically check and convert the input value to an instance of IPv4Addre
any Pydantic model uses it. This ensures that any field marked with this type is not just
an IPv4Address in form, but also valid according to the rules defined in ipv4_validator.
"""
# Define a custom port validator
Port: Final[Annotated] = Annotated[int, BeforeValidator(lambda n: 0 <= n <= 65535)]
"""Validates that network ports lie in the appropriate range of [0,65535]."""
# Define a custom IP protocol validator
PROTOCOL_LOOKUP: dict[str, str] = dict(
NONE="none",
TCP="tcp",
UDP="udp",
ICMP="icmp",
)
"""
Lookup table used for compatibility with PrimAITE <= 3.3. Configs with the capitalised protocol names are converted
to lowercase at runtime.
"""
VALID_PROTOCOLS = ["none", "tcp", "udp", "icmp"]
"""Supported protocols."""
def protocol_validator(v: Any) -> str:
"""
Validate that IP Protocols are chosen from the list of supported IP Protocols.
The protocol list is dynamic because plugins are able to extend it, therefore it is necessary to use this custom
validator instead of being able to specify a union of string literals.
"""
if v in PROTOCOL_LOOKUP:
return PROTOCOL_LOOKUP(v)
if v in VALID_PROTOCOLS:
return v
raise ValueError(f"{v} is not a valid IP Protocol. It must be one of the following: {VALID_PROTOCOLS}")
IPProtocol: Final[Annotated] = Annotated[str, BeforeValidator(protocol_validator)]
"""Validates that IP Protocols used in the simulation belong to the list of supported protocols."""