#2769 - add actions and tests for terminal

This commit is contained in:
Marek Wolan
2024-08-11 23:24:29 +01:00
parent f92a57cfc4
commit 3df55a708d
6 changed files with 253 additions and 139 deletions

View File

@@ -1101,15 +1101,14 @@ class NodeSessionsRemoteLoginAction(AbstractAction):
def form_request(self, node_id: str, username: str, password: str, remote_ip: str) -> RequestFormat: def form_request(self, node_id: str, username: str, password: str, remote_ip: str) -> RequestFormat:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" """Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
# TODO: change this so it creates a remote connection using terminal rather than a local remote login
node_name = self.manager.get_node_name_by_idx(node_id) node_name = self.manager.get_node_name_by_idx(node_id)
return [ return [
"network", "network",
"node", "node",
node_name, node_name,
"service", "service",
"UserSessionManager", "Terminal",
"remote_login", "ssh_to_remote",
username, username,
password, password,
remote_ip, remote_ip,
@@ -1122,11 +1121,21 @@ class NodeSessionsRemoteLogoutAction(AbstractAction):
def __init__(self, manager: "ActionManager", **kwargs) -> None: def __init__(self, manager: "ActionManager", **kwargs) -> None:
super().__init__(manager=manager) super().__init__(manager=manager)
def form_request(self, node_id: str, remote_session_id: str) -> RequestFormat: def form_request(self, node_id: str, remote_ip: str) -> RequestFormat:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" """Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
# TODO: change this so it destroys a remote connection using terminal rather than a local remote login
node_name = self.manager.get_node_name_by_idx(node_id) node_name = self.manager.get_node_name_by_idx(node_id)
return ["network", "node", node_name, "service", "UserSessionManager", "remote_logout", remote_session_id] return ["network", "node", node_name, "service", "Terminal", "remote_logoff", remote_ip]
class NodeSendRemoteCommandAction(AbstractAction):
"""Action which sends a terminal command to a remote node via SSH."""
def __init__(self, manager: "ActionManager", **kwargs) -> None:
super().__init__(manager=manager)
def form_request(self, node_id: int, remote_ip: str, command: RequestFormat) -> RequestFormat:
node_name = self.manager.get_node_name_by_idx(node_id)
return ["network", "node", node_name, "service", "Terminal", "send_remote_command", remote_ip, command]
class ActionManager: class ActionManager:
@@ -1180,9 +1189,10 @@ class ActionManager:
"CONFIGURE_DATABASE_CLIENT": ConfigureDatabaseClientAction, "CONFIGURE_DATABASE_CLIENT": ConfigureDatabaseClientAction,
"CONFIGURE_RANSOMWARE_SCRIPT": ConfigureRansomwareScriptAction, "CONFIGURE_RANSOMWARE_SCRIPT": ConfigureRansomwareScriptAction,
"CONFIGURE_DOSBOT": ConfigureDoSBotAction, "CONFIGURE_DOSBOT": ConfigureDoSBotAction,
"NODE_ACCOUNTS_CHANGEPASSWORD": NodeAccountsChangePasswordAction, "NODE_ACCOUNTS_CHANGE_PASSWORD": NodeAccountsChangePasswordAction,
"NODE_SESSIONS_REMOTE_LOGIN": NodeSessionsRemoteLoginAction, "SSH_TO_REMOTE": NodeSessionsRemoteLoginAction,
"NODE_SESSIONS_REMOTE_LOGOUT": NodeSessionsRemoteLogoutAction, "SSH_LOGOUT_LOGOUT": NodeSessionsRemoteLogoutAction,
"NODE_SEND_REMOTE_COMMAND": NodeSendRemoteCommandAction,
} }
"""Dictionary which maps action type strings to the corresponding action class.""" """Dictionary which maps action type strings to the corresponding action class."""

View File

@@ -92,7 +92,7 @@ class LocalTerminalConnection(TerminalClientConnection):
if not self.is_active: if not self.is_active:
self.parent_terminal.sys_log.warning("Connection inactive, cannot execute") self.parent_terminal.sys_log.warning("Connection inactive, cannot execute")
return None return None
return self.parent_terminal.execute(command, connection_id=self.connection_uuid) return self.parent_terminal.execute(command)
class RemoteTerminalConnection(TerminalClientConnection): class RemoteTerminalConnection(TerminalClientConnection):
@@ -162,22 +162,36 @@ class Terminal(Service):
def _init_request_manager(self) -> RequestManager: def _init_request_manager(self) -> RequestManager:
"""Initialise Request manager.""" """Initialise Request manager."""
rm = super()._init_request_manager() rm = super()._init_request_manager()
rm.add_request( # rm.add_request(
"send", # "send",
request_type=RequestType(func=lambda request, context: RequestResponse.from_bool(self.send())), # request_type=RequestType(func=lambda request, context: RequestResponse.from_bool(self.send())),
) # )
def _login(request: RequestFormat, context: Dict) -> RequestResponse: # def _login(request: RequestFormat, context: Dict) -> RequestResponse:
login = self._process_local_login(username=request[0], password=request[1]) # login = self._process_local_login(username=request[0], password=request[1])
if login: # if login:
return RequestResponse( # return RequestResponse(
status="success", # status="success",
data={ # data={
"ip_address": login.ip_address, # "ip_address": login.ip_address,
}, # },
) # )
else: # else:
return RequestResponse(status="failure", data={"reason": "Invalid login credentials"}) # return RequestResponse(status="failure", data={"reason": "Invalid login credentials"})
#
# rm.add_request(
# "Login",
# request_type=RequestType(func=_login),
# )
# def _logoff(request: RequestFormat, context: Dict) -> RequestResponse:
# """Logoff from connection."""
# connection_uuid = request[0]
# self.parent.user_session_manager.local_logout(connection_uuid)
# self._disconnect(connection_uuid)
# return RequestResponse(status="success", data={})
#
# rm.add_request("Logoff", request_type=RequestType(func=_logoff))
def _remote_login(request: RequestFormat, context: Dict) -> RequestResponse: def _remote_login(request: RequestFormat, context: Dict) -> RequestResponse:
login = self._send_remote_login(username=request[0], password=request[1], ip_address=request[2]) login = self._send_remote_login(username=request[0], password=request[1], ip_address=request[2])
@@ -191,10 +205,34 @@ class Terminal(Service):
else: else:
return RequestResponse(status="failure", data={}) 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={},
)
else:
return RequestResponse(
status="failure",
data={"reason": "No remote connection held."},
)
rm.add_request("remote_logoff", request_type=RequestType(func=_remote_logoff))
def remote_execute_request(request: RequestFormat, context: Dict) -> RequestResponse: def remote_execute_request(request: RequestFormat, context: Dict) -> RequestResponse:
"""Execute an instruction.""" """Execute an instruction."""
command: str = request[0] ip_address: IPv4Address = IPv4Address(request[0])
ip_address: IPv4Address = IPv4Address(request[1]) command: str = request[1]
remote_connection = self._get_connection_from_ip(ip_address=ip_address) remote_connection = self._get_connection_from_ip(ip_address=ip_address)
if remote_connection: if remote_connection:
outcome = remote_connection.execute(command) outcome = remote_connection.execute(command)
@@ -209,30 +247,11 @@ class Terminal(Service):
data={}, data={},
) )
def _logoff(request: RequestFormat, context: Dict) -> RequestResponse:
"""Logoff from connection."""
connection_uuid = request[0]
self.parent.user_session_manager.local_logout(connection_uuid)
self._disconnect(connection_uuid)
return RequestResponse(status="success", data={})
rm.add_request( rm.add_request(
"Login", "send_remote_command",
request_type=RequestType(func=_login),
)
rm.add_request(
"Remote Login",
request_type=RequestType(func=_remote_login),
)
rm.add_request(
"Execute",
request_type=RequestType(func=remote_execute_request), request_type=RequestType(func=remote_execute_request),
) )
rm.add_request("Logoff", request_type=RequestType(func=_logoff))
return rm return rm
def execute(self, command: List[Any]) -> Optional[RequestResponse]: def execute(self, command: List[Any]) -> Optional[RequestResponse]:
@@ -280,13 +299,9 @@ class Terminal(Service):
if self.operating_state != ServiceOperatingState.RUNNING: if self.operating_state != ServiceOperatingState.RUNNING:
self.sys_log.warning(f"{self.name}: Cannot login as service is not running.") self.sys_log.warning(f"{self.name}: Cannot login as service is not running.")
return None return None
connection_request_id = str(uuid4())
self._client_connection_requests[connection_request_id] = None
if ip_address: if ip_address:
# Assuming that if IP is passed we are connecting to remote # Assuming that if IP is passed we are connecting to remote
return self._send_remote_login( return self._send_remote_login(username=username, password=password, ip_address=ip_address)
username=username, password=password, ip_address=ip_address, connection_request_id=connection_request_id
)
else: else:
return self._process_local_login(username=username, password=password) return self._process_local_login(username=username, password=password)
@@ -320,32 +335,24 @@ class Terminal(Service):
username: str, username: str,
password: str, password: str,
ip_address: IPv4Address, ip_address: IPv4Address,
connection_request_id: str, connection_request_id: Optional[str] = None,
is_reattempt: bool = False, is_reattempt: bool = False,
) -> Optional[RemoteTerminalConnection]: ) -> Optional[RemoteTerminalConnection]:
"""Send a remote login attempt and connect to Node. """Send a remote login attempt and connect to Node.
:param: username: Username used to connect to the remote node. :param: username: Username used to connect to the remote node.
:type: username: str :type: username: str
:param: password: Password used to connect to the remote node :param: password: Password used to connect to the remote node
:type: password: str :type: password: str
:param: ip_address: Target Node IP address for login attempt. :param: ip_address: Target Node IP address for login attempt.
:type: ip_address: IPv4Address :type: ip_address: IPv4Address
:param: connection_request_id: Connection Request ID, if not provided, a new one is generated
:param: connection_request_id: Connection Request ID :type: connection_request_id: Optional[str]
:type: connection_request_id: str
:param: is_reattempt: True if the request has been reattempted. Default False. :param: is_reattempt: True if the request has been reattempted. Default False.
:type: is_reattempt: Optional[bool] :type: is_reattempt: Optional[bool]
:return: RemoteTerminalConnection: Connection Object for sending further commands if successful, else False. :return: RemoteTerminalConnection: Connection Object for sending further commands if successful, else False.
""" """
self.sys_log.info( connection_request_id = connection_request_id or str(uuid4())
f"{self.name}: Sending Remote login attempt to {ip_address}. Connection_id is {connection_request_id}"
)
if is_reattempt: if is_reattempt:
valid_connection_request = self._validate_client_connection_request(connection_id=connection_request_id) valid_connection_request = self._validate_client_connection_request(connection_id=connection_request_id)
if valid_connection_request: if valid_connection_request:
@@ -360,6 +367,9 @@ class Terminal(Service):
self.sys_log.warning(f"{self.name}: Remote connection to {ip_address} declined.") self.sys_log.warning(f"{self.name}: Remote connection to {ip_address} declined.")
return None 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 transport_message: SSHTransportMessage = SSHTransportMessage.SSH_MSG_USERAUTH_REQUEST
connection_message: SSHConnectionMessage = SSHConnectionMessage.SSH_MSG_CHANNEL_DATA connection_message: SSHConnectionMessage = SSHConnectionMessage.SSH_MSG_CHANNEL_DATA
user_details: SSHUserCredentials = SSHUserCredentials(username=username, password=password) user_details: SSHUserCredentials = SSHUserCredentials(username=username, password=password)

View File

@@ -458,9 +458,10 @@ def game_and_agent():
{"type": "HOST_NIC_DISABLE"}, {"type": "HOST_NIC_DISABLE"},
{"type": "NETWORK_PORT_ENABLE"}, {"type": "NETWORK_PORT_ENABLE"},
{"type": "NETWORK_PORT_DISABLE"}, {"type": "NETWORK_PORT_DISABLE"},
{"type": "NODE_ACCOUNTS_CHANGEPASSWORD"}, {"type": "NODE_ACCOUNTS_CHANGE_PASSWORD"},
{"type": "NODE_SESSIONS_REMOTE_LOGIN"}, {"type": "SSH_TO_REMOTE"},
{"type": "NODE_SESSIONS_REMOTE_LOGOUT"}, {"type": "SSH_LOGOUT_LOGOUT"},
{"type": "NODE_SEND_REMOTE_COMMAND"},
] ]
action_space = ActionManager( action_space = ActionManager(

View File

@@ -0,0 +1,165 @@
# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
from typing import Tuple
import pytest
from primaite.game.agent.interface import ProxyAgent
from primaite.game.game import PrimaiteGame
from primaite.simulator.network.hardware.base import UserManager
from primaite.simulator.network.hardware.nodes.host.computer import Computer
from primaite.simulator.network.hardware.nodes.host.server import Server
from primaite.simulator.system.services.service import ServiceOperatingState
from primaite.simulator.system.services.terminal.terminal import RemoteTerminalConnection
@pytest.fixture
def game_and_agent_fixture(game_and_agent):
"""Create a game with a simple agent that can be controlled by the tests."""
game, agent = game_and_agent
client_1: Computer = game.simulation.network.get_node_by_hostname("client_1")
client_1.start_up_duration = 3
return (game, agent)
def test_remote_login(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAgent]):
game, agent = game_and_agent_fixture
server_1: Server = game.simulation.network.get_node_by_hostname("server_1")
client_1 = game.simulation.network.get_node_by_hostname("client_1")
# create a new user account on server_1 that will be logged into remotely
server_1_usm: UserManager = server_1.software_manager.software["UserManager"]
server_1_usm.add_user("user123", "password", is_admin=True)
action = (
"SSH_TO_REMOTE",
{
"node_id": 0,
"username": "user123",
"password": "password",
"remote_ip": str(server_1.network_interface[1].ip_address),
},
)
agent.store_action(action)
game.step()
assert agent.history[-1].response.status == "success"
connection_established = False
for conn_str, conn_obj in client_1.terminal.connections.items():
conn_obj: RemoteTerminalConnection
if conn_obj.ip_address == server_1.network_interface[1].ip_address:
connection_established = True
if not connection_established:
pytest.fail("Remote SSH connection could not be established")
def test_remote_login_wrong_password(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAgent]):
game, agent = game_and_agent_fixture
server_1: Server = game.simulation.network.get_node_by_hostname("server_1")
client_1 = game.simulation.network.get_node_by_hostname("client_1")
# create a new user account on server_1 that will be logged into remotely
server_1_usm: UserManager = server_1.software_manager.software["UserManager"]
server_1_usm.add_user("user123", "password", is_admin=True)
action = (
"SSH_TO_REMOTE",
{
"node_id": 0,
"username": "user123",
"password": "wrong_password",
"remote_ip": str(server_1.network_interface[1].ip_address),
},
)
agent.store_action(action)
game.step()
assert agent.history[-1].response.status == "failure"
connection_established = False
for conn_str, conn_obj in client_1.terminal.connections.items():
conn_obj: RemoteTerminalConnection
if conn_obj.ip_address == server_1.network_interface[1].ip_address:
connection_established = True
if connection_established:
pytest.fail("Remote SSH connection was established despite wrong password")
def test_remote_login_change_password(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAgent]):
game, agent = game_and_agent_fixture
server_1: Server = game.simulation.network.get_node_by_hostname("server_1")
client_1 = game.simulation.network.get_node_by_hostname("client_1")
# create a new user account on server_1 that will be logged into remotely
server_1_um: UserManager = server_1.software_manager.software["UserManager"]
server_1_um.add_user("user123", "password", is_admin=True)
action = (
"NODE_ACCOUNTS_CHANGE_PASSWORD",
{
"node_id": 1, # server_1
"username": "user123",
"current_password": "password",
"new_password": "different_password",
},
)
agent.store_action(action)
game.step()
assert agent.history[-1].response.status == "success"
assert server_1_um.users["user123"].password == "different_password"
def test_change_password_logs_out_user(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAgent]):
game, agent = game_and_agent_fixture
server_1: Server = game.simulation.network.get_node_by_hostname("server_1")
client_1 = game.simulation.network.get_node_by_hostname("client_1")
# create a new user account on server_1 that will be logged into remotely
server_1_usm: UserManager = server_1.software_manager.software["UserManager"]
server_1_usm.add_user("user123", "password", is_admin=True)
# Log in remotely
action = (
"SSH_TO_REMOTE",
{
"node_id": 0,
"username": "user123",
"password": "password",
"remote_ip": str(server_1.network_interface[1].ip_address),
},
)
agent.store_action(action)
game.step()
# Change password
action = (
"NODE_ACCOUNTS_CHANGE_PASSWORD",
{
"node_id": 1, # server_1
"username": "user123",
"current_password": "password",
"new_password": "different_password",
},
)
agent.store_action(action)
game.step()
# Assert that the user cannot execute an action
# TODO: should the db conn object get destroyed on both nodes? or is that not realistic?
action = (
"NODE_SEND_REMOTE_COMMAND",
{
"node_id": 0,
"remote_ip": server_1.network_interface[1].ip_address,
"command": ["file_system", "create", "file", "folder123", "doggo.pdf", False],
},
)
agent.store_action(action)
game.step()
assert server_1.file_system.get_folder("folder123") is None
assert server_1.file_system.get_file("folder123", "doggo.pdf") is None

View File

@@ -1,49 +0,0 @@
# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
from primaite.simulator.network.hardware.nodes.host.computer import Computer
def test_remote_logon(game_and_agent):
"""Test that the remote session login action works."""
game, agent = game_and_agent
client_1: Computer = game.simulation.network.get_node_by_hostname("client_1")
client_1.user_manager.add_user(username="test_user", password="password", bypass_can_perform_action=True)
action = (
"NODE_SESSIONS_REMOTE_LOGIN",
{"node_id": 0, "username": "test_user", "password": "password", "remote_ip": "10.0.2.2"},
)
agent.store_action(action)
game.step()
assert len(client_1.user_session_manager.remote_sessions) == 1
def test_remote_logoff(game_and_agent):
"""Test that the remote session logout action works."""
game, agent = game_and_agent
client_1: Computer = game.simulation.network.get_node_by_hostname("client_1")
client_1.user_manager.add_user(username="test_user", password="password", bypass_can_perform_action=True)
action = (
"NODE_SESSIONS_REMOTE_LOGIN",
{"node_id": 0, "username": "test_user", "password": "password"},
)
agent.store_action(action)
game.step()
assert len(client_1.user_session_manager.remote_sessions) == 1
remote_session_id = client_1.user_session_manager.remote_sessions[0].uuid
action = (
"NODE_SESSIONS_REMOTE_LOGOUT",
{"node_id": 0, "remote_session_id": remote_session_id},
)
agent.store_action(action)
game.step()
assert len(client_1.user_session_manager.remote_sessions) == 0

View File

@@ -1,23 +0,0 @@
# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
from primaite.simulator.network.hardware.nodes.host.computer import Computer
def test_remote_logon(game_and_agent):
"""Test that the remote session login action works."""
game, agent = game_and_agent
client_1: Computer = game.simulation.network.get_node_by_hostname("client_1")
client_1.user_manager.add_user(username="test_user", password="password", bypass_can_perform_action=True)
user = next((user for user in client_1.user_manager.users.values() if user.username == "test_user"), None)
assert user.password == "password"
action = (
"NODE_ACCOUNTS_CHANGEPASSWORD",
{"node_id": 0, "username": user.username, "current_password": user.password, "new_password": "test_pass"},
)
agent.store_action(action)
game.step()
assert user.password == "test_pass"