Merge remote-tracking branch 'origin/4.0.0a1-dev' into feature/2913_Extensible_Rewards
This commit is contained in:
Marek Wolan
@@ -1 +1,338 @@
|
||||
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
|
||||
from ipaddress import IPv4Address
|
||||
from typing import Dict, Optional
|
||||
|
||||
from primaite import getLogger
|
||||
from primaite.interface.request import RequestFormat, RequestResponse
|
||||
from primaite.simulator.core import RequestManager, RequestType
|
||||
from primaite.simulator.file_system.file_system import File
|
||||
from primaite.simulator.network.protocols.ftp import FTPCommand, FTPPacket, FTPStatusCode
|
||||
from primaite.simulator.system.core.software_manager import SoftwareManager
|
||||
from primaite.simulator.system.services.ftp.ftp_service import FTPServiceABC
|
||||
from primaite.utils.validation.ip_protocol import PROTOCOL_LOOKUP
|
||||
from primaite.utils.validation.port import Port, PORT_LOOKUP
|
||||
|
||||
_LOGGER = getLogger(__name__)
|
||||
|
||||
|
||||
class FTPClient(FTPServiceABC):
|
||||
"""
|
||||
A class for simulating an FTP client service.
|
||||
|
||||
This class inherits from the `Service` class and provides methods to emulate FTP
|
||||
RFC 959: https://datatracker.ietf.org/doc/html/rfc959
|
||||
"""
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
kwargs["name"] = "FTPClient"
|
||||
kwargs["port"] = PORT_LOOKUP["FTP"]
|
||||
kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"]
|
||||
super().__init__(**kwargs)
|
||||
self.start()
|
||||
|
||||
def _init_request_manager(self) -> RequestManager:
|
||||
"""
|
||||
Initialise the request manager.
|
||||
|
||||
More information in user guide and docstring for SimComponent._init_request_manager.
|
||||
"""
|
||||
rm = super()._init_request_manager()
|
||||
|
||||
def _send_data_request(request: RequestFormat, context: Dict) -> RequestResponse:
|
||||
"""
|
||||
Request for sending data via the ftp_client using the request options parameters.
|
||||
|
||||
:param request: Request with one element containing a dict of parameters for the send method.
|
||||
:type request: RequestFormat
|
||||
:param context: additional context for resolving this action, currently unused
|
||||
:type context: dict
|
||||
:return: RequestResponse object with a success code reflecting whether the configuration could be applied.
|
||||
:rtype: RequestResponse
|
||||
"""
|
||||
dest_ip = request[-1].get("dest_ip_address")
|
||||
dest_ip = None if dest_ip is None else IPv4Address(dest_ip)
|
||||
|
||||
# Missing FTP Options results is an automatic failure.
|
||||
src_folder = request[-1].get("src_folder_name", None)
|
||||
src_file_name = request[-1].get("src_file_name", None)
|
||||
dest_folder = request[-1].get("dest_folder_name", None)
|
||||
dest_file_name = request[-1].get("dest_file_name", None)
|
||||
|
||||
if not self.file_system.access_file(folder_name=src_folder, file_name=src_file_name):
|
||||
self.sys_log.debug(
|
||||
f"{self.name}: Received a FTP Request to transfer file: {src_file_name} to Remote IP: {dest_ip}."
|
||||
)
|
||||
return RequestResponse(
|
||||
status="failure",
|
||||
data={
|
||||
"reason": "Unable to locate given file on local file system. Perhaps given options are invalid?"
|
||||
},
|
||||
)
|
||||
|
||||
return RequestResponse.from_bool(
|
||||
self.send_file(
|
||||
dest_ip_address=dest_ip,
|
||||
src_folder_name=src_folder,
|
||||
src_file_name=src_file_name,
|
||||
dest_folder_name=dest_folder,
|
||||
dest_file_name=dest_file_name,
|
||||
)
|
||||
)
|
||||
|
||||
rm.add_request("send", request_type=RequestType(func=_send_data_request)),
|
||||
return rm
|
||||
|
||||
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 client service is down, return error
|
||||
if not self._can_perform_action():
|
||||
payload.status_code = FTPStatusCode.ERROR
|
||||
return payload
|
||||
|
||||
self.sys_log.info(f"{self.name}: Received FTP {payload.ftp_command.name} {payload.ftp_command_args}")
|
||||
|
||||
# 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_LOOKUP["FTP"],
|
||||
session_id: Optional[str] = None,
|
||||
is_reattempt: Optional[bool] = False,
|
||||
) -> bool:
|
||||
"""
|
||||
Connects the client to a given FTP server.
|
||||
|
||||
: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]
|
||||
:param: is_reattempt: Set to True if attempt to connect to FTP Server has been attempted. Default False.
|
||||
:type: is_reattempt: Optional[bool]
|
||||
"""
|
||||
# make sure the service is running before attempting
|
||||
if not self._can_perform_action():
|
||||
return False
|
||||
|
||||
# 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_LOOKUP["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}"
|
||||
)
|
||||
self.add_connection(connection_id="server_connection", session_id=session_id)
|
||||
return True
|
||||
else:
|
||||
if is_reattempt:
|
||||
# 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}"
|
||||
)
|
||||
return False
|
||||
else:
|
||||
# try again
|
||||
self._connect_to_server(
|
||||
dest_ip_address=dest_ip_address, dest_port=dest_port, session_id=session_id, is_reattempt=True
|
||||
)
|
||||
else:
|
||||
self.sys_log.warning(f"{self.name}: Unable to send FTPPacket")
|
||||
return False
|
||||
|
||||
def _disconnect_from_server(
|
||||
self, dest_ip_address: Optional[IPv4Address] = None, dest_port: Optional[Port] = PORT_LOOKUP["FTP"]
|
||||
) -> bool:
|
||||
"""
|
||||
Connects the client from a given FTP server.
|
||||
|
||||
: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]
|
||||
:param: is_reattempt: Set to True if attempt to disconnect from FTP Server has been attempted. Default False.
|
||||
:type: is_reattempt: Optional[bool]
|
||||
"""
|
||||
# send a disconnect request payload to FTP server
|
||||
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
|
||||
)
|
||||
return payload.status_code == FTPStatusCode.OK
|
||||
|
||||
def send_file(
|
||||
self,
|
||||
dest_ip_address: IPv4Address,
|
||||
src_folder_name: str,
|
||||
src_file_name: str,
|
||||
dest_folder_name: str,
|
||||
dest_file_name: str,
|
||||
dest_port: Optional[Port] = PORT_LOOKUP["FTP"],
|
||||
session_id: Optional[str] = None,
|
||||
) -> bool:
|
||||
"""
|
||||
Send a file to a target IP address.
|
||||
|
||||
The function checks if the file exists in the FTP Client host.
|
||||
The STOR command is then sent to the FTP Server.
|
||||
|
||||
:param: dest_ip_address: The IP address of the machine that hosts the FTP Server.
|
||||
:type: dest_ip_address: IPv4Address
|
||||
|
||||
:param: src_folder_name: The name of the folder that contains the file to send to the FTP Server.
|
||||
:type: src_folder_name: str
|
||||
|
||||
:param: src_file_name: The name of the file to send to the FTP Server.
|
||||
:type: src_file_name: str
|
||||
|
||||
: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_port: The open port of the machine that hosts the FTP Server. Default is Port["FTP"].
|
||||
:type: dest_port: Optional[Port]
|
||||
|
||||
:param: session_id: The id of the session
|
||||
:type: session_id: Optional[str]
|
||||
"""
|
||||
# check if the file to transfer exists on the client
|
||||
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.warning(f"Unable to send file that does not exist: {src_folder_name}/{src_file_name}")
|
||||
return False
|
||||
|
||||
# check if FTP is currently connected to IP
|
||||
self._connect_to_server(dest_ip_address=dest_ip_address, dest_port=dest_port)
|
||||
|
||||
if not len(self.connections):
|
||||
return False
|
||||
else:
|
||||
self.sys_log.info(f"Sending file {src_folder_name}/{src_file_name} to {str(dest_ip_address)}")
|
||||
# send STOR request
|
||||
if self._send_data(
|
||||
file=file_to_transfer,
|
||||
dest_folder_name=dest_folder_name,
|
||||
dest_file_name=dest_file_name,
|
||||
dest_ip_address=dest_ip_address,
|
||||
dest_port=dest_port,
|
||||
):
|
||||
return self._disconnect_from_server(dest_ip_address=dest_ip_address, dest_port=dest_port)
|
||||
|
||||
return False
|
||||
|
||||
def request_file(
|
||||
self,
|
||||
dest_ip_address: IPv4Address,
|
||||
src_folder_name: str,
|
||||
src_file_name: str,
|
||||
dest_folder_name: str,
|
||||
dest_file_name: str,
|
||||
dest_port: Optional[Port] = PORT_LOOKUP["FTP"],
|
||||
) -> bool:
|
||||
"""
|
||||
Request a file from a target IP address.
|
||||
|
||||
Sends a RETR command to the FTP Server.
|
||||
|
||||
:param: dest_ip_address: The IP address of the machine that hosts the FTP Server.
|
||||
:type: dest_ip_address: IPv4Address
|
||||
|
||||
:param: src_folder_name: The name of the folder that contains the file to send to the FTP Server.
|
||||
:type: src_folder_name: str
|
||||
|
||||
:param: src_file_name: The name of the file to send to the FTP Server.
|
||||
:type: src_file_name: str
|
||||
|
||||
: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_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)
|
||||
|
||||
if not len(self.connections):
|
||||
return False
|
||||
else:
|
||||
# send retrieve request
|
||||
payload: FTPPacket = FTPPacket(
|
||||
ftp_command=FTPCommand.RETR,
|
||||
ftp_command_args={
|
||||
"src_folder_name": src_folder_name,
|
||||
"src_file_name": src_file_name,
|
||||
"dest_file_name": dest_file_name,
|
||||
"dest_folder_name": dest_folder_name,
|
||||
},
|
||||
)
|
||||
self.sys_log.info(f"Requesting file {src_folder_name}/{src_file_name} from {str(dest_ip_address)}")
|
||||
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
|
||||
)
|
||||
|
||||
# the payload should have ok status code
|
||||
if payload.status_code == FTPStatusCode.OK:
|
||||
self.sys_log.info(f"{self.name}: File {src_folder_name}/{src_file_name} found in FTP server.")
|
||||
return True
|
||||
else:
|
||||
self.sys_log.error(f"{self.name}: File {src_folder_name}/{src_file_name} does not exist in FTP server")
|
||||
return False
|
||||
|
||||
def receive(self, payload: FTPPacket, session_id: Optional[str] = None, **kwargs) -> bool:
|
||||
"""
|
||||
Receives a payload from the SessionManager.
|
||||
|
||||
:param: payload: FTPPacket payload.
|
||||
:type: payload: FTPPacket
|
||||
|
||||
:param: session_id: ID of the session. Optional.
|
||||
:type: session_id: Optional[str]
|
||||
"""
|
||||
if not isinstance(payload, FTPPacket):
|
||||
self.sys_log.warning(f"{self.name}: Payload is not an FTP packet")
|
||||
self.sys_log.debug(f"{self.name}: {payload}")
|
||||
return False
|
||||
|
||||
"""
|
||||
Ignore ftp payload if status code is None.
|
||||
|
||||
This helps prevent an FTP request loop - FTP client and servers can exist on
|
||||
the same node.
|
||||
"""
|
||||
if not self._can_perform_action():
|
||||
return False
|
||||
|
||||
if payload.status_code is None:
|
||||
self.sys_log.error(f"FTP Server could not be found - Error Code: {FTPStatusCode.NOT_FOUND.value}")
|
||||
return False
|
||||
|
||||
# if PORT succeeded, add the connection as an active connection list
|
||||
if payload.ftp_command is FTPCommand.PORT and payload.status_code is FTPStatusCode.OK:
|
||||
self.add_connection(connection_id=session_id, session_id=session_id)
|
||||
|
||||
# if QUIT succeeded, remove the session from active connection list
|
||||
if payload.ftp_command is FTPCommand.QUIT and payload.status_code is FTPStatusCode.OK:
|
||||
self.terminate_connection(connection_id=session_id)
|
||||
|
||||
self.sys_log.info(f"{self.name}: Received FTP Response {payload.ftp_command.name} {payload.status_code.value}")
|
||||
|
||||
self._process_ftp_command(payload=payload, session_id=session_id)
|
||||
return True
|
||||
|
||||
@@ -1 +1,545 @@
|
||||
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import abstractmethod
|
||||
from datetime import datetime
|
||||
from ipaddress import IPv4Address
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
from uuid import uuid4
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from primaite.interface.request import RequestFormat, RequestResponse
|
||||
from primaite.simulator.core import RequestManager, RequestType
|
||||
from primaite.simulator.network.protocols.ssh import (
|
||||
SSHConnectionMessage,
|
||||
SSHPacket,
|
||||
SSHTransportMessage,
|
||||
SSHUserCredentials,
|
||||
)
|
||||
from primaite.simulator.system.core.software_manager import SoftwareManager
|
||||
from primaite.simulator.system.services.service import Service, ServiceOperatingState
|
||||
from primaite.utils.validation.ip_protocol import PROTOCOL_LOOKUP
|
||||
from primaite.utils.validation.port import PORT_LOOKUP
|
||||
|
||||
|
||||
# TODO 2824: Since remote terminal connections and remote user sessions are the same thing, we could refactor
|
||||
# the terminal to leverage the user session manager's list. This way we avoid potential bugs and code ducplication
|
||||
class TerminalClientConnection(BaseModel):
|
||||
"""
|
||||
TerminalClientConnection Class.
|
||||
|
||||
This class is used to record current User Connections to the Terminal class.
|
||||
"""
|
||||
|
||||
parent_terminal: Terminal
|
||||
"""The parent Node that this connection was created on."""
|
||||
|
||||
ssh_session_id: str = None
|
||||
"""Session ID that connection is linked to, used for sending commands via session manager."""
|
||||
|
||||
connection_uuid: str = None
|
||||
"""Connection UUID"""
|
||||
|
||||
connection_request_id: str = None
|
||||
"""Connection request ID"""
|
||||
|
||||
time: datetime = None
|
||||
"""Timestamp connection was created."""
|
||||
|
||||
ip_address: IPv4Address
|
||||
"""Source IP of Connection"""
|
||||
|
||||
is_active: bool = True
|
||||
"""Flag to state whether the connection is active or not"""
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"{self.__class__.__name__}(connection_id: '{self.connection_uuid}, ip_address: {self.ip_address}')"
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return self.__str__()
|
||||
|
||||
def __getitem__(self, key: Any) -> Any:
|
||||
return getattr(self, key)
|
||||
|
||||
@property
|
||||
def client(self) -> Optional[Terminal]:
|
||||
"""The Terminal that holds this connection."""
|
||||
return self.parent_terminal
|
||||
|
||||
def disconnect(self) -> bool:
|
||||
"""Disconnect the session."""
|
||||
return self.parent_terminal._disconnect(connection_uuid=self.connection_uuid)
|
||||
|
||||
@abstractmethod
|
||||
def execute(self, command: Any) -> bool:
|
||||
"""Execute a given command."""
|
||||
pass
|
||||
|
||||
|
||||
class LocalTerminalConnection(TerminalClientConnection):
|
||||
"""
|
||||
LocalTerminalConnectionClass.
|
||||
|
||||
This class represents a local terminal when connected.
|
||||
"""
|
||||
|
||||
ip_address: str = "Local Connection"
|
||||
|
||||
def execute(self, command: Any) -> Optional[RequestResponse]:
|
||||
"""Execute a given command on local Terminal."""
|
||||
if self.parent_terminal.operating_state != ServiceOperatingState.RUNNING:
|
||||
self.parent_terminal.sys_log.warning("Cannot process command as system not running")
|
||||
return None
|
||||
if not self.is_active:
|
||||
self.parent_terminal.sys_log.warning("Connection inactive, cannot execute")
|
||||
return None
|
||||
return self.parent_terminal.execute(command)
|
||||
|
||||
|
||||
class RemoteTerminalConnection(TerminalClientConnection):
|
||||
"""
|
||||
RemoteTerminalConnection Class.
|
||||
|
||||
This class acts as broker between the terminal and remote.
|
||||
|
||||
"""
|
||||
|
||||
def execute(self, command: Any) -> bool:
|
||||
"""Execute a given command on the remote Terminal."""
|
||||
if self.parent_terminal.operating_state != ServiceOperatingState.RUNNING:
|
||||
self.parent_terminal.sys_log.warning("Cannot process command as system not running")
|
||||
return False
|
||||
if not self.is_active:
|
||||
self.parent_terminal.sys_log.warning("Connection inactive, cannot execute")
|
||||
return False
|
||||
# Send command to remote terminal to process.
|
||||
|
||||
transport_message: SSHTransportMessage = SSHTransportMessage.SSH_MSG_SERVICE_REQUEST
|
||||
connection_message: SSHConnectionMessage = SSHConnectionMessage.SSH_MSG_CHANNEL_DATA
|
||||
|
||||
payload: SSHPacket = SSHPacket(
|
||||
transport_message=transport_message,
|
||||
connection_message=connection_message,
|
||||
connection_request_uuid=self.connection_request_id,
|
||||
connection_uuid=self.connection_uuid,
|
||||
ssh_command=command,
|
||||
)
|
||||
|
||||
return self.parent_terminal.send(payload=payload, session_id=self.ssh_session_id)
|
||||
|
||||
|
||||
class Terminal(Service):
|
||||
"""Class used to simulate a generic terminal service. Can be interacted with by other terminals via SSH."""
|
||||
|
||||
_client_connection_requests: Dict[str, Optional[Union[str, TerminalClientConnection]]] = {}
|
||||
"""Dictionary of connect requests made to remote nodes."""
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
kwargs["name"] = "Terminal"
|
||||
kwargs["port"] = PORT_LOOKUP["SSH"]
|
||||
kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"]
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def describe_state(self) -> Dict:
|
||||
"""
|
||||
Produce a dictionary describing the current state of this object.
|
||||
|
||||
Please see :py:meth:`primaite.simulator.core.SimComponent.describe_state` for a more detailed explanation.
|
||||
|
||||
:return: Current state of this object and child objects.
|
||||
:rtype: Dict
|
||||
"""
|
||||
state = super().describe_state()
|
||||
return state
|
||||
|
||||
def show(self, markdown: bool = False):
|
||||
"""
|
||||
Display the remote connections to this terminal instance in tabular format.
|
||||
|
||||
:param markdown: Whether to display the table in Markdown format or not. Default is `False`.
|
||||
"""
|
||||
self.show_connections(markdown=markdown)
|
||||
|
||||
def _init_request_manager(self) -> RequestManager:
|
||||
"""Initialise Request manager."""
|
||||
rm = super()._init_request_manager()
|
||||
|
||||
def _remote_login(request: RequestFormat, context: Dict) -> RequestResponse:
|
||||
login = self._send_remote_login(username=request[0], password=request[1], ip_address=request[2])
|
||||
if login:
|
||||
return RequestResponse(
|
||||
status="success",
|
||||
data={
|
||||
"ip_address": str(login.ip_address),
|
||||
"username": request[0],
|
||||
},
|
||||
)
|
||||
else:
|
||||
return RequestResponse(status="failure", data={})
|
||||
|
||||
rm.add_request(
|
||||
"ssh_to_remote",
|
||||
request_type=RequestType(func=_remote_login),
|
||||
)
|
||||
|
||||
def _remote_logoff(request: RequestFormat, context: Dict) -> RequestResponse:
|
||||
"""Logoff from remote connection."""
|
||||
ip_address = IPv4Address(request[0])
|
||||
remote_connection = self._get_connection_from_ip(ip_address=ip_address)
|
||||
if remote_connection:
|
||||
outcome = self._disconnect(remote_connection.connection_uuid)
|
||||
if outcome:
|
||||
return RequestResponse(status="success", data={})
|
||||
|
||||
return RequestResponse(status="failure", data={})
|
||||
|
||||
rm.add_request("remote_logoff", request_type=RequestType(func=_remote_logoff))
|
||||
|
||||
def remote_execute_request(request: RequestFormat, context: Dict) -> RequestResponse:
|
||||
"""Execute an instruction."""
|
||||
ip_address: IPv4Address = IPv4Address(request[0])
|
||||
command: str = request[1]["command"]
|
||||
remote_connection = self._get_connection_from_ip(ip_address=ip_address)
|
||||
if remote_connection:
|
||||
outcome = remote_connection.execute(command)
|
||||
if outcome:
|
||||
return RequestResponse(
|
||||
status="success",
|
||||
data={},
|
||||
)
|
||||
else:
|
||||
return RequestResponse(
|
||||
status="failure",
|
||||
data={},
|
||||
)
|
||||
|
||||
rm.add_request(
|
||||
"send_remote_command",
|
||||
request_type=RequestType(func=remote_execute_request),
|
||||
)
|
||||
|
||||
return rm
|
||||
|
||||
def execute(self, command: List[Any]) -> Optional[RequestResponse]:
|
||||
"""Execute a passed ssh command via the request manager."""
|
||||
return self.parent.apply_request(command)
|
||||
|
||||
def _get_connection_from_ip(self, ip_address: IPv4Address) -> Optional[RemoteTerminalConnection]:
|
||||
"""Find Remote Terminal Connection from a given IP."""
|
||||
for connection in self._connections.values():
|
||||
if connection.ip_address == ip_address:
|
||||
return connection
|
||||
|
||||
def _create_local_connection(self, connection_uuid: str, session_id: str) -> TerminalClientConnection:
|
||||
"""Create a new connection object and amend to list of active connections.
|
||||
|
||||
:param connection_uuid: Connection ID of the new local connection
|
||||
:param session_id: Session ID of the new local connection
|
||||
:return: TerminalClientConnection object
|
||||
"""
|
||||
new_connection = LocalTerminalConnection(
|
||||
parent_terminal=self,
|
||||
connection_uuid=connection_uuid,
|
||||
ssh_session_id=session_id,
|
||||
time=datetime.now(),
|
||||
)
|
||||
self._connections[connection_uuid] = new_connection
|
||||
self._client_connection_requests[connection_uuid] = new_connection
|
||||
|
||||
return new_connection
|
||||
|
||||
def login(
|
||||
self, username: str, password: str, ip_address: Optional[IPv4Address] = None
|
||||
) -> Optional[TerminalClientConnection]:
|
||||
"""Login to the terminal. Will attempt a remote login if ip_address is given, else local.
|
||||
|
||||
:param: username: Username used to connect to the remote node.
|
||||
:type: username: str
|
||||
|
||||
:param: password: Password used to connect to the remote node
|
||||
:type: password: str
|
||||
|
||||
:param: ip_address: Target Node IP address for login attempt. If None, login is assumed local.
|
||||
:type: ip_address: Optional[IPv4Address]
|
||||
"""
|
||||
if self.operating_state != ServiceOperatingState.RUNNING:
|
||||
self.sys_log.warning(f"{self.name}: Cannot login as service is not running.")
|
||||
return None
|
||||
if ip_address:
|
||||
# Assuming that if IP is passed we are connecting to remote
|
||||
return self._send_remote_login(username=username, password=password, ip_address=ip_address)
|
||||
else:
|
||||
return self._process_local_login(username=username, password=password)
|
||||
|
||||
def _process_local_login(self, username: str, password: str) -> Optional[TerminalClientConnection]:
|
||||
"""Local session login to terminal.
|
||||
|
||||
:param username: Username for login.
|
||||
:param password: Password for login.
|
||||
:return: boolean, True if successful, else False
|
||||
"""
|
||||
# TODO: Un-comment this when UserSessionManager is merged.
|
||||
connection_uuid = self.parent.user_session_manager.local_login(username=username, password=password)
|
||||
if connection_uuid:
|
||||
self.sys_log.info(f"{self.name}: Login request authorised, connection uuid: {connection_uuid}")
|
||||
# Add new local session to list of connections and return
|
||||
return self._create_local_connection(connection_uuid=connection_uuid, session_id="Local_Connection")
|
||||
else:
|
||||
self.sys_log.warning(f"{self.name}: Login failed, incorrect Username or Password")
|
||||
return None
|
||||
|
||||
def _validate_client_connection_request(self, connection_id: str) -> bool:
|
||||
"""Check that client_connection_id is valid."""
|
||||
return connection_id in self._client_connection_requests
|
||||
|
||||
def _check_client_connection(self, connection_id: str) -> bool:
|
||||
"""Check that client_connection_id is valid."""
|
||||
if not self.parent.user_session_manager.validate_remote_session_uuid(connection_id):
|
||||
self._disconnect(connection_id)
|
||||
return False
|
||||
return connection_id in self._connections
|
||||
|
||||
def _send_remote_login(
|
||||
self,
|
||||
username: str,
|
||||
password: str,
|
||||
ip_address: IPv4Address,
|
||||
connection_request_id: Optional[str] = None,
|
||||
is_reattempt: bool = False,
|
||||
) -> Optional[RemoteTerminalConnection]:
|
||||
"""Send a remote login attempt and connect to Node.
|
||||
|
||||
:param: username: Username used to connect to the remote node.
|
||||
:type: username: str
|
||||
:param: password: Password used to connect to the remote node
|
||||
:type: password: str
|
||||
:param: ip_address: Target Node IP address for login attempt.
|
||||
:type: ip_address: IPv4Address
|
||||
:param: connection_request_id: Connection Request ID, if not provided, a new one is generated
|
||||
:type: connection_request_id: Optional[str]
|
||||
:param: is_reattempt: True if the request has been reattempted. Default False.
|
||||
:type: is_reattempt: Optional[bool]
|
||||
:return: RemoteTerminalConnection: Connection Object for sending further commands if successful, else False.
|
||||
"""
|
||||
connection_request_id = connection_request_id or str(uuid4())
|
||||
if is_reattempt:
|
||||
valid_connection_request = self._validate_client_connection_request(connection_id=connection_request_id)
|
||||
if valid_connection_request:
|
||||
remote_terminal_connection = self._client_connection_requests.pop(connection_request_id)
|
||||
if isinstance(remote_terminal_connection, RemoteTerminalConnection):
|
||||
self.sys_log.info(f"{self.name}: Remote Connection to {ip_address} authorised.")
|
||||
return remote_terminal_connection
|
||||
else:
|
||||
self.sys_log.warning(f"{self.name}: Connection request {connection_request_id} declined")
|
||||
return None
|
||||
else:
|
||||
self.sys_log.warning(f"{self.name}: Remote connection to {ip_address} declined.")
|
||||
return None
|
||||
|
||||
self.sys_log.info(
|
||||
f"{self.name}: Sending Remote login attempt to {ip_address}. Connection_id is {connection_request_id}"
|
||||
)
|
||||
transport_message: SSHTransportMessage = SSHTransportMessage.SSH_MSG_USERAUTH_REQUEST
|
||||
connection_message: SSHConnectionMessage = SSHConnectionMessage.SSH_MSG_CHANNEL_DATA
|
||||
user_details: SSHUserCredentials = SSHUserCredentials(username=username, password=password)
|
||||
|
||||
payload_contents = {
|
||||
"type": "login_request",
|
||||
"username": username,
|
||||
"password": password,
|
||||
"connection_request_id": connection_request_id,
|
||||
}
|
||||
|
||||
payload: SSHPacket = SSHPacket(
|
||||
payload=payload_contents,
|
||||
transport_message=transport_message,
|
||||
connection_message=connection_message,
|
||||
user_account=user_details,
|
||||
connection_request_uuid=connection_request_id,
|
||||
)
|
||||
|
||||
software_manager: SoftwareManager = self.software_manager
|
||||
software_manager.send_payload_to_session_manager(
|
||||
payload=payload, dest_ip_address=ip_address, dest_port=self.port
|
||||
)
|
||||
return self._send_remote_login(
|
||||
username=username,
|
||||
password=password,
|
||||
ip_address=ip_address,
|
||||
is_reattempt=True,
|
||||
connection_request_id=connection_request_id,
|
||||
)
|
||||
|
||||
def _create_remote_connection(
|
||||
self, connection_id: str, connection_request_id: str, session_id: str, source_ip: str
|
||||
) -> None:
|
||||
"""Create a new TerminalClientConnection Object.
|
||||
|
||||
:param: connection_request_id: Connection Request ID
|
||||
:type: connection_request_id: str
|
||||
|
||||
:param: session_id: Session ID of connection.
|
||||
:type: session_id: str
|
||||
"""
|
||||
client_connection = RemoteTerminalConnection(
|
||||
parent_terminal=self,
|
||||
ssh_session_id=session_id,
|
||||
connection_uuid=connection_id,
|
||||
ip_address=source_ip,
|
||||
connection_request_id=connection_request_id,
|
||||
time=datetime.now(),
|
||||
)
|
||||
self._connections[connection_id] = client_connection
|
||||
self._client_connection_requests[connection_request_id] = client_connection
|
||||
|
||||
def receive(self, session_id: str, payload: Union[SSHPacket, Dict], **kwargs) -> bool:
|
||||
"""
|
||||
Receive a payload from the Software Manager.
|
||||
|
||||
:param payload: A payload to receive.
|
||||
:param session_id: The session id the payload relates to.
|
||||
:return: True.
|
||||
"""
|
||||
source_ip = kwargs["frame"].ip.src_ip_address
|
||||
self.sys_log.info(f"{self.name}: Received payload: {payload}. Source: {source_ip}")
|
||||
if isinstance(payload, SSHPacket):
|
||||
if payload.transport_message == SSHTransportMessage.SSH_MSG_USERAUTH_REQUEST:
|
||||
# validate & add connection
|
||||
# TODO: uncomment this as part of 2781
|
||||
username = payload.user_account.username
|
||||
password = payload.user_account.password
|
||||
connection_id = self.parent.user_session_manager.remote_login(
|
||||
username=username, password=password, remote_ip_address=source_ip
|
||||
)
|
||||
if connection_id:
|
||||
connection_request_id = payload.connection_request_uuid
|
||||
self.sys_log.info(f"{self.name}: Connection authorised, session_id: {session_id}")
|
||||
self._create_remote_connection(
|
||||
connection_id=connection_id,
|
||||
connection_request_id=connection_request_id,
|
||||
session_id=session_id,
|
||||
source_ip=source_ip,
|
||||
)
|
||||
|
||||
transport_message = SSHTransportMessage.SSH_MSG_USERAUTH_SUCCESS
|
||||
connection_message = SSHConnectionMessage.SSH_MSG_CHANNEL_DATA
|
||||
|
||||
payload_contents = {
|
||||
"type": "login_success",
|
||||
"username": username,
|
||||
"password": password,
|
||||
"connection_request_id": connection_request_id,
|
||||
"connection_id": connection_id,
|
||||
}
|
||||
payload: SSHPacket = SSHPacket(
|
||||
payload=payload_contents,
|
||||
transport_message=transport_message,
|
||||
connection_message=connection_message,
|
||||
connection_request_uuid=connection_request_id,
|
||||
connection_uuid=connection_id,
|
||||
)
|
||||
|
||||
software_manager: SoftwareManager = self.software_manager
|
||||
software_manager.send_payload_to_session_manager(
|
||||
payload=payload, dest_port=self.port, session_id=session_id
|
||||
)
|
||||
elif payload.transport_message == SSHTransportMessage.SSH_MSG_USERAUTH_SUCCESS:
|
||||
self.sys_log.info(f"{self.name}: Login Successful")
|
||||
self._create_remote_connection(
|
||||
connection_id=payload.connection_uuid,
|
||||
connection_request_id=payload.connection_request_uuid,
|
||||
session_id=session_id,
|
||||
source_ip=source_ip,
|
||||
)
|
||||
|
||||
elif payload.transport_message == SSHTransportMessage.SSH_MSG_SERVICE_REQUEST:
|
||||
# Requesting a command to be executed
|
||||
self.sys_log.info(f"{self.name}: Received command to execute")
|
||||
command = payload.ssh_command
|
||||
valid_connection = self._check_client_connection(payload.connection_uuid)
|
||||
if valid_connection:
|
||||
remote_session = self.software_manager.node.user_session_manager.remote_sessions.get(
|
||||
payload.connection_uuid
|
||||
)
|
||||
remote_session.last_active_step = self.software_manager.node.user_session_manager.current_timestep
|
||||
self.execute(command)
|
||||
return True
|
||||
else:
|
||||
self.sys_log.error(
|
||||
f"{self.name}: Connection UUID:{payload.connection_uuid} is not valid. Rejecting Command."
|
||||
)
|
||||
|
||||
if isinstance(payload, dict) and payload.get("type"):
|
||||
if payload["type"] == "disconnect":
|
||||
connection_id = payload["connection_id"]
|
||||
valid_id = self._check_client_connection(connection_id)
|
||||
if valid_id:
|
||||
self.sys_log.info(f"{self.name}: Received disconnect command for {connection_id=} from remote.")
|
||||
self._disconnect(payload["connection_id"])
|
||||
self.parent.user_session_manager.remote_logout(remote_session_id=connection_id)
|
||||
else:
|
||||
self.sys_log.error(f"{self.name}: No Active connection held for received connection ID.")
|
||||
|
||||
if payload["type"] == "user_timeout":
|
||||
connection_id = payload["connection_id"]
|
||||
valid_id = connection_id in self._connections
|
||||
if valid_id:
|
||||
connection = self._connections.pop(connection_id)
|
||||
connection.is_active = False
|
||||
self.sys_log.info(f"{self.name}: Connection {connection_id} disconnected due to inactivity.")
|
||||
else:
|
||||
self.sys_log.error(f"{self.name}: Connection {connection_id} is invalid.")
|
||||
|
||||
return True
|
||||
|
||||
def _disconnect(self, connection_uuid: str) -> bool:
|
||||
"""Disconnect connection.
|
||||
|
||||
:param connection_uuid: Connection ID that we want to disconnect.
|
||||
:return True if successful, False otherwise.
|
||||
"""
|
||||
# TODO: Handle the possibility of attempting to disconnect
|
||||
if not self._connections:
|
||||
self.sys_log.warning(f"{self.name}: No remote connection present")
|
||||
return False
|
||||
|
||||
connection = self._connections.pop(connection_uuid, None)
|
||||
if not connection:
|
||||
return False
|
||||
connection.is_active = False
|
||||
|
||||
if isinstance(connection, RemoteTerminalConnection):
|
||||
# Send disconnect command via software manager
|
||||
session_id = connection.ssh_session_id
|
||||
|
||||
software_manager: SoftwareManager = self.software_manager
|
||||
software_manager.send_payload_to_session_manager(
|
||||
payload={"type": "disconnect", "connection_id": connection_uuid},
|
||||
dest_port=self.port,
|
||||
session_id=session_id,
|
||||
)
|
||||
self.sys_log.info(f"{self.name}: Disconnected {connection_uuid}")
|
||||
return True
|
||||
|
||||
elif isinstance(connection, LocalTerminalConnection):
|
||||
self.parent.user_session_manager.local_logout()
|
||||
return True
|
||||
|
||||
def send(
|
||||
self, payload: SSHPacket, dest_ip_address: Optional[IPv4Address] = None, session_id: Optional[str] = None
|
||||
) -> bool:
|
||||
"""
|
||||
Send a payload out from the Terminal.
|
||||
|
||||
:param payload: The payload to be sent.
|
||||
:param dest_up_address: The IP address of the payload destination.
|
||||
"""
|
||||
if self.operating_state != ServiceOperatingState.RUNNING:
|
||||
self.sys_log.warning(f"{self.name}: Cannot send commands when Operating state is {self.operating_state}!")
|
||||
return False
|
||||
|
||||
self.sys_log.debug(f"{self.name}: Sending payload: {payload}")
|
||||
return super().send(
|
||||
payload=payload, dest_ip_address=dest_ip_address, dest_port=self.port, session_id=session_id
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user