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

2
src/primaite/notebooks/Command-&-Control-E2E-Demonstration.ipynb
View File

@@ -1783,7 +1783,7 @@
"from primaite.simulator.network.transmission.network_layer import IPProtocol\n",
"from primaite.simulator.network.transmission.transport_layer import Port\n",
"# As we're configuring via the PrimAITE API we need to pass the actual IPProtocol/Port (Agents leverage the simulation via the game layer and thus can pass strings).\n",
"c2_beacon.configure(c2_server_ip_address=\"192.168.10.21\", masquerade_protocol=IPProtocol.UDP, masquerade_port=Port.DNS)\n",
"c2_beacon.configure(c2_server_ip_address=\"192.168.10.21\", masquerade_protocol=IPProtocol["UDP"], masquerade_port=Port["DNS"])\n",
"c2_beacon.establish()\n",
"c2_beacon.show()"
]

2
src/primaite/simulator/_package_data/create-simulation_demo.ipynb
View File

@@ -182,7 +182,7 @@
"metadata": {},
"outputs": [],
"source": [
"mspaint = MSPaint(name = \"mspaint\", health_state_actual=SoftwareHealthState.GOOD, health_state_visible=SoftwareHealthState.GOOD, criticality=SoftwareCriticality.MEDIUM, port=Port.HTTP, protocol = IPProtocol.NONE,operating_state=ApplicationOperatingState.RUNNING,execution_control_status='manual', file_system=FileSystem(sys_log=SysLog(hostname=\"Test\"), sim_root=Path(__name__).parent),)"
"mspaint = MSPaint(name = \"mspaint\", health_state_actual=SoftwareHealthState.GOOD, health_state_visible=SoftwareHealthState.GOOD, criticality=SoftwareCriticality.MEDIUM, port=Port["HTTP"], protocol = IPProtocol["NONE"],operating_state=ApplicationOperatingState.RUNNING,execution_control_status='manual', file_system=FileSystem(sys_log=SysLog(hostname=\"Test\"), sim_root=Path(__name__).parent),)"
]
},
{

2
src/primaite/simulator/_package_data/network_simulator_demo.ipynb
View File

@@ -537,7 +537,7 @@
"from primaite.simulator.network.hardware.nodes.network.router import ACLAction\n",
"network.get_node_by_hostname(\"router_1\").acl.add_rule(\n",
" action=ACLAction.DENY,\n",
" protocol=IPProtocol.ICMP,\n",
" protocol=IPProtocol["ICMP"],\n",
" src_ip_address=\"192.168.10.22\",\n",
" position=1\n",
")"

82
src/primaite/simulator/network/airspace.py
View File

@@ -3,10 +3,12 @@ from __future__ import annotations
from abc import ABC, abstractmethod
from enum import Enum
from typing import Any, Dict, List
from typing import Any, ClassVar, Dict, List, Type
from pydantic._internal._generics import PydanticGenericMetadata
from typing_extensions import Unpack
from prettytable import MARKDOWN, PrettyTable
from pydantic import BaseModel, Field
from pydantic import BaseModel, ConfigDict, Field
from primaite import getLogger
from primaite.simulator.network.hardware.base import Layer3Interface, NetworkInterface, WiredNetworkInterface
@@ -40,51 +42,29 @@ def format_hertz(hertz: float, format_terahertz: bool = False, decimals: int = 3
else: # Hertz
return format_str.format(hertz) + " Hz"
AirSpaceFrequencyRegistry: Dict[str,Dict] = {
"WIFI_2_4" : {'frequency': 2.4e9, 'data_rate_bps':100_000_000.0},
"WIFI_5" : {'frequency': 5e9, 'data_rate_bps':500_000_000.0},
}
class AirSpaceFrequency(Enum):
"""Enumeration representing the operating frequencies for wireless communications."""
def register_frequency(freq_name: str, freq_hz: int, data_rate_bps: int) -> None:
if freq_name in AirSpaceFrequencyRegistry:
raise RuntimeError(f"Cannot register new frequency {freq_name} because it's already registered.")
AirSpaceFrequencyRegistry.update({freq_name:{'frequency': freq_hz, 'data_rate_bps':data_rate_bps}})
WIFI_2_4 = 2.4e9
"""WiFi 2.4 GHz. Known for its extensive range and ability to penetrate solid objects effectively."""
WIFI_5 = 5e9
"""WiFi 5 GHz. Known for its higher data transmission speeds and reduced interference from other devices."""
def maximum_data_rate_mbps(frequency_name:str) -> float:
"""
Retrieves the maximum data transmission rate in megabits per second (Mbps).
def __str__(self) -> str:
hertz_str = format_hertz(hertz=self.value)
if self == AirSpaceFrequency.WIFI_2_4:
return f"WiFi {hertz_str}"
if self == AirSpaceFrequency.WIFI_5:
return f"WiFi {hertz_str}"
return "Unknown Frequency"
This is derived by converting the maximum data rate from bits per second, as defined
in `maximum_data_rate_bps`, to megabits per second.
@property
def maximum_data_rate_bps(self) -> float:
"""
Retrieves the maximum data transmission rate in bits per second (bps).
:return: The maximum data rate in megabits per second.
"""
return AirSpaceFrequencyRegistry[frequency_name]['data_rate_bps']
return data_rate / 1_000_000.0
The maximum rates are predefined for frequencies.:
- WIFI 2.4 supports 100,000,000 bps
- WIFI 5 supports 500,000,000 bps
:return: The maximum data rate in bits per second.
"""
if self == AirSpaceFrequency.WIFI_2_4:
return 100_000_000.0 # 100 Megabits per second
if self == AirSpaceFrequency.WIFI_5:
return 500_000_000.0 # 500 Megabits per second
return 0.0
@property
def maximum_data_rate_mbps(self) -> float:
"""
Retrieves the maximum data transmission rate in megabits per second (Mbps).
This is derived by converting the maximum data rate from bits per second, as defined
in `maximum_data_rate_bps`, to megabits per second.
:return: The maximum data rate in megabits per second.
"""
return self.maximum_data_rate_bps / 1_000_000.0
class AirSpace(BaseModel):
@@ -97,13 +77,13 @@ class AirSpace(BaseModel):
"""
wireless_interfaces: Dict[str, WirelessNetworkInterface] = Field(default_factory=lambda: {})
wireless_interfaces_by_frequency: Dict[AirSpaceFrequency, List[WirelessNetworkInterface]] = Field(
wireless_interfaces_by_frequency: Dict[int, List[WirelessNetworkInterface]] = Field(
default_factory=lambda: {}
)
bandwidth_load: Dict[AirSpaceFrequency, float] = Field(default_factory=lambda: {})
frequency_max_capacity_mbps_: Dict[AirSpaceFrequency, float] = Field(default_factory=lambda: {})
bandwidth_load: Dict[int, float] = Field(default_factory=lambda: {})
frequency_max_capacity_mbps_: Dict[int, float] = Field(default_factory=lambda: {})
def get_frequency_max_capacity_mbps(self, frequency: AirSpaceFrequency) -> float:
def get_frequency_max_capacity_mbps(self, frequency: str) -> float:
"""
Retrieves the maximum data transmission capacity for a specified frequency.
@@ -117,9 +97,9 @@ class AirSpace(BaseModel):
"""
if frequency in self.frequency_max_capacity_mbps_:
return self.frequency_max_capacity_mbps_[frequency]
return frequency.maximum_data_rate_mbps
return maximum_data_rate_mbps(frequency)
def set_frequency_max_capacity_mbps(self, cfg: Dict[AirSpaceFrequency, float]):
def set_frequency_max_capacity_mbps(self, cfg: Dict[int, float]):
"""
Sets custom maximum data transmission capacities for multiple frequencies.
@@ -150,7 +130,7 @@ class AirSpace(BaseModel):
load_percent = load / maximum_capacity if maximum_capacity > 0 else 0.0
if load_percent > 1.0:
load_percent = 1.0
table.add_row([format_hertz(frequency.value), f"{load_percent:.0%}", f"{maximum_capacity:.3f}"])
table.add_row([format_hertz(frequency), f"{load_percent:.0%}", f"{maximum_capacity:.3f}"])
print(table)
def show_wireless_interfaces(self, markdown: bool = False):
@@ -182,7 +162,7 @@ class AirSpace(BaseModel):
interface.mac_address,
interface.ip_address if hasattr(interface, "ip_address") else None,
interface.subnet_mask if hasattr(interface, "subnet_mask") else None,
format_hertz(interface.frequency.value),
format_hertz(interface.frequency),
f"{interface.speed:.3f}",
status,
]
@@ -298,7 +278,7 @@ class WirelessNetworkInterface(NetworkInterface, ABC):
"""
airspace: AirSpace
frequency: AirSpaceFrequency = AirSpaceFrequency.WIFI_2_4
frequency: str = "WIFI_2_4"
def enable(self):
"""Attempt to enable the network interface."""
@@ -430,7 +410,7 @@ class IPWirelessNetworkInterface(WirelessNetworkInterface, Layer3Interface, ABC)
# Update the state with information from Layer3Interface
state.update(Layer3Interface.describe_state(self))
state["frequency"] = self.frequency.value
state["frequency"] = self.frequency
return state

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

@@ -98,8 +98,8 @@ 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["ARP"], dst_port=Port["ARP"], position=22)
router.acl.add_rule(action=ACLAction.PERMIT, protocol=IPProtocol["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)

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

@@ -203,16 +203,16 @@ class NetworkInterface(SimComponent, ABC):
# Initialise basic frame data variables
direction = "inbound" if inbound else "outbound" # Direction of the traffic
ip_address = str(frame.ip.src_ip_address if inbound else frame.ip.dst_ip_address) # Source or destination IP
protocol = frame.ip.protocol.name # Network protocol used in the frame
protocol = frame.ip.protocol # Network protocol used in the frame
# Initialise port variable; will be determined based on protocol type
port = None
# Determine the source or destination port based on the protocol (TCP/UDP)
if frame.tcp:
port = frame.tcp.src_port.value if inbound else frame.tcp.dst_port.value
port = frame.tcp.src_port if inbound else frame.tcp.dst_port
elif frame.udp:
port = frame.udp.src_port.value if inbound else frame.udp.dst_port.value
port = frame.udp.src_port if inbound else frame.udp.dst_port
# Convert frame payload to string for keyword checking
frame_str = str(frame.payload)
@@ -274,20 +274,20 @@ class NetworkInterface(SimComponent, ABC):
# Identify the protocol and port from the frame
if frame.tcp:
protocol = IPProtocol.TCP
protocol = IPProtocol["TCP"]
port = frame.tcp.dst_port
elif frame.udp:
protocol = IPProtocol.UDP
protocol = IPProtocol["UDP"]
port = frame.udp.dst_port
elif frame.icmp:
protocol = IPProtocol.ICMP
protocol = IPProtocol["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 != IPProtocol["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 +843,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["NONE"]
kwargs["protocol"] = IPProtocol["NONE"]
super().__init__(**kwargs)
self.start()
@@ -1166,8 +1166,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["NONE"]
kwargs["protocol"] = IPProtocol["NONE"]
super().__init__(**kwargs)
self.start()
@@ -1312,7 +1312,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["SSH"],
dest_ip_address=session.remote_ip_address,
)
@@ -1845,8 +1845,9 @@ class Node(SimComponent):
table.align = "l"
table.title = f"{self.hostname} Open Ports"
for port in self.software_manager.get_open_ports():
if port.value > 0:
table.add_row([port.value, port.name])
if port > 0:
# TODO: do a reverse lookup for port name, or change this to only show port int
table.add_row([port, port])
print(table.get_string(sortby="Port"))
@property

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

@@ -58,8 +58,8 @@ class Firewall(Router):
>>> # Permit HTTP traffic to the DMZ
>>> firewall.dmz_inbound_acl.add_rule(
... action=ACLAction.PERMIT,
... protocol=IPProtocol.TCP,
... dst_port=Port.HTTP,
... protocol=IPProtocol["TCP"],
... dst_port=Port["HTTP"],
... src_ip_address="0.0.0.0",
... src_wildcard_mask="0.0.0.0",
... dst_ip_address="172.16.0.0",

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

@@ -7,7 +7,7 @@ from ipaddress import IPv4Address, IPv4Network
from typing import Any, ClassVar, Dict, List, Optional, Tuple, Union
from prettytable import MARKDOWN, PrettyTable
from pydantic import validate_call
from pydantic import field_validator, validate_call
from primaite.interface.request import RequestResponse
from primaite.simulator.core import RequestManager, RequestType, SimComponent
@@ -106,7 +106,7 @@ class ACLRule(SimComponent):
:ivar ACLAction action: Specifies whether to `PERMIT` or `DENY` the traffic that matches the rule conditions.
The default action is `DENY`.
:ivar Optional[IPProtocol] protocol: The network protocol (e.g., TCP, UDP, ICMP) to match. If `None`, the rule
:ivar Optional[str] protocol: The network protocol (e.g., TCP, UDP, ICMP) to match. If `None`, the rule
applies to all protocols.
:ivar Optional[IPV4Address] src_ip_address: The source IP address to match. If combined with `src_wildcard_mask`,
it specifies the start of an IP range.
@@ -116,20 +116,33 @@ class ACLRule(SimComponent):
`dst_wildcard_mask`, it specifies the start of an IP range.
:ivar Optional[IPv4Address] dst_wildcard_mask: The wildcard mask for the destination IP address, defining the
range of addresses to match.
:ivar Optional[Port] src_port: The source port number to match. Relevant for TCP/UDP protocols.
:ivar Optional[Port] dst_port: The destination port number to match. Relevant for TCP/UDP protocols.
:ivar Optional[int] src_port: The source port number to match. Relevant for TCP/UDP protocols.
:ivar Optional[int] dst_port: The destination port number to match. Relevant for TCP/UDP protocols.
"""
action: ACLAction = ACLAction.DENY
protocol: Optional[IPProtocol] = None
protocol: Optional[str] = None
src_ip_address: Optional[IPV4Address] = None
src_wildcard_mask: Optional[IPV4Address] = None
dst_ip_address: Optional[IPV4Address] = None
dst_wildcard_mask: Optional[IPV4Address] = None
src_port: Optional[Port] = None
dst_port: Optional[Port] = None
src_port: Optional[int] = None
dst_port: Optional[int] = None
match_count: int = 0
@field_validator('protocol', mode='before')
def protocol_valid(cls, val:Optional[str]) -> Optional[str]:
if val is not None:
assert val in IPProtocol.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]:
if val is not None:
assert val in Port.values(), f"Cannot create ACL rule with invalid port {val}"
return val
def __str__(self) -> str:
rule_strings = []
for key, value in self.model_dump(exclude={"uuid", "request_manager"}).items():
@@ -149,13 +162,13 @@ class ACLRule(SimComponent):
"""
state = super().describe_state()
state["action"] = self.action.value
state["protocol"] = self.protocol.name if self.protocol else None
state["protocol"] = self.protocol if self.protocol else None
state["src_ip_address"] = str(self.src_ip_address) if self.src_ip_address else None
state["src_wildcard_mask"] = str(self.src_wildcard_mask) if self.src_wildcard_mask else None
state["src_port"] = self.src_port.name if self.src_port else None
state["src_port"] = self.src_port if self.src_port else None
state["dst_ip_address"] = str(self.dst_ip_address) if self.dst_ip_address else None
state["dst_wildcard_mask"] = str(self.dst_wildcard_mask) if self.dst_wildcard_mask else None
state["dst_port"] = self.dst_port.name if self.dst_port else None
state["dst_port"] = self.dst_port if self.dst_port else None
state["match_count"] = self.match_count
return state
@@ -265,7 +278,7 @@ class AccessControlList(SimComponent):
>>> acl = AccessControlList()
>>> acl.add_rule(
... action=ACLAction.PERMIT,
... protocol=IPProtocol.TCP,
... protocol=IPProtocol["TCP"],
... src_ip_address="192.168.1.0",
... src_wildcard_mask="0.0.0.255",
... dst_ip_address="192.168.2.0",
@@ -323,13 +336,13 @@ class AccessControlList(SimComponent):
func=lambda request, context: RequestResponse.from_bool(
self.add_rule(
action=ACLAction[request[0]],
protocol=None if request[1] == "ALL" else IPProtocol[request[1]],
protocol=None if request[1] == "ALL" else request[1],
src_ip_address=None if request[2] == "ALL" else IPv4Address(request[2]),
src_wildcard_mask=None if request[3] == "NONE" else IPv4Address(request[3]),
src_port=None if request[4] == "ALL" else Port[request[4]],
src_port=None if request[4] == "ALL" else request[4],
dst_ip_address=None if request[5] == "ALL" else IPv4Address(request[5]),
dst_wildcard_mask=None if request[6] == "NONE" else IPv4Address(request[6]),
dst_port=None if request[7] == "ALL" else Port[request[7]],
dst_port=None if request[7] == "ALL" else request[7],
position=int(request[8]),
)
)
@@ -377,13 +390,13 @@ class AccessControlList(SimComponent):
def add_rule(
self,
action: ACLAction = ACLAction.DENY,
protocol: Optional[IPProtocol] = None,
protocol: Optional[str] = None,
src_ip_address: Optional[IPV4Address] = None,
src_wildcard_mask: Optional[IPV4Address] = None,
dst_ip_address: Optional[IPV4Address] = None,
dst_wildcard_mask: Optional[IPV4Address] = None,
src_port: Optional[Port] = None,
dst_port: Optional[Port] = None,
src_port: Optional[int] = None,
dst_port: Optional[int] = None,
position: int = 0,
) -> bool:
"""
@@ -399,11 +412,11 @@ class AccessControlList(SimComponent):
>>> router = Router("router")
>>> router.add_rule(
... action=ACLAction.DENY,
... protocol=IPProtocol.TCP,
... protocol=IPProtocol["TCP"],
... src_ip_address="192.168.1.0",
... src_wildcard_mask="0.0.0.255",
... dst_ip_address="10.10.10.5",
... dst_port=Port.SSH,
... dst_port=Port["SSH"],
... position=5
... )
>>> # This permits SSH traffic from the 192.168.1.0/24 subnet to the 10.10.10.5 server.
@@ -411,10 +424,10 @@ class AccessControlList(SimComponent):
>>> # Then if we want to allow a specific IP address from this subnet to SSH into the server
>>> router.add_rule(
... action=ACLAction.PERMIT,
... protocol=IPProtocol.TCP,
... protocol=IPProtocol["TCP"],
... src_ip_address="192.168.1.25",
... dst_ip_address="10.10.10.5",
... dst_port=Port.SSH,
... dst_port=Port["SSH"],
... position=4
... )
@@ -485,11 +498,11 @@ class AccessControlList(SimComponent):
def get_relevant_rules(
self,
protocol: IPProtocol,
protocol: str,
src_ip_address: Union[str, IPv4Address],
src_port: Port,
src_port: int,
dst_ip_address: Union[str, IPv4Address],
dst_port: Port,
dst_port: int,
) -> List[ACLRule]:
"""
Get the list of relevant rules for a packet with given properties.
@@ -552,13 +565,13 @@ class AccessControlList(SimComponent):
[
index,
rule.action.name if rule.action else "ANY",
rule.protocol.name if rule.protocol else "ANY",
rule.protocol if rule.protocol else "ANY",
rule.src_ip_address if rule.src_ip_address else "ANY",
rule.src_wildcard_mask if rule.src_wildcard_mask else "ANY",
f"{rule.src_port.value} ({rule.src_port.name})" if rule.src_port else "ANY",
f"{rule.src_port} ({rule.src_port})" if rule.src_port else "ANY",
rule.dst_ip_address if rule.dst_ip_address else "ANY",
rule.dst_wildcard_mask if rule.dst_wildcard_mask else "ANY",
f"{rule.dst_port.value} ({rule.dst_port.name})" if rule.dst_port else "ANY",
f"{rule.dst_port} ({rule.dst_port})" if rule.dst_port else "ANY",
rule.match_count,
]
)
@@ -1088,17 +1101,17 @@ class RouterSessionManager(SessionManager):
def resolve_outbound_transmission_details(
self,
dst_ip_address: Optional[Union[IPv4Address, IPv4Network]] = None,
src_port: Optional[Port] = None,
dst_port: Optional[Port] = None,
protocol: Optional[IPProtocol] = None,
src_port: Optional[int] = None,
dst_port: Optional[int] = None,
protocol: Optional[str] = None,
session_id: Optional[str] = None,
) -> Tuple[
Optional[RouterInterface],
Optional[str],
IPv4Address,
Optional[Port],
Optional[Port],
Optional[IPProtocol],
Optional[int],
Optional[int],
Optional[str],
bool,
]:
"""
@@ -1118,19 +1131,19 @@ class RouterSessionManager(SessionManager):
treats the transmission as a broadcast to that network. Optional.
:type dst_ip_address: Optional[Union[IPv4Address, IPv4Network]]
:param src_port: The source port number for the transmission. Optional.
:type src_port: Optional[Port]
:type src_port: Optional[int]
:param dst_port: The destination port number for the transmission. Optional.
:type dst_port: Optional[Port]
:type dst_port: Optional[int]
:param protocol: The IP protocol to be used for the transmission. Optional.
:type protocol: Optional[IPProtocol]
:type protocol: Optional[str]
:param session_id: The session ID associated with the transmission. If provided, the session details override
other parameters. Optional.
:type session_id: Optional[str]
:return: A tuple containing the resolved outbound network interface, destination MAC address, destination IP
address, source port, destination port, protocol, and a boolean indicating whether the transmission is a
broadcast.
:rtype: Tuple[Optional[RouterInterface], Optional[str], IPv4Address, Optional[Port], Optional[Port],
Optional[IPProtocol], bool]
:rtype: Tuple[Optional[RouterInterface], Optional[str], IPv4Address, Optional[int], Optional[int],
Optional[str], bool]
"""
if dst_ip_address and not isinstance(dst_ip_address, (IPv4Address, IPv4Network)):
dst_ip_address = IPv4Address(dst_ip_address)
@@ -1257,8 +1270,8 @@ 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["ARP"], dst_port=Port["ARP"], position=22)
self.acl.add_rule(action=ACLAction.PERMIT, protocol=IPProtocol["ICMP"], position=23)
def setup_for_episode(self, episode: int):
"""
@@ -1357,9 +1370,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 == IPProtocol["TCP"]:
dst_port = frame.tcp.dst_port
elif frame.ip.protocol == IPProtocol.UDP:
elif frame.ip.protocol == IPProtocol["UDP"]:
dst_port = frame.udp.dst_port
if self.ip_is_router_interface(dst_ip_address) and (

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

@@ -116,7 +116,7 @@ class WirelessRouter(Router):
>>> wireless_router.configure_wireless_access_point(
... ip_address="10.10.10.1",
... subnet_mask="255.255.255.0"
... frequency=AirSpaceFrequency.WIFI_2_4
... frequency=AirSpaceFrequency["WIFI_2_4"]
... )
"""
@@ -153,7 +153,7 @@ class WirelessRouter(Router):
self,
ip_address: IPV4Address,
subnet_mask: IPV4Address,
frequency: Optional[AirSpaceFrequency] = AirSpaceFrequency.WIFI_2_4,
frequency: Optional[int] = AirSpaceFrequency["WIFI_2_4"],
):
"""
Configures a wireless access point (WAP).
@@ -168,10 +168,10 @@ class WirelessRouter(Router):
:param subnet_mask: The subnet mask associated with the IP address
:param frequency: The operating frequency of the wireless access point, defined by the AirSpaceFrequency
enum. This determines the frequency band (e.g., 2.4 GHz or 5 GHz) the access point will use for wireless
communication. Default is AirSpaceFrequency.WIFI_2_4.
communication. Default is AirSpaceFrequency["WIFI_2_4"].
"""
if not frequency:
frequency = AirSpaceFrequency.WIFI_2_4
frequency = AirSpaceFrequency["WIFI_2_4"]
self.sys_log.info("Configuring wireless access point")
self.wireless_access_point.disable() # Temporarily disable the WAP for reconfiguration

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

@@ -79,9 +79,9 @@ 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["ARP"], dst_port=Port["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=IPProtocol["ICMP"], position=23)
return network
@@ -271,23 +271,23 @@ 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["ARP"], dst_port=Port["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=IPProtocol["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["POSTGRES_SERVER"], dst_port=Port["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["DNS"], dst_port=Port["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["FTP"], dst_port=Port["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["HTTP"], dst_port=Port["HTTP"], position=3)
return network

4
src/primaite/simulator/network/protocols/masquerade.py
View File

@@ -8,9 +8,9 @@ from primaite.simulator.network.protocols.packet import DataPacket
class MasqueradePacket(DataPacket):
"""Represents an generic malicious packet that is masquerading as another protocol."""
masquerade_protocol: Enum # The 'Masquerade' protocol that is currently in use
masquerade_protocol: str # The 'Masquerade' protocol that is currently in use
masquerade_port: Enum # The 'Masquerade' port that is currently in use
masquerade_port: int # The 'Masquerade' port that is currently in use
class C2Packet(MasqueradePacket):

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

@@ -70,15 +70,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 == IPProtocol["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 == IPProtocol["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 == IPProtocol["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 +165,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["ARP"]
@property
def is_icmp(self) -> bool:

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

@@ -9,25 +9,32 @@ from primaite.utils.validators import IPV4Address
_LOGGER = getLogger(__name__)
class IPProtocol(Enum):
"""
Enum representing transport layer protocols in IP header.
IPProtocol : dict[str, str] = dict(
NONE = "none",
TCP = "tcp",
UDP = "udp",
ICMP = "icmp",
)
.. _List of IPProtocols:
"""
# class IPProtocol(Enum):
# """
# Enum representing transport layer protocols in IP header.
NONE = "none"
"""Placeholder for a non-protocol."""
TCP = "tcp"
"""Transmission Control Protocol."""
UDP = "udp"
"""User Datagram Protocol."""
ICMP = "icmp"
"""Internet Control Message Protocol."""
# .. _List of IPProtocols:
# """
def model_dump(self) -> str:
"""Return as JSON-serialisable string."""
return self.name
# 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):
@@ -81,7 +88,7 @@ class IPPacket(BaseModel):
>>> ip_packet = IPPacket(
... src_ip_address=IPv4Address('192.168.0.1'),
... dst_ip_address=IPv4Address('10.0.0.1'),
... protocol=IPProtocol.TCP,
... protocol=IPProtocol["TCP"],
... ttl=64,
... precedence=Precedence.CRITICAL
... )
@@ -91,7 +98,7 @@ class IPPacket(BaseModel):
"Source IP address."
dst_ip_address: IPV4Address
"Destination IP address."
protocol: IPProtocol = IPProtocol.TCP
protocol: str = IPProtocol["TCP"]
"IPProtocol."
ttl: int = 64
"Time to Live (TTL) for the packet."

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

@@ -5,76 +5,112 @@ from typing import List, Union
from pydantic import BaseModel
class Port(Enum):
"""
Enumeration of common known TCP/UDP ports used by protocols for operation of network applications.
Port: dict[str, int] = dict(
UNUSED = -1,
NONE = 0,
WOL = 9,
FTP_DATA = 20,
FTP = 21,
SSH = 22,
SMTP = 25,
DNS = 53,
HTTP = 80,
POP3 = 110,
SFTP = 115,
NTP = 123,
IMAP = 143,
SNMP = 161,
SNMP_TRAP = 162,
ARP = 219,
LDAP = 389,
HTTPS = 443,
SMB = 445,
IPP = 631,
SQL_SERVER = 1433,
MYSQL = 3306,
RDP = 3389,
RTP = 5004,
RTP_ALT = 5005,
DNS_ALT = 5353,
HTTP_ALT = 8080,
HTTPS_ALT = 8443,
POSTGRES_SERVER = 5432,
)
.. _List of Ports:
"""
# class Port():
# def __getattr__()
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."
# class Port(Enum):
# """
# Enumeration of common known TCP/UDP ports used by protocols for operation of network applications.
def model_dump(self) -> str:
"""Return a json-serialisable string."""
return self.name
# .. _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):
@@ -87,13 +123,13 @@ class UDPHeader(BaseModel):
:Example:
>>> udp_header = UDPHeader(
... src_port=Port.HTTP_ALT,
... dst_port=Port.HTTP,
... src_port=Port["HTTP_ALT"],
... dst_port=Port["HTTP"],
... )
"""
src_port: Union[Port, int]
dst_port: Union[Port, int]
src_port: int
dst_port: int
class TCPFlags(Enum):
@@ -126,12 +162,12 @@ class TCPHeader(BaseModel):
:Example:
>>> tcp_header = TCPHeader(
... src_port=Port.HTTP_ALT,
... dst_port=Port.HTTP,
... src_port=Port["HTTP_ALT"],
... dst_port=Port["HTTP"],
... flags=[TCPFlags.SYN, TCPFlags.ACK]
... )
"""
src_port: Port
dst_port: Port
src_port: int
dst_port: int
flags: List[TCPFlags] = [TCPFlags.SYN]

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

@@ -90,8 +90,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["POSTGRES_SERVER"]
kwargs["protocol"] = IPProtocol["TCP"]
super().__init__(**kwargs)
def _init_request_manager(self) -> RequestManager:

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

@@ -24,8 +24,8 @@ class PortScanPayload(SimComponent):
"""
ip_address: IPV4Address
port: Port
protocol: IPProtocol
port: int
protocol: str
request: bool = True
def describe_state(self) -> Dict:
@@ -37,8 +37,8 @@ class PortScanPayload(SimComponent):
"""
state = super().describe_state()
state["ip_address"] = str(self.ip_address)
state["port"] = self.port.value
state["protocol"] = self.protocol.value
state["port"] = self.port
state["protocol"] = self.protocol
state["request"] = self.request
return state
@@ -64,8 +64,8 @@ class NMAP(Application, identifier="NMAP"):
def __init__(self, **kwargs):
kwargs["name"] = "NMAP"
kwargs["port"] = Port.NONE
kwargs["protocol"] = IPProtocol.NONE
kwargs["port"] = Port["NONE"]
kwargs["protocol"] = IPProtocol["NONE"]
super().__init__(**kwargs)
def _can_perform_network_action(self) -> bool:
@@ -218,14 +218,14 @@ class NMAP(Application, identifier="NMAP"):
print(table.get_string(sortby="IP Address"))
return active_nodes
def _determine_port_scan_type(self, target_ip_addresses: List[IPV4Address], target_ports: List[Port]) -> str:
def _determine_port_scan_type(self, target_ip_addresses: List[IPV4Address], target_ports: List[int]) -> str:
"""
Determine the type of port scan based on the number of target IP addresses and ports.
:param target_ip_addresses: The list of target IP addresses.
:type target_ip_addresses: List[IPV4Address]
:param target_ports: The list of target ports.
:type target_ports: List[Port]
:type target_ports: List[int]
:return: The type of port scan.
:rtype: str
@@ -238,8 +238,8 @@ class NMAP(Application, identifier="NMAP"):
def _check_port_open_on_ip_address(
self,
ip_address: IPv4Address,
port: Port,
protocol: IPProtocol,
port: int,
protocol: str,
is_re_attempt: bool = False,
port_scan_uuid: Optional[str] = None,
) -> bool:
@@ -251,7 +251,7 @@ class NMAP(Application, identifier="NMAP"):
:param port: The target port.
:type port: Port
:param protocol: The protocol used for the port scan.
:type protocol: IPProtocol
:type protocol: str
:param is_re_attempt: Flag indicating if this is a reattempt. Defaults to False.
:type is_re_attempt: bool
:param port_scan_uuid: The UUID of the port scan payload. Defaults to None.
@@ -272,8 +272,8 @@ class NMAP(Application, identifier="NMAP"):
payload = PortScanPayload(ip_address=ip_address, port=port, protocol=protocol)
self._active_port_scans[payload.uuid] = payload
self.sys_log.info(
f"{self.name}: Sending port scan request over {payload.protocol.name} on port {payload.port.value} "
f"({payload.port.name}) to {payload.ip_address}"
f"{self.name}: Sending port scan request over {payload.protocol} on port {payload.port} "
f"({payload.port}) to {payload.ip_address}"
)
self.software_manager.send_payload_to_session_manager(
payload=payload, dest_ip_address=ip_address, src_port=port, dest_port=port, ip_protocol=protocol
@@ -295,8 +295,8 @@ class NMAP(Application, identifier="NMAP"):
self._active_port_scans.pop(payload.uuid)
self._port_scan_responses[payload.uuid] = payload
self.sys_log.info(
f"{self.name}: Received port scan response from {payload.ip_address} on port {payload.port.value} "
f"({payload.port.name}) over {payload.protocol.name}"
f"{self.name}: Received port scan response from {payload.ip_address} on port {payload.port} "
f"({payload.port}) over {payload.protocol}"
)
def _process_port_scan_request(self, payload: PortScanPayload, session_id: str) -> None:
@@ -311,8 +311,8 @@ class NMAP(Application, identifier="NMAP"):
if self.software_manager.check_port_is_open(port=payload.port, protocol=payload.protocol):
payload.request = False
self.sys_log.info(
f"{self.name}: Responding to port scan request for port {payload.port.value} "
f"({payload.port.name}) over {payload.protocol.name}",
f"{self.name}: Responding to port scan request for port {payload.port} "
f"({payload.port}) over {payload.protocol}",
)
self.software_manager.send_payload_to_session_manager(payload=payload, session_id=session_id)
@@ -320,20 +320,20 @@ class NMAP(Application, identifier="NMAP"):
def port_scan(
self,
target_ip_address: Union[IPV4Address, List[IPV4Address], IPv4Network, List[IPv4Network]],
target_protocol: Optional[Union[IPProtocol, List[IPProtocol]]] = None,
target_port: Optional[Union[Port, List[Port]]] = None,
target_protocol: Optional[Union[str, List[str]]] = None,
target_port: Optional[Union[int, List[int]]] = None,
show: bool = True,
json_serializable: bool = False,
) -> Dict[IPv4Address, Dict[IPProtocol, List[Port]]]:
) -> Dict[IPv4Address, Dict[str, List[int]]]:
"""
Perform a port scan on the target IP address(es).
:param target_ip_address: The target IP address(es) or network(s) for the port scan.
:type target_ip_address: Union[IPV4Address, List[IPV4Address], IPv4Network, List[IPv4Network]]
:param target_protocol: The protocol(s) to use for the port scan. Defaults to None, which includes TCP and UDP.
:type target_protocol: Optional[Union[IPProtocol, List[IPProtocol]]]
:type target_protocol: Optional[Union[str, List[str]]]
:param target_port: The port(s) to scan. Defaults to None, which includes all valid ports.
:type target_port: Optional[Union[Port, List[Port]]]
:type target_port: Optional[Union[int, List[int]]]
:param show: Flag indicating whether to display the scan results. Defaults to True.
:type show: bool
:param json_serializable: Flag indicating whether the return value should be JSON serializable. Defaults to
@@ -341,19 +341,19 @@ class NMAP(Application, identifier="NMAP"):
:type json_serializable: bool
:return: A dictionary mapping IP addresses to protocols and lists of open ports.
:rtype: Dict[IPv4Address, Dict[IPProtocol, List[Port]]]
:rtype: Dict[IPv4Address, Dict[str, List[int]]]
"""
ip_addresses = self._explode_ip_address_network_array(target_ip_address)
if isinstance(target_port, Port):
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 if port not in {Port["NONE"], Port["UNUSED"]}]
if isinstance(target_protocol, IPProtocol):
if isinstance(target_protocol, str):
target_protocol = [target_protocol]
elif target_protocol is None:
target_protocol = [IPProtocol.TCP, IPProtocol.UDP]
target_protocol = [IPProtocol["TCP"], IPProtocol["UDP"]]
scan_type = self._determine_port_scan_type(list(ip_addresses), target_port)
active_ports = {}
@@ -372,10 +372,10 @@ class NMAP(Application, identifier="NMAP"):
if port_open:
if show:
table.add_row([ip_address, port.value, port.name, protocol.name])
table.add_row([ip_address, port, port, protocol])
_ip_address = ip_address if not json_serializable else str(ip_address)
_protocol = protocol if not json_serializable else protocol.value
_port = port if not json_serializable else port.value
_protocol = protocol if not json_serializable else protocol
_port = port if not json_serializable else port
if _ip_address not in active_ports:
active_ports[_ip_address] = dict()
if _protocol not in active_ports[_ip_address]:
@@ -390,12 +390,12 @@ class NMAP(Application, identifier="NMAP"):
def network_service_recon(
self,
target_ip_address: Union[IPV4Address, List[IPV4Address], IPv4Network, List[IPv4Network]],
target_protocol: Optional[Union[IPProtocol, List[IPProtocol]]] = None,
target_port: Optional[Union[Port, List[Port]]] = None,
target_protocol: Optional[Union[str, List[str]]] = None,
target_port: Optional[Union[int, List[int]]] = None,
show: bool = True,
show_online_only: bool = True,
json_serializable: bool = False,
) -> Dict[IPv4Address, Dict[IPProtocol, List[Port]]]:
) -> Dict[IPv4Address, Dict[str, List[int]]]:
"""
Perform a network service reconnaissance which includes a ping scan followed by a port scan.
@@ -408,9 +408,9 @@ class NMAP(Application, identifier="NMAP"):
:param target_ip_address: The target IP address(es) or network(s) for the port scan.
:type target_ip_address: Union[IPV4Address, List[IPV4Address], IPv4Network, List[IPv4Network]]
:param target_protocol: The protocol(s) to use for the port scan. Defaults to None, which includes TCP and UDP.
:type target_protocol: Optional[Union[IPProtocol, List[IPProtocol]]]
:type target_protocol: Optional[Union[str, List[str]]]
:param target_port: The port(s) to scan. Defaults to None, which includes all valid ports.
:type target_port: Optional[Union[Port, List[Port]]]
:type target_port: Optional[Union[int, List[int]]]
:param show: Flag indicating whether to display the scan results. Defaults to True.
:type show: bool
:param show_online_only: Flag indicating whether to show only the online hosts. Defaults to True.
@@ -420,7 +420,7 @@ class NMAP(Application, identifier="NMAP"):
:type json_serializable: bool
:return: A dictionary mapping IP addresses to protocols and lists of open ports.
:rtype: Dict[IPv4Address, Dict[IPProtocol, List[Port]]]
:rtype: Dict[IPv4Address, Dict[str, List[int]]]
"""
ping_scan_results = self.ping_scan(
target_ip_address=target_ip_address, show=show, show_online_only=show_online_only, json_serializable=False

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

@@ -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: IPProtocol = Field(default=IPProtocol.TCP)
masquerade_protocol: str = Field(default=IPProtocol["TCP"])
"""The currently chosen protocol that the C2 traffic is masquerading as. Defaults as TCP."""
masquerade_port: Port = Field(default=Port.HTTP)
masquerade_port: int = Field(default=Port["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["HTTP"], Port["FTP"], Port["DNS"]}
kwargs["port"] = Port["NONE"]
kwargs["protocol"] = IPProtocol["TCP"]
super().__init__(**kwargs)
@property
@@ -367,7 +367,7 @@ class AbstractC2(Application, identifier="AbstractC2"):
:rtype: bool
"""
# Validating that they are valid Enums.
if not isinstance(payload.masquerade_port, Port) or not isinstance(payload.masquerade_protocol, IPProtocol):
if not isinstance(payload.masquerade_port, int) or not isinstance(payload.masquerade_protocol, str):
self.sys_log.warning(
f"{self.name}: Received invalid Masquerade Values within Keep Alive."
f"Port: {payload.masquerade_port} Protocol: {payload.masquerade_protocol}."
@@ -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["HTTP"]
self.c2_config.masquerade_protocol = IPProtocol["TCP"]
@abstractmethod
def _confirm_remote_connection(self, timestep: int) -> bool:

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

@@ -130,8 +130,8 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"):
self,
c2_server_ip_address: IPv4Address = None,
keep_alive_frequency: int = 5,
masquerade_protocol: Enum = IPProtocol.TCP,
masquerade_port: Enum = Port.HTTP,
masquerade_protocol: str = IPProtocol["TCP"],
masquerade_port: int = Port["HTTP"],
) -> bool:
"""
Configures the C2 beacon to communicate with the C2 server.

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

@@ -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["NONE"]
kwargs["protocol"] = IPProtocol["NONE"]
super().__init__(**kwargs)
self._db_connection: Optional[DatabaseClientConnection] = None

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

@@ -35,7 +35,7 @@ class DoSBot(DatabaseClient, identifier="DoSBot"):
target_ip_address: Optional[IPv4Address] = None
"""IP address of the target service."""
target_port: Optional[Port] = None
target_port: Optional[int] = None
"""Port of the target service."""
payload: Optional[str] = None
@@ -94,7 +94,7 @@ class DoSBot(DatabaseClient, identifier="DoSBot"):
def configure(
self,
target_ip_address: IPv4Address,
target_port: Optional[Port] = Port.POSTGRES_SERVER,
target_port: Optional[int] = Port["POSTGRES_SERVER"],
payload: Optional[str] = None,
repeat: bool = False,
port_scan_p_of_success: float = 0.1,
@@ -105,7 +105,7 @@ class DoSBot(DatabaseClient, identifier="DoSBot"):
Configure the Denial of Service bot.
:param: target_ip_address: The IP address of the Node containing the target service.
:param: target_port: The port of the target service. Optional - Default is `Port.HTTP`
:param: target_port: The port of the target service. Optional - Default is `Port["HTTP"]`
:param: payload: The payload the DoS Bot will throw at the target service. Optional - Default is `None`
:param: repeat: If True, the bot will maintain the attack. Optional - Default is `True`
:param: port_scan_p_of_success: The chance of the port scan being successful. Optional - Default is 0.1 (10%)

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

@@ -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["NONE"]
kwargs["protocol"] = IPProtocol["NONE"]
super().__init__(**kwargs)
self._db_connection: Optional[DatabaseClientConnection] = None

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

@@ -43,10 +43,10 @@ class WebBrowser(Application, identifier="WebBrowser"):
def __init__(self, **kwargs):
kwargs["name"] = "WebBrowser"
kwargs["protocol"] = IPProtocol.TCP
kwargs["protocol"] = IPProtocol["TCP"]
# default for web is port 80
if kwargs.get("port") is None:
kwargs["port"] = Port.HTTP
kwargs["port"] = Port["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["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[Port] = Port.HTTP,
dest_port: Optional[int] = Port["HTTP"],
session_id: Optional[str] = None,
**kwargs,
) -> bool:

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

@@ -34,14 +34,14 @@ class Session(SimComponent):
:param connected: A flag indicating whether the session is connected.
"""
protocol: IPProtocol
protocol: str
with_ip_address: IPv4Address
src_port: Optional[Port]
dst_port: Optional[Port]
src_port: Optional[int]
dst_port: Optional[int]
connected: bool = False
@classmethod
def from_session_key(cls, session_key: Tuple[IPProtocol, IPv4Address, Optional[Port], Optional[Port]]) -> Session:
def from_session_key(cls, session_key: Tuple[str, IPv4Address, Optional[int], Optional[int]]) -> Session:
"""
Create a Session instance from a session key tuple.
@@ -77,7 +77,7 @@ class SessionManager:
def __init__(self, sys_log: SysLog):
self.sessions_by_key: Dict[
Tuple[IPProtocol, IPv4Address, IPv4Address, Optional[Port], Optional[Port]], Session
Tuple[str, IPv4Address, IPv4Address, Optional[int], Optional[int]], Session
] = {}
self.sessions_by_uuid: Dict[str, Session] = {}
self.sys_log: SysLog = sys_log
@@ -103,7 +103,7 @@ class SessionManager:
@staticmethod
def _get_session_key(
frame: Frame, inbound_frame: bool = True
) -> Tuple[IPProtocol, IPv4Address, Optional[Port], Optional[Port]]:
) -> Tuple[str, IPv4Address, Optional[int], Optional[int]]:
"""
Extracts the session key from the given frame.
@@ -111,15 +111,15 @@ class SessionManager:
- IPProtocol: The transport protocol (e.g. TCP, UDP, ICMP).
- IPv4Address: The source IP address.
- IPv4Address: The destination IP address.
- Optional[Port]: The source port number (if applicable).
- Optional[Port]: The destination port number (if applicable).
- Optional[int]: The source port number (if applicable).
- Optional[int]: The destination port number (if applicable).
:param frame: The network frame from which to extract the session key.
:return: A tuple containing the session key.
"""
protocol = frame.ip.protocol
with_ip_address = frame.ip.src_ip_address
if protocol == IPProtocol.TCP:
if protocol == IPProtocol["TCP"]:
if inbound_frame:
src_port = frame.tcp.src_port
dst_port = frame.tcp.dst_port
@@ -127,7 +127,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 == IPProtocol["UDP"]:
if inbound_frame:
src_port = frame.udp.src_port
dst_port = frame.udp.dst_port
@@ -167,17 +167,17 @@ class SessionManager:
def resolve_outbound_transmission_details(
self,
dst_ip_address: Optional[Union[IPv4Address, IPv4Network]] = None,
src_port: Optional[Port] = None,
dst_port: Optional[Port] = None,
protocol: Optional[IPProtocol] = None,
src_port: Optional[int] = None,
dst_port: Optional[int] = None,
protocol: Optional[str] = None,
session_id: Optional[str] = None,
) -> Tuple[
Optional["NetworkInterface"],
Optional[str],
IPv4Address,
Optional[Port],
Optional[Port],
Optional[IPProtocol],
Optional[int],
Optional[int],
Optional[str],
bool,
]:
"""
@@ -196,19 +196,19 @@ class SessionManager:
treats the transmission as a broadcast to that network. Optional.
:type dst_ip_address: Optional[Union[IPv4Address, IPv4Network]]
:param src_port: The source port number for the transmission. Optional.
:type src_port: Optional[Port]
:type src_port: Optional[int]
:param dst_port: The destination port number for the transmission. Optional.
:type dst_port: Optional[Port]
:type dst_port: Optional[int]
:param protocol: The IP protocol to be used for the transmission. Optional.
:type protocol: Optional[IPProtocol]
:type protocol: Optional[str]
:param session_id: The session ID associated with the transmission. If provided, the session details override
other parameters. Optional.
:type session_id: Optional[str]
:return: A tuple containing the resolved outbound network interface, destination MAC address, destination IP
address, source port, destination port, protocol, and a boolean indicating whether the transmission is a
broadcast.
:rtype: Tuple[Optional["NetworkInterface"], Optional[str], IPv4Address, Optional[Port], Optional[Port],
Optional[IPProtocol], bool]
:rtype: Tuple[Optional["NetworkInterface"], Optional[str], IPv4Address, Optional[int], Optional[int],
Optional[str], bool]
"""
if dst_ip_address and not isinstance(dst_ip_address, (IPv4Address, IPv4Network)):
dst_ip_address = IPv4Address(dst_ip_address)
@@ -259,10 +259,10 @@ class SessionManager:
self,
payload: Any,
dst_ip_address: Optional[Union[IPv4Address, IPv4Network]] = None,
src_port: Optional[Port] = None,
dst_port: Optional[Port] = None,
src_port: Optional[int] = None,
dst_port: Optional[int] = None,
session_id: Optional[str] = None,
ip_protocol: IPProtocol = IPProtocol.TCP,
ip_protocol: str = IPProtocol["TCP"],
icmp_packet: Optional[ICMPPacket] = None,
) -> Union[Any, None]:
"""
@@ -286,7 +286,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 = IPProtocol["UDP"]
else:
vals = self.resolve_outbound_transmission_details(
dst_ip_address=dst_ip_address,
@@ -311,26 +311,26 @@ class SessionManager:
if not outbound_network_interface or not dst_mac_address:
return False
if not (src_port or dst_port):
if src_port is None and dst_port is None:
raise ValueError(
"Failed to resolve src or dst port. Have you sent the port from the service or application?"
)
tcp_header = None
udp_header = None
if ip_protocol == IPProtocol.TCP:
if ip_protocol == IPProtocol["TCP"]:
tcp_header = TCPHeader(
src_port=dst_port,
dst_port=dst_port,
)
elif ip_protocol == IPProtocol.UDP:
elif ip_protocol == IPProtocol["UDP"]:
udp_header = UDPHeader(
src_port=dst_port,
dst_port=dst_port,
)
# TODO: Only create IP packet if not ARP
# ip_packet = None
# if dst_port != Port.ARP:
# if dst_port != Port["ARP"]:
# IPPacket(
# src_ip_address=outbound_network_interface.ip_address,
# dst_ip_address=dst_ip_address,
@@ -387,7 +387,7 @@ class SessionManager:
elif frame.udp:
dst_port = frame.udp.dst_port
elif frame.icmp:
dst_port = Port.NONE
dst_port = Port["NONE"]
self.software_manager.receive_payload_from_session_manager(
payload=frame.payload,
port=dst_port,
@@ -413,5 +413,5 @@ class SessionManager:
table.align = "l"
table.title = f"{self.sys_log.hostname} Session Manager"
for session in self.sessions_by_key.values():
table.add_row([session.dst_ip_address, session.dst_port.value, session.protocol.name])
table.add_row([session.dst_ip_address, session.dst_port, session.protocol])
print(table)

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

@@ -52,7 +52,7 @@ class SoftwareManager:
self.session_manager = session_manager
self.software: Dict[str, Union[Service, Application]] = {}
self._software_class_to_name_map: Dict[Type[IOSoftware], str] = {}
self.port_protocol_mapping: Dict[Tuple[Port, IPProtocol], Union[Service, Application]] = {}
self.port_protocol_mapping: Dict[Tuple[int, str], Union[Service, Application]] = {}
self.sys_log: SysLog = sys_log
self.file_system: FileSystem = file_system
self.dns_server: Optional[IPv4Address] = dns_server
@@ -67,7 +67,7 @@ class SoftwareManager:
"""Provides access to the ICMP service instance, if installed."""
return self.software.get("ICMP") # noqa
def get_open_ports(self) -> List[Port]:
def get_open_ports(self) -> List[int]:
"""
Get a list of open ports.
@@ -81,7 +81,7 @@ class SoftwareManager:
open_ports += list(software.listen_on_ports)
return open_ports
def check_port_is_open(self, port: Port, protocol: IPProtocol) -> bool:
def check_port_is_open(self, port: int, protocol: str) -> bool:
"""
Check if a specific port is open and running a service using the specified protocol.
@@ -93,7 +93,7 @@ class SoftwareManager:
:param port: The port to check.
:type port: Port
:param protocol: The protocol to check (e.g., TCP, UDP).
:type protocol: IPProtocol
:type protocol: str
:return: True if the port is open and a service is running on it using the specified protocol, False otherwise.
:rtype: bool
"""
@@ -189,9 +189,9 @@ class SoftwareManager:
self,
payload: Any,
dest_ip_address: Optional[Union[IPv4Address, IPv4Network]] = None,
src_port: Optional[Port] = None,
dest_port: Optional[Port] = None,
ip_protocol: IPProtocol = IPProtocol.TCP,
src_port: Optional[int] = None,
dest_port: Optional[int] = None,
ip_protocol: str = IPProtocol["TCP"],
session_id: Optional[str] = None,
) -> bool:
"""
@@ -219,8 +219,8 @@ class SoftwareManager:
def receive_payload_from_session_manager(
self,
payload: Any,
port: Port,
protocol: IPProtocol,
port: int,
protocol: str,
session_id: str,
from_network_interface: "NIC",
frame: Frame,
@@ -275,8 +275,8 @@ class SoftwareManager:
software_type,
software.operating_state.name,
software.health_state_actual.name,
software.port.value if software.port != Port.NONE else None,
software.protocol.value,
software.port if software.port != Port["NONE"] else None,
software.protocol,
]
)
print(table)

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

@@ -26,8 +26,8 @@ class ARP(Service):
def __init__(self, **kwargs):
kwargs["name"] = "ARP"
kwargs["port"] = Port.ARP
kwargs["protocol"] = IPProtocol.UDP
kwargs["port"] = Port["ARP"]
kwargs["protocol"] = IPProtocol["UDP"]
super().__init__(**kwargs)
def describe_state(self) -> Dict:

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

@@ -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["POSTGRES_SERVER"]
kwargs["protocol"] = IPProtocol["TCP"]
super().__init__(**kwargs)
self._create_db_file()

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

@@ -22,11 +22,11 @@ class DNSClient(Service):
def __init__(self, **kwargs):
kwargs["name"] = "DNSClient"
kwargs["port"] = Port.DNS
kwargs["port"] = Port["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"] = IPProtocol["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["DNS"]
)
# recursively re-call the function passing is_reattempt=True
@@ -110,7 +110,7 @@ class DNSClient(Service):
payload: DNSPacket,
session_id: Optional[str] = None,
dest_ip_address: Optional[IPv4Address] = None,
dest_port: Optional[Port] = None,
dest_port: Optional[int] = None,
**kwargs,
) -> bool:
"""

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

@@ -21,11 +21,11 @@ class DNSServer(Service):
def __init__(self, **kwargs):
kwargs["name"] = "DNSServer"
kwargs["port"] = Port.DNS
kwargs["port"] = Port["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"] = IPProtocol["TCP"]
super().__init__(**kwargs)
self.start()

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

@@ -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["FTP"]
kwargs["protocol"] = IPProtocol["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[Port] = Port.FTP,
dest_port: Optional[int] = Port["FTP"],
session_id: Optional[str] = None,
is_reattempt: Optional[bool] = False,
) -> bool:
@@ -114,7 +114,7 @@ class FTPClient(FTPServiceABC):
:param: dest_ip_address: IP address of the FTP server the client needs to connect to. Optional.
:type: dest_ip_address: Optional[IPv4Address]
:param: dest_port: Port of the FTP server the client needs to connect to. Optional.
:type: dest_port: Optional[Port]
:type: dest_port: Optional[int]
:param: is_reattempt: Set to True if attempt to connect to FTP Server has been attempted. Default False.
:type: is_reattempt: Optional[bool]
"""
@@ -124,13 +124,13 @@ 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["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:
self.sys_log.info(
f"{self.name}: Successfully connected to FTP Server "
f"{dest_ip_address} via port {payload.ftp_command_args.value}"
f"{dest_ip_address} via port {payload.ftp_command_args}"
)
self.add_connection(connection_id="server_connection", session_id=session_id)
return True
@@ -139,7 +139,7 @@ class FTPClient(FTPServiceABC):
# reattempt failed
self.sys_log.warning(
f"{self.name}: Unable to connect to FTP Server "
f"{dest_ip_address} via port {payload.ftp_command_args.value}"
f"{dest_ip_address} via port {payload.ftp_command_args}"
)
return False
else:
@@ -152,7 +152,7 @@ class FTPClient(FTPServiceABC):
return False
def _disconnect_from_server(
self, dest_ip_address: Optional[IPv4Address] = None, dest_port: Optional[Port] = Port.FTP
self, dest_ip_address: Optional[IPv4Address] = None, dest_port: Optional[int] = Port["FTP"]
) -> bool:
"""
Connects the client from a given FTP server.
@@ -160,7 +160,7 @@ class FTPClient(FTPServiceABC):
:param: dest_ip_address: IP address of the FTP server the client needs to disconnect from. Optional.
:type: dest_ip_address: Optional[IPv4Address]
:param: dest_port: Port of the FTP server the client needs to disconnect from. Optional.
:type: dest_port: Optional[Port]
:type: dest_port: Optional[int]
:param: is_reattempt: Set to True if attempt to disconnect from FTP Server has been attempted. Default False.
:type: is_reattempt: Optional[bool]
"""
@@ -179,7 +179,7 @@ class FTPClient(FTPServiceABC):
src_file_name: str,
dest_folder_name: str,
dest_file_name: str,
dest_port: Optional[Port] = Port.FTP,
dest_port: Optional[int] = Port["FTP"],
session_id: Optional[str] = None,
) -> bool:
"""
@@ -203,8 +203,8 @@ class FTPClient(FTPServiceABC):
:param: dest_file_name: The name of the file to be saved on the FTP Server.
:type: dest_file_name: str
:param: dest_port: The open port of the machine that hosts the FTP Server. Default is Port.FTP.
:type: dest_port: Optional[Port]
:param: dest_port: The open port of the machine that hosts the FTP Server. Default is Port["FTP"].
:type: dest_port: Optional[int]
:param: session_id: The id of the session
:type: session_id: Optional[str]
@@ -241,7 +241,7 @@ class FTPClient(FTPServiceABC):
src_file_name: str,
dest_folder_name: str,
dest_file_name: str,
dest_port: Optional[Port] = Port.FTP,
dest_port: Optional[int] = Port["FTP"],
) -> bool:
"""
Request a file from a target IP address.
@@ -263,8 +263,8 @@ class FTPClient(FTPServiceABC):
:param: dest_file_name: The name of the file to be saved on the FTP Server.
:type: dest_file_name: str
:param: dest_port: The open port of the machine that hosts the FTP Server. Default is Port.FTP.
:type: dest_port: Optional[Port]
:param: dest_port: The open port of the machine that hosts the FTP Server. Default is Port["FTP"].
:type: dest_port: Optional[int]
"""
# check if FTP is currently connected to IP
self._connect_to_server(dest_ip_address=dest_ip_address, dest_port=dest_port)

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

@@ -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["FTP"]
kwargs["protocol"] = IPProtocol["TCP"]
super().__init__(**kwargs)
self.start()
@@ -52,7 +52,7 @@ class FTPServer(FTPServiceABC):
# process server specific commands, otherwise call super
if payload.ftp_command == FTPCommand.PORT:
# check that the port is valid
if isinstance(payload.ftp_command_args, Port) and payload.ftp_command_args.value in range(0, 65535):
if isinstance(payload.ftp_command_args, int) and (0 <= payload.ftp_command_args < 65535):
# return successful connection
self.add_connection(connection_id=session_id, session_id=session_id)
payload.status_code = FTPStatusCode.OK

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

@@ -78,7 +78,7 @@ class FTPServiceABC(Service, ABC):
dest_folder_name: str,
dest_file_name: str,
dest_ip_address: Optional[IPv4Address] = None,
dest_port: Optional[Port] = None,
dest_port: Optional[int] = None,
session_id: Optional[str] = None,
is_response: bool = False,
) -> bool:
@@ -97,8 +97,8 @@ class FTPServiceABC(Service, ABC):
:param: dest_ip_address: The IP address of the machine that hosts the FTP Server.
:type: dest_ip_address: Optional[IPv4Address]
:param: dest_port: The open port of the machine that hosts the FTP Server. Default is Port.FTP.
:type: dest_port: Optional[Port]
:param: dest_port: The open port of the machine that hosts the FTP Server. Default is Port["FTP"].
:type: dest_port: Optional[int]
:param: session_id: session ID linked to the FTP Packet. Optional.
:type: session_id: Optional[str]
@@ -168,7 +168,7 @@ class FTPServiceABC(Service, ABC):
payload: FTPPacket,
session_id: Optional[str] = None,
dest_ip_address: Optional[IPv4Address] = None,
dest_port: Optional[Port] = None,
dest_port: Optional[int] = None,
**kwargs,
) -> bool:
"""

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

@@ -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["NONE"]
kwargs["protocol"] = IPProtocol["ICMP"]
super().__init__(**kwargs)
def describe_state(self) -> Dict:

4
src/primaite/simulator/system/services/icmp/router_icmp.py
View File

@@ -36,13 +36,13 @@
# self.sys_log.info(f"Received echo request from {frame.ip.src_ip_address}")
# target_mac_address = self.arp.get_arp_cache_mac_address(frame.ip.src_ip_address)
# src_nic = self.arp.get_arp_cache_network_interface(frame.ip.src_ip_address)
# tcp_header = TCPHeader(src_port=Port.ARP, dst_port=Port.ARP)
# tcp_header = TCPHeader(src_port=Port["ARP"], dst_port=Port["ARP"])
#
# # Network Layer
# ip_packet = IPPacket(
# src_ip_address=network_interface.ip_address,
# dst_ip_address=frame.ip.src_ip_address,
# protocol=IPProtocol.ICMP,
# protocol=IPProtocol["ICMP"],
# )
# # Data Link Layer
# ethernet_header = EthernetHeader(

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

@@ -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["NTP"]
kwargs["protocol"] = IPProtocol["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: Port = Port.NTP,
dest_port: int = Port["NTP"],
**kwargs,
) -> bool:
"""Requests NTP data from NTP server.

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

@@ -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["NTP"]
kwargs["protocol"] = IPProtocol["UDP"]
super().__init__(**kwargs)
self.start()

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

@@ -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["SSH"]
kwargs["protocol"] = IPProtocol["TCP"]
super().__init__(**kwargs)
def describe_state(self) -> Dict:

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

@@ -49,10 +49,10 @@ class WebServer(Service):
def __init__(self, **kwargs):
kwargs["name"] = "WebServer"
kwargs["protocol"] = IPProtocol.TCP
kwargs["protocol"] = IPProtocol["TCP"]
# default for web is port 80
if kwargs.get("port") is None:
kwargs["port"] = Port.HTTP
kwargs["port"] = Port["HTTP"]
super().__init__(**kwargs)
self._install_web_files()
@@ -145,7 +145,7 @@ class WebServer(Service):
payload: HttpResponsePacket,
session_id: Optional[str] = None,
dest_ip_address: Optional[IPv4Address] = None,
dest_port: Optional[Port] = None,
dest_port: Optional[int] = None,
**kwargs,
) -> bool:
"""

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

@@ -251,11 +251,11 @@ class IOSoftware(Software):
"Indicates if the software uses TCP protocol for communication. Default is True."
udp: bool = True
"Indicates if the software uses UDP protocol for communication. Default is True."
port: Port
port: int
"The port to which the software is connected."
listen_on_ports: Set[Port] = Field(default_factory=set)
listen_on_ports: Set[int] = Field(default_factory=set)
"The set of ports to listen on."
protocol: IPProtocol
protocol: str
"The IP Protocol the Software operates on."
_connections: Dict[str, Dict] = {}
"Active connections."
@@ -277,7 +277,7 @@ class IOSoftware(Software):
"max_sessions": self.max_sessions,
"tcp": self.tcp,
"udp": self.udp,
"port": self.port.value,
"port": self.port,
}
)
return state
@@ -386,8 +386,8 @@ class IOSoftware(Software):
payload: Any,
session_id: Optional[str] = None,
dest_ip_address: Optional[Union[IPv4Address, IPv4Network]] = None,
dest_port: Optional[Port] = None,
ip_protocol: IPProtocol = IPProtocol.TCP,
dest_port: Optional[int] = None,
ip_protocol: str = IPProtocol["TCP"],
**kwargs,
) -> bool:
"""