diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 91230171..4d38f840 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -3,7 +3,11 @@ repos: rev: v4.4.0 hooks: - id: check-yaml - exclude: scenario_with_placeholders/ + exclude: | + (?x)^( + .*/scenario_with_placeholders/.*| + .*/mini_scenario_with_simulation_variation/.* + )$ - id: end-of-file-fixer - id: trailing-whitespace - id: check-added-large-files 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 new file mode 100644 index 00000000..b4457a28 --- /dev/null +++ b/src/primaite/config/_package_data/mini_scenario_with_simulation_variation/base_scenario.yaml @@ -0,0 +1,86 @@ +game: + max_episode_length: 128 + ports: [] + protocols: [] + +agents: + - ref: RL_Agent + type: ProxyAgent + observation_space: null + action_space: + action_list: + - type: DONOTHING + - type: NODE_SHUTDOWN + - type: NODE_STARTUP + - type: HOST_NIC_ENABLE + - type: HOST_NIC_DISABLE + action_map: + 0: + action: DONOTHING + options: {} + 1: + action: NODE_SHUTDOWN + options: + node_id: 0 + 2: + action: NODE_SHUTDOWN + options: + node_id: 1 + 3: + action: NODE_STARTUP + options: + node_id: 0 + 4: + action: NODE_STARTUP + options: + node_id: 1 + 5: + action: HOST_NIC_DISABLE + options: + node_id: 0 + nic_id: 0 + 6: + action: HOST_NIC_DISABLE + options: + node_id: 1 + nic_id: 0 + 7: + action: HOST_NIC_ENABLE + options: + node_id: 0 + nic_id: 0 + 8: + action: HOST_NIC_ENABLE + options: + node_id: 1 + nic_id: 0 + options: + nodes: + - node_name: client_1 + - node_name: server + max_folders_per_node: 0 + max_files_per_folder: 0 + max_services_per_node: 0 + max_nics_per_node: 1 + max_acl_rules: 0 + ip_list: + - 192.168.1.2 + - 192.168.1.3 + reward_function: + reward_components: [] + +simulation: + network: + nodes: + - hostname: server + type: server + ip_address: 192.168.1.2 + subnet_mask: 255.255.255.0 + default_gateway: 192.168.1.1 + services: *server_services + - hostname: client_1 + type: computer + ip_address: 192.168.1.3 + subnet_mask: 255.255.255.0 + default_gateway: 192.168.1.1 + applications: *client_applications diff --git a/src/primaite/config/_package_data/mini_scenario_with_simulation_variation/schedule.yaml b/src/primaite/config/_package_data/mini_scenario_with_simulation_variation/schedule.yaml new file mode 100644 index 00000000..59da7646 --- /dev/null +++ b/src/primaite/config/_package_data/mini_scenario_with_simulation_variation/schedule.yaml @@ -0,0 +1,6 @@ +base_scenario: base_scenario.yaml +schedule: + 0: + - simulation_variant_1.yaml + 1: + - simulation_variant_2.yaml diff --git a/src/primaite/config/_package_data/mini_scenario_with_simulation_variation/simulation_variant_1.yaml b/src/primaite/config/_package_data/mini_scenario_with_simulation_variation/simulation_variant_1.yaml new file mode 100644 index 00000000..3e27cc27 --- /dev/null +++ b/src/primaite/config/_package_data/mini_scenario_with_simulation_variation/simulation_variant_1.yaml @@ -0,0 +1,5 @@ +server_services: &server_services + - type: DatabaseService + +client_applications: &client_applications + - type: DatabaseClient diff --git a/src/primaite/config/_package_data/mini_scenario_with_simulation_variation/simulation_variant_2.yaml b/src/primaite/config/_package_data/mini_scenario_with_simulation_variation/simulation_variant_2.yaml new file mode 100644 index 00000000..207e0c73 --- /dev/null +++ b/src/primaite/config/_package_data/mini_scenario_with_simulation_variation/simulation_variant_2.yaml @@ -0,0 +1,5 @@ +server_services: &server_services + - type: FTPServer + +client_applications: &client_applications + - type: RansomwareScript diff --git a/src/primaite/notebooks/Using-Episode-Schedules.ipynb b/src/primaite/notebooks/Using-Episode-Schedules.ipynb index b0669472..b6ca16ce 100644 --- a/src/primaite/notebooks/Using-Episode-Schedules.ipynb +++ b/src/primaite/notebooks/Using-Episode-Schedules.ipynb @@ -352,6 +352,101 @@ "source": [ "env.reset(); # semicolon suppresses jupyter outputting the observation space.\n" ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Other uses\n", + "\n", + "Since the episode schedules make use of yaml aliases and anchors, it's possible to use them in any part of the config, not just agent definitions. For instance, we can vary the simulation setup by changing what software is installed on hosts, how that software is configured, or even change the nodes themselves.\n", + "\n", + "As an example, we will vary what software is installed on nodes in a basic test network." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "mini_scenario_path = PRIMAITE_PATHS.user_config_path / \"example_config/mini_scenario_with_simulation_variation\"\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Let's open the base scenario to see the placeholders. client_1 and server both have placeholders in the software installed on them. The server has a placeholder called `*server_services` and the client has `*client_applications`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "with open(mini_scenario_path/\"base_scenario.yaml\") as f:\n", + " print(f.read())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "In the 0th episode, simulation_variant_1.yaml is loaded in and the server gets a `DatabaseService`, while client_1 gets `DatabaseClient`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "with open(mini_scenario_path/\"simulation_variant_1.yaml\") as f:\n", + " print(f.read())" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "env = PrimaiteGymEnv(env_config=mini_scenario_path)\n", + "print(f\"Episode: {env.episode_counter}\")\n", + "env.game.simulation.network.get_node_by_hostname('server').software_manager.show()\n", + "env.game.simulation.network.get_node_by_hostname('client_1').software_manager.show()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "In the 1st episode, `simulation_variant_2.yaml` is loaded in, therefore the server gets a `FTPServer` and client_1 gets a `RansomwareScript`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "with open(mini_scenario_path/\"simulation_variant_2.yaml\") as f:\n", + " print(f.read())" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "env.reset()\n", + "print(f\"Episode: {env.episode_counter}\")\n", + "env.game.simulation.network.get_node_by_hostname('server').software_manager.show()\n", + "env.game.simulation.network.get_node_by_hostname('client_1').software_manager.show()" + ] } ], "metadata": {