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

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]