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

Make notebook for varying red agent behaviour in uc2

Browse Source
This commit is contained in:
Marek Wolan
2024-03-04 21:04:27 +00:00
parent a222a8c58f
commit 758f892b74
5 changed files with 453 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
docs/index.rst
View File

@@ -109,6 +109,7 @@ Head over to the :ref:`getting-started` page to install and setup PrimAITE!
source/game_layer
source/config
source/environment
source/customising_scenarios
.. toctree::
:caption: Developer information:

4
docs/source/customising_scenarios.rst Normal file
View File

@@ -0,0 +1,4 @@
Customising Agents
******************
For an example of how to customise red agent behaviour in the Data Manipulation scenario, please refer to the notebook ``Data-Manipulation-Customising-Red-Agent.ipynb``.

3
src/primaite/config/_package_data/example_config.yaml
View File

@@ -134,9 +134,6 @@ agents:
action_list:
- type: DONOTHING
- type: NODE_APPLICATION_EXECUTE
- type: NODE_FILE_DELETE
- type: NODE_FILE_CORRUPT
- type: NODE_OS_SCAN
options:
nodes:
- node_name: client_1

444
src/primaite/notebooks/Data-Manipulation-Customising-Red-Agent.ipynb Normal file
View File

@@ -0,0 +1,444 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Customising red agents\n",
"\n",
"This notebook will go over some examples of how red agent behaviour can be varied by changing its configuration parameters.\n",
"\n",
"First, let's load the standard Data Manipulation config file, and see what the red agent does.\n",
"\n",
"*(For a full explanation of the Data Manipulation scenario, check out the notebook `Data-Manipulation-E2E-Demonstration.ipynb`)*"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Imports\n",
"\n",
"from primaite.config.load import data_manipulation_config_path\n",
"from primaite.session.environment import PrimaiteGymEnv\n",
"import yaml\n",
"from pprint import pprint"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def make_cfg_have_flat_obs(cfg):\n",
" for agent in cfg['agents']:\n",
" if agent['type'] == \"ProxyAgent\":\n",
" agent['agent_settings']['flatten_obs'] = False"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"with open(data_manipulation_config_path(), 'r') as f:\n",
" cfg = yaml.safe_load(f)\n",
" make_cfg_have_flat_obs(cfg)\n",
"\n",
"env = PrimaiteGymEnv(game_config = cfg)\n",
"obs, info = env.reset()\n",
"print('env created successfully')"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def friendly_output_red_action(info):\n",
" # parse the info dict form step output and write out what the red agent is doing\n",
" red_info = info['agent_actions']['data_manipulation_attacker']\n",
" red_action = red_info[0]\n",
" if red_action == 'DONOTHING':\n",
" red_str = 'DO NOTHING'\n",
" elif red_action == 'NODE_APPLICATION_EXECUTE':\n",
" client = \"client 1\" if red_info[1]['node_id'] == 0 else \"client 2\"\n",
" red_str = f\"ATTACK from {client}\"\n",
" return red_str"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"By default, the red agent can start on client 1 or client 2. It starts its attack on a random step between 20 and 30, and it repeats its attack every 15-25 steps.\n",
"\n",
"It also has a 20% chance to fail to perform the port scan, and a 20% chance to fail launching the SQL attack. However it will continue where it left off after a failed step. I.e. if lucky, it can perform the port scan and SQL attack on the first try. If the port scan works, but the sql attack fails the first time it tries to attack, the next time it will not need to port scan again, it can go straight to trying to use SQL attack again."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"for step in range(35):\n",
" step_num = env.game.step_counter\n",
" obs, reward, terminated, truncated, info = env.step(0)\n",
" red = friendly_output_red_action(info)\n",
" print(f\"step: {step_num:3}, Red action: {friendly_output_red_action(info)}, Blue reward:{reward:.2f}\" )"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Since the agent does nothing most of the time, let's only print the steps where it performs an attack."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"env.reset()\n",
"for step in range(100):\n",
" step_num = env.game.step_counter\n",
" obs, reward, terminated, truncated, info = env.step(0)\n",
" red = friendly_output_red_action(info)\n",
" if red.startswith(\"ATTACK\"):\n",
" print(f\"step: {step_num:3}, Red action: {friendly_output_red_action(info)}, Blue reward:{reward:.2f}\" )"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Red Configuration\n",
"\n",
"There are two important parts of the YAML config for varying red agent behaviour.\n",
"\n",
"#### The red agent settings\n",
"Here is an annotated config for the red agent in the data manipulation scenario.\n",
"```yaml\n",
" - ref: data_manipulation_attacker # name of agent\n",
" team: RED # not used, just for human reference\n",
" type: RedDatabaseCorruptingAgent # 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: UC2RedObservation\n",
" options:\n",
" nodes: {}\n",
"\n",
" action_space:\n",
"\n",
" # The agent has two action choices, either do nothing, or execute a pre-scripted attack by using \n",
" action_list:\n",
" - type: DONOTHING\n",
" - type: NODE_APPLICATION_EXECUTE\n",
"\n",
" # The agent has access to the DataManipulationBoth on clients 1 and 2.\n",
" options:\n",
" nodes:\n",
" - node_name: client_1 # The network should have a node called client_1\n",
" applications:\n",
" - application_name: DataManipulationBot # The node client_1 should have DataManipulationBot configured on it\n",
" - node_name: client_2 # The network should have a node called client_2\n",
" applications:\n",
" - application_name: DataManipulationBot # The node client_2 should have DataManipulationBot configured on it\n",
"\n",
" # not important\n",
" max_folders_per_node: 1\n",
" max_files_per_folder: 1\n",
" max_services_per_node: 1\n",
"\n",
" # red agent does not need a reward function\n",
" reward_function:\n",
" reward_components:\n",
" - type: DUMMY\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",
"```\n",
"\n",
"#### The settings of the red agent's attack application\n",
"The red agent uses an application called `DataManipulationBot` which leverages a node's `DatabaseClient` to send a malicious SQL query to the database server. Here's an annotated example of how this is configured in the yaml *(with impertinent config items omitted)*:\n",
"```yaml\n",
"simulation:\n",
" network:\n",
" nodes:\n",
" - ref: client_1\n",
" hostname: client_1\n",
" type: computer\n",
" ip_address: 192.168.10.21\n",
" subnet_mask: 255.255.255.0\n",
" default_gateway: 192.168.10.1\n",
" \n",
" # \n",
" applications:\n",
" - ref: data_manipulation_bot\n",
" type: DataManipulationBot\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",
" - ref: client_1_database_client\n",
" type: DatabaseClient # 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",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Removing randomness from attack timing\n",
"\n",
"We can make the attacks happen at completely predictable intervals if we set the variance parameter of the red agent to 0."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"change = yaml.safe_load(\"\"\"\n",
"start_settings:\n",
" start_step: 25\n",
" frequency: 20\n",
" variance: 0\n",
"\"\"\")\n",
"\n",
"with open(data_manipulation_config_path(), 'r') as f:\n",
" cfg = yaml.safe_load(f)\n",
" for agent in cfg['agents']:\n",
" if agent['ref'] == \"data_manipulation_attacker\":\n",
" agent['agent_settings'] = change\n",
"\n",
"env = PrimaiteGymEnv(game_config = cfg)\n",
"env.reset()\n",
"for step in range(100):\n",
" step_num = env.game.step_counter\n",
" obs, reward, terminated, truncated, info = env.step(0)\n",
" red = friendly_output_red_action(info)\n",
" if red.startswith(\"ATTACK\"):\n",
" print(f\"step: {step_num:3}, Red action: {friendly_output_red_action(info)}, Blue reward:{reward:.2f}\" )"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### Making the start node always the same\n",
"\n",
"Normally, the agent randomly chooses between the nodes in its action space to send attacks from:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Open the config without changing anyhing\n",
"with open(data_manipulation_config_path(), 'r') as f:\n",
" cfg = yaml.safe_load(f)\n",
"\n",
"env = PrimaiteGymEnv(game_config = cfg)\n",
"env.reset()\n",
"for ep in range(12):\n",
" env.reset()\n",
" for step in range(31):\n",
" step_num = env.game.step_counter\n",
" obs, reward, terminated, truncated, info = env.step(0)\n",
" red = friendly_output_red_action(info)\n",
" if red.startswith(\"ATTACK\"):\n",
" print(f\"Episode: {ep:2}, step: {step_num:3}, Red action: {friendly_output_red_action(info)}\" )"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We can make the agent always start on a node of our choice letting that be the only node in the agent's action space."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"change = yaml.safe_load(\"\"\"\n",
"action_space:\n",
" action_list:\n",
" - type: DONOTHING\n",
" - type: NODE_APPLICATION_EXECUTE\n",
" options:\n",
" nodes:\n",
" - node_name: client_1\n",
" applications:\n",
" - application_name: DataManipulationBot\n",
" max_folders_per_node: 1\n",
" max_files_per_folder: 1\n",
" max_services_per_node: 1\n",
"\"\"\")\n",
"\n",
"with open(data_manipulation_config_path(), 'r') as f:\n",
" cfg = yaml.safe_load(f)\n",
" for agent in cfg['agents']:\n",
" if agent['ref'] == \"data_manipulation_attacker\":\n",
" agent.update(change)\n",
"\n",
"env = PrimaiteGymEnv(game_config = cfg)\n",
"env.reset()\n",
"for ep in range(12):\n",
" env.reset()\n",
" for step in range(31):\n",
" step_num = env.game.step_counter\n",
" obs, reward, terminated, truncated, info = env.step(0)\n",
" red = friendly_output_red_action(info)\n",
" if red.startswith(\"ATTACK\"):\n",
" print(f\"Episode: {ep:2}, step: {step_num:3}, Red action: {friendly_output_red_action(info)}\" )"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### Make the attack less likely to succeed.\n",
"\n",
"We can change the success probabilities within the data manipulation bot application. When the attack succeeds, the reward goes down.\n",
"\n",
"Setting the probabilities to 1.0 means the attack always succeeds - the reward will always drop\n",
"\n",
"Setting the probabilities to 0.0 means the attack always fails - the reward will never drop."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Make attack always succeed.\n",
"change = yaml.safe_load(\"\"\"\n",
" applications:\n",
" - ref: data_manipulation_bot\n",
" type: DataManipulationBot\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",
" - ref: client_1_web_browser\n",
" type: WebBrowser\n",
" options:\n",
" target_url: http://arcd.com/users/\n",
" - ref: client_1_database_client\n",
" type: DatabaseClient\n",
" options:\n",
" db_server_ip: 192.168.1.14\n",
"\"\"\")\n",
"\n",
"with open(data_manipulation_config_path(), 'r') as f:\n",
" cfg = yaml.safe_load(f)\n",
" cfg['simulation']['network']\n",
" for node in cfg['simulation']['network']['nodes']:\n",
" if node['ref'] in ['client_1', 'client_2']:\n",
" node['applications'] = change['applications']\n",
"\n",
"env = PrimaiteGymEnv(game_config = cfg)\n",
"env.reset()\n",
"for ep in range(5):\n",
" env.reset()\n",
" for step in range(36):\n",
" step_num = env.game.step_counter\n",
" obs, reward, terminated, truncated, info = env.step(0)\n",
" red = friendly_output_red_action(info)\n",
" if step_num == 35:\n",
" print(f\"Episode: {ep:2}, step: {step_num:3}, Reward: {reward:.2f}\" )"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Make attack always fail.\n",
"change = yaml.safe_load(\"\"\"\n",
" applications:\n",
" - ref: data_manipulation_bot\n",
" type: DataManipulationBot\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",
" - ref: client_1_web_browser\n",
" type: WebBrowser\n",
" options:\n",
" target_url: http://arcd.com/users/\n",
" - ref: client_1_database_client\n",
" type: DatabaseClient\n",
" options:\n",
" db_server_ip: 192.168.1.14\n",
"\"\"\")\n",
"\n",
"with open(data_manipulation_config_path(), 'r') as f:\n",
" cfg = yaml.safe_load(f)\n",
" cfg['simulation']['network']\n",
" for node in cfg['simulation']['network']['nodes']:\n",
" if node['ref'] in ['client_1', 'client_2']:\n",
" node['applications'] = change['applications']\n",
"\n",
"env = PrimaiteGymEnv(game_config = cfg)\n",
"env.reset()\n",
"for ep in range(5):\n",
" env.reset()\n",
" for step in range(36):\n",
" step_num = env.game.step_counter\n",
" obs, reward, terminated, truncated, info = env.step(0)\n",
" red = friendly_output_red_action(info)\n",
" if step_num == 35:\n",
" print(f\"Episode: {ep:2}, step: {step_num:3}, Reward: {reward:.2f}\" )"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "venv",
"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,
"nbformat_minor": 2
}

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

@@ -6,10 +6,13 @@ import gymnasium
from gymnasium.core import ActType, ObsType
from ray.rllib.env.multi_agent_env import MultiAgentEnv
from primaite import getLogger
from primaite.game.agent.interface import ProxyAgent
from primaite.game.game import PrimaiteGame
from primaite.simulator import SIM_OUTPUT
_LOGGER = getLogger(__name__)
class PrimaiteGymEnv(gymnasium.Env):
"""
@@ -75,7 +78,7 @@ class PrimaiteGymEnv(gymnasium.Env):
def reset(self, seed: Optional[int] = None) -> Tuple[ObsType, Dict[str, Any]]:
"""Reset the environment."""
print(
_LOGGER.info(
f"Resetting environment, episode {self.episode_counter}, "
f"avg. reward: {self.agent.reward_function.total_reward}"
)