#2713 - eod commit. Initial RequestManager Test implemented, along with an initial setup of the additional Request Manager methods.

This commit is contained in:
Charlie Crane
2024-07-26 16:56:03 +01:00
parent 978e2c5a52
commit 0ac1c6702c
3 changed files with 205 additions and 19 deletions

View File

@@ -2,9 +2,10 @@
from __future__ import annotations
from ipaddress import IPv4Address
from typing import Dict, List, Optional
from typing import Any, Dict, List, Optional
from uuid import uuid4
from prettytable import MARKDOWN, PrettyTable
from pydantic import BaseModel
from primaite.interface.request import RequestFormat, RequestResponse
@@ -35,17 +36,12 @@ class TerminalClientConnection(BaseModel):
is_active: bool = True
"""Flag to state whether the connection is still active or not."""
_dest_ip_address: IPv4Address
dest_ip_address: IPv4Address = None
"""Destination IP address of connection"""
_connection_uuid: str = None
"""Connection UUID"""
@property
def dest_ip_address(self) -> Optional[IPv4Address]:
"""Destination IP Address."""
return self._dest_ip_address
@property
def client(self) -> Optional[Terminal]:
"""The Terminal that holds this connection."""
@@ -95,6 +91,21 @@ class Terminal(Service):
"""Apply Terminal Request."""
return super().apply_request(request, context)
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`.
"""
table = PrettyTable(["Connection ID", "IP_Address", "Active"])
if markdown:
table.set_style(MARKDOWN)
table.align = "l"
table.title = f"{self.sys_log.hostname} {self.name} Remote Connections"
for connection_id, connection in self.remote_connection.items():
table.add_row([connection_id, connection.dest_ip_address, connection.is_active])
print(table.get_string(sortby="Connection ID"))
def _init_request_manager(self) -> RequestManager:
"""Initialise Request manager."""
_login_valid = Terminal._LoginValidator(terminal=self)
@@ -106,12 +117,52 @@ class Terminal(Service):
func=lambda request, context: RequestResponse.from_bool(self.send()), validator=_login_valid
),
)
return rm
def _validate_login(self) -> bool:
"""Validate login credentials are valid."""
# return self.parent.UserSessionManager.validate_remote_session_uuid(self.connection_uuid)
return True
def _login(request: List[Any], context: Any) -> RequestResponse:
login = self._process_local_login(username=request[0], password=request[1])
if login == True:
return RequestResponse(status="success", data={})
else:
return RequestResponse(status="failure", data={})
def _remote_login(request: List[Any], context: Any) -> RequestResponse:
self._process_remote_login(username=request[0], password=request[1], ip_address=request[2])
if self.is_connected:
return RequestResponse(status="success", data={})
else:
return RequestResponse(status="failure", data={})
def _execute(request: List[Any], context: Any) -> RequestResponse:
"""Execute an instruction."""
command: str = request[0]
self.execute(command)
return RequestResponse(status="success", data={})
def _logoff() -> RequestResponse:
"""Logoff from connection."""
self.parent.UserSessionManager.logoff(self.connection_uuid)
self.disconnect(self.connection_uuid)
return RequestResponse(status="success")
rm.add_request(
"Login",
request_type=RequestType(func=_login),
)
rm.add_request(
"Remote Login",
request_type=RequestType(func=_remote_login),
)
rm.add_request(
"Execute",
request_type=RequestType(func=_execute),
)
rm.add_request("Logoff", request_type=RequestType(func=_logoff))
return rm
class _LoginValidator(RequestPermissionValidator):
"""
@@ -155,7 +206,8 @@ class Terminal(Service):
def _process_local_login(self, username: str, password: str) -> bool:
"""Local session login to terminal."""
# self.connection_uuid = self.parent.UserSessionManager.login(username=username, password=password)
self.connection_uuid = str(uuid4())
self.connection_uuid = str(uuid4()) # TODO: Remove following merging of UserSessionManager.
self.is_connected = True
if self.connection_uuid:
self.sys_log.info(f"Login request authorised, connection uuid: {self.connection_uuid}")
return True
@@ -187,7 +239,7 @@ class Terminal(Service):
self.sys_log.info(f"Sending UserAuth request to UserSessionManager, username={username}, password={password}")
# connection_uuid = self.parent.UserSessionManager.remote_login(username=username, password=password)
connection_uuid = str(uuid4())
self.is_connected = True
if connection_uuid:
# Send uuid to remote
self.sys_log.info(
@@ -203,14 +255,14 @@ class Terminal(Service):
sender_ip_address=self.parent.network_interface[1].ip_address,
target_ip_address=payload.sender_ip_address,
)
self.send(payload=return_payload, dest_ip_address=return_payload.target_ip_address)
self.remote_connection[connection_uuid] = TerminalClientConnection(
parent_node=self.software_manager.node,
_dest_ip_address=payload.sender_ip_address,
dest_ip_address=payload.sender_ip_address,
connection_uuid=connection_uuid,
)
self.send(payload=return_payload, dest_ip_address=return_payload.target_ip_address)
return True
else:
# UserSessionManager has returned None
@@ -246,6 +298,9 @@ class Terminal(Service):
self.is_connected = True
return True
elif payload.transport_message == SSHTransportMessage.SSH_MSG_SERVICE_REQUEST:
return self.execute(command=payload.payload)
else:
self.sys_log.warning("Encounter unexpected message type, rejecting connection")
return False
@@ -254,6 +309,14 @@ class Terminal(Service):
# %% Outbound
def execute(self, command: List[Any]) -> bool:
"""Execute a passed ssh command via the request manager."""
if command[0] == "install":
self.parent.software_manager.software.install(command[1])
return True
# TODO: Expand as necessary
def _disconnect(self, dest_ip_address: IPv4Address) -> bool:
"""Disconnect from the remote."""
if not self.is_connected:
@@ -266,7 +329,7 @@ class Terminal(Service):
software_manager: SoftwareManager = self.software_manager
software_manager.send_payload_to_session_manager(
payload={"type": "disconnect", "connection_id": self.remote_connection._connection_uuid},
payload={"type": "disconnect", "connection_id": self.connection_uuid},
dest_ip_address=dest_ip_address,
dest_port=self.port,
)