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

Merge remote-tracking branch 'origin/dev' into feature/2445-make-observation-thresholds-configurable

Browse Source
This commit is contained in:
Czar Echavez
2024-09-25 09:04:18 +01:00
parent 4f172634d0 b0363a6b48
commit 171dd83f2f
21 changed files with 679 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
CHANGELOG.md
View File

@@ -5,13 +5,18 @@ All notable changes to this project will be documented in this file.
The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
## [3.3.0] - 2024-09-04
## [3.4.0]
### Added
- Log observation space data by episode and step.
- Added `show_history` method to Agents, allowing you to view actions taken by an agent per step. By default, `DONOTHING` actions are omitted.
- New ``NODE_SEND_LOCAL_COMMAND`` action implemented which grants agents the ability to execute commands locally. (Previously limited to remote only)
### Changed
- ACL's are no longer applied to layer-2 traffic.
- Random number seed values are recorded in simulation/seed.log if the seed is set in the config file
or `generate_seed_value` is set to `true`.
- ARP .show() method will now include the port number associated with each entry.
- Added `services_requires_scan` and `applications_requires_scan` to agent observation space config to allow the agents to be able to see actual health states of services and applications without requiring scans (Default `True`, set to `False` to allow agents to see actual health state without scanning).

1
docs/source/configuration/agents.rst
View File

@@ -177,3 +177,4 @@ If ``True``, gymnasium flattening will be performed on the observation space bef
-----------------
Agents will record their action log for each step. This is a summary of what the agent did, along with response information from requests within the simulation.
A summary of the actions taken by the agent can be viewed using the `show_history()` function. By default, this will display all actions taken apart from ``DONOTHING``.

2
docs/source/request_system.rst
View File

@@ -2,6 +2,8 @@
© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
.. _request_system:
Request System
**************

119
docs/source/simulation_components/system/services/terminal.rst
View File

@@ -23,6 +23,14 @@ Key capabilities
- Simulates common Terminal processes/commands.
- Leverages the Service base class for install/uninstall, status tracking etc.
Usage
"""""
- Pre-Installs on any `Node` component (with the exception of `Switches`).
- Terminal Clients connect, execute commands and disconnect from remote nodes.
- Ensures that users are logged in to the component before executing any commands.
- Service runs on SSH port 22 by default.
- Enables Agents to send commands both remotely and locally.
Implementation
""""""""""""""
@@ -30,19 +38,112 @@ Implementation
- Manages remote connections in a dictionary by session ID.
- Processes commands, forwarding to the ``RequestManager`` or ``SessionManager`` where appropriate.
- Extends Service class.
- A detailed guide on the implementation and functionality of the Terminal class can be found in the "Terminal-Processing" jupyter notebook.
A detailed guide on the implementation and functionality of the Terminal class can be found in the "Terminal-Processing" jupyter notebook.
Command Format
^^^^^^^^^^^^^^
Terminals implement their commands through leveraging the pre-existing :ref:`request_system`.
Due to this Terminals will only accept commands passed within the ``RequestFormat``.
:py:class:`primaite.game.interface.RequestFormat`
For example, ``terminal`` command actions when used in ``yaml`` format are formatted as follows:
.. code-block:: yaml
command:
- "file_system"
- "create"
- "file"
- "downloads"
- "cat.png"
- "False
This is then loaded from yaml into a dictionary containing the terminal command:
.. code-block:: python
{"command":["file_system", "create", "file", "downloads", "cat.png", "False"]}
Which is then passed to the ``Terminals`` Request Manager to be executed.
Game Layer Usage (Agents)
========================
The below code examples demonstrate how to use terminal related actions in yaml files.
yaml
""""
``NODE_SEND_LOCAL_COMMAND``
"""""""""""""""""""""""""""
Agents can execute local commands without needing to perform a separate remote login action (``SSH_TO_REMOTE``).
.. code-block:: yaml
...
...
action: NODE_SEND_LOCAL_COMMAND
options:
node_id: 0
username: admin
password: admin
command: # Example command - Creates a file called 'cat.png' in the downloads folder.
- "file_system"
- "create"
- "file"
- "downloads"
- "cat.png"
- "False"
Usage
"""""
``SSH_TO_REMOTE``
"""""""""""""""""
- Pre-Installs on all ``Nodes`` (with the exception of ``Switches``).
- Terminal Clients connect, execute commands and disconnect from remote nodes.
- Ensures that users are logged in to the component before executing any commands.
- Service runs on SSH port 22 by default.
Agents are able to use the terminal to login into remote nodes via ``SSH`` which allows for agents to execute commands on remote hosts.
.. code-block:: yaml
...
...
action: SSH_TO_REMOTE
options:
node_id: 0
username: admin
password: admin
remote_ip: 192.168.0.10 # Example Ip Address. (The remote host's IP that will be used by ssh)
``NODE_SEND_REMOTE_COMMAND``
""""""""""""""""""""""""""""
After remotely logging into another host, an agent can use the ``NODE_SEND_REMOTE_COMMAND`` to execute commands across the network remotely.
.. code-block:: yaml
...
...
action: NODE_SEND_REMOTE_COMMAND
options:
node_id: 0
remote_ip: 192.168.0.10
command:
- "file_system"
- "create"
- "file"
- "downloads"
- "cat.png"
- "False"
Simulation Layer Usage
======================
Usage
=====
The below code examples demonstrate how to create a terminal, a remote terminal, and how to send a basic application install command to a remote node.

2
src/primaite/VERSION
View File

@@ -1 +1 @@
3.3.0
3.4.0-dev

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

@@ -1298,6 +1298,28 @@ class NodeSendRemoteCommandAction(AbstractAction):
]
class NodeSendLocalCommandAction(AbstractAction):
"""Action which sends a terminal command using a local terminal session."""
def __init__(self, manager: "ActionManager", **kwargs) -> None:
super().__init__(manager=manager)
def form_request(self, node_id: int, username: str, password: str, command: RequestFormat) -> RequestFormat:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
node_name = self.manager.get_node_name_by_idx(node_id)
return [
"network",
"node",
node_name,
"service",
"Terminal",
"send_local_command",
username,
password,
{"command": command},
]
class TerminalC2ServerAction(AbstractAction):
"""Action which causes the C2 Server to send a command to the C2 Beacon to execute the terminal command passed."""
@@ -1406,6 +1428,7 @@ class ActionManager:
"SSH_TO_REMOTE": NodeSessionsRemoteLoginAction,
"SESSIONS_REMOTE_LOGOFF": NodeSessionsRemoteLogoutAction,
"NODE_SEND_REMOTE_COMMAND": NodeSendRemoteCommandAction,
"NODE_SEND_LOCAL_COMMAND": NodeSendLocalCommandAction,
}
"""Dictionary which maps action type strings to the corresponding action class."""

32
src/primaite/game/agent/interface.py
View File

@@ -4,6 +4,7 @@ from abc import ABC, abstractmethod
from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING
from gymnasium.core import ActType, ObsType
from prettytable import PrettyTable
from pydantic import BaseModel, model_validator
from primaite.game.agent.actions import ActionManager
@@ -126,6 +127,37 @@ class AbstractAgent(ABC):
self.history: List[AgentHistoryItem] = []
self.logger = AgentLog(agent_name)
def add_agent_action(self, item: AgentHistoryItem, table: PrettyTable) -> PrettyTable:
"""Update the given table with information from given AgentHistoryItem."""
node, application = "unknown", "unknown"
if (node_id := item.parameters.get("node_id")) is not None:
node = self.action_manager.node_names[node_id]
if (application_id := item.parameters.get("application_id")) is not None:
application = self.action_manager.application_names[node_id][application_id]
if (application_name := item.parameters.get("application_name")) is not None:
application = application_name
table.add_row([item.timestep, item.action, node, application, item.response.status])
return table
def show_history(self, ignored_actions: Optional[list] = None):
"""
Print an agent action provided it's not the DONOTHING action.
:param ignored_actions: OPTIONAL: List of actions to be ignored when displaying the history.
If not provided, defaults to ignore DONOTHING actions.
"""
if not ignored_actions:
ignored_actions = ["DONOTHING"]
table = PrettyTable()
table.field_names = ["Step", "Action", "Node", "Application", "Response"]
print(f"Actions for '{self.agent_name}':")
for item in self.history:
if item.action in ignored_actions:
pass
else:
table = self.add_agent_action(item=item, table=table)
print(table)
def update_observation(self, state: Dict) -> ObsType:
"""
Convert a state from the simulator into an observation for the agent using the observation space.

8
src/primaite/game/agent/observations/nic_observations.py
View File

@@ -1,19 +1,23 @@
# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
from __future__ import annotations
from typing import Dict, Optional
from typing import ClassVar, Dict, Optional
from gymnasium import spaces
from gymnasium.core import ObsType
from primaite.game.agent.observations.observations import AbstractObservation, WhereType
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
from primaite.simulator.network.nmne import NMNEConfig
from primaite.simulator.network.transmission.transport_layer import Port
class NICObservation(AbstractObservation, identifier="NETWORK_INTERFACE"):
"""Status information about a network interface within the simulation environment."""
capture_nmne: ClassVar[bool] = NMNEConfig().capture_nmne
"A Boolean specifying whether malicious network events should be captured."
class ConfigSchema(AbstractObservation.ConfigSchema):
"""Configuration schema for NICObservation."""
@@ -182,7 +186,7 @@ class NICObservation(AbstractObservation, identifier="NETWORK_INTERFACE"):
for port in self.monitored_traffic[protocol]:
obs["TRAFFIC"][protocol][Port[port].value] = {"inbound": 0, "outbound": 0}
if self.include_nmne:
if self.capture_nmne and self.include_nmne:
obs.update({"NMNE": {}})
direction_dict = nic_state["nmne"].get("direction", {})
inbound_keywords = direction_dict.get("inbound", {}).get("keywords", {})

4
src/primaite/game/game.py
View File

@@ -9,6 +9,7 @@ from pydantic import BaseModel, ConfigDict
from primaite import DEFAULT_BANDWIDTH, getLogger
from primaite.game.agent.actions import ActionManager
from primaite.game.agent.interface import AbstractAgent, AgentSettings, ProxyAgent
from primaite.game.agent.observations import NICObservation
from primaite.game.agent.observations.observation_manager import ObservationManager
from primaite.game.agent.rewards import RewardFunction, SharedReward
from primaite.game.agent.scripted_agents.data_manipulation_bot import DataManipulationAgent
@@ -80,6 +81,8 @@ class PrimaiteGameOptions(BaseModel):
seed: int = None
"""Random number seed for RNGs."""
generate_seed_value: bool = False
"""Internally generated seed value."""
max_episode_length: int = 256
"""Maximum number of episodes for the PrimAITE game."""
ports: List[str]
@@ -277,6 +280,7 @@ class PrimaiteGame:
links_cfg = network_config.get("links", [])
# Set the NMNE capture config
NetworkInterface.nmne_config = NMNEConfig(**network_config.get("nmne_config", {}))
NICObservation.capture_nmne = NMNEConfig(**network_config.get("nmne_config", {})).capture_nmne
for node_cfg in nodes_cfg:
n_type = node_cfg["type"]

10
src/primaite/notebooks/Command-&-Control-E2E-Demonstration.ipynb
View File

@@ -1800,6 +1800,16 @@
"\n",
"display_obs_diffs(tcp_c2_obs, udp_c2_obs, blue_config_env.game.step_counter)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"\n",
"env.game.agents[\"CustomC2Agent\"].show_history()"
]
}
],
"metadata": {

11
src/primaite/notebooks/Data-Manipulation-E2E-Demonstration.ipynb
View File

@@ -675,6 +675,15 @@
" print(f\"step: {env.game.step_counter}, Red action: {info['agent_actions']['data_manipulation_attacker'].action}, Blue reward:{reward:.2f}\" )"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"env.game.agents[\"data_manipulation_attacker\"].show_history()"
]
},
{
"cell_type": "markdown",
"metadata": {},
@@ -708,7 +717,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.12"
"version": "3.10.11"
}
},
"nbformat": 4,

45
src/primaite/notebooks/Getting-Information-Out-Of-PrimAITE.ipynb
View File

@@ -144,6 +144,49 @@
"PRIMAITE_CONFIG[\"developer_mode\"][\"enabled\"] = was_enabled\n",
"PRIMAITE_CONFIG[\"developer_mode\"][\"output_sys_logs\"] = was_syslogs_enabled"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Viewing Agent history\n",
"\n",
"It's possible to view the actions carried out by an agent for a given training session using the `show_history()` method. By default, this will be all actions apart from DONOTHING actions."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"with open(data_manipulation_config_path(), 'r') as f:\n",
" cfg = yaml.safe_load(f)\n",
"\n",
"env = PrimaiteGymEnv(env_config=cfg)\n",
"\n",
"# Run the training session to generate some resultant data.\n",
"for i in range(100):\n",
" env.step(0)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Calling `.show_history()` should show us when the Data Manipulation used the `NODE_APPLICATION_EXECUTE` action."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"attacker = env.game.agents[\"data_manipulation_attacker\"]\n",
"\n",
"attacker.show_history()"
]
}
],
"metadata": {
@@ -162,7 +205,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.8"
"version": "3.10.11"
}
},
"nbformat": 4,

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

@@ -9,6 +9,13 @@
"© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Simulation Layer Implementation."
]
},
{
"cell_type": "markdown",
"metadata": {},
@@ -198,6 +205,271 @@
"source": [
"computer_b.user_session_manager.show(include_historic=True, include_session_id=True)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Game Layer Implementation\n",
"\n",
"This notebook section will detail the implementation of how the game layer utilises the terminal to support different agent actions.\n",
"\n",
"The ``Terminal`` is used in a variety of different ways in the game layer. Specifically, the terminal is leveraged to implement the following actions:\n",
"\n",
"\n",
"| Game Layer Action | Simulation Layer |\n",
"|-----------------------------------|--------------------------|\n",
"| ``NODE_SEND_LOCAL_COMMAND`` | Uses the given user credentials, creates a ``LocalTerminalSession`` and executes the given command and returns the ``RequestResponse``.\n",
"| ``SSH_TO_REMOTE`` | Uses the given user credentials and remote IP to create a ``RemoteTerminalSession``.\n",
"| ``NODE_SEND_REMOTE_COMMAND`` | Uses the given remote IP to locate the correct ``RemoteTerminalSession``, executes the given command and returns the ``RequestsResponse``."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Game Layer Setup\n",
"\n",
"Similar to other notebooks, the next code cells create a custom proxy agent to demonstrate how these commands can be leveraged by agents in the ``UC2`` network environment.\n",
"\n",
"If you're unfamiliar with ``UC2`` then please refer to the [UC2-E2E-Demo notebook for further reference](./Data-Manipulation-E2E-Demonstration.ipynb)."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import yaml\n",
"from primaite.config.load import data_manipulation_config_path\n",
"from primaite.session.environment import PrimaiteGymEnv"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"custom_terminal_agent = \"\"\"\n",
" - ref: CustomC2Agent\n",
" team: RED\n",
" type: ProxyAgent\n",
" observation_space: null\n",
" action_space:\n",
" action_list:\n",
" - type: DONOTHING\n",
" - type: NODE_SEND_LOCAL_COMMAND\n",
" - type: SSH_TO_REMOTE\n",
" - type: NODE_SEND_REMOTE_COMMAND\n",
" options:\n",
" nodes:\n",
" - node_name: client_1\n",
" max_folders_per_node: 1\n",
" max_files_per_folder: 1\n",
" max_services_per_node: 2\n",
" max_nics_per_node: 8\n",
" max_acl_rules: 10\n",
" ip_list:\n",
" - 192.168.1.21\n",
" - 192.168.1.14\n",
" wildcard_list:\n",
" - 0.0.0.1\n",
" action_map:\n",
" 0:\n",
" action: DONOTHING\n",
" options: {}\n",
" 1:\n",
" action: NODE_SEND_LOCAL_COMMAND\n",
" options:\n",
" node_id: 0\n",
" username: admin\n",
" password: admin\n",
" command:\n",
" - file_system\n",
" - create\n",
" - file\n",
" - downloads\n",
" - dog.png\n",
" - False\n",
" 2:\n",
" action: SSH_TO_REMOTE\n",
" options:\n",
" node_id: 0\n",
" username: admin\n",
" password: admin\n",
" remote_ip: 192.168.10.22\n",
" 3:\n",
" action: NODE_SEND_REMOTE_COMMAND\n",
" options:\n",
" node_id: 0\n",
" remote_ip: 192.168.10.22\n",
" command:\n",
" - file_system\n",
" - create\n",
" - file\n",
" - downloads\n",
" - cat.png\n",
" - False\n",
" reward_function:\n",
" reward_components:\n",
" - type: DUMMY\n",
"\"\"\"\n",
"custom_terminal_agent_yaml = yaml.safe_load(custom_terminal_agent)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"with open(data_manipulation_config_path()) as f:\n",
" cfg = yaml.safe_load(f)\n",
" # removing all agents & adding the custom agent.\n",
" cfg['agents'] = {}\n",
" cfg['agents'] = custom_terminal_agent_yaml\n",
" \n",
"env = PrimaiteGymEnv(env_config=cfg)\n",
"\n",
"client_1: Computer = env.game.simulation.network.get_node_by_hostname(\"client_1\")\n",
"client_2: Computer = env.game.simulation.network.get_node_by_hostname(\"client_2\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Terminal Action | ``NODE_SEND_LOCAL_COMMAND`` \n",
"\n",
"The yaml snippet below shows all the relevant agent options for this action:\n",
"\n",
"```yaml\n",
"\n",
" action_space:\n",
" action_list:\n",
" ...\n",
" - type: NODE_SEND_LOCAL_COMMAND\n",
" ...\n",
" options:\n",
" nodes: # Node List\n",
" - node_name: client_1\n",
" ...\n",
" ...\n",
" action_map:\n",
" 1:\n",
" action: NODE_SEND_LOCAL_COMMAND\n",
" options:\n",
" node_id: 0 # Index 0 at the node list.\n",
" username: admin\n",
" password: admin\n",
" command:\n",
" - file_system\n",
" - create\n",
" - file\n",
" - downloads\n",
" - dog.png\n",
" - False\n",
"```"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"env.step(1)\n",
"client_1.file_system.show(full=True)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Terminal Action | ``SSH_TO_REMOTE`` \n",
"\n",
"The yaml snippet below shows all the relevant agent options for this action:\n",
"\n",
"```yaml\n",
"\n",
" action_space:\n",
" action_list:\n",
" ...\n",
" - type: SSH_TO_REMOTE\n",
" ...\n",
" options:\n",
" nodes: # Node List\n",
" - node_name: client_1\n",
" ...\n",
" ...\n",
" action_map:\n",
" 2:\n",
" action: SSH_TO_REMOTE\n",
" options:\n",
" node_id: 0 # Index 0 at the node list.\n",
" username: admin\n",
" password: admin\n",
" remote_ip: 192.168.10.22 # client_2's ip address.\n",
"```"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"env.step(2)\n",
"client_2.session_manager.show()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Terminal Action | ``NODE_SEND_REMOTE_COMMAND``\n",
"\n",
"The yaml snippet below shows all the relevant agent options for this action:\n",
"\n",
"```yaml\n",
"\n",
" action_space:\n",
" action_list:\n",
" ...\n",
" - type: NODE_SEND_REMOTE_COMMAND\n",
" ...\n",
" options:\n",
" nodes: # Node List\n",
" - node_name: client_1\n",
" ...\n",
" ...\n",
" action_map:\n",
" 1:\n",
" action: NODE_SEND_REMOTE_COMMAND\n",
" options:\n",
" node_id: 0 # Index 0 at the node list.\n",
" remote_ip: 192.168.10.22\n",
" commands:\n",
" - file_system\n",
" - create\n",
" - file\n",
" - downloads\n",
" - cat.png\n",
" - False\n",
"```"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"env.step(3)\n",
"client_2.file_system.show(full=True)"
]
}
],
"metadata": {
@@ -216,7 +488,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.11"
"version": "3.10.12"
}
},
"nbformat": 4,

30
src/primaite/session/environment.py
View File

@@ -26,14 +26,26 @@ except ModuleNotFoundError:
_LOGGER.debug("Torch not available for importing")
def set_random_seed(seed: int) -> Union[None, int]:
def set_random_seed(seed: int, generate_seed_value: bool) -> Union[None, int]:
"""
Set random number generators.
If seed is None or -1 and generate_seed_value is True randomly generate a
seed value.
If seed is > -1 and generate_seed_value is True ignore the latter and use
the provide seed value.
:param seed: int
:param generate_seed_value: bool
:return: None or the int representing the seed used.
"""
if seed is None or seed == -1:
return None
if generate_seed_value:
rng = np.random.default_rng()
# 2**32-1 is highest value for python RNG seed.
seed = int(rng.integers(low=0, high=2**32 - 1))
else:
return None
elif seed < -1:
raise ValueError("Invalid random number seed")
# Seed python RNG
@@ -50,6 +62,13 @@ def set_random_seed(seed: int) -> Union[None, int]:
return seed
def log_seed_value(seed: int):
"""Log the selected seed value to file."""
path = SIM_OUTPUT.path / "seed.log"
with open(path, "w") as file:
file.write(f"Seed value = {seed}")
class PrimaiteGymEnv(gymnasium.Env):
"""
Thin wrapper env to provide agents with a gymnasium API.
@@ -65,7 +84,8 @@ class PrimaiteGymEnv(gymnasium.Env):
"""Object that returns a config corresponding to the current episode."""
self.seed = self.episode_scheduler(0).get("game", {}).get("seed")
"""Get RNG seed from config file. NB: Must be before game instantiation."""
self.seed = set_random_seed(self.seed)
self.generate_seed_value = self.episode_scheduler(0).get("game", {}).get("generate_seed_value")
self.seed = set_random_seed(self.seed, self.generate_seed_value)
self.io = PrimaiteIO.from_config(self.episode_scheduler(0).get("io_settings", {}))
"""Handles IO for the environment. This produces sys logs, agent logs, etc."""
self.game: PrimaiteGame = PrimaiteGame.from_config(self.episode_scheduler(0))
@@ -79,6 +99,8 @@ class PrimaiteGymEnv(gymnasium.Env):
_LOGGER.info(f"PrimaiteGymEnv RNG seed = {self.seed}")
log_seed_value(self.seed)
def action_masks(self) -> np.ndarray:
"""
Return the action mask for the agent.
@@ -146,7 +168,7 @@ class PrimaiteGymEnv(gymnasium.Env):
f"avg. reward: {self.agent.reward_function.total_reward}"
)
if seed is not None:
set_random_seed(seed)
set_random_seed(seed, self.generate_seed_value)
self.total_reward_per_episode[self.episode_counter] = self.agent.reward_function.total_reward
if self.io.settings.save_agent_actions:

2
src/primaite/simulator/system/core/session_manager.py
View File

@@ -413,5 +413,5 @@ class SessionManager:
table.align = "l"
table.title = f"{self.sys_log.hostname} Session Manager"
for session in self.sessions_by_key.values():
table.add_row([session.dst_ip_address, session.dst_port.value, session.protocol.name])
table.add_row([session.with_ip_address, session.dst_port.value, session.protocol.name])
print(table)

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

@@ -208,17 +208,37 @@ class Terminal(Service):
status="success",
data={},
)
else:
return RequestResponse(
status="failure",
data={},
)
return RequestResponse(
status="failure",
data={},
)
rm.add_request(
"send_remote_command",
request_type=RequestType(func=remote_execute_request),
)
def local_execute_request(request: RequestFormat, context: Dict) -> RequestResponse:
"""Executes a command using a local terminal session."""
command: str = request[2]["command"]
local_connection = self._process_local_login(username=request[0], password=request[1])
if local_connection:
outcome = local_connection.execute(command)
if outcome:
return RequestResponse(
status="success",
data={"reason": outcome},
)
return RequestResponse(
status="success",
data={"reason": "Local Terminal failed to resolve command. Potentially invalid credentials?"},
)
rm.add_request(
"send_local_command",
request_type=RequestType(func=local_execute_request),
)
return rm
def execute(self, command: List[Any]) -> Optional[RequestResponse]:

1
tests/conftest.py
View File

@@ -466,6 +466,7 @@ def game_and_agent():
{"type": "SSH_TO_REMOTE"},
{"type": "SESSIONS_REMOTE_LOGOFF"},
{"type": "NODE_SEND_REMOTE_COMMAND"},
{"type": "NODE_SEND_LOCAL_COMMAND"},
]
action_space = ActionManager(

52
tests/integration_tests/game_layer/actions/test_terminal_actions.py
View File

@@ -164,3 +164,55 @@ def test_change_password_logs_out_user(game_and_agent_fixture: Tuple[PrimaiteGam
assert server_1.file_system.get_folder("folder123") is None
assert server_1.file_system.get_file("folder123", "doggo.pdf") is None
def test_local_terminal(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAgent]):
game, agent = game_and_agent_fixture
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
client_1_usm: UserManager = client_1.software_manager.software["UserManager"]
client_1_usm.add_user("user123", "password", is_admin=True)
action = (
"NODE_SEND_LOCAL_COMMAND",
{
"node_id": 0,
"username": "user123",
"password": "password",
"command": ["file_system", "create", "file", "folder123", "doggo.pdf", False],
},
)
agent.store_action(action)
game.step()
assert client_1.file_system.get_folder("folder123")
assert client_1.file_system.get_file("folder123", "doggo.pdf")
# Change password
action = (
"NODE_ACCOUNTS_CHANGE_PASSWORD",
{
"node_id": 0, # server_1
"username": "user123",
"current_password": "password",
"new_password": "different_password",
},
)
agent.store_action(action)
game.step()
action = (
"NODE_SEND_LOCAL_COMMAND",
{
"node_id": 0,
"username": "user123",
"password": "password",
"command": ["file_system", "create", "file", "folder123", "cat.pdf", False],
},
)
agent.store_action(action)
game.step()
assert client_1.file_system.get_file("folder123", "cat.pdf") is None
client_1.session_manager.show()

8
tests/integration_tests/game_layer/observations/test_nic_observations.py
View File

@@ -77,6 +77,14 @@ def test_nic(simulation):
nic_obs = NICObservation(where=["network", "nodes", pc.hostname, "NICs", 1], include_nmne=True)
# The Simulation object created by the fixture also creates the
# NICObservation class with the NICObservation.capture_nmnme class variable
# set to False. Under normal (non-test) circumstances this class variable
# is set from a config file such as data_manipulation.yaml. So although
# capture_nmne is set to True in the NetworkInterface class it's still False
# in the NICObservation class so we set it now.
nic_obs.capture_nmne = True
# Set the NMNE configuration to capture DELETE/ENCRYPT queries as MNEs
nmne_config = {
"capture_nmne": True, # Enable the capture of MNEs

22
tests/integration_tests/game_layer/test_RNG_seed.py
View File

@@ -7,6 +7,7 @@ import yaml
from primaite.config.load import data_manipulation_config_path
from primaite.game.agent.interface import AgentHistoryItem
from primaite.session.environment import PrimaiteGymEnv
from primaite.simulator import SIM_OUTPUT
@pytest.fixture()
@@ -33,6 +34,11 @@ def test_rng_seed_set(create_env):
assert a == b
# Check that seed log file was created.
path = SIM_OUTPUT.path / "seed.log"
with open(path, "r") as file:
assert file
def test_rng_seed_unset(create_env):
"""Test with no RNG seed."""
@@ -48,3 +54,19 @@ def test_rng_seed_unset(create_env):
b = [item.timestep for item in env.game.agents["client_2_green_user"].history if item.action != "DONOTHING"]
assert a != b
def test_for_generated_seed():
"""
Show that setting generate_seed_value to true producess a valid seed.
"""
with open(data_manipulation_config_path(), "r") as f:
cfg = yaml.safe_load(f)
cfg["game"]["generate_seed_value"] = True
PrimaiteGymEnv(env_config=cfg)
path = SIM_OUTPUT.path / "seed.log"
with open(path, "r") as file:
data = file.read()
assert data.split(" ")[3] != None

22
tests/integration_tests/network/test_capture_nmne.py
View File

@@ -1,5 +1,11 @@
# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
from itertools import product
import yaml
from primaite.config.load import data_manipulation_config_path
from primaite.game.agent.observations.nic_observations import NICObservation
from primaite.session.environment import PrimaiteGymEnv
from primaite.simulator.network.container import Network
from primaite.simulator.network.hardware.nodes.host.host_node import NIC
from primaite.simulator.network.hardware.nodes.host.server import Server
@@ -277,3 +283,19 @@ def test_capture_nmne_observations(uc2_network: Network):
assert web_nic_obs["outbound"] == expected_nmne
assert db_nic_obs["inbound"] == expected_nmne
uc2_network.apply_timestep(timestep=0)
def test_nmne_parameter_settings():
"""
Check that the four permutations of the values of capture_nmne and
include_nmne work as expected.
"""
with open(data_manipulation_config_path(), "r") as f:
cfg = yaml.safe_load(f)
DEFENDER = 3
for capture, include in product([True, False], [True, False]):
cfg["simulation"]["network"]["nmne_config"]["capture_nmne"] = capture
cfg["agents"][DEFENDER]["observation_space"]["options"]["components"][0]["options"]["include_nmne"] = include
PrimaiteGymEnv(env_config=cfg)