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

Add second green agent and make rewards based on webbrowser

Browse Source
This commit is contained in:
Marek Wolan
2024-02-06 14:42:59 +00:00
parent cfc254753f
commit 5b5b750d4d
7 changed files with 202 additions and 464 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
CHANGELOG.md
View File

@@ -6,6 +6,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
## [Unreleased]
- Changed the data manipulation scenario to include a second green agent on client 1.
- Refactored actions and observations to be configurable via object name, instead of UUID.
- Fixed a bug where ACL rules were not resetting on episode reset.
- Fixed a bug where blue agent's ACL actions were being applied against the wrong IP addresses

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

@@ -60,6 +60,31 @@ agents:
frequency: 4
variance: 3
- ref: client_1_green_user
team: GREEN
type: GreenWebBrowsingAgent
observation_space:
type: UC2GreenObservation
action_space:
action_list:
- type: DONOTHING
- type: NODE_APPLICATION_EXECUTE
options:
nodes:
- node_name: client_1
applications:
- application_name: WebBrowser
max_folders_per_node: 1
max_files_per_folder: 1
max_services_per_node: 1
max_applications_per_node: 1
reward_function:
reward_components:
- type: DUMMY
- ref: client_1_data_manipulation_red_bot
team: RED
type: RedDatabaseCorruptingAgent
@@ -112,7 +137,7 @@ agents:
- service_name: DNSServer
- node_hostname: web_server
services:
- service_name: web_server_web_service
- service_name: WebServer
- node_hostname: database_server
folders:
- folder_name: database
@@ -241,25 +266,25 @@ agents:
action: "NODE_FILE_SCAN"
options:
node_id: 2
folder_id: 1
folder_id: 0
file_id: 0
10:
action: "NODE_FILE_CHECKHASH"
options:
node_id: 2
folder_id: 1
folder_id: 0
file_id: 0
11:
action: "NODE_FILE_DELETE"
options:
node_id: 2
folder_id: 1
folder_id: 0
file_id: 0
12:
action: "NODE_FILE_REPAIR"
options:
node_id: 2
folder_id: 1
folder_id: 0
file_id: 0
13:
action: "NODE_SERVICE_PATCH"
@@ -270,22 +295,22 @@ agents:
action: "NODE_FOLDER_SCAN"
options:
node_id: 2
folder_id: 1
folder_id: 0
15:
action: "NODE_FOLDER_CHECKHASH"
options:
node_id: 2
folder_id: 1
folder_id: 0
16:
action: "NODE_FOLDER_REPAIR"
options:
node_id: 2
folder_id: 1
folder_id: 0
17:
action: "NODE_FOLDER_RESTORE"
options:
node_id: 2
folder_id: 1
folder_id: 0
18:
action: "NODE_OS_SCAN"
options:
@@ -302,7 +327,7 @@ agents:
action: "NODE_RESET"
options:
node_id: 5
22: # "ACL: ADDRULE - Block outgoing traffic from client 1" (not supported in Primaite)
22: # "ACL: ADDRULE - Block outgoing traffic from client 1"
action: "NETWORK_ACL_ADDRULE"
options:
position: 1
@@ -312,7 +337,7 @@ agents:
source_port_id: 1
dest_port_id: 1
protocol_id: 1
23: # "ACL: ADDRULE - Block outgoing traffic from client 2" (not supported in Primaite)
23: # "ACL: ADDRULE - Block outgoing traffic from client 2"
action: "NETWORK_ACL_ADDRULE"
options:
position: 2
@@ -497,6 +522,8 @@ agents:
- folder_name: database
files:
- file_name: database.db
services:
- service_name: DatabaseService
- node_name: backup_server
- node_name: security_suite
- node_name: client_1
@@ -529,18 +556,19 @@ agents:
reward_function:
reward_components:
- type: DATABASE_FILE_INTEGRITY
weight: 0.5
weight: 0.34
options:
node_hostname: database_server
folder_name: database
file_name: database.db
- type: WEB_SERVER_404_PENALTY
weight: 0.5
- type: WEBPAGE_UNAVAILABLE_PENALTY
weight: 0.33
options:
node_hostname: web_server
service_name: WebServer
node_hostname: client_1
- type: WEBPAGE_UNAVAILABLE_PENALTY
weight: 0.33
options:
node_hostname: client_2
agent_settings:
@@ -682,6 +710,10 @@ simulation:
data_manipulation_p_of_success: 0.8
payload: "DELETE"
server_ip: 192.168.1.14
- ref: client_1_web_browser
type: WebBrowser
options:
target_url: http://arcd.com/users/
services:
- ref: client_1_dns_client
type: DNSClient

72
src/primaite/game/agent/rewards.py
View File

@@ -26,16 +26,13 @@ the structure:
```
"""
from abc import abstractmethod
from typing import Dict, List, Tuple, Type, TYPE_CHECKING
from typing import Dict, List, Tuple, Type
from primaite import getLogger
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
_LOGGER = getLogger(__name__)
if TYPE_CHECKING:
from primaite.game.game import PrimaiteGame
class AbstractReward:
"""Base class for reward function components."""
@@ -47,13 +44,11 @@ class AbstractReward:
@classmethod
@abstractmethod
def from_config(cls, config: dict, game: "PrimaiteGame") -> "AbstractReward":
def from_config(cls, config: dict) -> "AbstractReward":
"""Create a reward function component from a config dictionary.
:param config: dict of options for the reward component's constructor
:type config: dict
:param game: Reference to the PrimAITE Game object
:type game: PrimaiteGame
:return: The reward component.
:rtype: AbstractReward
"""
@@ -68,13 +63,11 @@ class DummyReward(AbstractReward):
return 0.0
@classmethod
def from_config(cls, config: dict, game: "PrimaiteGame") -> "DummyReward":
def from_config(cls, config: dict) -> "DummyReward":
"""Create a reward function component from a config dictionary.
:param config: dict of options for the reward component's constructor. Should be empty.
:type config: dict
:param game: Reference to the PrimAITE Game object
:type game: PrimaiteGame
"""
return cls()
@@ -126,13 +119,11 @@ class DatabaseFileIntegrity(AbstractReward):
return 0
@classmethod
def from_config(cls, config: Dict, game: "PrimaiteGame") -> "DatabaseFileIntegrity":
def from_config(cls, config: Dict) -> "DatabaseFileIntegrity":
"""Create a reward function component from a config dictionary.
:param config: dict of options for the reward component's constructor
:type config: Dict
:param game: Reference to the PrimAITE Game object
:type game: PrimaiteGame
:return: The reward component.
:rtype: DatabaseFileIntegrity
"""
@@ -179,13 +170,11 @@ class WebServer404Penalty(AbstractReward):
return 0.0
@classmethod
def from_config(cls, config: Dict, game: "PrimaiteGame") -> "WebServer404Penalty":
def from_config(cls, config: Dict) -> "WebServer404Penalty":
"""Create a reward function component from a config dictionary.
:param config: dict of options for the reward component's constructor
:type config: Dict
:param game: Reference to the PrimAITE Game object
:type game: PrimaiteGame
:return: The reward component.
:rtype: WebServer404Penalty
"""
@@ -202,6 +191,50 @@ class WebServer404Penalty(AbstractReward):
return cls(node_hostname=node_hostname, service_name=service_name)
class WebpageUnavailablePenalty(AbstractReward):
"""Penalises the agent when the web browser fails to fetch a webpage."""
def __init__(self, node_hostname: str) -> None:
"""
Initialise the reward component.
:param node_hostname: Hostname of the node which has the web browser.
:type node_hostname: str
"""
self._node = node_hostname
self.location_in_state = ["network", "nodes", node_hostname, "applications", "WebBrowser"]
def calculate(self, state: Dict) -> float:
"""
Calculate the reward based on current simulation state.
:param state: The current state of the simulation.
:type state: Dict
"""
web_browser_state = access_from_nested_dict(state, self.location_in_state)
if web_browser_state is NOT_PRESENT_IN_STATE or "history" not in web_browser_state:
_LOGGER.info(
"Web browser reward could not be calculated because the web browser history on node",
f"{self._node} was not reported in the simulation state. Returning 0.0",
)
return 0.0 # 0 if the web browser cannot be found
if not web_browser_state["history"]:
return 0.0 # 0 if no requests have been attempted yet
outcome = web_browser_state["history"][-1]["outcome"]
if outcome == "PENDING":
return 0.0 # 0 if a request was attempted but not yet resolved
elif outcome == 200:
return 1.0 # 1 for successful request
else: # includes failure codes and SERVER_UNREACHABLE
return -1.0 # -1 for failure
@classmethod
def from_config(cls, config: dict) -> AbstractReward:
"""Build the reward component object from config."""
node_hostname = config.get("node_hostname")
return cls(node_hostname=node_hostname)
class RewardFunction:
"""Manages the reward function for the agent."""
@@ -209,6 +242,7 @@ class RewardFunction:
"DUMMY": DummyReward,
"DATABASE_FILE_INTEGRITY": DatabaseFileIntegrity,
"WEB_SERVER_404_PENALTY": WebServer404Penalty,
"WEBPAGE_UNAVAILABLE_PENALTY": WebpageUnavailablePenalty,
}
def __init__(self):
@@ -243,13 +277,11 @@ class RewardFunction:
return self.current_reward
@classmethod
def from_config(cls, config: Dict, game: "PrimaiteGame") -> "RewardFunction":
def from_config(cls, config: Dict) -> "RewardFunction":
"""Create a reward function from a config dictionary.
:param config: dict of options for the reward manager's constructor
:type config: Dict
:param game: Reference to the PrimAITE Game object
:type game: PrimaiteGame
:return: The reward manager.
:rtype: RewardFunction
"""
@@ -259,6 +291,6 @@ class RewardFunction:
rew_type = rew_component_cfg["type"]
weight = rew_component_cfg.get("weight", 1.0)
rew_class = cls.__rew_class_identifiers[rew_type]
rew_instance = rew_class.from_config(config=rew_component_cfg.get("options", {}), game=game)
rew_instance = rew_class.from_config(config=rew_component_cfg.get("options", {}))
new.register_component(component=rew_instance, weight=weight)
return new

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

@@ -345,7 +345,7 @@ class PrimaiteGame:
action_space = ActionManager.from_config(game, action_space_cfg)
# CREATE REWARD FUNCTION
rew_function = RewardFunction.from_config(reward_function_cfg, game=game)
rew_function = RewardFunction.from_config(reward_function_cfg)
# OTHER AGENT SETTINGS
agent_settings = AgentSettings.from_config(agent_cfg.get("agent_settings"))

449
src/primaite/notebooks/uc2_demo.ipynb
View File

@@ -46,7 +46,7 @@
"source": [
"## Green agent\n",
"\n",
"The green agent is logged onto client 2. It sometimes uses the web browser on client 2 to navigate to `http://arcd.com/users`. The web server replies with a status code 200 if the data is available on the database or 404 if not available."
"There are green agents is logged onto client 1 and client 2. They use the web browser to navigate to `http://arcd.com/users`. The web server replies with a status code 200 if the data is available on the database or 404 if not available."
]
},
{
@@ -68,7 +68,7 @@
"source": [
"## Blue agent\n",
"\n",
"The blue agent can view the entire network, but the health statuses of components are not updated until a scan is performed. The blue agent should restore the database file from backup after it was compromised. It can also prevent further attacks by blocking client 1 from reaching the database server. This can be done by removing client 1's network connection or adding ACL rules on the router to stop the packets from arriving."
"The blue agent can view the entire network, but the health statuses of components are not updated until a scan is performed. The blue agent should restore the database file from backup after it was compromised. It can also prevent further attacks by blocking client 1 from sending the malicious SQL query to the database server. This can be done by removing implementing an ACL rule on the router."
]
},
{
@@ -100,9 +100,9 @@
"The red agent does not use information about the state of the network to decide its action.\n",
"\n",
"### Green\n",
"The green agent sits on client 2 and uses the web browser application to send requests to the web server. The schedule of the green agent is currently random, meaning it will request webpage with a 50% probability, and do nothing with a 50% probability.\n",
"The green agents use the web browser application to send requests to the web server. The schedule of each green agent is currently random, meaning it will request webpage with a 50% probability, and do nothing with a 50% probability.\n",
"\n",
"When the green agent is blocked from accessing the data through the webpage, this incurs a negative reward to the RL defender."
"When a green agent is blocked from accessing the data through the webpage, this incurs a negative reward to the RL defender."
]
},
{
@@ -295,7 +295,7 @@
"- `28-37`: Remove ACL rules 1-10\n",
"- `42`: Disconnect client 1 from the network\n",
"\n",
"The other actions will either have no effect or will negatively impact the network, so the blue agent should avoid taking other actions, and learn about these actions."
"The other actions will either have no effect or will negatively impact the network, so the blue agent should avoid taking them."
]
},
{
@@ -306,8 +306,8 @@
"\n",
"The blue agent's reward is calculated using two measures:\n",
"1. Whether the database file is in a good state (+1 for good, -1 for corrupted, 0 for any other state)\n",
"2. Whether the green agent's most recent webpage request was successful (+1 for a `200` return code, -1 for a `404` return code and 0 otherwise).\n",
"These two components are averaged to get the final reward.\n"
"2. Whether each green agents' most recent webpage request was successful (+1 for a `200` return code, -1 for a `404` return code and 0 otherwise).\n",
"The file status reward and the two green-agent-related reward are averaged to get a total step reward.\n"
]
},
{
@@ -326,7 +326,7 @@
},
{
"cell_type": "code",
"execution_count": 1,
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
@@ -336,20 +336,9 @@
},
{
"cell_type": "code",
"execution_count": 2,
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"/home/cade/repos/PrimAITE/venv/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
" from .autonotebook import tqdm as notebook_tqdm\n",
"2024-01-25 14:43:32,056\tINFO util.py:159 -- Missing packages: ['ipywidgets']. Run `pip install -U ipywidgets`, then restart the notebook server for rich notebook output.\n",
"2024-01-25 14:43:35,213\tINFO util.py:159 -- Missing packages: ['ipywidgets']. Run `pip install -U ipywidgets`, then restart the notebook server for rich notebook output.\n"
]
}
],
"outputs": [],
"source": [
"# Imports\n",
"from primaite.config.load import example_config_path\n",
@@ -370,134 +359,9 @@
},
{
"cell_type": "code",
"execution_count": 3,
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Resetting environment, episode 0, avg. reward: 0.0\n",
"env created successfully\n",
"{'ACL': {1: {'dest_node_id': 0,\n",
" 'dest_port': 0,\n",
" 'permission': 0,\n",
" 'position': 0,\n",
" 'protocol': 0,\n",
" 'source_node_id': 0,\n",
" 'source_port': 0},\n",
" 2: {'dest_node_id': 0,\n",
" 'dest_port': 0,\n",
" 'permission': 0,\n",
" 'position': 1,\n",
" 'protocol': 0,\n",
" 'source_node_id': 0,\n",
" 'source_port': 0},\n",
" 3: {'dest_node_id': 0,\n",
" 'dest_port': 0,\n",
" 'permission': 0,\n",
" 'position': 2,\n",
" 'protocol': 0,\n",
" 'source_node_id': 0,\n",
" 'source_port': 0},\n",
" 4: {'dest_node_id': 0,\n",
" 'dest_port': 0,\n",
" 'permission': 0,\n",
" 'position': 3,\n",
" 'protocol': 0,\n",
" 'source_node_id': 0,\n",
" 'source_port': 0},\n",
" 5: {'dest_node_id': 0,\n",
" 'dest_port': 0,\n",
" 'permission': 0,\n",
" 'position': 4,\n",
" 'protocol': 0,\n",
" 'source_node_id': 0,\n",
" 'source_port': 0},\n",
" 6: {'dest_node_id': 0,\n",
" 'dest_port': 0,\n",
" 'permission': 0,\n",
" 'position': 5,\n",
" 'protocol': 0,\n",
" 'source_node_id': 0,\n",
" 'source_port': 0},\n",
" 7: {'dest_node_id': 0,\n",
" 'dest_port': 0,\n",
" 'permission': 0,\n",
" 'position': 6,\n",
" 'protocol': 0,\n",
" 'source_node_id': 0,\n",
" 'source_port': 0},\n",
" 8: {'dest_node_id': 0,\n",
" 'dest_port': 0,\n",
" 'permission': 0,\n",
" 'position': 7,\n",
" 'protocol': 0,\n",
" 'source_node_id': 0,\n",
" 'source_port': 0},\n",
" 9: {'dest_node_id': 0,\n",
" 'dest_port': 0,\n",
" 'permission': 0,\n",
" 'position': 8,\n",
" 'protocol': 0,\n",
" 'source_node_id': 0,\n",
" 'source_port': 0},\n",
" 10: {'dest_node_id': 0,\n",
" 'dest_port': 0,\n",
" 'permission': 0,\n",
" 'position': 9,\n",
" 'protocol': 0,\n",
" 'source_node_id': 0,\n",
" 'source_port': 0}},\n",
" 'ICS': 0,\n",
" 'LINKS': {1: {'PROTOCOLS': {'ALL': 1}},\n",
" 2: {'PROTOCOLS': {'ALL': 1}},\n",
" 3: {'PROTOCOLS': {'ALL': 1}},\n",
" 4: {'PROTOCOLS': {'ALL': 1}},\n",
" 5: {'PROTOCOLS': {'ALL': 1}},\n",
" 6: {'PROTOCOLS': {'ALL': 1}},\n",
" 7: {'PROTOCOLS': {'ALL': 1}},\n",
" 8: {'PROTOCOLS': {'ALL': 1}},\n",
" 9: {'PROTOCOLS': {'ALL': 1}},\n",
" 10: {'PROTOCOLS': {'ALL': 1}}},\n",
" 'NODES': {1: {'FOLDERS': {1: {'FILES': {1: {'health_status': 0}},\n",
" 'health_status': 0}},\n",
" 'NICS': {1: {'nic_status': 1}, 2: {'nic_status': 0}},\n",
" 'SERVICES': {1: {'health_status': 0, 'operating_status': 1}},\n",
" 'operating_status': 1},\n",
" 2: {'FOLDERS': {1: {'FILES': {1: {'health_status': 0}},\n",
" 'health_status': 0}},\n",
" 'NICS': {1: {'nic_status': 1}, 2: {'nic_status': 0}},\n",
" 'SERVICES': {1: {'health_status': 0, 'operating_status': 1}},\n",
" 'operating_status': 1},\n",
" 3: {'FOLDERS': {1: {'FILES': {1: {'health_status': 1}},\n",
" 'health_status': 1}},\n",
" 'NICS': {1: {'nic_status': 1}, 2: {'nic_status': 0}},\n",
" 'SERVICES': {1: {'health_status': 0, 'operating_status': 0}},\n",
" 'operating_status': 1},\n",
" 4: {'FOLDERS': {1: {'FILES': {1: {'health_status': 0}},\n",
" 'health_status': 0}},\n",
" 'NICS': {1: {'nic_status': 1}, 2: {'nic_status': 0}},\n",
" 'SERVICES': {1: {'health_status': 0, 'operating_status': 0}},\n",
" 'operating_status': 1},\n",
" 5: {'FOLDERS': {1: {'FILES': {1: {'health_status': 0}},\n",
" 'health_status': 0}},\n",
" 'NICS': {1: {'nic_status': 1}, 2: {'nic_status': 0}},\n",
" 'SERVICES': {1: {'health_status': 0, 'operating_status': 0}},\n",
" 'operating_status': 1},\n",
" 6: {'FOLDERS': {1: {'FILES': {1: {'health_status': 0}},\n",
" 'health_status': 0}},\n",
" 'NICS': {1: {'nic_status': 1}, 2: {'nic_status': 0}},\n",
" 'SERVICES': {1: {'health_status': 0, 'operating_status': 0}},\n",
" 'operating_status': 1},\n",
" 7: {'FOLDERS': {1: {'FILES': {1: {'health_status': 0}},\n",
" 'health_status': 0}},\n",
" 'NICS': {1: {'nic_status': 1}, 2: {'nic_status': 0}},\n",
" 'SERVICES': {1: {'health_status': 0, 'operating_status': 0}},\n",
" 'operating_status': 1}}}\n"
]
}
],
"outputs": [],
"source": [
"# create the env\n",
"with open(example_config_path(), 'r') as f:\n",
@@ -523,48 +387,9 @@
},
{
"cell_type": "code",
"execution_count": 4,
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"step: 1, Red action: DONOTHING, Blue reward:0.5\n",
"step: 2, Red action: DONOTHING, Blue reward:0.5\n",
"step: 3, Red action: DONOTHING, Blue reward:0.5\n",
"step: 4, Red action: DONOTHING, Blue reward:0.5\n",
"step: 5, Red action: DONOTHING, Blue reward:1.0\n",
"step: 6, Red action: DONOTHING, Blue reward:1.0\n",
"step: 7, Red action: DONOTHING, Blue reward:1.0\n",
"step: 8, Red action: DONOTHING, Blue reward:1.0\n",
"step: 9, Red action: DONOTHING, Blue reward:1.0\n",
"step: 10, Red action: DONOTHING, Blue reward:1.0\n",
"step: 11, Red action: DONOTHING, Blue reward:1.0\n",
"step: 12, Red action: DONOTHING, Blue reward:1.0\n",
"step: 13, Red action: DONOTHING, Blue reward:1.0\n",
"step: 14, Red action: DONOTHING, Blue reward:1.0\n",
"step: 15, Red action: DONOTHING, Blue reward:1.0\n",
"step: 16, Red action: DONOTHING, Blue reward:1.0\n",
"step: 17, Red action: DONOTHING, Blue reward:1.0\n",
"step: 18, Red action: DONOTHING, Blue reward:1.0\n",
"step: 19, Red action: DONOTHING, Blue reward:1.0\n",
"step: 20, Red action: DONOTHING, Blue reward:1.0\n",
"step: 21, Red action: DONOTHING, Blue reward:1.0\n",
"step: 22, Red action: NODE_APPLICATION_EXECUTE, Blue reward:0.0\n",
"step: 23, Red action: DONOTHING, Blue reward:0.0\n",
"step: 24, Red action: DONOTHING, Blue reward:0.0\n",
"step: 25, Red action: DONOTHING, Blue reward:0.0\n",
"step: 26, Red action: DONOTHING, Blue reward:-1.0\n",
"step: 27, Red action: DONOTHING, Blue reward:-1.0\n",
"step: 28, Red action: DONOTHING, Blue reward:-1.0\n",
"step: 29, Red action: DONOTHING, Blue reward:-1.0\n",
"step: 30, Red action: DONOTHING, Blue reward:-1.0\n",
"step: 31, Red action: DONOTHING, Blue reward:-1.0\n",
"step: 32, Red action: DONOTHING, Blue reward:-1.0\n"
]
}
],
"outputs": [],
"source": [
"for step in range(32):\n",
" obs, reward, terminated, truncated, info = env.step(0)\n",
@@ -580,44 +405,9 @@
},
{
"cell_type": "code",
"execution_count": 5,
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"{1: {'FOLDERS': {1: {'FILES': {1: {'health_status': 0}}, 'health_status': 0}},\n",
" 'NICS': {1: {'nic_status': 1}, 2: {'nic_status': 0}},\n",
" 'SERVICES': {1: {'health_status': 0, 'operating_status': 1}},\n",
" 'operating_status': 1},\n",
" 2: {'FOLDERS': {1: {'FILES': {1: {'health_status': 0}}, 'health_status': 0}},\n",
" 'NICS': {1: {'nic_status': 1}, 2: {'nic_status': 0}},\n",
" 'SERVICES': {1: {'health_status': 0, 'operating_status': 1}},\n",
" 'operating_status': 1},\n",
" 3: {'FOLDERS': {1: {'FILES': {1: {'health_status': 1}}, 'health_status': 1}},\n",
" 'NICS': {1: {'nic_status': 1}, 2: {'nic_status': 0}},\n",
" 'SERVICES': {1: {'health_status': 0, 'operating_status': 0}},\n",
" 'operating_status': 1},\n",
" 4: {'FOLDERS': {1: {'FILES': {1: {'health_status': 0}}, 'health_status': 0}},\n",
" 'NICS': {1: {'nic_status': 1}, 2: {'nic_status': 0}},\n",
" 'SERVICES': {1: {'health_status': 0, 'operating_status': 0}},\n",
" 'operating_status': 1},\n",
" 5: {'FOLDERS': {1: {'FILES': {1: {'health_status': 0}}, 'health_status': 0}},\n",
" 'NICS': {1: {'nic_status': 1}, 2: {'nic_status': 0}},\n",
" 'SERVICES': {1: {'health_status': 0, 'operating_status': 0}},\n",
" 'operating_status': 1},\n",
" 6: {'FOLDERS': {1: {'FILES': {1: {'health_status': 0}}, 'health_status': 0}},\n",
" 'NICS': {1: {'nic_status': 1}, 2: {'nic_status': 0}},\n",
" 'SERVICES': {1: {'health_status': 0, 'operating_status': 0}},\n",
" 'operating_status': 1},\n",
" 7: {'FOLDERS': {1: {'FILES': {1: {'health_status': 0}}, 'health_status': 0}},\n",
" 'NICS': {1: {'nic_status': 1}, 2: {'nic_status': 0}},\n",
" 'SERVICES': {1: {'health_status': 0, 'operating_status': 0}},\n",
" 'operating_status': 1}}\n"
]
}
],
"outputs": [],
"source": [
"pprint(obs['NODES'])"
]
@@ -631,44 +421,9 @@
},
{
"cell_type": "code",
"execution_count": 6,
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"{1: {'FOLDERS': {1: {'FILES': {1: {'health_status': 0}}, 'health_status': 0}},\n",
" 'NICS': {1: {'nic_status': 1}, 2: {'nic_status': 0}},\n",
" 'SERVICES': {1: {'health_status': 0, 'operating_status': 1}},\n",
" 'operating_status': 1},\n",
" 2: {'FOLDERS': {1: {'FILES': {1: {'health_status': 0}}, 'health_status': 0}},\n",
" 'NICS': {1: {'nic_status': 1}, 2: {'nic_status': 0}},\n",
" 'SERVICES': {1: {'health_status': 3, 'operating_status': 1}},\n",
" 'operating_status': 1},\n",
" 3: {'FOLDERS': {1: {'FILES': {1: {'health_status': 2}}, 'health_status': 1}},\n",
" 'NICS': {1: {'nic_status': 1}, 2: {'nic_status': 0}},\n",
" 'SERVICES': {1: {'health_status': 0, 'operating_status': 0}},\n",
" 'operating_status': 1},\n",
" 4: {'FOLDERS': {1: {'FILES': {1: {'health_status': 0}}, 'health_status': 0}},\n",
" 'NICS': {1: {'nic_status': 1}, 2: {'nic_status': 0}},\n",
" 'SERVICES': {1: {'health_status': 0, 'operating_status': 0}},\n",
" 'operating_status': 1},\n",
" 5: {'FOLDERS': {1: {'FILES': {1: {'health_status': 0}}, 'health_status': 0}},\n",
" 'NICS': {1: {'nic_status': 1}, 2: {'nic_status': 0}},\n",
" 'SERVICES': {1: {'health_status': 0, 'operating_status': 0}},\n",
" 'operating_status': 1},\n",
" 6: {'FOLDERS': {1: {'FILES': {1: {'health_status': 0}}, 'health_status': 0}},\n",
" 'NICS': {1: {'nic_status': 1}, 2: {'nic_status': 0}},\n",
" 'SERVICES': {1: {'health_status': 0, 'operating_status': 0}},\n",
" 'operating_status': 1},\n",
" 7: {'FOLDERS': {1: {'FILES': {1: {'health_status': 0}}, 'health_status': 0}},\n",
" 'NICS': {1: {'nic_status': 1}, 2: {'nic_status': 0}},\n",
" 'SERVICES': {1: {'health_status': 0, 'operating_status': 0}},\n",
" 'operating_status': 1}}\n"
]
}
],
"outputs": [],
"source": [
"obs, reward, terminated, truncated, info = env.step(9) # scan database file\n",
"obs, reward, terminated, truncated, info = env.step(1) # scan webapp service\n",
@@ -692,24 +447,14 @@
},
{
"cell_type": "code",
"execution_count": 7,
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"step: 35\n",
"Red action: DONOTHING\n",
"Green action: NODE_APPLICATION_EXECUTE\n",
"Blue reward:-1.0\n"
]
}
],
"outputs": [],
"source": [
"obs, reward, terminated, truncated, info = env.step(13) # patch the database\n",
"print(f\"step: {env.game.step_counter}\")\n",
"print(f\"Red action: {info['agent_actions']['client_1_data_manipulation_red_bot'][0]}\" )\n",
"print(f\"Green action: {info['agent_actions']['client_1_green_user'][0]}\" )\n",
"print(f\"Green action: {info['agent_actions']['client_2_green_user'][0]}\" )\n",
"print(f\"Blue reward:{reward}\" )"
]
@@ -727,25 +472,15 @@
},
{
"cell_type": "code",
"execution_count": 8,
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"step: 36\n",
"Red action: DONOTHING\n",
"Green action: NODE_APPLICATION_EXECUTE\n",
"Blue reward:0.0\n"
]
}
],
"outputs": [],
"source": [
"obs, reward, terminated, truncated, info = env.step(0) # patch the database\n",
"print(f\"step: {env.game.step_counter}\")\n",
"print(f\"Red action: {info['agent_actions']['client_1_data_manipulation_red_bot'][0]}\" )\n",
"print(f\"Green action: {info['agent_actions']['client_2_green_user'][0]}\" )\n",
"print(f\"Green action: {info['agent_actions']['client_1_green_user'][0]}\" )\n",
"print(f\"Blue reward:{reward}\" )"
]
},
@@ -758,48 +493,9 @@
},
{
"cell_type": "code",
"execution_count": 9,
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"step: 37, Red action: DONOTHING, Blue reward:0.0\n",
"step: 38, Red action: DONOTHING, Blue reward:0.0\n",
"step: 39, Red action: DONOTHING, Blue reward:1.0\n",
"step: 40, Red action: DONOTHING, Blue reward:1.0\n",
"step: 41, Red action: DONOTHING, Blue reward:1.0\n",
"step: 42, Red action: DONOTHING, Blue reward:1.0\n",
"step: 43, Red action: DONOTHING, Blue reward:1.0\n",
"step: 44, Red action: DONOTHING, Blue reward:1.0\n",
"step: 45, Red action: DONOTHING, Blue reward:1.0\n",
"step: 46, Red action: NODE_APPLICATION_EXECUTE, Blue reward:1.0\n",
"step: 47, Red action: DONOTHING, Blue reward:1.0\n",
"step: 48, Red action: DONOTHING, Blue reward:1.0\n",
"step: 49, Red action: DONOTHING, Blue reward:1.0\n",
"step: 50, Red action: DONOTHING, Blue reward:1.0\n",
"step: 51, Red action: DONOTHING, Blue reward:1.0\n",
"step: 52, Red action: DONOTHING, Blue reward:1.0\n",
"step: 53, Red action: DONOTHING, Blue reward:1.0\n",
"step: 54, Red action: DONOTHING, Blue reward:1.0\n",
"step: 55, Red action: DONOTHING, Blue reward:1.0\n",
"step: 56, Red action: DONOTHING, Blue reward:1.0\n",
"step: 57, Red action: DONOTHING, Blue reward:1.0\n",
"step: 58, Red action: DONOTHING, Blue reward:1.0\n",
"step: 59, Red action: DONOTHING, Blue reward:1.0\n",
"step: 60, Red action: DONOTHING, Blue reward:1.0\n",
"step: 61, Red action: DONOTHING, Blue reward:1.0\n",
"step: 62, Red action: DONOTHING, Blue reward:1.0\n",
"step: 63, Red action: DONOTHING, Blue reward:1.0\n",
"step: 64, Red action: DONOTHING, Blue reward:1.0\n",
"step: 65, Red action: DONOTHING, Blue reward:1.0\n",
"step: 66, Red action: DONOTHING, Blue reward:1.0\n",
"step: 67, Red action: DONOTHING, Blue reward:1.0\n",
"step: 68, Red action: DONOTHING, Blue reward:1.0\n"
]
}
],
"outputs": [],
"source": [
"env.step(13) # Patch the database\n",
"print(f\"step: {env.game.step_counter}, Red action: {info['agent_actions']['client_1_data_manipulation_red_bot'][0]}, Blue reward:{reward}\" )\n",
@@ -826,101 +522,14 @@
"Let's also have a look at the ACL observation to verify our new ACL rule at position 5."
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"{1: {'position': 0,\n",
" 'permission': 0,\n",
" 'source_node_id': 0,\n",
" 'source_port': 0,\n",
" 'dest_node_id': 0,\n",
" 'dest_port': 0,\n",
" 'protocol': 0},\n",
" 2: {'position': 1,\n",
" 'permission': 0,\n",
" 'source_node_id': 0,\n",
" 'source_port': 0,\n",
" 'dest_node_id': 0,\n",
" 'dest_port': 0,\n",
" 'protocol': 0},\n",
" 3: {'position': 2,\n",
" 'permission': 0,\n",
" 'source_node_id': 0,\n",
" 'source_port': 0,\n",
" 'dest_node_id': 0,\n",
" 'dest_port': 0,\n",
" 'protocol': 0},\n",
" 4: {'position': 3,\n",
" 'permission': 0,\n",
" 'source_node_id': 0,\n",
" 'source_port': 0,\n",
" 'dest_node_id': 0,\n",
" 'dest_port': 0,\n",
" 'protocol': 0},\n",
" 5: {'position': 4,\n",
" 'permission': 2,\n",
" 'source_node_id': 7,\n",
" 'source_port': 1,\n",
" 'dest_node_id': 4,\n",
" 'dest_port': 1,\n",
" 'protocol': 3},\n",
" 6: {'position': 5,\n",
" 'permission': 0,\n",
" 'source_node_id': 0,\n",
" 'source_port': 0,\n",
" 'dest_node_id': 0,\n",
" 'dest_port': 0,\n",
" 'protocol': 0},\n",
" 7: {'position': 6,\n",
" 'permission': 0,\n",
" 'source_node_id': 0,\n",
" 'source_port': 0,\n",
" 'dest_node_id': 0,\n",
" 'dest_port': 0,\n",
" 'protocol': 0},\n",
" 8: {'position': 7,\n",
" 'permission': 0,\n",
" 'source_node_id': 0,\n",
" 'source_port': 0,\n",
" 'dest_node_id': 0,\n",
" 'dest_port': 0,\n",
" 'protocol': 0},\n",
" 9: {'position': 8,\n",
" 'permission': 0,\n",
" 'source_node_id': 0,\n",
" 'source_port': 0,\n",
" 'dest_node_id': 0,\n",
" 'dest_port': 0,\n",
" 'protocol': 0},\n",
" 10: {'position': 9,\n",
" 'permission': 0,\n",
" 'source_node_id': 0,\n",
" 'source_port': 0,\n",
" 'dest_node_id': 0,\n",
" 'dest_port': 0,\n",
" 'protocol': 0}}"
]
},
"execution_count": 10,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"obs['ACL']"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
"source": [
"obs['ACL']"
]
}
],
"metadata": {

34
src/primaite/simulator/system/applications/web_browser.py
View File

@@ -1,7 +1,9 @@
from ipaddress import IPv4Address
from typing import Dict, Optional
from typing import Dict, List, Literal, Optional, Union
from urllib.parse import urlparse
from pydantic import BaseModel, ConfigDict
from primaite import getLogger
from primaite.simulator.core import RequestManager, RequestType
from primaite.simulator.network.protocols.http import (
@@ -33,6 +35,9 @@ class WebBrowser(Application):
latest_response: Optional[HttpResponsePacket] = None
"""Keeps track of the latest HTTP response."""
history: List["BrowserHistoryItem"] = []
"""Keep a log of visited websites and information about the visit, such as response code."""
def __init__(self, **kwargs):
kwargs["name"] = "WebBrowser"
kwargs["protocol"] = IPProtocol.TCP
@@ -71,7 +76,7 @@ class WebBrowser(Application):
:return: A dictionary capturing the current state of the WebBrowser and its child objects.
"""
state = super().describe_state()
state["last_response_status_code"] = self.latest_response.status_code if self.latest_response else None
state["history"] = [hist_item.state() for hist_item in self.history]
return state
def reset_component_for_episode(self, episode: int):
@@ -119,7 +124,8 @@ class WebBrowser(Application):
# create HTTPRequest payload
payload = HttpRequestPacket(request_method=HttpRequestMethod.GET, request_url=url)
# send request
# send request - As part of the self.send call, a response will be received and stored in the
# self.latest_response variable
if self.send(
payload=payload,
dest_ip_address=self.domain_name_ip_address,
@@ -129,9 +135,11 @@ class WebBrowser(Application):
f"{self.name}: Received HTTP {payload.request_method.name} "
f"Response {payload.request_url} - {self.latest_response.status_code.value}"
)
self.history.append(WebBrowser.BrowserHistoryItem(url=url, outcome=self.latest_response.status_code))
return self.latest_response.status_code is HttpStatusCode.OK
else:
self.sys_log.error(f"Error sending Http Packet {str(payload)}")
self.history.append(WebBrowser.BrowserHistoryItem(url=url, outcome="SERVER_UNREACHABLE"))
return False
def send(
@@ -172,3 +180,23 @@ class WebBrowser(Application):
self.sys_log.info(f"{self.name}: Received HTTP {payload.status_code.value}")
self.latest_response = payload
return True
class BrowserHistoryItem(BaseModel):
"""Simple representation of browser history, used for tracking success of web requests to calculate rewards."""
model_config = ConfigDict(extra="forbid")
"""Error if incorrect specification."""
url: str
"""The URL that was attempted to be fetched by the browser"""
outcome: Union[HttpStatusCode, Literal["PENDING", "SERVER_UNREACHABLE"]] = "PENDING"
"""HTTP response code that was received, or PENDING if a response was not yet received."""
def state(self) -> Dict:
"""Return the contents of this dataclass as a dict for use with describe_state method."""
if isinstance(self.outcome, HttpStatusCode):
outcome = self.outcome.value
else:
outcome = self.outcome
return {"url": self.url, "outcome": outcome}

40
tests/integration_tests/system/test_web_client_server_and_database.py
View File

@@ -97,12 +97,48 @@ def web_client_web_server_database(example_network) -> Tuple[Computer, Server, S
assert dns_client.check_domain_exists("arcd.com")
assert db_client.connect()
return computer, web_server, db_server
return example_network, computer, web_server, db_server
def test_web_client_requests_users(web_client_web_server_database):
computer, web_server, db_server = web_client_web_server_database
_, computer, _, _ = web_client_web_server_database
web_browser: WebBrowser = computer.software_manager.software.get("WebBrowser")
assert web_browser.get_webpage()
class TestWebBrowserHistory:
def test_populating_history(self, web_client_web_server_database):
network, computer, _, _ = web_client_web_server_database
web_browser: WebBrowser = computer.software_manager.software.get("WebBrowser")
assert web_browser.history == []
web_browser.get_webpage()
assert len(web_browser.history) == 1
web_browser.get_webpage()
assert len(web_browser.history) == 2
assert web_browser.history[-1].outcome == 200
router = network.get_node_by_hostname("router_1")
router.acl.add_rule(action=ACLAction.DENY, src_port=Port.HTTP, dst_port=Port.HTTP, position=0)
assert not web_browser.get_webpage()
assert len(web_browser.history) == 3
assert web_browser.history[-1].outcome == 404
def test_history_in_state(self, web_client_web_server_database):
network, computer, _, _ = web_client_web_server_database
web_browser: WebBrowser = computer.software_manager.software.get("WebBrowser")
state = computer.describe_state()
assert "history" in state["applications"]["WebBrowser"]
assert len(state["applications"]["WebBrowser"]["history"]) == 0
web_browser.get_webpage()
router = network.get_node_by_hostname("router_1")
router.acl.add_rule(action=ACLAction.DENY, src_port=Port.HTTP, dst_port=Port.HTTP, position=0)
web_browser.get_webpage()
state = computer.describe_state()
assert state["applications"]["WebBrowser"]["history"][0]["outcome"] == 200
assert state["applications"]["WebBrowser"]["history"][1]["outcome"] == 404