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

#2769 - add actions and tests for terminal

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files

28
src/primaite/game/agent/actions.py
View File

@@ -1101,15 +1101,14 @@ class NodeSessionsRemoteLoginAction(AbstractAction):
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."""
# 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)
return [
"network",
"node",
node_name,
"service",
"UserSessionManager",
"remote_login",
"Terminal",
"ssh_to_remote",
username,
password,
remote_ip,
@@ -1122,11 +1121,21 @@ class NodeSessionsRemoteLogoutAction(AbstractAction):
def __init__(self, manager: "ActionManager", **kwargs) -> None:
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."""
# 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)
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:
@@ -1180,9 +1189,10 @@ class ActionManager:
"CONFIGURE_DATABASE_CLIENT": ConfigureDatabaseClientAction,
"CONFIGURE_RANSOMWARE_SCRIPT": ConfigureRansomwareScriptAction,
"CONFIGURE_DOSBOT": ConfigureDoSBotAction,
"NODE_ACCOUNTS_CHANGEPASSWORD": NodeAccountsChangePasswordAction,
"NODE_SESSIONS_REMOTE_LOGIN": NodeSessionsRemoteLoginAction,
"NODE_SESSIONS_REMOTE_LOGOUT": NodeSessionsRemoteLogoutAction,
"NODE_ACCOUNTS_CHANGE_PASSWORD": NodeAccountsChangePasswordAction,
"SSH_TO_REMOTE": NodeSessionsRemoteLoginAction,
"SSH_LOGOUT_LOGOUT": NodeSessionsRemoteLogoutAction,
"NODE_SEND_REMOTE_COMMAND": NodeSendRemoteCommandAction,
}
"""Dictionary which maps action type strings to the corresponding action class."""

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

@@ -92,7 +92,7 @@ class LocalTerminalConnection(TerminalClientConnection):
if not self.is_active:
self.parent_terminal.sys_log.warning("Connection inactive, cannot execute")
return None
return self.parent_terminal.execute(command, connection_id=self.connection_uuid)
return self.parent_terminal.execute(command)
class RemoteTerminalConnection(TerminalClientConnection):
@@ -162,22 +162,36 @@ class Terminal(Service):
def _init_request_manager(self) -> RequestManager:
"""Initialise Request manager."""
rm = super()._init_request_manager()
rm.add_request(
"send",
request_type=RequestType(func=lambda request, context: RequestResponse.from_bool(self.send())),
)
# rm.add_request(
# "send",
# request_type=RequestType(func=lambda request, context: RequestResponse.from_bool(self.send())),
# )
def _login(request: RequestFormat, context: Dict) -> RequestResponse:
login = self._process_local_login(username=request[0], password=request[1])
if login:
return RequestResponse(
status="success",
data={
"ip_address": login.ip_address,
},
)
else:
return RequestResponse(status="failure", data={"reason": "Invalid login credentials"})
# def _login(request: RequestFormat, context: Dict) -> RequestResponse:
# login = self._process_local_login(username=request[0], password=request[1])
# if login:
# return RequestResponse(
# status="success",
# data={
# "ip_address": login.ip_address,
# },
# )
# else:
# 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:
login = self._send_remote_login(username=request[0], password=request[1], ip_address=request[2])
@@ -191,10 +205,34 @@ class Terminal(Service):
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={},
)
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:
"""Execute an instruction."""
command: str = request[0]
ip_address: IPv4Address = IPv4Address(request[1])
ip_address: IPv4Address = IPv4Address(request[0])
command: str = request[1]
remote_connection = self._get_connection_from_ip(ip_address=ip_address)
if remote_connection:
outcome = remote_connection.execute(command)
@@ -209,30 +247,11 @@ class Terminal(Service):
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(
"Login",
request_type=RequestType(func=_login),
)
rm.add_request(
"Remote Login",
request_type=RequestType(func=_remote_login),
)
rm.add_request(
"Execute",
"send_remote_command",
request_type=RequestType(func=remote_execute_request),
)
rm.add_request("Logoff", request_type=RequestType(func=_logoff))
return rm
def execute(self, command: List[Any]) -> Optional[RequestResponse]:
@@ -280,13 +299,9 @@ class Terminal(Service):
if self.operating_state != ServiceOperatingState.RUNNING:
self.sys_log.warning(f"{self.name}: Cannot login as service is not running.")
return None
connection_request_id = str(uuid4())
self._client_connection_requests[connection_request_id] = 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, connection_request_id=connection_request_id
)
return self._send_remote_login(username=username, password=password, ip_address=ip_address)
else:
return self._process_local_login(username=username, password=password)
@@ -320,32 +335,24 @@ class Terminal(Service):
username: str,
password: str,
ip_address: IPv4Address,
connection_request_id: str,
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
:type: connection_request_id: str
: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.
"""
self.sys_log.info(
f"{self.name}: Sending Remote login attempt to {ip_address}. Connection_id is {connection_request_id}"
)
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:
@@ -360,6 +367,9 @@ class Terminal(Service):
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)

7
tests/conftest.py
View File

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

165
tests/integration_tests/game_layer/actions/test_terminal_actions.py Normal file
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

49
tests/integration_tests/game_layer/actions/user_account_actions/test_remote_user_account_actions.py
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

23
tests/integration_tests/game_layer/actions/user_account_actions/test_user_account_change_password.py
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"