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