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

#2706 - Updated terminal.receive() to work with SSHPacket class, fixed some tests and updated RemoteTerminalConnection to hold Source_IP for easier reading

Browse Source
This commit is contained in:
Charlie Crane
2024-08-05 15:44:52 +01:00
parent 972b0b9712
commit 4fe9753fcf
5 changed files with 222 additions and 103 deletions
Download Patch File Download Diff File Expand all files Collapse all files

103
src/primaite/notebooks/Terminal-Processing.ipynb
View File

@@ -15,7 +15,7 @@
"source": [
"This notebook serves as a guide on the functionality and use of the new Terminal simulation component.\n",
"\n",
"By default, the Terminal will come pre-installed on any simulation component which inherits from `HostNode` (Computer, Server, Printer), and simulates the Secure Shell (SSH) protocol as the communication method."
"The Terminal service comes pre-installed on most Nodes (The exception being Switches, as these are currently dumb). "
]
},
{
@@ -27,15 +27,9 @@
"from primaite.simulator.system.services.terminal.terminal import Terminal\n",
"from primaite.simulator.network.container import Network\n",
"from primaite.simulator.network.hardware.nodes.host.computer import Computer\n",
"from primaite.simulator.system.applications.red_applications.ransomware_script import RansomwareScript"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from primaite.simulator.system.applications.red_applications.ransomware_script import RansomwareScript\n",
"from primaite.simulator.system.services.terminal.terminal import RemoteTerminalConnection\n",
"\n",
"def basic_network() -> Network:\n",
" \"\"\"Utility function for creating a default network to demonstrate Terminal functionality\"\"\"\n",
" network = Network()\n",
@@ -51,9 +45,9 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"The terminal can be accessed from a `HostNode` via the `software_manager` as demonstrated below. \n",
"The terminal can be accessed from a `Node` via the `software_manager` as demonstrated below. \n",
"\n",
"In the example, we have a basic network consisting of two computers "
"In the example, we have a basic network consisting of two computers, connected to form a basic network."
]
},
{
@@ -66,15 +60,17 @@
"computer_a: Computer = network.get_node_by_hostname(\"node_a\")\n",
"terminal_a: Terminal = computer_a.software_manager.software.get(\"Terminal\")\n",
"computer_b: Computer = network.get_node_by_hostname(\"node_b\")\n",
"terminal_b: Terminal = computer_b.software_manager.software.get(\"Terminal\")\n"
"terminal_b: Terminal = computer_b.software_manager.software.get(\"Terminal\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"To be able to send commands from `node_a` to `node_b`, you will need to `login` to `node_b` first, using valid user credentials. In the example below, we are logging in to the 'admin' account on `node_b`. \n",
"If you are not logged in, any commands sent will be rejected."
"To be able to send commands from `node_a` to `node_b`, you will need to `login` to `node_b` first, using valid user credentials. In the example below, we are remotely logging in to the 'admin' account on `node_b`, from `node_a`. \n",
"If you are not logged in, any commands sent will be rejected by the remote.\n",
"\n",
"Remote Logins return a RemoteTerminalConnection object, which can be used for sending commands to the remote node. "
]
},
{
@@ -84,10 +80,14 @@
"outputs": [],
"source": [
"# Login to the remote (node_b) from local (node_a)\n",
"from primaite.simulator.system.services.terminal.terminal import RemoteTerminalConnection\n",
"\n",
"\n",
"term_a_term_b_remote_connection: RemoteTerminalConnection = terminal_a.login(username=\"admin\", password=\"Admin123!\", ip_address=computer_b.network_interface[1].ip_address)"
"term_a_term_b_remote_connection: RemoteTerminalConnection = terminal_a.login(username=\"admin\", password=\"Admin123!\", ip_address=\"192.168.0.11\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"You can view all active connections to a terminal through use of the `show()` method"
]
},
{
@@ -96,7 +96,14 @@
"metadata": {},
"outputs": [],
"source": [
"computer_b.software_manager.show()"
"terminal_b.show()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The new connection object allows us to forward commands to be executed on the target node. The example below demonstrates how you can remotely install an application on the target node."
]
},
{
@@ -105,7 +112,6 @@
"metadata": {},
"outputs": [],
"source": [
"print(type(term_a_term_b_remote_connection))\n",
"term_a_term_b_remote_connection.execute([\"software_manager\", \"application\", \"install\", \"RansomwareScript\"])"
]
},
@@ -122,7 +128,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"You can view all remote connections to a terminal through use of the `show()` method"
"The code block below demonstrates how the Terminal class allows the user of `terminal_a`, on `computer_a`, to send a command to `computer_b` to create a downloads folder. \n"
]
},
{
@@ -131,23 +137,11 @@
"metadata": {},
"outputs": [],
"source": [
"terminal_b.show()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The Terminal can be used to send requests to install new software. The code block below demonstrates how the Terminal class allows the user of `terminal_a`, on `computer_a`, to send a command to `computer_b` to install the `RansomwareScript` application. \n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The below example shows how you can send a command via the terminal to create a folder on the target Node.\n",
"# Display the current state of the file system on computer_b\n",
"computer_b.file_system.show()\n",
"\n",
"Here, we send a command to `computer_b` to create a new folder titled \"Downloads\"."
"# Send command\n",
"term_a_term_b_remote_connection.execute([\"file_system\", \"create\", \"folder\", \"downloads\"])"
]
},
{
@@ -156,6 +150,39 @@
"source": [
"The resultant call to `computer_b.file_system.show()` shows that the new folder has been created."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"computer_b.file_system.show()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"When finished, the connection can be closed by calling the `disconnect` function of the Remote Client object"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Display active connection\n",
"terminal_a.show()\n",
"terminal_b.show()\n",
"\n",
"term_a_term_b_remote_connection.disconnect()\n",
"\n",
"terminal_a.show()\n",
"\n",
"terminal_b.show()"
]
}
],
"metadata": {

186
src/primaite/simulator/system/services/terminal/terminal.py
View File

@@ -2,7 +2,7 @@
from __future__ import annotations
from ipaddress import IPv4Address
from typing import Any, Dict, List, Optional
from typing import Any, Dict, List, Optional, Union
from uuid import uuid4
from prettytable import MARKDOWN, PrettyTable
@@ -10,7 +10,12 @@ from pydantic import BaseModel
from primaite.interface.request import RequestResponse
from primaite.simulator.core import RequestManager, RequestType
from primaite.simulator.network.protocols.ssh import SSHPacket
from primaite.simulator.network.protocols.ssh import (
SSHConnectionMessage,
SSHPacket,
SSHTransportMessage,
SSHUserCredentials,
)
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
@@ -33,6 +38,9 @@ class TerminalClientConnection(BaseModel):
connection_uuid: str = None
"""Connection UUID"""
connection_request_id: str = None
"""Connection request ID"""
@property
def client(self) -> Optional[Terminal]:
"""The Terminal that holds this connection."""
@@ -51,6 +59,9 @@ class RemoteTerminalConnection(TerminalClientConnection):
"""
source_ip: IPv4Address
"""Source IP of Connection"""
def execute(self, command: Any) -> bool:
"""Execute a given command on the remote Terminal."""
if self.parent_terminal.operating_state != ServiceOperatingState.RUNNING:
@@ -88,13 +99,13 @@ class Terminal(Service):
:param markdown: Whether to display the table in Markdown format or not. Default is `False`.
"""
table = PrettyTable(["Connection ID", "Session_ID"])
table = PrettyTable(["Connection ID", "Connection request ID", "Source IP"])
if markdown:
table.set_style(MARKDOWN)
table.align = "l"
table.title = f"{self.sys_log.hostname} {self.name} Connections"
for connection_id, connection in self._connections.items():
table.add_row([connection_id, connection.session_id])
table.add_row([connection_id, connection.connection_request_id, connection.source_ip])
print(table.get_string(sortby="Connection ID"))
def _init_request_manager(self) -> RequestManager:
@@ -130,7 +141,7 @@ class Terminal(Service):
connection_uuid = request[0]
# TODO: Uncomment this when UserSessionManager merged.
# self.parent.UserSessionManager.logoff(connection_uuid)
self.disconnect(connection_uuid)
self._disconnect(connection_uuid)
return RequestResponse(status="success", data={})
@@ -157,8 +168,13 @@ class Terminal(Service):
"""Execute a passed ssh command via the request manager."""
return self.parent.apply_request(command)
def _create_local_connection(self, connection_uuid: str, session_id: str) -> RemoteTerminalConnection:
"""Create a new connection object and amend to list of active connections."""
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 = TerminalClientConnection(
parent_terminal=self,
connection_uuid=connection_uuid,
@@ -172,7 +188,17 @@ class Terminal(Service):
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."""
"""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("Cannot login as service is not running.")
return None
@@ -199,8 +225,10 @@ class Terminal(Service):
if connection_uuid:
self.sys_log.info(f"Login request authorised, connection uuid: {connection_uuid}")
# Add new local session to list of connections
self._create_local_connection(connection_uuid=connection_uuid, session_id="")
return TerminalClientConnection(parent_terminal=self, session_id="", connection_uuid=connection_uuid)
self._create_local_connection(connection_uuid=connection_uuid, session_id="Local_Connection")
return TerminalClientConnection(
parent_terminal=self, session_id="Local_Connection", connection_uuid=connection_uuid
)
else:
self.sys_log.warning("Login failed, incorrect Username or Password")
return None
@@ -217,7 +245,26 @@ class Terminal(Service):
connection_request_id: str,
is_reattempt: bool = False,
) -> Optional[RemoteTerminalConnection]:
"""Process a remote login attempt."""
"""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
:type: connection_request_id: 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.
"""
self.sys_log.info(f"Sending Remote login attempt to {ip_address}. Connection_id is {connection_request_id}")
if is_reattempt:
valid_connection = self._check_client_connection(connection_id=connection_request_id)
@@ -229,12 +276,24 @@ class Terminal(Service):
self.sys_log.warning(f"{self.name}: Remote connection to {ip_address} declined.")
return None
payload = {
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,
)
software_manager: SoftwareManager = self.software_manager
software_manager.send_payload_to_session_manager(
payload=payload, dest_ip_address=ip_address, dest_port=self.port
@@ -247,15 +306,28 @@ class Terminal(Service):
connection_request_id=connection_request_id,
)
def _create_remote_connection(self, connection_id: str, connection_request_id: str, session_id: str) -> None:
"""Create a new TerminalClientConnection Object."""
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, session_id=session_id, connection_uuid=connection_id
parent_terminal=self,
session_id=session_id,
connection_uuid=connection_id,
source_ip=source_ip,
connection_request_id=connection_request_id,
)
self._connections[connection_id] = client_connection
self._client_connection_requests[connection_request_id] = client_connection
def receive(self, session_id: str, payload: Any, **kwargs) -> bool:
def receive(self, session_id: str, payload: Union[SSHPacket, Dict, List], **kwargs) -> bool:
"""
Receive a payload from the Software Manager.
@@ -263,42 +335,62 @@ class Terminal(Service):
:param session_id: The session id the payload relates to.
:return: True.
"""
self.sys_log.info(f"Received payload: {payload}")
if isinstance(payload, dict) and payload.get("type"):
if payload["type"] == "login_request":
# add connection
connection_request_id = payload["connection_request_id"]
username = payload["username"]
password = payload["password"]
print(f"Connection ID is: {connection_request_id}")
self.sys_log.info(f"Connection authorised, session_id: {session_id}")
source_ip = kwargs["from_network_interface"].ip_address
self.sys_log.info(f"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
# connection_id = self.parent.UserSessionManager.login(username=username, password=password)
connection_id = str(uuid4())
if connection_id:
connection_request_id = payload.connection_request_uuid
username = payload.user_account.username
password = payload.user_account.password
print(f"Connection ID is: {connection_request_id}")
self.sys_log.info(f"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("Login Successful")
self._create_remote_connection(
connection_id=connection_request_id,
connection_request_id=payload["connection_request_id"],
connection_id=payload.connection_uuid,
connection_request_id=payload.connection_request_uuid,
session_id=session_id,
)
payload = {
"type": "login_success",
"username": username,
"password": password,
"connection_request_id": connection_request_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["type"] == "login_success":
self.sys_log.info(f"Login was successful! session_id is:{session_id}")
connection_request_id = payload["connection_request_id"]
self._create_remote_connection(
connection_id=connection_request_id,
session_id=session_id,
connection_request_id=connection_request_id,
source_ip=source_ip,
)
elif payload["type"] == "disconnect":
if isinstance(payload, dict) and payload.get("type"):
if payload["type"] == "disconnect":
connection_id = payload["connection_id"]
self.sys_log.info(f"{self.name}: Received disconnect command for {connection_id=} from the server")
self.sys_log.info(f"{self.name}: Received disconnect command for {connection_id=} from remote.")
self._disconnect(payload["connection_id"])
if isinstance(payload, list):