#2711 - Added in remote_login and process_login methods. Minor updates to make pydantic happy. Starting to flesh out functionality of Terminal Service in more detail

This commit is contained in:
Charlie Crane
2024-07-02 16:47:39 +01:00
parent bd05f4d4e8
commit ebf6e7a90e
2 changed files with 73 additions and 72 deletions

View File

@@ -60,11 +60,11 @@ class SSHConnectionMessage(IntEnum):
class SSHPacket(DataPacket): class SSHPacket(DataPacket):
"""Represents an SSHPacket.""" """Represents an SSHPacket."""
transport_message: SSHTransportMessage transport_message: SSHTransportMessage = None
connection_message: SSHConnectionMessage connection_message: SSHConnectionMessage = None
ssh_command: Optional[any] = None # This is the request string ssh_command: Optional[str] = None # This is the request string
ssh_output: Optional[RequestResponse] = None # The Request Manager's returned RequestResponse ssh_output: Optional[RequestResponse] = None # The Request Manager's returned RequestResponse

View File

@@ -8,8 +8,7 @@ from uuid import uuid4
from primaite.interface.request import RequestFormat, RequestResponse from primaite.interface.request import RequestFormat, RequestResponse
from primaite.simulator.core import RequestManager, RequestPermissionValidator from primaite.simulator.core import RequestManager, RequestPermissionValidator
from primaite.simulator.network.protocols.icmp import ICMPPacket from primaite.simulator.network.protocols.icmp import ICMPPacket
from primaite.simulator.network.protocols.ssh import SSHConnectionMessage, SSHPacket, SSHTransportMessage
# from primaite.simulator.network.protocols.ssh import SSHPacket, SSHTransportMessage, SSHConnectionMessage
from primaite.simulator.network.transmission.network_layer import IPProtocol from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.system.services.service import Service, ServiceOperatingState from primaite.simulator.system.services.service import Service, ServiceOperatingState
@@ -21,19 +20,22 @@ class Terminal(Service):
user_account: Optional[str] = None user_account: Optional[str] = None
"The User Account used for login" "The User Account used for login"
connected: bool = False is_connected: bool = False
"Boolean Value for whether connected" "Boolean Value for whether connected"
connection_uuid: Optional[str] = None connection_uuid: Optional[str] = None
"Uuid for connection requests" "Uuid for connection requests"
operating_state: ServiceOperatingState = ServiceOperatingState.INSTALLING
"""Service Operating State""" # Install at start ??? Maybe ???
def __init__(self, **kwargs): def __init__(self, **kwargs):
kwargs["name"] = "Terminal" kwargs["name"] = "Terminal"
kwargs["port"] = Port.SSH kwargs["port"] = Port.SSH
kwargs["protocol"] = IPProtocol.TCP kwargs["protocol"] = IPProtocol.TCP
super().__init__(**kwargs) super().__init__(**kwargs)
self.operating_state = ServiceOperatingState.RUNNING # self.operating_state = ServiceOperatingState.RUNNING
class _LoginValidator(RequestPermissionValidator): class _LoginValidator(RequestPermissionValidator):
""" """
@@ -46,34 +48,22 @@ class Terminal(Service):
def __call__(self, request: RequestFormat, context: Dict) -> bool: def __call__(self, request: RequestFormat, context: Dict) -> bool:
"""Return whether the login credentials are valid.""" """Return whether the login credentials are valid."""
pass # TODO: Expand & Implement logic when we have User Accounts.
if self.terminal.is_connected:
return True
else:
self.terminal.sys_log.error("terminal is not logged in.")
@property @property
def fail_message(self) -> str: def fail_message(self) -> str:
"""Message that is reported when a request is rejected by this validator.""" """Message that is reported when a request is rejected by this validator."""
return ( return f"Cannot perform request on Terminal '{self.terminal.name}' because login credentials are invalid"
f"Cannot perform request on Terminal '{self.terminal.hostname}' because login credentials are invalid"
)
def _validate_login(self) -> bool: def _validate_login(self) -> bool:
"""Validate login credentials when receiving commands.""" """Validate login credentials when receiving commands."""
# TODO: Implement # TODO: Implement
return True return True
def receive_payload_from_software_manager(
self,
payload: Any,
dst_ip_address: Optional[Union[IPv4Address, IPv4Network]] = None,
src_port: Optional[Port] = None,
dst_port: Optional[Port] = None,
session_id: Optional[str] = None,
ip_protocol: IPProtocol = IPProtocol.TCP,
icmp_packet: Optional[ICMPPacket] = None,
connection_id: Optional[str] = None,
) -> Union[Any, None]:
"""Receive Software Manager Payload."""
self._validate_login()
def _init_request_manager(self) -> RequestManager: def _init_request_manager(self) -> RequestManager:
"""Initialise Request manager.""" """Initialise Request manager."""
# _login_is_valid = Terminal._LoginValidator(terminal=self) # _login_is_valid = Terminal._LoginValidator(terminal=self)
@@ -101,7 +91,7 @@ class Terminal(Service):
""" """
state = super().describe_state() state = super().describe_state()
# TBD # TBD
state.update({"hostname": self.hostname}) state.update({"hostname": self.name})
return state return state
def execute(self, command: Any, request: Any) -> Optional[RequestResponse]: def execute(self, command: Any, request: Any) -> Optional[RequestResponse]:
@@ -133,58 +123,69 @@ class Terminal(Service):
# %% # %%
# def _ssh_process_login(self, user_account: dict, **kwargs) -> SSHPacket: def _ssh_process_login(self, user_account: dict, **kwargs) -> SSHPacket:
# """Processes the login attempt. Returns a SSHPacket which either rejects the login or accepts it.""" """Processes the login attempt. Returns a SSHPacket which either rejects the login or accepts it."""
# # we assume that the login fails unless we meet all the criteria. # we assume that the login fails unless we meet all the criteria.
# transport_message = SSHTransportMessage.SSH_MSG_USERAUTH_FAILURE transport_message = SSHTransportMessage.SSH_MSG_USERAUTH_FAILURE
# connection_message = SSHConnectionMessage.SSH_MSG_CHANNEL_OPEN_FAILED connection_message = SSHConnectionMessage.SSH_MSG_CHANNEL_OPEN_FAILED
# # operating state validation here(if overwhelmed) # operating state validation here(if overwhelmed)
# # Hard coded at current - replace with another method to handle local accounts. # Hard coded at current - replace with another method to handle local accounts.
# if user_account == f"{self.user_name:} placeholder, {self.password:} placeholder": # hardcoded if user_account == f"{self.user_name:} placeholder, {self.password:} placeholder": # hardcoded
# connection_id = self._generate_connection_id() connection_id = self._generate_connection_id()
# if not self.add_connection(self, connection_id="ssh_connection", session_id=self.session_id): if not self.add_connection(self, connection_id="ssh_connection", session_id=self.session_id):
# self.sys_log.warning(f"{self.name}: Connect request for {self.src_ip} declined. self.sys_log.warning(
# Service is at capacity.") f"{self.name}: Connect request for {self.src_ip} declined. Service is at capacity."
# ... )
# else: else:
# self.sys_log.info(f"{self.name}: Connect request for {connection_id=} authorised") self.sys_log.info(f"{self.name}: Connect request for {connection_id=} authorised")
# transport_message = SSHTransportMessage.SSH_MSG_USERAUTH_SUCCESS transport_message = SSHTransportMessage.SSH_MSG_USERAUTH_SUCCESS
# connection_message = SSHConnectionMessage.SSH_MSG_CHANNEL_OPEN_CONFIRMATION connection_message = SSHConnectionMessage.SSH_MSG_CHANNEL_OPEN_CONFIRMATION
self.is_connected = True
# payload: SSHPacket = SSHPacket(transport_message = transport_message, connection_message = connection_message) payload: SSHPacket = SSHPacket(transport_message=transport_message, connection_message=connection_message)
# return payload return payload
# %% # %%
# Copy + Paste from Terminal Wiki # Copy + Paste from Terminal Wiki
# def ssh_remote_login(self, dest_ip_address = IPv4Address, user_account: Optional[dict] = None) -> bool: def ssh_remote_login(self, dest_ip_address: IPv4Address, user_account: Optional[dict] = None) -> bool:
# if user_account: """Remote login to terminal via SSH."""
# # Setting default creds (Best to use this until we have more clarification on the specifics of user accounts) if user_account:
# self.user_account = {self.user_name:"placeholder", self.password:"placeholder"} # Setting default creds (Best to use this until we have more clarification around user accounts)
self.user_account = {self.user_name: "placeholder", self.password: "placeholder"}
# Implement SSHPacket class
payload: SSHPacket = SSHPacket(
transport_message=SSHTransportMessage.SSH_MSG_USERAUTH_REQUEST,
connection_message=SSHConnectionMessage.SSH_MSG_CHANNEL_OPEN,
user_account=user_account,
)
if self.send(payload=payload, dest_ip_address=dest_ip_address):
if payload.connection_message == SSHTransportMessage.SSH_MSG_USERAUTH_SUCCESS:
self.sys_log.info(f"{self.name} established an ssh connection with {dest_ip_address}")
# Need to confirm if self.uuid is correct.
self.add_connection(self, connection_id=self.uuid, session_id=self.session_id)
return True
else:
self.sys_log.error("Payload type incorrect, Login Failed")
return False
else:
self.sys_log.error("Incorrect credentials provided. Login Failed.")
return False
# # Implement SSHPacket class
# payload: SSHPacket = SSHPacket(transport_message= SSHTransportMessage.SSH_MSG_USERAUTH_REQUEST,
# connection_message= SSHConnectionMessage.SSH_MSG_CHANNEL_OPEN,
# user_account=user_account)
# if self.send(payload=payload,dest_ip_address=dest_ip_address):
# if payload.connection_message == SSHTransportMessage.SSH_MSG_USERAUTH_SUCCESS:
# self.sys_log.info(f"{self.name} established an ssh connection with {dest_ip_address}")
# # Need to confirm if self.uuid is correct.
# self.add_connection(self, connection_id=self.uuid, session_id=self.session_id)
# return True
# else:
# self.sys_log.error("Payload type incorrect, Login Failed")
# return False
# else:
# self.sys_log.error("Incorrect credentials provided. Login Failed.")
# return False
# %% # %%
def connect(self, **kwargs): def receive_payload_from_software_manager(
"""Send connect request.""" self,
self._connect(self, **kwargs) payload: Any,
dst_ip_address: Optional[Union[IPv4Address, IPv4Network]] = None,
def _connect(self): src_port: Optional[Port] = None,
"""Do something.""" dst_port: Optional[Port] = None,
pass session_id: Optional[str] = None,
ip_protocol: IPProtocol = IPProtocol.TCP,
icmp_packet: Optional[ICMPPacket] = None,
connection_id: Optional[str] = None,
) -> Union[Any, None]:
"""Receive Software Manager Payload."""
self._validate_login()