From 2c234ab67abbc18ca0dc22cb28773b2e2ae1c0e6 Mon Sep 17 00:00:00 2001 From: "Czar.Echavez" Date: Fri, 22 Sep 2023 15:38:01 +0100 Subject: [PATCH] #1916: Setting up a connected states + added tests + error states for if service is interacted with when not running --- .../simulator/network/protocols/ftp.py | 4 +- .../system/services/ftp/ftp_client.py | 40 +++++++++++++++--- .../system/services/ftp/ftp_server.py | 41 ++++++++++++++++++- .../system/services/ftp/ftp_service.py | 7 ---- .../_simulator/_system/_services/test_ftp.py | 41 +++++++++++++++++++ 5 files changed, 117 insertions(+), 16 deletions(-) diff --git a/src/primaite/simulator/network/protocols/ftp.py b/src/primaite/simulator/network/protocols/ftp.py index ab277045..91080219 100644 --- a/src/primaite/simulator/network/protocols/ftp.py +++ b/src/primaite/simulator/network/protocols/ftp.py @@ -1,5 +1,5 @@ from enum import Enum -from typing import Any +from typing import Any, Optional from primaite.simulator.network.protocols.packet import DataPacket @@ -48,7 +48,7 @@ class FTPPacket(DataPacket): ftp_command: FTPCommand """Command type of the packet.""" - ftp_command_args: Any + ftp_command_args: Optional[Any] = None """Arguments for command.""" status_code: FTPStatusCode = None diff --git a/src/primaite/simulator/system/services/ftp/ftp_client.py b/src/primaite/simulator/system/services/ftp/ftp_client.py index 8e93df1b..0e8f3dce 100644 --- a/src/primaite/simulator/system/services/ftp/ftp_client.py +++ b/src/primaite/simulator/system/services/ftp/ftp_client.py @@ -7,6 +7,7 @@ from primaite.simulator.network.transmission.network_layer import IPProtocol from primaite.simulator.network.transmission.transport_layer import Port from primaite.simulator.system.core.software_manager import SoftwareManager from primaite.simulator.system.services.ftp.ftp_service import FTPServiceABC +from primaite.simulator.system.services.service import ServiceOperatingState class FTPClient(FTPServiceABC): @@ -27,6 +28,15 @@ class FTPClient(FTPServiceABC): super().__init__(**kwargs) self.start() + def _process_ftp_command(self, payload: FTPPacket, session_id: Optional[str] = None, **kwargs) -> FTPPacket: + # if server is down, return error + if self.operating_state != ServiceOperatingState.RUNNING: + payload.status_code = FTPStatusCode.ERROR + return payload + + # process client specific commands, otherwise call super + return super()._process_ftp_command(payload=payload, session_id=session_id, **kwargs) + def _connect_to_server( self, dest_ip_address: Optional[IPv4Address] = None, dest_port: Optional[Port] = Port.FTP ) -> bool: @@ -54,12 +64,18 @@ class FTPClient(FTPServiceABC): return payload.status_code == FTPStatusCode.OK def _disconnect_from_server( - self, - ftp_server_ip_address: Optional[IPv4Address] = None, + self, dest_ip_address: Optional[IPv4Address] = None, dest_port: Optional[Port] = Port.FTP ) -> bool: # send a disconnect request payload to FTP server - # return true if connected successfully else false - self.connected = False + payload: FTPPacket = FTPPacket(ftp_command=FTPCommand.QUIT) + software_manager: SoftwareManager = self.software_manager + software_manager.send_payload_to_session_manager( + payload=payload, dest_ip_address=dest_ip_address, dest_port=dest_port + ) + if payload.status_code == FTPStatusCode.OK: + self.connected = False + return True + return False def send_file( self, @@ -72,6 +88,10 @@ class FTPClient(FTPServiceABC): is_reattempt: Optional[bool] = False, ) -> bool: """Send a file to a target IP address.""" + # if service is not running, return error + if self.operating_state != ServiceOperatingState.RUNNING: + self.sys_log.error(f"FTPClient not running for {self.sys_log.hostname}") + return False file_to_transfer: File = self.file_system.get_file(folder_name=src_folder_name, file_name=src_file_name) if not file_to_transfer: self.sys_log.error(f"Unable to send file that does not exist: {src_folder_name}/{src_file_name}") @@ -98,7 +118,7 @@ class FTPClient(FTPServiceABC): ) else: # send STOR request - return self._send_data( + self._send_data( file=file_to_transfer, dest_folder_name=dest_folder_name, dest_file_name=dest_file_name, @@ -106,6 +126,9 @@ class FTPClient(FTPServiceABC): dest_port=dest_port, ) + # send disconnect + return self._disconnect_from_server(dest_ip_address=dest_ip_address, dest_port=dest_port) + def request_file( self, dest_ip_address: IPv4Address, @@ -117,6 +140,10 @@ class FTPClient(FTPServiceABC): is_reattempt: Optional[bool] = False, ) -> bool: """Request a file from a target IP address.""" + # if service is not running, return error + if self.operating_state != ServiceOperatingState.RUNNING: + self.sys_log.error(f"FTPClient not running for {self.sys_log.hostname}") + return False # check if FTP is currently connected to IP self.connected = self._connect_to_server( dest_ip_address=dest_ip_address, @@ -152,6 +179,9 @@ class FTPClient(FTPServiceABC): payload=payload, dest_ip_address=dest_ip_address, dest_port=dest_port ) + # send disconnect + self._disconnect_from_server(dest_ip_address=dest_ip_address, dest_port=dest_port) + # the payload should have ok status code if payload.status_code == FTPStatusCode.OK: self.sys_log.info(f"File {src_folder_name}/{src_file_name} found in FTP server.") diff --git a/src/primaite/simulator/system/services/ftp/ftp_server.py b/src/primaite/simulator/system/services/ftp/ftp_server.py index c575479a..6371d53a 100644 --- a/src/primaite/simulator/system/services/ftp/ftp_server.py +++ b/src/primaite/simulator/system/services/ftp/ftp_server.py @@ -1,9 +1,12 @@ -from typing import Any, Optional +from ipaddress import IPv4Address +from typing import Any, Dict, Optional -from primaite.simulator.network.protocols.ftp import FTPPacket +from primaite.simulator.network.protocols.ftp import FTPCommand, FTPPacket, FTPStatusCode from primaite.simulator.network.transmission.network_layer import IPProtocol from primaite.simulator.network.transmission.transport_layer import Port +from primaite.simulator.system.core.session_manager import Session from primaite.simulator.system.services.ftp.ftp_service import FTPServiceABC +from primaite.simulator.system.services.service import ServiceOperatingState class FTPServer(FTPServiceABC): @@ -17,6 +20,9 @@ class FTPServer(FTPServiceABC): server_password: Optional[str] = None """Password needed to connect to FTP server. Default is None.""" + connections: Dict[str, IPv4Address] = {} + """Current active connections to the FTP server.""" + def __init__(self, **kwargs): kwargs["name"] = "FTPServer" kwargs["port"] = Port.FTP @@ -24,6 +30,37 @@ class FTPServer(FTPServiceABC): super().__init__(**kwargs) self.start() + def _get_session_details(self, session_id: str) -> Session: + """ + Returns the Session object from the given session id. + + :param: session_id: ID of the session that needs details retrieved + """ + return self.software_manager.session_manager.sessions_by_uuid[session_id] + + def _process_ftp_command(self, payload: FTPPacket, session_id: Optional[str] = None, **kwargs) -> FTPPacket: + # if server is down, return error + if self.operating_state != ServiceOperatingState.RUNNING: + payload.status_code = FTPStatusCode.ERROR + return payload + + # 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): + # return successful connection + session_details = self._get_session_details(session_id) + self.connections[session_id] = session_details.with_ip_address + payload.status_code = FTPStatusCode.OK + return payload + + if payload.ftp_command == FTPCommand.QUIT: + session_details = self._get_session_details(session_id) + self.connections.pop(session_id) + payload.status_code = FTPStatusCode.OK + + return super()._process_ftp_command(payload=payload, session_id=session_id, **kwargs) + def receive(self, payload: Any, session_id: Optional[str] = None, **kwargs) -> bool: """Receives a payload from the SessionManager.""" if not isinstance(payload, FTPPacket): diff --git a/src/primaite/simulator/system/services/ftp/ftp_service.py b/src/primaite/simulator/system/services/ftp/ftp_service.py index 6214c510..a41d647c 100644 --- a/src/primaite/simulator/system/services/ftp/ftp_service.py +++ b/src/primaite/simulator/system/services/ftp/ftp_service.py @@ -17,13 +17,6 @@ class FTPServiceABC(Service, ABC): """ def _process_ftp_command(self, payload: FTPPacket, session_id: Optional[str] = None, **kwargs) -> FTPPacket: - # handle PORT request - 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): - # return successful connection - payload.status_code = FTPStatusCode.OK - # handle STOR request if payload.ftp_command == FTPCommand.STOR: # check that the file is created in the computed hosting the FTP server diff --git a/tests/unit_tests/_primaite/_simulator/_system/_services/test_ftp.py b/tests/unit_tests/_primaite/_simulator/_system/_services/test_ftp.py index 64013207..ea563a88 100644 --- a/tests/unit_tests/_primaite/_simulator/_system/_services/test_ftp.py +++ b/tests/unit_tests/_primaite/_simulator/_system/_services/test_ftp.py @@ -3,6 +3,7 @@ from ipaddress import IPv4Address import pytest from primaite.simulator.network.hardware.base import Node +from primaite.simulator.network.protocols.ftp import FTPCommand, FTPPacket from primaite.simulator.network.transmission.network_layer import IPProtocol from primaite.simulator.network.transmission.transport_layer import Port from primaite.simulator.system.services.ftp.ftp_client import FTPClient @@ -39,3 +40,43 @@ def test_create_ftp_client(ftp_client): assert ftp_client_service.name is "FTPClient" assert ftp_client_service.port is Port.FTP assert ftp_client_service.protocol is IPProtocol.TCP + + +def test_ftp_server_store_file(ftp_server): + """Test to make sure the FTP Server knows how to deal with request responses.""" + assert ftp_server.file_system.get_file(folder_name="downloads", file_name="file.txt") is None + + response: FTPPacket = FTPPacket( + ftp_command=FTPCommand.STOR, + ftp_command_args={ + "dest_folder_name": "downloads", + "dest_file_name": "file.txt", + "file_size": 24, + }, + packet_payload_size=24, + ) + + ftp_server_service: FTPServer = ftp_server.software_manager.software["FTPServer"] + ftp_server_service.receive(response) + + assert ftp_server.file_system.get_file(folder_name="downloads", file_name="file.txt") + + +def test_ftp_client_store_file(ftp_client): + """Test to make sure the FTP Client knows how to deal with request responses.""" + assert ftp_client.file_system.get_file(folder_name="downloads", file_name="file.txt") is None + + response: FTPPacket = FTPPacket( + ftp_command=FTPCommand.STOR, + ftp_command_args={ + "dest_folder_name": "downloads", + "dest_file_name": "file.txt", + "file_size": 24, + }, + packet_payload_size=24, + ) + + ftp_client_service: FTPClient = ftp_client.software_manager.software["FTPClient"] + ftp_client_service.receive(response) + + assert ftp_client.file_system.get_file(folder_name="downloads", file_name="file.txt")