#1916: Setting up a connected states + added tests + error states for if service is interacted with when not running

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

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,
+        self, dest_ip_address: Optional[IPv4Address] = None, dest_port: Optional[Port] = Port.FTP
-        ftp_server_ip_address: Optional[IPv4Address] = None,
     ) -> bool:
         # send a disconnect request payload to FTP server
-        # return true if connected successfully else false
+        payload: FTPPacket = FTPPacket(ftp_command=FTPCommand.QUIT)
-        self.connected = False
+        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.")

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

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

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")