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/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:
"""