PrimAITE/src/primaite/simulator/system/services/ftp/ftp_service.py

from abc import ABC
from ipaddress import IPv4Address
from typing import Dict, Optional

from primaite.simulator.file_system.file_system import File
from primaite.simulator.network.protocols.ftp import FTPCommand, FTPPacket, FTPStatusCode
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.system.services.service import Service


class FTPServiceABC(Service, ABC):
    """
    Abstract Base Class for FTP Client and Service.

    Contains shared methods between both classes.
    """

    def describe_state(self) -> Dict:
        """Returns a Dict of the FTPService state."""
        return super().describe_state()

    def _process_ftp_command(self, payload: FTPPacket, session_id: Optional[str] = None, **kwargs) -> FTPPacket:
        """
        Process the command in the FTP Packet.

        :param: payload: The FTP Packet to process
        :type: payload: FTPPacket
        :param: session_id: session ID linked to the FTP Packet. Optional.
        :type: session_id: Optional[str]
        """
        if payload.ftp_command is not None:
            self.sys_log.info(f"Received FTP {payload.ftp_command.name} command.")

        # handle STOR request
        if payload.ftp_command == FTPCommand.STOR:
            # check that the file is created in the computed hosting the FTP server
            if self._store_data(payload=payload):
                payload.status_code = FTPStatusCode.OK

        if payload.ftp_command == FTPCommand.RETR:
            if self._retrieve_data(payload=payload, session_id=session_id):
                payload.status_code = FTPStatusCode.OK

        return payload

    def _store_data(self, payload: FTPPacket) -> bool:
        """
        Stores the data in the FTP Service's host machine.

        :param: payload: The FTP Packet that contains the file data
        :type: FTPPacket
        """
        try:
            file_name = payload.ftp_command_args["dest_file_name"]
            folder_name = payload.ftp_command_args["dest_folder_name"]
            file_size = payload.ftp_command_args["file_size"]
            health_status = payload.ftp_command_args["health_status"]
            file = self.file_system.create_file(
                file_name=file_name,
                folder_name=folder_name,
                size=file_size,
            )
            file.health_status = health_status
            self.sys_log.info(
                f"{self.name}: Created item in {self.sys_log.hostname}: {payload.ftp_command_args['dest_folder_name']}/"
                f"{payload.ftp_command_args['dest_file_name']}"
            )
            # file should exist
            return self.file_system.get_file(file_name=file_name, folder_name=folder_name) is not None
        except Exception as e:
            self.sys_log.error(f"Unable to create file in {self.sys_log.hostname}: {e}")
            return False

    def _send_data(
        self,
        file: File,
        dest_folder_name: str,
        dest_file_name: str,
        dest_ip_address: Optional[IPv4Address] = None,
        dest_port: Optional[Port] = None,
        session_id: Optional[str] = None,
        is_response: bool = False,
    ) -> bool:
        """
        Sends data from the host FTP Service's machine to another FTP Service's host machine.

        :param: file: File to send to the target FTP Service.
        :type: file: File

        :param: dest_folder_name: The name of the folder where the file will be stored in the FTP Server.
        :type: dest_folder_name: str

        :param: dest_file_name: The name of the file to be saved on the FTP Server.
        :type: dest_file_name: str

        :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: session_id: session ID linked to the FTP Packet. Optional.
        :type: session_id: Optional[str]

        :param: is_response: is true if the data being sent is in response to a request. Default False.
        :type: is_response: bool
        """
        # send STOR request
        payload: FTPPacket = FTPPacket(
            ftp_command=FTPCommand.STOR,
            ftp_command_args={
                "dest_folder_name": dest_folder_name,
                "dest_file_name": dest_file_name,
                "file_size": file.sim_size,
                "health_status": file.health_status,
            },
            packet_payload_size=file.sim_size,
            status_code=FTPStatusCode.OK if is_response else None,
        )
        self.sys_log.info(f"{self.name}: Sending file {file.folder.name}/{file.name}")
        response = self.send(
            payload=payload, dest_ip_address=dest_ip_address, dest_port=dest_port, session_id=session_id
        )

        if response and payload.status_code == FTPStatusCode.OK:
            return True

        return False

    def _retrieve_data(self, payload: FTPPacket, session_id: Optional[str] = None) -> bool:
        """
        Handle the transfer of data from Server to Client.

        :param: payload: The FTP Packet that contains the file data
        :type: FTPPacket
        """
        try:
            # find the file
            file_name = payload.ftp_command_args["src_file_name"]
            folder_name = payload.ftp_command_args["src_folder_name"]
            dest_folder_name = payload.ftp_command_args["dest_folder_name"]
            dest_file_name = payload.ftp_command_args["dest_file_name"]
            retrieved_file: File = self.file_system.get_file(folder_name=folder_name, file_name=file_name)

            # if file does not exist, return an error
            if not retrieved_file:
                self.sys_log.error(
                    f"File  {payload.ftp_command_args['dest_folder_name']}/"
                    f"{payload.ftp_command_args['dest_file_name']} does not exist in {self.sys_log.hostname}"
                )
                return False
            else:
                # send requested data
                return self._send_data(
                    file=retrieved_file,
                    dest_file_name=dest_file_name,
                    dest_folder_name=dest_folder_name,
                    session_id=session_id,
                    is_response=True,
                )
        except Exception as e:
            self.sys_log.error(f"Unable to retrieve file from {self.sys_log.hostname}: {e}")
            return False

    def send(
        self,
        payload: FTPPacket,
        session_id: Optional[str] = None,
        dest_ip_address: Optional[IPv4Address] = None,
        dest_port: Optional[Port] = None,
        **kwargs,
    ) -> bool:
        """
        Sends a payload to the SessionManager.

        :param payload: The payload to be sent.
        :param dest_ip_address: The ip address of the payload destination.
        :param dest_port: The port of the payload destination.
        :param session_id: The Session ID the payload is to originate from. Optional.

        :return: True if successful, False otherwise.
        """
        self.sys_log.info(f"{self.name}: Sending FTP {payload.ftp_command.name} {payload.ftp_command_args}")

        return super().send(
            payload=payload, dest_ip_address=dest_ip_address, dest_port=dest_port, session_id=session_id, **kwargs
        )