diff --git a/.azure/azure-ci-build-pipeline.yaml b/.azure/azure-ci-build-pipeline.yaml index 624c9ca4..b6f24777 100644 --- a/.azure/azure-ci-build-pipeline.yaml +++ b/.azure/azure-ci-build-pipeline.yaml @@ -113,10 +113,8 @@ stages: - script: | pytest --nbmake -n=auto src/primaite/notebooks --junit-xml=./notebook-tests/notebooks.xml notebooks_exit_code=$? - pytest --nbmake -n=auto src/primaite/simulator/_package_data --junit-xml=./notebook-tests/package-notebooks.xml - package_notebooks_exit_code=$? - # Fail step if either of these do not have exit code 0 - if [ $notebooks_exit_code -ne 0 ] || [ $package_notebooks_exit_code -ne 0 ]; then + # Fail step if exit code not equal to 0 + if [ $notebooks_exit_code -ne 0 ]; then exit 1 fi displayName: 'Run notebooks on Linux and macOS' @@ -126,11 +124,8 @@ stages: - script: | pytest --nbmake -n=auto src/primaite/notebooks --junit-xml=./notebook-tests/notebooks.xml set notebooks_exit_code=%ERRORLEVEL% - pytest --nbmake -n=auto src/primaite/simulator/_package_data --junit-xml=./notebook-tests/package-notebooks.xml - set package_notebooks_exit_code=%ERRORLEVEL% - rem Fail step if either of these do not have exit code 0 + rem Fail step if exit code not equal to 0 if %notebooks_exit_code% NEQ 0 exit /b 1 - if %package_notebooks_exit_code% NEQ 0 exit /b 1 displayName: 'Run notebooks on Windows' condition: eq(variables['Agent.OS'], 'Windows_NT') diff --git a/MANIFEST.in b/MANIFEST.in index 2ac7b306..51ae4ddf 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -1,3 +1,2 @@ include src/primaite/setup/_package_data/primaite_config.yaml include src/primaite/config/_package_data/*.yaml -include src/primaite/simulator/_package_data/*.ipynb diff --git a/src/primaite/config/_package_data/basic_lan_network_example.yaml b/src/primaite/config/_package_data/basic_lan_network_example.yaml index 9490ff00..9996be84 100644 --- a/src/primaite/config/_package_data/basic_lan_network_example.yaml +++ b/src/primaite/config/_package_data/basic_lan_network_example.yaml @@ -1,3 +1,6 @@ +metadata: + version: 3.0 + game: ports: - ARP diff --git a/src/primaite/config/_package_data/client_server_p2p_network_example.yaml b/src/primaite/config/_package_data/client_server_p2p_network_example.yaml index 798dd318..1a9fca98 100644 --- a/src/primaite/config/_package_data/client_server_p2p_network_example.yaml +++ b/src/primaite/config/_package_data/client_server_p2p_network_example.yaml @@ -1,3 +1,6 @@ +metadata: + version: 3.0 + game: ports: - ARP diff --git a/src/primaite/config/_package_data/data_manipulation.yaml b/src/primaite/config/_package_data/data_manipulation.yaml index 4705c135..7d9fdf36 100644 --- a/src/primaite/config/_package_data/data_manipulation.yaml +++ b/src/primaite/config/_package_data/data_manipulation.yaml @@ -1,3 +1,6 @@ +metadata: + version: 3.0 + io_settings: save_agent_actions: true save_step_metadata: false @@ -81,7 +84,7 @@ agents: action: node-application-execute options: node_name: client_1 - application_name: web-browser + application_name: database-client reward_function: reward_components: diff --git a/src/primaite/config/_package_data/data_manipulation_marl.yaml b/src/primaite/config/_package_data/data_manipulation_marl.yaml index 0263edb0..71507acb 100644 --- a/src/primaite/config/_package_data/data_manipulation_marl.yaml +++ b/src/primaite/config/_package_data/data_manipulation_marl.yaml @@ -1,3 +1,6 @@ +metadata: + version: 3.0 + io_settings: save_agent_actions: true save_step_metadata: false diff --git a/src/primaite/config/_package_data/mini_scenario_with_simulation_variation/base_scenario.yaml b/src/primaite/config/_package_data/mini_scenario_with_simulation_variation/base_scenario.yaml index 4c7734d8..2ea18867 100644 --- a/src/primaite/config/_package_data/mini_scenario_with_simulation_variation/base_scenario.yaml +++ b/src/primaite/config/_package_data/mini_scenario_with_simulation_variation/base_scenario.yaml @@ -1,3 +1,6 @@ +metadata: + version: 3.0 + game: max_episode_length: 128 ports: [] diff --git a/src/primaite/config/_package_data/multi_lan_internet_network_example.yaml b/src/primaite/config/_package_data/multi_lan_internet_network_example.yaml index cbd1c01e..deaef3bd 100644 --- a/src/primaite/config/_package_data/multi_lan_internet_network_example.yaml +++ b/src/primaite/config/_package_data/multi_lan_internet_network_example.yaml @@ -1,3 +1,6 @@ +metadata: + version: 3.0 + game: ports: - ARP diff --git a/src/primaite/config/_package_data/scenario_with_placeholders/scenario.yaml b/src/primaite/config/_package_data/scenario_with_placeholders/scenario.yaml index 1a96e3a0..4ec3d257 100644 --- a/src/primaite/config/_package_data/scenario_with_placeholders/scenario.yaml +++ b/src/primaite/config/_package_data/scenario_with_placeholders/scenario.yaml @@ -1,3 +1,6 @@ +metadata: + version: 3.0 + io_settings: save_agent_actions: true save_step_metadata: false diff --git a/src/primaite/game/agent/interface.py b/src/primaite/game/agent/interface.py index 0547973f..1fef14ef 100644 --- a/src/primaite/game/agent/interface.py +++ b/src/primaite/game/agent/interface.py @@ -77,7 +77,7 @@ class AbstractAgent(BaseModel, ABC): config: ConfigSchema = Field(default_factory=lambda: AbstractAgent.ConfigSchema()) - logger: AgentLog = AgentLog(agent_name="Abstract_Agent") + logger: AgentLog = None history: List[AgentHistoryItem] = [] action_manager: ActionManager = Field(default_factory=lambda: ActionManager()) @@ -86,6 +86,11 @@ class AbstractAgent(BaseModel, ABC): _registry: ClassVar[Dict[str, Type[AbstractAgent]]] = {} + def __init__(self, **kwargs): + """Initialise and setup agent logger.""" + super().__init__(**kwargs) + self.logger: AgentLog = AgentLog(agent_name=kwargs["config"]["ref"]) + def __init_subclass__(cls, discriminator: Optional[str] = None, **kwargs: Any) -> None: super().__init_subclass__(**kwargs) if discriminator is None: diff --git a/src/primaite/notebooks/Action-masking.ipynb b/src/primaite/notebooks/Action-masking.ipynb index 7fde0a49..74504878 100644 --- a/src/primaite/notebooks/Action-masking.ipynb +++ b/src/primaite/notebooks/Action-masking.ipynb @@ -6,7 +6,7 @@ "source": [ "# Action Masking\n", "\n", - "© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n", + "© Crown-owned copyright 2025, Defence Science and Technology Laboratory UK\n", "\n", "PrimAITE environments support action masking. The action mask shows which of the agent's actions are applicable with the current environment state. For example, a node can only be turned on if it is currently turned off." ] @@ -204,7 +204,7 @@ ], "metadata": { "kernelspec": { - "display_name": ".venv", + "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, @@ -218,7 +218,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.11" + "version": "3.10.12" } }, "nbformat": 4, diff --git a/src/primaite/notebooks/Command-and-Control-E2E-Demonstration.ipynb b/src/primaite/notebooks/Command-and-Control-E2E-Demonstration.ipynb index 52499ea6..882c3429 100644 --- a/src/primaite/notebooks/Command-and-Control-E2E-Demonstration.ipynb +++ b/src/primaite/notebooks/Command-and-Control-E2E-Demonstration.ipynb @@ -6,11 +6,20 @@ "source": [ "# Command and Control Application Suite E2E Demonstration\n", "\n", - "© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n", + "© Crown-owned copyright 2025, Defence Science and Technology Laboratory UK\n", "\n", "This notebook demonstrates the current implementation of the command and control (C2) server and beacon applications in primAITE." ] }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "!primaite setup" + ] + }, { "cell_type": "code", "execution_count": null, @@ -50,12 +59,12 @@ "custom_c2_agent = \"\"\"\n", " - ref: CustomC2Agent\n", " team: RED\n", - " type: proxy-a.gent\n", + " type: ProxyAgent\n", "\n", " action_space:\n", " action_map:\n", " 0:\n", - " action: do_nothing\n", + " action: do-nothing\n", " options: {}\n", " 1:\n", " action: node-application-install\n", @@ -65,35 +74,34 @@ " 2:\n", " action: configure-c2-beacon\n", " options:\n", - " node_name: web_server\n", + " node_id: 0\n", " config:\n", " c2_server_ip_address: 192.168.10.21\n", " keep_alive_frequency:\n", " masquerade_protocol:\n", " masquerade_port:\n", " 3:\n", - " action: node_application_execute\n", + " action: node-application-execute\n", " options:\n", " node_name: web_server\n", " application_name: c2-beacon\n", " 4:\n", " action: c2-server-terminal-command\n", " options:\n", - " node_id: 1\n", + " node_name: client_1\n", " ip_address:\n", - " account:\n", - " username: admin\n", - " password: admin\n", + " username: admin\n", + " password: admin\n", " commands:\n", " -\n", " - software_manager\n", " - application\n", " - install\n", - " - RansomwareScript\n", + " - ransomware-script\n", " 5:\n", " action: c2-server-ransomware-configure\n", " options:\n", - " node_name: client_1\n", + " node_id: 1\n", " config:\n", " server_ip_address: 192.168.1.14\n", " payload: ENCRYPT\n", @@ -105,9 +113,8 @@ " target_folder_name: \"database\"\n", " exfiltration_folder_name: \"spoils\"\n", " target_ip_address: 192.168.1.14\n", - " account:\n", - " username: admin\n", - " password: admin\n", + " username: admin\n", + " password: admin\n", "\n", " 7:\n", " action: c2-server-ransomware-launch\n", @@ -116,7 +123,7 @@ " 8:\n", " action: configure-c2-beacon\n", " options:\n", - " node_name: web_server\n", + " node_id: 0\n", " config:\n", " c2_server_ip_address: 192.168.10.21\n", " keep_alive_frequency: 10\n", @@ -125,12 +132,16 @@ " 9:\n", " action: configure-c2-beacon\n", " options:\n", - " node_name: web_server\n", + " node_id: 0\n", " config:\n", " c2_server_ip_address: 192.168.10.22\n", " keep_alive_frequency:\n", " masquerade_protocol:\n", " masquerade_port:\n", + "\n", + " reward_function:\n", + " reward_components:\n", + " - type: DUMMY\n", "\"\"\"\n", "c2_agent_yaml = yaml.safe_load(custom_c2_agent)" ] @@ -171,7 +182,7 @@ "source": [ "client_1: Computer = env.game.simulation.network.get_node_by_hostname(\"client_1\")\n", "client_1.software_manager.install(C2Server)\n", - "c2_server: C2Server = client_1.software_manager.software[\"C2Server\"]\n", + "c2_server: C2Server = client_1.software_manager.software[\"c2-server\"]\n", "c2_server.run()\n", "client_1.software_manager.show()" ] @@ -208,7 +219,7 @@ " ...\n", " action_map:\n", " 1:\n", - " action: node_application_install \n", + " action: node-application-install \n", " options:\n", " node_id: 0 # Index 0 at the node list.\n", " application_name: c2-beacon\n", @@ -237,23 +248,18 @@ "The yaml snippet below shows all the relevant agent options for this action:\n", "\n", "```yaml\n", + "\n", " action_space:\n", - " options:\n", - " nodes: # Node List\n", - " - node_name: web_server\n", - " ...\n", - " ...\n", " action_map:\n", - " ...\n", " 2:\n", " action: configure-c2-beacon\n", " options:\n", - " node_id: 0 # Node Index\n", - " config: # Further information about these config options can be found at the bottom of this notebook.\n", - " c2_server_ip_address: 192.168.10.21\n", - " keep_alive_frequency:\n", - " masquerade_protocol:\n", - " masquerade_port:\n", + " node_name: web_server\n", + " c2_server_ip_address: 192.168.10.21 # Further information about these config options can be found at the bottom of this notebook.\n", + " keep_alive_frequency:\n", + " masquerade_protocol:\n", + " masquerade_port:\n", + "\n", "```" ] }, @@ -275,7 +281,7 @@ "source": [ "### **Command and Control** | C2 Beacon Actions | node_application_execute\n", "\n", - "The final action is ``node_application_execute`` which is used to establish a connection for the C2 application. This action can be called by the Red Agent via action ``3`` in it's action map. \n", + "The final action is ``node-application-execute`` which is used to establish a connection for the C2 application. This action can be called by the Red Agent via action ``3`` in it's action map. \n", "\n", "The yaml snippet below shows all the relevant agent options for this action:\n", "\n", @@ -285,16 +291,15 @@ " nodes: # Node List\n", " - node_name: web_server\n", " applications: \n", - " - application_name: c2-beacon\n", + " - application_name: C2Beacon\n", " ...\n", " ...\n", " action_map:\n", - " ...\n", " 3:\n", " action: node-application-execute\n", " options:\n", - " node_id: 0\n", - " application_id: 0\n", + " node_name: web_server\n", + " application_name: c2-beacon\n", "```" ] }, @@ -333,15 +338,15 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "### **Command and Control** | C2 Server Actions | C2_SERVER_TERMINAL_COMMAND\n", + "### **Command and Control** | C2 Server Actions | c2-server-terminal-command\n", "\n", - "The C2 Server's terminal action: ``C2_SERVER_TERMINAL_COMMAND`` is indexed at ``4`` in it's action map. \n", + "The C2 Server's terminal action: ``c2-server-terminal-command`` is indexed at ``4`` in it's action map. \n", "\n", "This action leverages the terminal service that is installed by default on all nodes to grant red agents a lot more configurability. If you're unfamiliar with terminals then it's recommended that you refer to the ``Terminal Processing`` notebook.\n", "\n", "It's worth noting that an additional benefit a red agent has when using the terminal service via the C2 Server is that you can execute multiple commands in one action. \n", "\n", - "In this notebook, the ``C2_SERVER_TERMINAL_COMMAND`` is used to install a RansomwareScript application on the ``web_server`` node.\n", + "In this notebook, the ``c2-server-terminal-command`` is used to install a RansomwareScript application on the ``web_server`` node.\n", "\n", "The yaml snippet below shows all the relevant agent options for this action:\n", "\n", @@ -352,15 +357,14 @@ " ...\n", " - node_name: client_1\n", " applications: \n", - " - application_name: c2-server\n", + " - application_name: C2Server\n", " ...\n", " action_map:\n", " 4:\n", " action: c2-server-terminal-command\n", " options:\n", - " node_id: 1\n", + " node_name: client_1\n", " ip_address:\n", - " account:\n", " username: admin\n", " password: admin\n", " commands:\n", @@ -368,7 +372,7 @@ " - software_manager\n", " - application\n", " - install\n", - " - RansomwareScript\n", + " - ransomware-script\n", "```" ] }, @@ -409,16 +413,15 @@ " ...\n", " - node_name: client_1\n", " applications: \n", - " - application_name: c2-server\n", + " - application_name: C2Server\n", " ...\n", " action_map:\n", " 5:\n", " action: c2-server-ransomware-configure\n", " options:\n", - " node_id: 1\n", - " config:\n", - " server_ip_address: 192.168.1.14\n", - " payload: ENCRYPT\n", + " node_name: client_1\n", + " server_ip_address: 192.168.1.14\n", + " payload: ENCRYPT\n", "```\n" ] }, @@ -446,9 +449,9 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "### **Command and Control** | C2 Server Actions | c2_server_data_exfiltrate\n", + "### **Command and Control** | C2 Server Actions | c2-server-data-exfiltrate\n", "\n", - "The second to last action available is the ``c2_server_data_exfiltrate`` which is indexed as action ``6`` in the action map.\n", + "The second to last action available is the ``c2-server-data-exfiltrate`` which is indexed as action ``6`` in the action map.\n", "\n", "This action can be used to exfiltrate a target file on a remote node to the C2 Beacon and the C2 Server's host file system via the ``FTP`` services.\n", "\n", @@ -461,7 +464,7 @@ " ...\n", " - node_name: client_1\n", " applications: \n", - " - application_name: c2-server\n", + " - application_name: C2Server\n", " ...\n", " action_map:\n", " 6:\n", @@ -472,9 +475,8 @@ " target_folder_name: \"database\"\n", " exfiltration_folder_name: \"spoils\"\n", " target_ip_address: \"192.168.1.14\"\n", - " account:\n", - " username: \"admin\",\n", - " password: \"admin\"\n", + " username: \"admin\"\n", + " password: \"admin\"\n", "\n", "```" ] @@ -512,9 +514,9 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "### **Command and Control** | C2 Server Actions | c2_server_ransomware_launch\n", + "### **Command and Control** | C2 Server Actions | c2-server-ransomware-launch\n", "\n", - "Finally, the last available action is for the c2_server_ransomware_launch to start the ransomware script installed on the same node as the C2 beacon.\n", + "Finally, the last available action is for the c2-server-ransomware-launch to start the ransomware script installed on the same node as the C2 beacon.\n", "\n", "This action is indexed as action ``7``.\n", "\n", @@ -527,7 +529,7 @@ " ...\n", " - node_name: client_1\n", " applications: \n", - " - application_name: c2-server\n", + " - application_name: C2Server\n", " ...\n", " action_map:\n", " 7:\n", @@ -582,7 +584,7 @@ " type: custom\n", " options:\n", " components:\n", - " - type: nodes\n", + " - type: NODES\n", " label: NODES\n", " options:\n", " hosts:\n", @@ -660,17 +662,17 @@ " action_space:\n", " action_map:\n", " 0:\n", - " action: do_nothing\n", + " action: do-nothing\n", " options: {}\n", " 1:\n", " action: node-application-remove\n", " options:\n", - " node_name: web-server\n", + " node_id: 0\n", " application_name: C2Beacon\n", " 2:\n", " action: node-shutdown\n", " options:\n", - " node_name: web-server\n", + " node_id: 0\n", " 3:\n", " action: router-acl-add-rule\n", " options:\n", @@ -685,6 +687,36 @@ " source_wildcard_id: 0\n", " dest_wildcard_id: 0\n", "\n", + "\n", + " options:\n", + " nodes:\n", + " - node_name: web_server\n", + " applications:\n", + " - application_name: C2Beacon\n", + "\n", + " - node_name: database_server\n", + " folders:\n", + " - folder_name: database\n", + " files:\n", + " - file_name: database.db\n", + " services:\n", + " - service_name: DatabaseService\n", + " - node_name: router_1\n", + "\n", + " max_folders_per_node: 2\n", + " max_files_per_folder: 2\n", + " max_services_per_node: 2\n", + " max_nics_per_node: 8\n", + " max_acl_rules: 10\n", + " ip_list:\n", + " - 192.168.10.21\n", + " - 192.168.1.12\n", + " wildcard_list:\n", + " - 0.0.0.1\n", + " reward_function:\n", + " reward_components:\n", + " - type: DUMMY\n", + "\n", " agent_settings:\n", " flatten_obs: False\n", "\"\"\"\n", @@ -774,12 +806,12 @@ "\n", "# Installing the C2 Server.\n", "client_1.software_manager.install(C2Server)\n", - "c2_server: C2Server = client_1.software_manager.software[\"C2Server\"]\n", + "c2_server: C2Server = client_1.software_manager.software[\"c2-server\"]\n", "c2_server.run()\n", "\n", "# Installing the C2 Beacon.\n", "web_server.software_manager.install(C2Beacon)\n", - "c2_beacon: C2Beacon = web_server.software_manager.software[\"C2Beacon\"]\n", + "c2_beacon: C2Beacon = web_server.software_manager.software[\"c2-beacon\"]\n", "c2_beacon.configure(c2_server_ip_address=\"192.168.10.21\")\n", "c2_beacon.establish()" ] @@ -1335,18 +1367,18 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "As demonstrated earlier, red agents can use the ``configure_c2_beacon`` action to configure these settings mid episode through the configuration options:\n", + "As demonstrated earlier, red agents can use the ``configure-c2-beacon`` action to configure these settings mid episode through the configuration options:\n", "\n", "``` YAML\n", "...\n", - " action: configure-c2-beacon\n", + " action: configure_c2_beacon\n", " options:\n", " node_id: 0\n", " config:\n", " c2_server_ip_address: 192.168.10.21\n", " keep_alive_frequency: 10\n", - " masquerade_protocol: TCP\n", - " masquerade_port: DNS\n", + " masquerade_protocol: tcp\n", + " masquerade_port: dns\n", "```" ] }, diff --git a/src/primaite/notebooks/Data-Manipulation-Customising-Red-Agent.ipynb b/src/primaite/notebooks/Data-Manipulation-Customising-Red-Agent.ipynb index 2dbf750e..9ac1da9b 100644 --- a/src/primaite/notebooks/Data-Manipulation-Customising-Red-Agent.ipynb +++ b/src/primaite/notebooks/Data-Manipulation-Customising-Red-Agent.ipynb @@ -6,7 +6,7 @@ "source": [ "# Customising UC2 Red Agents\n", "\n", - "© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n", + "© Crown-owned copyright 2025, Defence Science and Technology Laboratory UK\n", "\n", "This notebook will go over some examples of how red agent behaviour can be varied by changing its configuration parameters.\n", "\n", @@ -79,8 +79,7 @@ " if red_action == 'do-nothing':\n", " red_str = 'DO NOTHING'\n", " elif red_action == 'node-application-execute':\n", - " client = \"client 1\" if red_info.parameters['node_id'] == 0 else \"client 2\"\n", - " red_str = f\"ATTACK from {client}\"\n", + " red_str = f\"ATTACK from {red_info.parameters['node_name']}\"\n", " return red_str" ] }, @@ -149,18 +148,13 @@ " team: RED # not used, just for human reference\n", " type: red-database-corrupting-agent # type of agent - this lets primaite know which agent class to use\n", "\n", - " # Since the agent does not need to react to what is happening in the environment, the observation space is empty.\n", - " observation_space:\n", - " type: uc2-red-observation # TODO: what\n", - " options:\n", - " nodes: {}\n", - "\n", " # These actions are passed to the RedDatabaseCorruptingAgent init method, they dictate the schedule of attacks\n", " agent_settings:\n", - " start_settings:\n", - " start_step: 25 # first attack at step 25\n", - " frequency: 20 # attacks will happen every 20 steps (on average)\n", - " variance: 5 # the timing of attacks will vary by up to 5 steps earlier or later\n", + " possible_start_nodes: [client_1, client_2] # List of clients the attack can start from\n", + " target_application: data-manipulation-bot\n", + " start_step: 25 # first attack at step 25\n", + " frequency: 20 # attacks will happen every 20 steps (on average)\n", + " variance: 5 # the timing of attacks will vary by up to 5 steps earlier or later\n", "```" ] }, @@ -180,8 +174,7 @@ "simulation:\n", " network:\n", " nodes:\n", - " - ref: client_1\n", - " hostname: client_1\n", + " - hostname: client_1\n", " type: computer\n", " ip_address: 192.168.10.21\n", " subnet_mask: 255.255.255.0\n", @@ -189,13 +182,13 @@ " \n", " # \n", " applications:\n", - " - type: data-manipulation-bot\n", + " - type: data-manipulation-bot\n", " options:\n", " port_scan_p_of_success: 0.8 # Probability that port scan is successful\n", " data_manipulation_p_of_success: 0.8 # Probability that SQL attack is successful\n", " payload: \"DELETE\" # The SQL query which causes the attack (this has to be DELETE)\n", " server_ip: 192.168.1.14 # IP address of server hosting the database\n", - " - type: database-client # Database client must be installed in order for DataManipulationBot to function\n", + " - type: database-client # Database client must be installed in order for DataManipulationBot to function\n", " options:\n", " db_server_ip: 192.168.1.14 # IP address of server hosting the database\n", "```" @@ -219,7 +212,8 @@ "outputs": [], "source": [ "change = yaml.safe_load(\"\"\"\n", - "start_settings:\n", + " possible_start_nodes: [client_1]\n", + " target_application: DataManipulationBot\n", " start_step: 25\n", " frequency: 20\n", " variance: 0\n", @@ -229,7 +223,9 @@ " cfg = yaml.safe_load(f)\n", " for agent in cfg['agents']:\n", " if agent['ref'] == \"data_manipulation_attacker\":\n", + " print(f\"{agent['agent_settings']=}\")\n", " agent['agent_settings'] = change\n", + " print(f\"{agent['agent_settings']=}\")\n", "\n", "env = PrimaiteGymEnv(env_config = cfg)\n", "env.reset()\n", @@ -286,9 +282,21 @@ "outputs": [], "source": [ "change = yaml.safe_load(\"\"\"\n", - "# TODO:\n", + " agent_settings:\n", + " possible_start_nodes: [client_1]\n", + " target_application: DataManipulationBot\n", + "\n", + " action_space:\n", + " action_map:\n", + " 0:\n", + " action: do-nothing\n", + " options: {}\n", + " 1:\n", + " action: node-application-execute\n", + " options:\n", + " node_name: client_1\n", + " application_name: DataManipulationBot\n", "\"\"\")\n", - "#TODO 2869 fix\n", "\n", "with open(data_manipulation_config_path(), 'r') as f:\n", " cfg = yaml.safe_load(f)\n", @@ -330,16 +338,19 @@ "# Make attack always succeed.\n", "change = yaml.safe_load(\"\"\"\n", " applications:\n", - " - type: data-manipulation-bot\n", + " - ref: data_manipulation_bot\n", + " type: data-manipulation-bot\n", " options:\n", " port_scan_p_of_success: 1.0\n", " data_manipulation_p_of_success: 1.0\n", " payload: \"DELETE\"\n", " server_ip: 192.168.1.14\n", - " - type: web-browser\n", + " - ref: client_1_web_browser\n", + " type: web-browser\n", " options:\n", " target_url: http://arcd.com/users/\n", - " - type: database-client\n", + " - ref: client_1_database_client\n", + " type: database-client\n", " options:\n", " db_server_ip: 192.168.1.14\n", "\"\"\")\n", @@ -372,16 +383,19 @@ "# Make attack always fail.\n", "change = yaml.safe_load(\"\"\"\n", " applications:\n", - " - type: data-manipulation-bot\n", + " - ref: data_manipulation_bot\n", + " type: data-manipulation-bot\n", " options:\n", " port_scan_p_of_success: 0.0\n", " data_manipulation_p_of_success: 0.0\n", " payload: \"DELETE\"\n", " server_ip: 192.168.1.14\n", - " - type: web-browser\n", + " - ref: client_1_web_browser\n", + " type: web-browser\n", " options:\n", " target_url: http://arcd.com/users/\n", - " - type: database-client\n", + " - ref: client_1_database_client\n", + " type: database-client\n", " options:\n", " db_server_ip: 192.168.1.14\n", "\"\"\")\n", @@ -408,7 +422,7 @@ ], "metadata": { "kernelspec": { - "display_name": ".venv", + "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, @@ -422,7 +436,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.11" + "version": "3.10.12" } }, "nbformat": 4, diff --git a/src/primaite/notebooks/Data-Manipulation-E2E-Demonstration.ipynb b/src/primaite/notebooks/Data-Manipulation-E2E-Demonstration.ipynb index 2070e03c..562f0f91 100644 --- a/src/primaite/notebooks/Data-Manipulation-E2E-Demonstration.ipynb +++ b/src/primaite/notebooks/Data-Manipulation-E2E-Demonstration.ipynb @@ -6,7 +6,7 @@ "source": [ "# Data Manipulation Scenario\n", "\n", - "© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK" + "© Crown-owned copyright 2025, Defence Science and Technology Laboratory UK" ] }, { @@ -429,9 +429,9 @@ " cfg = yaml.safe_load(f)\n", " # set success probability to 1.0 to avoid rerunning cells.\n", " cfg['simulation']['network']['nodes'][8]['applications'][0]['options']['data_manipulation_p_of_success'] = 1.0\n", - " cfg['simulation']['network']['nodes'][9]['applications'][0]['options']['data_manipulation_p_of_success'] = 1.0\n", + " cfg['simulation']['network']['nodes'][9]['applications'][1]['options']['data_manipulation_p_of_success'] = 1.0\n", " cfg['simulation']['network']['nodes'][8]['applications'][0]['options']['port_scan_p_of_success'] = 1.0\n", - " cfg['simulation']['network']['nodes'][9]['applications'][0]['options']['port_scan_p_of_success'] = 1.0\n", + " cfg['simulation']['network']['nodes'][9]['applications'][1]['options']['port_scan_p_of_success'] = 1.0\n", " # don't flatten observations so that we can see what is going on\n", " cfg['agents'][3]['agent_settings']['flatten_obs'] = False\n", "\n", @@ -458,10 +458,10 @@ " # parse the info dict form step output and write out what the red agent is doing\n", " red_info : AgentHistoryItem = info['agent_actions']['data_manipulation_attacker']\n", " red_action = red_info.action\n", - " if red_action == 'do_nothing':\n", + " if red_action == 'do-nothing':\n", " red_str = 'DO NOTHING'\n", - " elif red_action == 'node_application_execute':\n", - " client = \"client 1\" if red_info.parameters['node_id'] == 0 else \"client 2\"\n", + " elif red_action == 'node-application-execute':\n", + " client = \"client 1\" if red_info.parameters['node_name'] == 0 else \"client 2\"\n", " red_str = f\"ATTACK from {client}\"\n", " return red_str" ] @@ -600,7 +600,7 @@ "while abs(reward - 0.8) > 1e-5:\n", " obs, reward, terminated, truncated, info = env.step(0) # do nothing\n", " print(f\"step: {env.game.step_counter}, Red action: {info['agent_actions']['data_manipulation_attacker'].action}, Blue reward:{reward:.2f}\" )\n", - " if env.game.step_counter > 10000:\n", + " if env.game.step_counter > 2000:\n", " break # make sure there's no infinite loop if something went wrong" ] }, diff --git a/src/primaite/notebooks/Getting-Information-Out-Of-PrimAITE.ipynb b/src/primaite/notebooks/Getting-Information-Out-Of-PrimAITE.ipynb index 58573ac6..40250d0c 100644 --- a/src/primaite/notebooks/Getting-Information-Out-Of-PrimAITE.ipynb +++ b/src/primaite/notebooks/Getting-Information-Out-Of-PrimAITE.ipynb @@ -6,7 +6,7 @@ "source": [ "# Getting information out of PrimAITE\n", "\n", - "© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n" + "© Crown-owned copyright 2025, Defence Science and Technology Laboratory UK\n" ] }, { @@ -200,7 +200,7 @@ ], "metadata": { "kernelspec": { - "display_name": "venv", + "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, @@ -214,7 +214,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.11" + "version": "3.10.12" } }, "nbformat": 4, diff --git a/src/primaite/notebooks/Privilege-Escalation-and-Data-Loss-Example.ipynb b/src/primaite/notebooks/Privilege-Escalation-and-Data-Loss-Example.ipynb index fcda4dbd..deb38eea 100644 --- a/src/primaite/notebooks/Privilege-Escalation-and-Data-Loss-Example.ipynb +++ b/src/primaite/notebooks/Privilege-Escalation-and-Data-Loss-Example.ipynb @@ -6,7 +6,7 @@ "source": [ "# Simulating Privilege Escalation and Data Loss Using SSH and ACLs Manipulation\n", "\n", - "© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n", + "© Crown-owned copyright 2025, Defence Science and Technology Laboratory UK\n", "\n", "## Overview\n", "\n", @@ -51,6 +51,15 @@ "## The Scenario" ] }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "!primaite setup" + ] + }, { "cell_type": "code", "execution_count": null, @@ -68,7 +77,8 @@ "from primaite.simulator.network.hardware.nodes.host.server import Server\n", "from primaite.simulator.system.applications.database_client import DatabaseClient\n", "from primaite.simulator.system.applications.web_browser import WebBrowser\n", - "from primaite.simulator.system.services.database.database_service import DatabaseService" + "from primaite.simulator.system.services.database.database_service import DatabaseService\n", + "from primaite.simulator.network.hardware.nodes.network import firewall\n" ] }, { @@ -110,11 +120,11 @@ "outputs": [], "source": [ "some_tech_jnr_dev_pc: Computer = game.simulation.network.get_node_by_hostname(\"some_tech_jnr_dev_pc\")\n", - "some_tech_jnr_dev_db_client: DatabaseClient = some_tech_jnr_dev_pc.software_manager.software[\"DatabaseClient\"]\n", - "some_tech_jnr_dev_web_browser: WebBrowser = some_tech_jnr_dev_pc.software_manager.software[\"WebBrowser\"]\n", + "some_tech_jnr_dev_db_client: DatabaseClient = some_tech_jnr_dev_pc.software_manager.software[\"database-client\"]\n", + "some_tech_jnr_dev_web_browser: WebBrowser = some_tech_jnr_dev_pc.software_manager.software[\"web-browser\"]\n", "some_tech_rt: Router = game.simulation.network.get_node_by_hostname(\"some_tech_rt\")\n", "some_tech_db_srv: Server = game.simulation.network.get_node_by_hostname(\"some_tech_db_srv\")\n", - "some_tech_db_service: DatabaseService = some_tech_db_srv.software_manager.software[\"DatabaseService\"]\n", + "some_tech_db_service: DatabaseService = some_tech_db_srv.software_manager.software[\"database-service\"]\n", "some_tech_storage_srv: Server = game.simulation.network.get_node_by_hostname(\"some_tech_storage_srv\")\n", "some_tech_web_srv: Server = game.simulation.network.get_node_by_hostname(\"some_tech_web_srv\")" ] @@ -201,7 +211,7 @@ "source": [ "caos_action = [\n", " \"network\", \"node\", \"some_tech_jnr_dev_pc\", \n", - " \"service\", \"Terminal\", \"node_session_remote_login\", \"admin\", \"admin\", str(some_tech_storage_srv.network_interface[1].ip_address)\n", + " \"service\", \"terminal\", \"node-session-remote-login\", \"admin\", \"admin\", str(some_tech_storage_srv.network_interface[1].ip_address)\n", "]\n", "game.simulation.apply_request(caos_action)" ] @@ -223,7 +233,7 @@ }, "outputs": [], "source": [ - "caos_action = [\"network\", \"node\", \"some_tech_jnr_dev_pc\", \"application\", \"WebBrowser\", \"execute\"]\n", + "caos_action = [\"network\", \"node\", \"some_tech_jnr_dev_pc\", \"application\", \"web-browser\", \"execute\"]\n", "game.simulation.apply_request(caos_action)" ] }, @@ -246,7 +256,7 @@ "metadata": {}, "outputs": [], "source": [ - "game.get_sim_state()[\"network\"][\"nodes\"][\"some_tech_rt\"][\"services\"][\"UserSessionManager\"][\"active_remote_sessions\"]" + "game.get_sim_state()[\"network\"][\"nodes\"][\"some_tech_rt\"][\"services\"][\"user-session-manager\"][\"active_remote_sessions\"]" ] }, { @@ -259,7 +269,7 @@ "source": [ "caos_action = [\n", " \"network\", \"node\", \"some_tech_jnr_dev_pc\", \n", - " \"service\", \"Terminal\", \"node_session_remote_login\", \"admin\", \"admin\", str(some_tech_rt.network_interface[4].ip_address)\n", + " \"service\", \"terminal\", \"node-session-remote-login\", \"admin\", \"admin\", str(some_tech_rt.network_interface[4].ip_address)\n", "]\n", "game.simulation.apply_request(caos_action)" ] @@ -272,7 +282,7 @@ }, "outputs": [], "source": [ - "game.get_sim_state()[\"network\"][\"nodes\"][\"some_tech_rt\"][\"services\"][\"UserSessionManager\"][\"active_remote_sessions\"]" + "game.get_sim_state()[\"network\"][\"nodes\"][\"some_tech_rt\"][\"services\"][\"user-session-manager\"][\"active_remote_sessions\"]" ] }, { @@ -305,7 +315,7 @@ "source": [ "caos_action = [\n", " \"network\", \"node\", \"some_tech_jnr_dev_pc\", \n", - " \"service\", \"Terminal\", \"send_remote_command\", str(some_tech_rt.network_interface[4].ip_address),\n", + " \"service\", \"terminal\", \"send_remote_command\", str(some_tech_rt.network_interface[4].ip_address),\n", " {\n", " \"command\": [\n", " \"acl\", \"add_rule\", \"PERMIT\", \"TCP\",\n", @@ -358,7 +368,7 @@ "source": [ "caos_action = [\n", " \"network\", \"node\", \"some_tech_jnr_dev_pc\", \n", - " \"service\", \"Terminal\", \"remote_logoff\", str(some_tech_rt.network_interface[4].ip_address)\n", + " \"service\", \"terminal\", \"remote_logoff\", str(some_tech_rt.network_interface[4].ip_address)\n", "]\n", "game.simulation.apply_request(caos_action)" ] @@ -376,7 +386,7 @@ "metadata": {}, "outputs": [], "source": [ - "game.get_sim_state()[\"network\"][\"nodes\"][\"some_tech_rt\"][\"services\"][\"UserSessionManager\"][\"active_remote_sessions\"]" + "game.get_sim_state()[\"network\"][\"nodes\"][\"some_tech_rt\"][\"services\"][\"user-session-manager\"][\"active_remote_sessions\"]" ] }, { @@ -396,7 +406,7 @@ "source": [ "caos_action = [\n", " \"network\", \"node\", \"some_tech_jnr_dev_pc\", \n", - " \"service\", \"Terminal\", \"node_session_remote_login\", \"admin\", \"admin\", str(some_tech_storage_srv.network_interface[1].ip_address)\n", + " \"service\", \"terminal\", \"node-session-remote-login\", \"admin\", \"admin\", str(some_tech_storage_srv.network_interface[1].ip_address)\n", "]\n", "game.simulation.apply_request(caos_action)" ] @@ -411,7 +421,7 @@ "source": [ "caos_action = [\n", " \"network\", \"node\", \"some_tech_jnr_dev_pc\", \n", - " \"service\", \"Terminal\", \"send_remote_command\", str(some_tech_storage_srv.network_interface[1].ip_address),\n", + " \"service\", \"terminal\", \"send_remote_command\", str(some_tech_storage_srv.network_interface[1].ip_address),\n", " {\n", " \"command\": [\n", " \"file_system\", \"delete\", \"file\", db_backup_folder, \"database.db\"\n", @@ -466,7 +476,7 @@ }, "outputs": [], "source": [ - "caos_action = [\"network\", \"node\", \"some_tech_jnr_dev_pc\", \"application\", \"WebBrowser\", \"execute\"]\n", + "caos_action = [\"network\", \"node\", \"some_tech_jnr_dev_pc\", \"application\", \"web-browser\", \"execute\"]\n", "game.simulation.apply_request(caos_action)" ] }, @@ -525,7 +535,7 @@ }, "outputs": [], "source": [ - "caos_action = [\"network\", \"node\", \"some_tech_jnr_dev_pc\", \"application\", \"WebBrowser\", \"execute\"]\n", + "caos_action = [\"network\", \"node\", \"some_tech_jnr_dev_pc\", \"application\", \"web-browser\", \"execute\"]\n", "game.simulation.apply_request(caos_action)" ] }, diff --git a/src/primaite/notebooks/Requests-and-Responses.ipynb b/src/primaite/notebooks/Requests-and-Responses.ipynb index 9260b29d..c29d41dc 100644 --- a/src/primaite/notebooks/Requests-and-Responses.ipynb +++ b/src/primaite/notebooks/Requests-and-Responses.ipynb @@ -6,7 +6,7 @@ "source": [ "# Requests and Responses\n", "\n", - "© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n", + "© Crown-owned copyright 2025, Defence Science and Technology Laboratory UK\n", "\n", "Agents interact with the PrimAITE simulation via the Request system.\n" ] @@ -52,12 +52,17 @@ "outputs": [], "source": [ "sim = Simulation()\n", + "\n", "sim.network.add_node(\n", - " HostNode(\n", - " hostname=\"client\",\n", - " ip_address='10.0.0.1',\n", - " subnet_mask='255.255.255.0',\n", - " operating_state=NodeOperatingState.ON)\n", + " HostNode.from_config(\n", + " config = {\n", + " 'type': \"host-node\",\n", + " 'hostname': \"client\",\n", + " 'ip_address': '10.0.0.1',\n", + " 'subnet_mask': '255.255.255.0',\n", + " 'operating_state': \"ON\",\n", + " }\n", + " )\n", ")\n", "client = sim.network.get_node_by_hostname('client')\n" ] @@ -94,7 +99,7 @@ "outputs": [], "source": [ "response = sim.apply_request(\n", - " request=[\"network\", \"node\", \"client\", \"service\", \"DNSClient\", \"stop\"],\n", + " request=[\"network\", \"node\", \"client\", \"service\", \"dns-client\", \"stop\"],\n", " context={}\n", " )\n", "print(response)" @@ -138,7 +143,7 @@ "outputs": [], "source": [ "response = sim.apply_request(\n", - " request=[\"network\", \"node\", \"client\", \"service\", \"NonExistentApplication\", \"stop\"],\n", + " request=[\"network\", \"node\", \"client\", \"service\", \"non-existent-application\", \"stop\"],\n", " context={}\n", " )\n", "print(response)" @@ -201,7 +206,7 @@ "outputs": [], "source": [ "response = sim.apply_request(\n", - " request=[\"network\", \"node\", \"client\", \"service\", \"DNSClient\", \"start\"],\n", + " request=[\"network\", \"node\", \"client\", \"service\", \"dns-client\", \"start\"],\n", " context={}\n", " )\n", "print(response)" @@ -210,7 +215,7 @@ ], "metadata": { "kernelspec": { - "display_name": "venv", + "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, diff --git a/src/primaite/notebooks/Terminal-Processing.ipynb b/src/primaite/notebooks/Terminal-Processing.ipynb index 7c94d432..07d38791 100644 --- a/src/primaite/notebooks/Terminal-Processing.ipynb +++ b/src/primaite/notebooks/Terminal-Processing.ipynb @@ -6,7 +6,7 @@ "source": [ "# Terminal Processing\n", "\n", - "© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK" + "© Crown-owned copyright 2025, Defence Science and Technology Laboratory UK" ] }, { @@ -49,9 +49,26 @@ "def basic_network() -> Network:\n", " \"\"\"Utility function for creating a default network to demonstrate Terminal functionality\"\"\"\n", " network = Network()\n", - " node_a = Computer(hostname=\"node_a\", ip_address=\"192.168.0.10\", subnet_mask=\"255.255.255.0\", start_up_duration=0)\n", + " node_a = Computer.from_config(\n", + " config = {\n", + " \"type\": \"computer\",\n", + " \"hostname\": \"node_a\",\n", + " \"ip_address\": \"192.168.0.10\",\n", + " \"subnet_mask\": \"255.255.255.0\",\n", + " # \"startup_duration\": 0,\n", + " }\n", + " )\n", + " print(f\"{node_a=}\")\n", " node_a.power_on()\n", - " node_b = Computer(hostname=\"node_b\", ip_address=\"192.168.0.11\", subnet_mask=\"255.255.255.0\", start_up_duration=0)\n", + " node_b = Computer.from_config(\n", + " config = {\n", + " \"type\": \"computer\",\n", + " \"hostname\": \"node_b\",\n", + " \"ip_address\": \"192.168.0.11\",\n", + " \"subnet_mask\": \"255.255.255.0\",\n", + " # \"startup_duration\": 0,\n", + " }\n", + " )\n", " node_b.power_on()\n", " network.connect(node_a.network_interface[1], node_b.network_interface[1])\n", " return network" @@ -491,7 +508,7 @@ ], "metadata": { "kernelspec": { - "display_name": ".venv", + "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, @@ -505,7 +522,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.11" + "version": "3.10.12" } }, "nbformat": 4, diff --git a/src/primaite/notebooks/Training-an-RLLIB-MARL-System.ipynb b/src/primaite/notebooks/Training-an-RLLIB-MARL-System.ipynb index 76cab86a..87d9c377 100644 --- a/src/primaite/notebooks/Training-an-RLLIB-MARL-System.ipynb +++ b/src/primaite/notebooks/Training-an-RLLIB-MARL-System.ipynb @@ -6,7 +6,7 @@ "source": [ "# Train a Multi agent system using RLLIB\n", "\n", - "© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n", + "© Crown-owned copyright 2025, Defence Science and Technology Laboratory UK\n", "\n", "This notebook will demonstrate how to use the `PrimaiteRayMARLEnv` to train a very basic system with two PPO agents." ] @@ -40,6 +40,7 @@ "import ray\n", "from ray.rllib.algorithms.ppo import PPOConfig\n", "from primaite.session.ray_envs import PrimaiteRayMARLEnv\n", + "from primaite.game.agent.scripted_agents import probabilistic_agent\n", "\n", "with open(PRIMAITE_PATHS.user_config_path / 'example_config/data_manipulation_marl.yaml', 'r') as f:\n", " cfg = yaml.safe_load(f)\n", @@ -110,21 +111,9 @@ ], "metadata": { "kernelspec": { - "display_name": "venv", + "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.10.12" } }, "nbformat": 4, diff --git a/src/primaite/notebooks/Training-an-RLLib-Agent.ipynb b/src/primaite/notebooks/Training-an-RLLib-Agent.ipynb index 7252b046..79740bca 100644 --- a/src/primaite/notebooks/Training-an-RLLib-Agent.ipynb +++ b/src/primaite/notebooks/Training-an-RLLib-Agent.ipynb @@ -6,7 +6,7 @@ "source": [ "# Train a Single agent system using RLLib\n", "\n", - "© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n", + "© Crown-owned copyright 2025, Defence Science and Technology Laboratory UK\n", "\n", "This notebook will demonstrate how to use PrimaiteRayEnv to train a basic PPO agent." ] @@ -32,6 +32,8 @@ "from primaite.session.ray_envs import PrimaiteRayEnv\n", "import ray\n", "from ray.rllib.algorithms.ppo import PPOConfig\n", + "from primaite.game.agent.scripted_agents import probabilistic_agent\n", + "\n", "\n", "# If you get an error saying this config file doesn't exist, you may need to run `primaite setup` in your command line\n", "# to copy the files to your user data path.\n", @@ -104,21 +106,9 @@ ], "metadata": { "kernelspec": { - "display_name": ".venv", + "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.10.11" } }, "nbformat": 4, diff --git a/src/primaite/notebooks/Training-an-SB3-Agent.ipynb b/src/primaite/notebooks/Training-an-SB3-Agent.ipynb index 2b554475..1bec0183 100644 --- a/src/primaite/notebooks/Training-an-SB3-Agent.ipynb +++ b/src/primaite/notebooks/Training-an-SB3-Agent.ipynb @@ -6,7 +6,7 @@ "source": [ "# Training an SB3 Agent\n", "\n", - "© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n", + "© Crown-owned copyright 2025, Defence Science and Technology Laboratory UK\n", "\n", "This notebook will demonstrate how to use primaite to create and train a PPO agent, using a pre-defined configuration file." ] @@ -35,6 +35,7 @@ "source": [ "from primaite.game.game import PrimaiteGame\n", "from primaite.session.environment import PrimaiteGymEnv\n", + "from primaite.game.agent.scripted_agents import probabilistic_agent\n", "import yaml" ] }, @@ -177,7 +178,7 @@ ], "metadata": { "kernelspec": { - "display_name": ".venv", + "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, @@ -191,7 +192,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.11" + "version": "3.10.12" } }, "nbformat": 4, diff --git a/src/primaite/notebooks/Using-Episode-Schedules.ipynb b/src/primaite/notebooks/Using-Episode-Schedules.ipynb index d08ca67b..4bc49a70 100644 --- a/src/primaite/notebooks/Using-Episode-Schedules.ipynb +++ b/src/primaite/notebooks/Using-Episode-Schedules.ipynb @@ -6,7 +6,7 @@ "source": [ "# Using Episode Schedules\n", "\n", - "© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n", + "© Crown-owned copyright 2025, Defence Science and Technology Laboratory UK\n", "\n", "PrimAITE supports the ability to use different variations on a scenario at different episodes. This can be used to increase \n", "domain randomisation to prevent overfitting, or to set up curriculum learning to train agents to perform more complicated tasks.\n", @@ -48,6 +48,7 @@ "from primaite.session.environment import PrimaiteGymEnv\n", "from primaite import PRIMAITE_PATHS\n", "from prettytable import PrettyTable\n", + "from primaite.game.agent.scripted_agents import probabilistic_agent, data_manipulation_bot\n", "scenario_path = PRIMAITE_PATHS.user_config_path / \"example_config/scenario_with_placeholders\"" ] }, @@ -409,7 +410,7 @@ ], "metadata": { "kernelspec": { - "display_name": "venv", + "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, @@ -423,7 +424,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.11" + "version": "3.10.12" } }, "nbformat": 4, diff --git a/src/primaite/simulator/_package_data/create-simulation_demo.ipynb b/src/primaite/notebooks/create-simulation_demo.ipynb similarity index 80% rename from src/primaite/simulator/_package_data/create-simulation_demo.ipynb rename to src/primaite/notebooks/create-simulation_demo.ipynb index 690e7856..46f03be6 100644 --- a/src/primaite/simulator/_package_data/create-simulation_demo.ipynb +++ b/src/primaite/notebooks/create-simulation_demo.ipynb @@ -6,7 +6,7 @@ "source": [ "# Build a simulation using the Python API\n", "\n", - "© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n", + "© Crown-owned copyright 2025, Defence Science and Technology Laboratory UK\n", "\n", "Currently, this notebook manipulates the simulation by directly placing objects inside of the attributes of the network and domain. It should be refactored when proper methods exist for adding these objects." ] @@ -70,9 +70,23 @@ "metadata": {}, "outputs": [], "source": [ - "my_pc = Computer(hostname=\"Computer\", ip_address=\"192.168.1.10\", subnet_mask=\"255.255.255.0\")\n", + "my_pc = Computer.from_config(\n", + " config={\n", + " \"type\": \"computer\",\n", + " \"hostname\":\"pc_1\",\n", + " \"ip_address\":\"192.168.1.10\",\n", + " \"subnet_mask\":\"255.255.255.0\",\n", + " }\n", + " )\n", "net.add_node(my_pc)\n", - "my_server = Server(hostname=\"Server\", ip_address=\"192.168.1.11\", subnet_mask=\"255.255.255.0\")\n", + "my_server = Server.from_config(\n", + " config={\n", + " \"type\": \"server\",\n", + " \"hostname\":\"Server\",\n", + " \"ip_address\":\"192.168.1.11\",\n", + " \"subnet_mask\":\"255.255.255.0\"\n", + " }\n", + ")\n", "net.add_node(my_server)\n" ] }, @@ -99,7 +113,13 @@ "metadata": {}, "outputs": [], "source": [ - "my_switch = Switch(hostname=\"switch1\", num_ports=12)\n", + "my_switch = Switch.from_config(\n", + " config = {\n", + " \"type\":\"switch\",\n", + " \"hostname\":\"switch1\",\n", + " \"num_ports\":12\n", + " }\n", + ")\n", "net.add_node(my_switch)\n", "\n", "pc_nic = NIC(ip_address=\"130.1.1.1\", gateway=\"130.1.1.255\", subnet_mask=\"255.255.255.0\")\n", @@ -163,16 +183,30 @@ "metadata": {}, "outputs": [], "source": [ + "from pydantic import Field\n", + "\n", "from pathlib import Path\n", "from primaite.simulator.system.applications.application import Application, ApplicationOperatingState\n", "from primaite.simulator.system.software import SoftwareHealthState, SoftwareCriticality\n", "from primaite.simulator.file_system.file_system import FileSystem\n", "from primaite.utils.validation.ip_protocol import PROTOCOL_LOOKUP\n", "from primaite.utils.validation.port import PORT_LOOKUP\n", + "from primaite.simulator.system.core.sys_log import SysLog\n", "\n", "\n", "# no applications exist yet so we will create our own.\n", "class MSPaint(Application, discriminator=\"MSPaint\"):\n", + " class ConfigSchema(Application.ConfigSchema):\n", + " type: str = \"MSPaint\"\n", + "\n", + " config: ConfigSchema = Field(default_factory=lambda: MSPaint.ConfigSchema())\n", + "\n", + " def __init__(self, **kwargs):\n", + " kwargs[\"name\"] = \"MSPaint\"\n", + " kwargs[\"port\"] = PORT_LOOKUP[\"HTTP\"]\n", + " kwargs[\"protocol\"] = PROTOCOL_LOOKUP[\"NONE\"]\n", + " super().__init__(**kwargs)\n", + "\n", " def describe_state(self):\n", " return super().describe_state()" ] @@ -183,7 +217,8 @@ "metadata": {}, "outputs": [], "source": [ - "mspaint = MSPaint(name = \"mspaint\", health_state_actual=SoftwareHealthState.GOOD, health_state_visible=SoftwareHealthState.GOOD, criticality=SoftwareCriticality.MEDIUM, port=PORT_LOOKUP[\"HTTP\"], protocol = PROTOCOL_LOOKUP[\"NONE\"],operating_state=ApplicationOperatingState.RUNNING,execution_control_status='manual', file_system=FileSystem(sys_log=SysLog(hostname=\"Test\"), sim_root=Path(__name__).parent),)" + "my_pc.software_manager.install(MSPaint)\n", + "mspaint = my_pc.software_manager.software.get(\"MSPaint\")" ] }, { @@ -250,7 +285,7 @@ ], "metadata": { "kernelspec": { - "display_name": ".venv", + "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, @@ -264,7 +299,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.11" + "version": "3.10.12" } }, "nbformat": 4, diff --git a/src/primaite/notebooks/multi-processing.ipynb b/src/primaite/notebooks/multi-processing.ipynb index ad386f34..e56bf362 100644 --- a/src/primaite/notebooks/multi-processing.ipynb +++ b/src/primaite/notebooks/multi-processing.ipynb @@ -6,7 +6,7 @@ "source": [ "# Simple multi-processing demonstration\n", "\n", - "© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n", + "© Crown-owned copyright 2025, Defence Science and Technology Laboratory UK\n", "\n", "This notebook uses SubprocVecEnv from SB3." ] @@ -37,7 +37,6 @@ "from stable_baselines3 import PPO\n", "from stable_baselines3.common.utils import set_random_seed\n", "from stable_baselines3.common.vec_env import SubprocVecEnv\n", - "\n", "from primaite.session.environment import PrimaiteGymEnv\n" ] }, @@ -138,7 +137,7 @@ ], "metadata": { "kernelspec": { - "display_name": ".venv", + "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, @@ -152,7 +151,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.11" + "version": "3.10.12" } }, "nbformat": 4, diff --git a/src/primaite/simulator/_package_data/network_simulator_demo.ipynb b/src/primaite/notebooks/network_simulator_demo.ipynb similarity index 97% rename from src/primaite/simulator/_package_data/network_simulator_demo.ipynb rename to src/primaite/notebooks/network_simulator_demo.ipynb index 58c07fe5..be930ac0 100644 --- a/src/primaite/simulator/_package_data/network_simulator_demo.ipynb +++ b/src/primaite/notebooks/network_simulator_demo.ipynb @@ -7,7 +7,7 @@ "source": [ "# PrimAITE Router Simulation Demo\n", "\n", - "© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n", + "© Crown-owned copyright 2025, Defence Science and Technology Laboratory UK\n", "\n", "This demo uses a modified version of the ARCD Use Case 2 Network (seen below) to demonstrate the capabilities of the Network simulator in PrimAITE." ] @@ -650,21 +650,9 @@ ], "metadata": { "kernelspec": { - "display_name": ".venv", + "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.10.11" } }, "nbformat": 4, diff --git a/src/primaite/simulator/network/airspace.py b/src/primaite/simulator/network/airspace.py index 7ede0bb0..434940f8 100644 --- a/src/primaite/simulator/network/airspace.py +++ b/src/primaite/simulator/network/airspace.py @@ -449,10 +449,8 @@ class IPWirelessNetworkInterface(WirelessNetworkInterface, Layer3Interface, ABC) `default_gateway_hello` method is not defined, ignoring such errors to proceed without interruption. """ super().enable() - try: + if hasattr(self._connected_node, "default_gateway_hello"): self._connected_node.default_gateway_hello() - except AttributeError: - pass @abstractmethod def receive_frame(self, frame: Frame) -> bool: diff --git a/src/primaite/simulator/network/container.py b/src/primaite/simulator/network/container.py index 2e494910..b0426537 100644 --- a/src/primaite/simulator/network/container.py +++ b/src/primaite/simulator/network/container.py @@ -163,16 +163,6 @@ class Network(SimComponent): :param links: Include link details in the output. Defaults to True. :param markdown: Use Markdown style in table output. Defaults to False. """ - nodes_type_map = { - "Router": self.router_nodes, - "Firewall": self.firewall_nodes, - "Switch": self.switch_nodes, - "Server": self.server_nodes, - "Computer": self.computer_nodes, - "Printer": self.printer_nodes, - "Wireless Router": self.wireless_router_nodes, - } - if nodes: table = PrettyTable(["Node", "Type", "Operating State"]) if markdown: @@ -189,21 +179,20 @@ class Network(SimComponent): table.set_style(MARKDOWN) table.align = "l" table.title = "IP Addresses" - for nodes in nodes_type_map.values(): - for node in nodes: - for i, port in node.network_interface.items(): - if hasattr(port, "ip_address"): - if port.ip_address != IPv4Address("127.0.0.1"): - port_str = port.port_name if port.port_name else port.port_num - table.add_row( - [ - node.config.hostname, - port_str, - port.ip_address, - port.subnet_mask, - node.config.default_gateway, - ] - ) + for node in self.nodes.values(): + for i, port in node.network_interface.items(): + if hasattr(port, "ip_address"): + if port.ip_address != IPv4Address("127.0.0.1"): + port_str = port.port_name if port.port_name else port.port_num + table.add_row( + [ + node.config.hostname, + port_str, + port.ip_address, + port.subnet_mask, + node.config.default_gateway, + ] + ) print(table) if links: @@ -215,22 +204,21 @@ class Network(SimComponent): table.align = "l" table.title = "Links" links = list(self.links.values()) - for nodes in nodes_type_map.values(): - for node in nodes: - for link in links[::-1]: - if node in [link.endpoint_a.parent, link.endpoint_b.parent]: - table.add_row( - [ - link.endpoint_a.parent.config.hostname, - str(link.endpoint_a), - link.endpoint_b.parent.config.hostname, - str(link.endpoint_b), - link.is_up, - link.bandwidth, - link.current_load_percent, - ] - ) - links.remove(link) + for node in self.nodes.values(): + for link in links[::-1]: + if node in [link.endpoint_a.parent, link.endpoint_b.parent]: + table.add_row( + [ + link.endpoint_a.parent.config.hostname, + str(link.endpoint_a), + link.endpoint_b.parent.config.hostname, + str(link.endpoint_b), + link.is_up, + link.bandwidth, + link.current_load_percent, + ] + ) + links.remove(link) print(table) def clear_links(self): diff --git a/src/primaite/simulator/network/creation.py b/src/primaite/simulator/network/creation.py index 255b7bf5..089ed00d 100644 --- a/src/primaite/simulator/network/creation.py +++ b/src/primaite/simulator/network/creation.py @@ -155,7 +155,12 @@ class OfficeLANAdder(NetworkNodeAdder, discriminator="office-lan"): # Create a core switch if more than one edge switch is needed if num_of_switches > 1: core_switch = Switch.from_config( - config={"type": "switch", "hostname": f"switch_core_{config.lan_name}", "start_up_duration": 0} + config={ + "type": "switch", + "hostname": f"switch_core_{config.lan_name}", + "start_up_duration": 0, + "num_ports": 24, + } ) core_switch.power_on() network.add_node(core_switch) @@ -183,7 +188,12 @@ class OfficeLANAdder(NetworkNodeAdder, discriminator="office-lan"): switch_port = 0 switch_n = 1 switch = Switch.from_config( - config={"type": "switch", "hostname": f"switch_edge_{switch_n}_{config.lan_name}", "start_up_duration": 0} + config={ + "type": "switch", + "hostname": f"switch_edge_{switch_n}_{config.lan_name}", + "start_up_duration": 0, + "num_ports": 24, + } ) switch.power_on() network.add_node(switch) @@ -207,6 +217,7 @@ class OfficeLANAdder(NetworkNodeAdder, discriminator="office-lan"): "type": "switch", "hostname": f"switch_edge_{switch_n}_{config.lan_name}", "start_up_duration": 0, + "num_ports": 24, } ) switch.power_on() diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py index 50a7d2a4..9cc39848 100644 --- a/src/primaite/simulator/network/hardware/base.py +++ b/src/primaite/simulator/network/hardware/base.py @@ -639,10 +639,8 @@ class IPWiredNetworkInterface(WiredNetworkInterface, Layer3Interface, ABC): `default_gateway_hello` method is not defined, ignoring such errors to proceed without interruption. """ super().enable() - try: + if hasattr(self._connected_node, "default_gateway_hello"): self._connected_node.default_gateway_hello() - except AttributeError: - pass return True @abstractmethod @@ -1211,7 +1209,7 @@ class UserSessionManager(Service, discriminator="user-session-manager"): """Request should take the form [username, password, remote_ip_address].""" username, password, remote_ip_address = request response = RequestResponse.from_bool(self.remote_login(username, password, remote_ip_address)) - response.data = {"remote_hostname": self.parent.hostname, "username": username} + response.data = {"remote_hostname": self.parent.config.hostname, "username": username} return response rm.add_request("remote_login", RequestType(func=_remote_login)) @@ -1244,7 +1242,7 @@ class UserSessionManager(Service, discriminator="user-session-manager"): if markdown: table.set_style(MARKDOWN) table.align = "l" - table.title = f"{self.parent.hostname} User Sessions" + table.title = f"{self.parent.config.hostname} User Sessions" def _add_session_to_table(user_session: UserSession): """ diff --git a/src/primaite/simulator/network/hardware/nodes/host/host_node.py b/src/primaite/simulator/network/hardware/nodes/host/host_node.py index 86ac790c..fcd0eb80 100644 --- a/src/primaite/simulator/network/hardware/nodes/host/host_node.py +++ b/src/primaite/simulator/network/hardware/nodes/host/host_node.py @@ -334,7 +334,7 @@ class HostNode(Node, discriminator="host-node"): class ConfigSchema(Node.ConfigSchema): """Configuration Schema for HostNode class.""" - type: Literal["host-node"] + type: Literal["host-node"] = "host-node" hostname: str = "HostNode" subnet_mask: IPV4Address = "255.255.255.0" ip_address: IPV4Address @@ -388,7 +388,7 @@ class HostNode(Node, discriminator="host-node"): This method is invoked to ensure the host node can communicate with its default gateway, primarily to confirm network connectivity and populate the ARP cache with the gateway's MAC address. """ - if self.operating_state == NodeOperatingState.ON and self.default_gateway: + if self.operating_state == NodeOperatingState.ON and self.config.default_gateway: self.software_manager.arp.get_default_gateway_mac_address() def receive_frame(self, frame: Frame, from_network_interface: NIC): diff --git a/src/primaite/simulator/network/hardware/nodes/network/router.py b/src/primaite/simulator/network/hardware/nodes/network/router.py index 47ee3169..33b55b9d 100644 --- a/src/primaite/simulator/network/hardware/nodes/network/router.py +++ b/src/primaite/simulator/network/hardware/nodes/network/router.py @@ -1352,7 +1352,6 @@ class Router(NetworkNode, discriminator="router"): :return: A dictionary representing the current state. """ state = super().describe_state() - state["num_ports"] = self.config.num_ports state["acl"] = self.acl.describe_state() return state diff --git a/src/primaite/simulator/network/hardware/nodes/network/switch.py b/src/primaite/simulator/network/hardware/nodes/network/switch.py index 1f2bc135..6e5814d0 100644 --- a/src/primaite/simulator/network/hardware/nodes/network/switch.py +++ b/src/primaite/simulator/network/hardware/nodes/network/switch.py @@ -103,8 +103,8 @@ class Switch(NetworkNode, discriminator="switch"): type: Literal["switch"] = "switch" hostname: str = "Switch" - num_ports: int = 24 - "The number of ports on the switch. Default is 24." + num_ports: int = 8 + "The number of ports on the switch." config: ConfigSchema = Field(default_factory=lambda: Switch.ConfigSchema()) @@ -139,7 +139,6 @@ class Switch(NetworkNode, discriminator="switch"): """ state = super().describe_state() state["ports"] = {port_num: port.describe_state() for port_num, port in self.network_interface.items()} - state["num_ports"] = self.config.num_ports # redundant? state["mac_address_table"] = {mac: port.port_num for mac, port in self.mac_address_table.items()} return state diff --git a/src/primaite/simulator/network/networks.py b/src/primaite/simulator/network/networks.py index 5e677679..5d558e80 100644 --- a/src/primaite/simulator/network/networks.py +++ b/src/primaite/simulator/network/networks.py @@ -295,7 +295,7 @@ def arcd_uc2_network() -> Network: # Security Suite security_suite_cfg = { "type": "server", - "hostname": "backup_server", + "hostname": "security_suite", "ip_address": "192.168.1.110", "subnet_mask": "255.255.255.0", "default_gateway": "192.168.1.1", diff --git a/src/primaite/simulator/system/applications/red_applications/data_manipulation_bot.py b/src/primaite/simulator/system/applications/red_applications/data_manipulation_bot.py index 7bf7b203..17862df4 100644 --- a/src/primaite/simulator/system/applications/red_applications/data_manipulation_bot.py +++ b/src/primaite/simulator/system/applications/red_applications/data_manipulation_bot.py @@ -52,6 +52,7 @@ class DataManipulationBot(Application, discriminator="data-manipulation-bot"): payload: str = "DELETE" port_scan_p_of_success: float = 0.1 data_manipulation_p_of_success: float = 0.1 + repeat: bool = True config: "DataManipulationBot.ConfigSchema" = Field(default_factory=lambda: DataManipulationBot.ConfigSchema()) @@ -76,6 +77,7 @@ class DataManipulationBot(Application, discriminator="data-manipulation-bot"): self.payload = self.config.payload self.port_scan_p_of_success = self.config.port_scan_p_of_success self.data_manipulation_p_of_success = self.config.data_manipulation_p_of_success + self.repeat = self.config.repeat def describe_state(self) -> Dict: """ diff --git a/src/primaite/simulator/system/core/session_manager.py b/src/primaite/simulator/system/core/session_manager.py index d6617efa..26e3be79 100644 --- a/src/primaite/simulator/system/core/session_manager.py +++ b/src/primaite/simulator/system/core/session_manager.py @@ -415,5 +415,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, session.protocol]) + table.add_row([session.with_ip_address, session.dst_port, session.protocol]) print(table) diff --git a/src/primaite/simulator/system/services/arp/arp.py b/src/primaite/simulator/system/services/arp/arp.py index c6b687ce..3302041d 100644 --- a/src/primaite/simulator/system/services/arp/arp.py +++ b/src/primaite/simulator/system/services/arp/arp.py @@ -138,8 +138,8 @@ class ARP(Service, discriminator="arp"): break if use_default_gateway: - if self.software_manager.node.default_gateway: - target_ip_address = self.software_manager.node.default_gateway + if self.software_manager.node.config.default_gateway: + target_ip_address = self.software_manager.node.config.default_gateway else: return diff --git a/src/primaite/simulator/system/software.py b/src/primaite/simulator/system/software.py index 950f77c6..86b57818 100644 --- a/src/primaite/simulator/system/software.py +++ b/src/primaite/simulator/system/software.py @@ -82,7 +82,7 @@ class Software(SimComponent, ABC): """Configurable options for all software.""" model_config = ConfigDict(extra="forbid") - starting_health_state: SoftwareHealthState = SoftwareHealthState.UNUSED + starting_health_state: SoftwareHealthState = SoftwareHealthState.GOOD criticality: SoftwareCriticality = SoftwareCriticality.LOWEST fixing_duration: int = 2 diff --git a/tests/assets/configs/action_penalty.yaml b/tests/assets/configs/action_penalty.yaml index 9c117dfe..6700172e 100644 --- a/tests/assets/configs/action_penalty.yaml +++ b/tests/assets/configs/action_penalty.yaml @@ -1,3 +1,6 @@ +metadata: + version: 3.0 + io_settings: save_agent_actions: false save_step_metadata: false diff --git a/tests/assets/configs/bad_primaite_session.yaml b/tests/assets/configs/bad_primaite_session.yaml index 18b37466..b8551caf 100644 --- a/tests/assets/configs/bad_primaite_session.yaml +++ b/tests/assets/configs/bad_primaite_session.yaml @@ -1,3 +1,6 @@ +metadata: + version: 3.0 + game: ports: - ARP diff --git a/tests/assets/configs/basic_c2_setup.yaml b/tests/assets/configs/basic_c2_setup.yaml index ac2b026e..9b569b44 100644 --- a/tests/assets/configs/basic_c2_setup.yaml +++ b/tests/assets/configs/basic_c2_setup.yaml @@ -4,6 +4,9 @@ # | node_a |------| switch_1 |------| node_b | # -------------- -------------- -------------- # +metadata: + version: 3.0 + io_settings: save_step_metadata: false save_pcap_logs: true diff --git a/tests/assets/configs/basic_firewall.yaml b/tests/assets/configs/basic_firewall.yaml index 8108ecbe..26038270 100644 --- a/tests/assets/configs/basic_firewall.yaml +++ b/tests/assets/configs/basic_firewall.yaml @@ -4,6 +4,8 @@ # | client_1 |------| switch_1 |------| client_2 | # -------------- -------------- -------------- # +metadata: + version: 3.0 io_settings: save_step_metadata: false diff --git a/tests/assets/configs/basic_node_with_software_listening_ports.yaml b/tests/assets/configs/basic_node_with_software_listening_ports.yaml index 7b04196b..6372de54 100644 --- a/tests/assets/configs/basic_node_with_software_listening_ports.yaml +++ b/tests/assets/configs/basic_node_with_software_listening_ports.yaml @@ -1,3 +1,6 @@ +metadata: + version: 3.0 + io_settings: save_step_metadata: false save_pcap_logs: true diff --git a/tests/assets/configs/basic_node_with_users.yaml b/tests/assets/configs/basic_node_with_users.yaml index 064519dd..20331ff2 100644 --- a/tests/assets/configs/basic_node_with_users.yaml +++ b/tests/assets/configs/basic_node_with_users.yaml @@ -1,3 +1,6 @@ +metadata: + version: 3.0 + io_settings: save_step_metadata: false save_pcap_logs: true diff --git a/tests/assets/configs/basic_switched_network.yaml b/tests/assets/configs/basic_switched_network.yaml index 7ffe4c08..76b4dfb9 100644 --- a/tests/assets/configs/basic_switched_network.yaml +++ b/tests/assets/configs/basic_switched_network.yaml @@ -4,6 +4,9 @@ # | client_1 |------| switch_1 |------| client_2 | # -------------- -------------- -------------- # +metadata: + version: 3.0 + io_settings: save_step_metadata: false save_pcap_logs: true diff --git a/tests/assets/configs/data_manipulation.yaml b/tests/assets/configs/data_manipulation.yaml index e9464f0f..59f97644 100644 --- a/tests/assets/configs/data_manipulation.yaml +++ b/tests/assets/configs/data_manipulation.yaml @@ -1,3 +1,6 @@ +metadata: + version: 3.0 + io_settings: save_agent_actions: true save_step_metadata: false diff --git a/tests/assets/configs/dmz_network.yaml b/tests/assets/configs/dmz_network.yaml index 2aa472d8..0accb3b2 100644 --- a/tests/assets/configs/dmz_network.yaml +++ b/tests/assets/configs/dmz_network.yaml @@ -30,6 +30,9 @@ # | external_computer |------| switch_3 |------| external_server | # ----------------------- -------------- --------------------- # +metadata: + version: 3.0 + io_settings: save_step_metadata: false save_pcap_logs: true diff --git a/tests/assets/configs/eval_only_primaite_session.yaml b/tests/assets/configs/eval_only_primaite_session.yaml index 3403a7ff..6085b1e7 100644 --- a/tests/assets/configs/eval_only_primaite_session.yaml +++ b/tests/assets/configs/eval_only_primaite_session.yaml @@ -1,3 +1,6 @@ +metadata: + version: 3.0 + game: ports: - ARP diff --git a/tests/assets/configs/extended_config.yaml b/tests/assets/configs/extended_config.yaml index 8bf7d506..a58a9d4a 100644 --- a/tests/assets/configs/extended_config.yaml +++ b/tests/assets/configs/extended_config.yaml @@ -1,3 +1,6 @@ +metadata: + version: 3.0 + io_settings: save_agent_actions: true save_step_metadata: false diff --git a/tests/assets/configs/firewall_actions_network.yaml b/tests/assets/configs/firewall_actions_network.yaml index eb39aa1a..66470f5a 100644 --- a/tests/assets/configs/firewall_actions_network.yaml +++ b/tests/assets/configs/firewall_actions_network.yaml @@ -30,6 +30,9 @@ # | external_computer |------| switch_3 |------| external_server | # ----------------------- -------------- --------------------- # +metadata: + version: 3.0 + io_settings: save_step_metadata: false save_pcap_logs: true diff --git a/tests/assets/configs/fixing_duration_one_item.yaml b/tests/assets/configs/fixing_duration_one_item.yaml index ff3a6504..02b69e5c 100644 --- a/tests/assets/configs/fixing_duration_one_item.yaml +++ b/tests/assets/configs/fixing_duration_one_item.yaml @@ -4,6 +4,9 @@ # | client_1 |------| switch_1 |------| client_2 | # -------------- -------------- -------------- # +metadata: + version: 3.0 + io_settings: save_step_metadata: false save_pcap_logs: true diff --git a/tests/assets/configs/install_and_configure_apps.yaml b/tests/assets/configs/install_and_configure_apps.yaml index 6b80519c..35546902 100644 --- a/tests/assets/configs/install_and_configure_apps.yaml +++ b/tests/assets/configs/install_and_configure_apps.yaml @@ -1,3 +1,6 @@ +metadata: + version: 3.0 + io_settings: save_step_metadata: false save_pcap_logs: false diff --git a/tests/assets/configs/multi_agent_session.yaml b/tests/assets/configs/multi_agent_session.yaml index e99ea49b..de0cdad9 100644 --- a/tests/assets/configs/multi_agent_session.yaml +++ b/tests/assets/configs/multi_agent_session.yaml @@ -1,3 +1,6 @@ +metadata: + version: 3.0 + io_settings: save_agent_actions: false save_step_metadata: false diff --git a/tests/assets/configs/nmap_network_service_recon_red_agent_config.yaml b/tests/assets/configs/nmap_network_service_recon_red_agent_config.yaml index 9b49d466..f7b8431e 100644 --- a/tests/assets/configs/nmap_network_service_recon_red_agent_config.yaml +++ b/tests/assets/configs/nmap_network_service_recon_red_agent_config.yaml @@ -1,3 +1,6 @@ +metadata: + version: 3.0 + io_settings: save_step_metadata: false save_pcap_logs: true diff --git a/tests/assets/configs/nmap_ping_scan_red_agent_config.yaml b/tests/assets/configs/nmap_ping_scan_red_agent_config.yaml index e3cb33e1..112d7266 100644 --- a/tests/assets/configs/nmap_ping_scan_red_agent_config.yaml +++ b/tests/assets/configs/nmap_ping_scan_red_agent_config.yaml @@ -1,3 +1,6 @@ +metadata: + version: 3.0 + io_settings: save_step_metadata: false save_pcap_logs: true diff --git a/tests/assets/configs/nmap_port_scan_red_agent_config.yaml b/tests/assets/configs/nmap_port_scan_red_agent_config.yaml index 2692a3bc..acd5319a 100644 --- a/tests/assets/configs/nmap_port_scan_red_agent_config.yaml +++ b/tests/assets/configs/nmap_port_scan_red_agent_config.yaml @@ -1,3 +1,6 @@ +metadata: + version: 3.0 + io_settings: save_step_metadata: false save_pcap_logs: true diff --git a/tests/assets/configs/no_nodes_links_agents_network.yaml b/tests/assets/configs/no_nodes_links_agents_network.yaml index b20835bc..ed279b51 100644 --- a/tests/assets/configs/no_nodes_links_agents_network.yaml +++ b/tests/assets/configs/no_nodes_links_agents_network.yaml @@ -1,3 +1,6 @@ +metadata: + version: 3.0 + io_settings: save_step_metadata: false save_pcap_logs: true diff --git a/tests/assets/configs/scenario_with_placeholders/scenario.yaml b/tests/assets/configs/scenario_with_placeholders/scenario.yaml index ab8c968c..57fe59ab 100644 --- a/tests/assets/configs/scenario_with_placeholders/scenario.yaml +++ b/tests/assets/configs/scenario_with_placeholders/scenario.yaml @@ -1,3 +1,6 @@ +metadata: + version: 3.0 + io_settings: save_agent_actions: true save_step_metadata: false diff --git a/tests/assets/configs/shared_rewards.yaml b/tests/assets/configs/shared_rewards.yaml index 07256f01..5aeb99fa 100644 --- a/tests/assets/configs/shared_rewards.yaml +++ b/tests/assets/configs/shared_rewards.yaml @@ -1,3 +1,6 @@ +metadata: + version: 3.0 + io_settings: save_agent_actions: false save_step_metadata: false diff --git a/tests/assets/configs/software_fixing_duration.yaml b/tests/assets/configs/software_fixing_duration.yaml index 63409f2b..66ba6f18 100644 --- a/tests/assets/configs/software_fixing_duration.yaml +++ b/tests/assets/configs/software_fixing_duration.yaml @@ -4,6 +4,9 @@ # | client_1 |------| switch_1 |------| client_2 | # -------------- -------------- -------------- # +metadata: + version: 3.0 + io_settings: save_step_metadata: false save_pcap_logs: true diff --git a/tests/assets/configs/test_application_install.yaml b/tests/assets/configs/test_application_install.yaml index 0266dd00..8a292f83 100644 --- a/tests/assets/configs/test_application_install.yaml +++ b/tests/assets/configs/test_application_install.yaml @@ -1,3 +1,6 @@ +metadata: + version: 3.0 + io_settings: save_agent_actions: true save_step_metadata: false diff --git a/tests/assets/configs/test_primaite_session.yaml b/tests/assets/configs/test_primaite_session.yaml index 65d16a47..ad43732f 100644 --- a/tests/assets/configs/test_primaite_session.yaml +++ b/tests/assets/configs/test_primaite_session.yaml @@ -1,3 +1,6 @@ +metadata: + version: 3.0 + io_settings: save_agent_actions: true save_step_metadata: true diff --git a/tests/assets/configs/wireless_wan_network_config.yaml b/tests/assets/configs/wireless_wan_network_config.yaml index 40a9126c..45721d8a 100644 --- a/tests/assets/configs/wireless_wan_network_config.yaml +++ b/tests/assets/configs/wireless_wan_network_config.yaml @@ -1,3 +1,6 @@ +metadata: + version: 3.0 + game: max_episode_length: 256 ports: diff --git a/tests/assets/configs/wireless_wan_network_config_freq_max_override.yaml b/tests/assets/configs/wireless_wan_network_config_freq_max_override.yaml index 13ceee16..20e48a89 100644 --- a/tests/assets/configs/wireless_wan_network_config_freq_max_override.yaml +++ b/tests/assets/configs/wireless_wan_network_config_freq_max_override.yaml @@ -1,3 +1,6 @@ +metadata: + version: 3.0 + game: max_episode_length: 256 ports: diff --git a/tests/assets/configs/wireless_wan_network_config_freq_max_override_blocked.yaml b/tests/assets/configs/wireless_wan_network_config_freq_max_override_blocked.yaml index 6c52fd95..6342d1b1 100644 --- a/tests/assets/configs/wireless_wan_network_config_freq_max_override_blocked.yaml +++ b/tests/assets/configs/wireless_wan_network_config_freq_max_override_blocked.yaml @@ -1,3 +1,6 @@ +metadata: + version: 3.0 + game: max_episode_length: 256 ports: diff --git a/tests/integration_tests/game_layer/test_action_shapes.py b/tests/integration_tests/game_layer/test_action_shapes.py deleted file mode 100644 index 48500d8f..00000000 --- a/tests/integration_tests/game_layer/test_action_shapes.py +++ /dev/null @@ -1,21 +0,0 @@ -# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK -from typing import Tuple - -from primaite.game.agent.interface import ProxyAgent -from primaite.game.game import PrimaiteGame -from tests import TEST_ASSETS_ROOT - -FIREWALL_ACTIONS_NETWORK = TEST_ASSETS_ROOT / "configs/firewall_actions_network.yaml" - - -def test_router_acl_add_rule_action_shape(game_and_agent: Tuple[PrimaiteGame, ProxyAgent]): - """Test to check ROUTER_ADD_ACL_RULE has the expected action shape.""" - game, agent = game_and_agent - - # assert that the shape of the actions is correct - router_acl_add_rule_action = agent.action_manager.actions.get("ROUTER_ACL_ADDRULE") - assert router_acl_add_rule_action.shape.get("source_ip_id") == len(agent.action_manager.ip_address_list) - assert router_acl_add_rule_action.shape.get("dest_ip_id") == len(agent.action_manager.ip_address_list) - assert router_acl_add_rule_action.shape.get("source_port_id") == len(agent.action_manager.ports) - assert router_acl_add_rule_action.shape.get("dest_port_id") == len(agent.action_manager.ports) - assert router_acl_add_rule_action.shape.get("protocol_id") == len(agent.action_manager.protocols) diff --git a/tests/integration_tests/game_layer/test_actions.py b/tests/integration_tests/game_layer/test_actions.py index 390a86bf..764c207e 100644 --- a/tests/integration_tests/game_layer/test_actions.py +++ b/tests/integration_tests/game_layer/test_actions.py @@ -117,7 +117,7 @@ def test_router_acl_addrule_integration(game_and_agent: Tuple[PrimaiteGame, Prox server_1 = game.simulation.network.get_node_by_hostname("server_1") server_2 = game.simulation.network.get_node_by_hostname("server_2") router = game.simulation.network.get_node_by_hostname("router") - assert router.acl.num_rules == 3 + assert router.acl.num_rules == 4 assert client_1.ping("10.0.2.3") # client_1 can ping server_2 assert server_2.ping("10.0.1.2") # server_2 can ping client_1 @@ -140,7 +140,7 @@ def test_router_acl_addrule_integration(game_and_agent: Tuple[PrimaiteGame, Prox agent.store_action(action) game.step() - # 3: Check that the acl now has 5 rules, and that client 1 cannot ping server 2 + # 3: Check that the acl now has 6 rules, and that client 1 cannot ping server 2 assert router.acl.num_rules == 5 assert not client_1.ping("10.0.2.3") # Cannot ping server_2 assert client_1.ping("10.0.2.2") # Can ping server_1 @@ -167,8 +167,8 @@ def test_router_acl_addrule_integration(game_and_agent: Tuple[PrimaiteGame, Prox agent.store_action(action) game.step() - # 5: Check that the ACL now has 5 rules, but that server_1 can still ping server_2 - assert router.acl.num_rules == 5 + # 5: Check that the ACL now has 6 rules, but that server_1 can still ping server_2 + assert router.acl.num_rules == 6 assert server_1.ping("10.0.2.3") # Can ping server_2 @@ -199,7 +199,7 @@ def test_router_acl_removerule_integration(game_and_agent: Tuple[PrimaiteGame, P game.step() # 3: Check that the ACL now has 2 rules, and that client 1 cannot access example.com - assert router.acl.num_rules == 2 + assert router.acl.num_rules == 3 assert not browser.get_webpage() client_1.software_manager.software.get("dns-client").dns_cache.clear() assert client_1.ping("10.0.2.2") # pinging still works because ICMP is allowed diff --git a/tests/unit_tests/_primaite/_game/_agent/test_agent.py b/tests/unit_tests/_primaite/_game/_agent/test_agent.py index a2693591..b555f1b2 100644 --- a/tests/unit_tests/_primaite/_game/_agent/test_agent.py +++ b/tests/unit_tests/_primaite/_game/_agent/test_agent.py @@ -4,7 +4,7 @@ from primaite.game.agent.scripted_agents.random_agent import RandomAgent def test_creating_empty_agent(): - agent = RandomAgent() + agent = RandomAgent(config={"ref": "Empty Agent"}) assert len(agent.action_manager.action_map) == 0 assert isinstance(agent.observation_manager.obs, NullObservation) assert len(agent.reward_function.reward_components) == 0 diff --git a/tests/unit_tests/_primaite/_simulator/_network/_hardware/nodes/test_switch.py b/tests/unit_tests/_primaite/_simulator/_network/_hardware/nodes/test_switch.py index e45fe45d..94b1764d 100644 --- a/tests/unit_tests/_primaite/_simulator/_network/_hardware/nodes/test_switch.py +++ b/tests/unit_tests/_primaite/_simulator/_network/_hardware/nodes/test_switch.py @@ -17,4 +17,3 @@ def switch() -> Switch: def test_describe_state(switch): state = switch.describe_state() assert len(state.get("ports")) is 8 - assert state.get("num_ports") is 8 diff --git a/tests/unit_tests/_primaite/_simulator/_system/_applications/test_applications.py b/tests/unit_tests/_primaite/_simulator/_system/_applications/test_applications.py index dd29f18e..6cccad91 100644 --- a/tests/unit_tests/_primaite/_simulator/_system/_applications/test_applications.py +++ b/tests/unit_tests/_primaite/_simulator/_system/_applications/test_applications.py @@ -18,7 +18,7 @@ def test_scan(application): def test_run_application(application): assert application.operating_state == ApplicationOperatingState.CLOSED - assert application.health_state_actual == SoftwareHealthState.UNUSED + assert application.health_state_actual == SoftwareHealthState.GOOD application.run() assert application.operating_state == ApplicationOperatingState.RUNNING @@ -37,9 +37,9 @@ def test_close_application(application): def test_application_describe_states(application): assert application.operating_state == ApplicationOperatingState.CLOSED - assert application.health_state_actual == SoftwareHealthState.UNUSED + assert application.health_state_actual == SoftwareHealthState.GOOD - assert SoftwareHealthState.UNUSED.value == application.describe_state().get("health_state_actual") + assert SoftwareHealthState.GOOD.value == application.describe_state().get("health_state_actual") application.run() assert SoftwareHealthState.GOOD.value == application.describe_state().get("health_state_actual") diff --git a/tests/unit_tests/_primaite/_simulator/_system/_services/test_services.py b/tests/unit_tests/_primaite/_simulator/_system/_services/test_services.py index 5598e1a7..fe78aa65 100644 --- a/tests/unit_tests/_primaite/_simulator/_system/_services/test_services.py +++ b/tests/unit_tests/_primaite/_simulator/_system/_services/test_services.py @@ -22,7 +22,7 @@ def test_scan(service): def test_start_service(service): assert service.operating_state == ServiceOperatingState.STOPPED - assert service.health_state_actual == SoftwareHealthState.UNUSED + assert service.health_state_actual == SoftwareHealthState.GOOD service.start() assert service.operating_state == ServiceOperatingState.RUNNING @@ -43,7 +43,7 @@ def test_pause_and_resume_service(service): assert service.operating_state == ServiceOperatingState.STOPPED service.resume() assert service.operating_state == ServiceOperatingState.STOPPED - assert service.health_state_actual == SoftwareHealthState.UNUSED + assert service.health_state_actual == SoftwareHealthState.GOOD service.start() assert service.health_state_actual == SoftwareHealthState.GOOD @@ -58,11 +58,11 @@ def test_pause_and_resume_service(service): def test_restart(service): assert service.operating_state == ServiceOperatingState.STOPPED - assert service.health_state_actual == SoftwareHealthState.UNUSED + assert service.health_state_actual == SoftwareHealthState.GOOD service.restart() # Service is STOPPED. Restart will only work if the service was PAUSED or RUNNING assert service.operating_state == ServiceOperatingState.STOPPED - assert service.health_state_actual == SoftwareHealthState.UNUSED + assert service.health_state_actual == SoftwareHealthState.GOOD service.start() assert service.operating_state == ServiceOperatingState.RUNNING @@ -157,11 +157,11 @@ def test_service_fixing(service): def test_enable_disable(service): service.disable() assert service.operating_state == ServiceOperatingState.DISABLED - assert service.health_state_actual == SoftwareHealthState.UNUSED + assert service.health_state_actual == SoftwareHealthState.GOOD service.enable() assert service.operating_state == ServiceOperatingState.STOPPED - assert service.health_state_actual == SoftwareHealthState.UNUSED + assert service.health_state_actual == SoftwareHealthState.GOOD def test_overwhelm_service(service): diff --git a/tests/unit_tests/_primaite/_simulator/_system/test_software.py b/tests/unit_tests/_primaite/_simulator/_system/test_software.py index 8c39c41d..9ad0dbcb 100644 --- a/tests/unit_tests/_primaite/_simulator/_system/test_software.py +++ b/tests/unit_tests/_primaite/_simulator/_system/test_software.py @@ -39,6 +39,6 @@ def test_software_creation(software): def test_software_set_health_state(software): - assert software.health_state_actual == SoftwareHealthState.UNUSED - software.set_health_state(SoftwareHealthState.GOOD) assert software.health_state_actual == SoftwareHealthState.GOOD + software.set_health_state(SoftwareHealthState.COMPROMISED) + assert software.health_state_actual == SoftwareHealthState.COMPROMISED