oxbrbtx/PrimAITE
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Czar.Echavez
2023-09-22 15:38:01 +01:00
parent 58edb6d3e4
commit 2c234ab67a
5 changed files with 117 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
src/primaite/simulator/network/protocols/ftp.py
View File

@@ -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

40
src/primaite/simulator/system/services/ftp/ftp_client.py
View File

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

41
src/primaite/simulator/system/services/ftp/ftp_server.py
View File

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

7
src/primaite/simulator/system/services/ftp/ftp_service.py
View File

@@ -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

41
tests/unit_tests/_primaite/_simulator/_system/_services/test_ftp.py
View File

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