diff --git a/CHANGELOG.md b/CHANGELOG.md index fdb530e2..bc39a2b9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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 diff --git a/src/primaite/config/_package_data/example_config.yaml b/src/primaite/config/_package_data/example_config.yaml index 7a286931..700a0c18 100644 --- a/src/primaite/config/_package_data/example_config.yaml +++ b/src/primaite/config/_package_data/example_config.yaml @@ -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 diff --git a/src/primaite/game/agent/rewards.py b/src/primaite/game/agent/rewards.py index 8944a184..1a37b954 100644 --- a/src/primaite/game/agent/rewards.py +++ b/src/primaite/game/agent/rewards.py @@ -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 diff --git a/src/primaite/game/game.py b/src/primaite/game/game.py index 368d899a..8ecb365e 100644 --- a/src/primaite/game/game.py +++ b/src/primaite/game/game.py @@ -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")) diff --git a/src/primaite/notebooks/uc2_demo.ipynb b/src/primaite/notebooks/uc2_demo.ipynb index 679e8226..4e2e5e30 100644 --- a/src/primaite/notebooks/uc2_demo.ipynb +++ b/src/primaite/notebooks/uc2_demo.ipynb @@ -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": { diff --git a/src/primaite/simulator/system/applications/web_browser.py b/src/primaite/simulator/system/applications/web_browser.py index a5738d76..bf778031 100644 --- a/src/primaite/simulator/system/applications/web_browser.py +++ b/src/primaite/simulator/system/applications/web_browser.py @@ -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} diff --git a/tests/integration_tests/system/test_web_client_server_and_database.py b/tests/integration_tests/system/test_web_client_server_and_database.py index a4ef3d52..bc6dc7e6 100644 --- a/tests/integration_tests/system/test_web_client_server_and_database.py +++ b/tests/integration_tests/system/test_web_client_server_and_database.py @@ -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