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

Merged PR 267: Make web status rewards based on web browser instead of web server.

Browse Source 
## Summary
- added a second green agent on client 1
- added browser history
- added a reward component based on the most recent return code in browser history
- modified the UC2 notebook to adhere to new reward.
- (also there were some issues in example_config from a previous merge that I failed to notice prior)

## Test process
New unit tests. Ran uc2 notebook to check that it produces results aligning with what is described in the markdown cells.

## Checklist
- [x] PR is linked to a **work item**
- [x] **acceptance criteria** of linked ticket are met
- [x] performed **self-review** of the code
- [x] written **tests** for any new functionality added with this PR
- [x] updated the **documentation** if this PR changes or adds functionality
- [ ] written/updated **design docs** if this PR implements new functionality
- [x] updated the **change log**
- [x] ran **pre-commit** checks for code style
- [x] attended to any **TO-DOs** left in the code

Add second green agent and make rewards based on webbrowser

Related work items: #2247
This commit is contained in:
Marek Wolan
2024-02-08 15:19:24 +00:00
parent cfc254753f 9b350ddd6f
commit 1a2298ebfb
16 changed files with 561 additions and 739 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). and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
## [Unreleased] ## [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. - 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 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 - 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 frequency: 4
variance: 3 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 - ref: client_1_data_manipulation_red_bot
team: RED team: RED
type: RedDatabaseCorruptingAgent type: RedDatabaseCorruptingAgent
@@ -112,7 +137,7 @@ agents:
- service_name: DNSServer - service_name: DNSServer
- node_hostname: web_server - node_hostname: web_server
services: services:
- service_name: web_server_web_service - service_name: WebServer
- node_hostname: database_server - node_hostname: database_server
folders: folders:
- folder_name: database - folder_name: database
@@ -241,25 +266,25 @@ agents:
action: "NODE_FILE_SCAN" action: "NODE_FILE_SCAN"
options: options:
node_id: 2 node_id: 2
folder_id: 1 folder_id: 0
file_id: 0 file_id: 0
10: 10:
action: "NODE_FILE_CHECKHASH" action: "NODE_FILE_CHECKHASH"
options: options:
node_id: 2 node_id: 2
folder_id: 1 folder_id: 0
file_id: 0 file_id: 0
11: 11:
action: "NODE_FILE_DELETE" action: "NODE_FILE_DELETE"
options: options:
node_id: 2 node_id: 2
folder_id: 1 folder_id: 0
file_id: 0 file_id: 0
12: 12:
action: "NODE_FILE_REPAIR" action: "NODE_FILE_REPAIR"
options: options:
node_id: 2 node_id: 2
folder_id: 1 folder_id: 0
file_id: 0 file_id: 0
13: 13:
action: "NODE_SERVICE_PATCH" action: "NODE_SERVICE_PATCH"
@@ -270,22 +295,22 @@ agents:
action: "NODE_FOLDER_SCAN" action: "NODE_FOLDER_SCAN"
options: options:
node_id: 2 node_id: 2
folder_id: 1 folder_id: 0
15: 15:
action: "NODE_FOLDER_CHECKHASH" action: "NODE_FOLDER_CHECKHASH"
options: options:
node_id: 2 node_id: 2
folder_id: 1 folder_id: 0
16: 16:
action: "NODE_FOLDER_REPAIR" action: "NODE_FOLDER_REPAIR"
options: options:
node_id: 2 node_id: 2
folder_id: 1 folder_id: 0
17: 17:
action: "NODE_FOLDER_RESTORE" action: "NODE_FOLDER_RESTORE"
options: options:
node_id: 2 node_id: 2
folder_id: 1 folder_id: 0
18: 18:
action: "NODE_OS_SCAN" action: "NODE_OS_SCAN"
options: options:
@@ -302,7 +327,7 @@ agents:
action: "NODE_RESET" action: "NODE_RESET"
options: options:
node_id: 5 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" action: "NETWORK_ACL_ADDRULE"
options: options:
position: 1 position: 1
@@ -312,7 +337,7 @@ agents:
source_port_id: 1 source_port_id: 1
dest_port_id: 1 dest_port_id: 1
protocol_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" action: "NETWORK_ACL_ADDRULE"
options: options:
position: 2 position: 2
@@ -497,6 +522,8 @@ agents:
- folder_name: database - folder_name: database
files: files:
- file_name: database.db - file_name: database.db
services:
- service_name: DatabaseService
- node_name: backup_server - node_name: backup_server
- node_name: security_suite - node_name: security_suite
- node_name: client_1 - node_name: client_1
@@ -529,18 +556,19 @@ agents:
reward_function: reward_function:
reward_components: reward_components:
- type: DATABASE_FILE_INTEGRITY - type: DATABASE_FILE_INTEGRITY
weight: 0.5 weight: 0.34
options: options:
node_hostname: database_server node_hostname: database_server
folder_name: database folder_name: database
file_name: database.db file_name: database.db
- type: WEBPAGE_UNAVAILABLE_PENALTY
weight: 0.33
- type: WEB_SERVER_404_PENALTY
weight: 0.5
options: options:
node_hostname: web_server node_hostname: client_1
service_name: WebServer - type: WEBPAGE_UNAVAILABLE_PENALTY
weight: 0.33
options:
node_hostname: client_2
agent_settings: agent_settings:
@@ -682,6 +710,10 @@ simulation:
data_manipulation_p_of_success: 0.8 data_manipulation_p_of_success: 0.8
payload: "DELETE" payload: "DELETE"
server_ip: 192.168.1.14 server_ip: 192.168.1.14
- ref: client_1_web_browser
type: WebBrowser
options:
target_url: http://arcd.com/users/
services: services:
- ref: client_1_dns_client - ref: client_1_dns_client
type: DNSClient type: DNSClient

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

@@ -26,16 +26,13 @@ the structure:
``` ```
""" """
from abc import abstractmethod 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 import getLogger
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
_LOGGER = getLogger(__name__) _LOGGER = getLogger(__name__)
if TYPE_CHECKING:
from primaite.game.game import PrimaiteGame
class AbstractReward: class AbstractReward:
"""Base class for reward function components.""" """Base class for reward function components."""
@@ -47,13 +44,11 @@ class AbstractReward:
@classmethod @classmethod
@abstractmethod @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. """Create a reward function component from a config dictionary.
:param config: dict of options for the reward component's constructor :param config: dict of options for the reward component's constructor
:type config: dict :type config: dict
:param game: Reference to the PrimAITE Game object
:type game: PrimaiteGame
:return: The reward component. :return: The reward component.
:rtype: AbstractReward :rtype: AbstractReward
""" """
@@ -68,13 +63,13 @@ class DummyReward(AbstractReward):
return 0.0 return 0.0
@classmethod @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. """Create a reward function component from a config dictionary.
:param config: dict of options for the reward component's constructor. Should be empty. :param config: dict of options for the reward component's constructor. Should be empty.
:type config: dict :type config: dict
:param game: Reference to the PrimAITE Game object :return: The reward component.
:type game: PrimaiteGame :rtype: DummyReward
""" """
return cls() return cls()
@@ -126,13 +121,11 @@ class DatabaseFileIntegrity(AbstractReward):
return 0 return 0
@classmethod @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. """Create a reward function component from a config dictionary.
:param config: dict of options for the reward component's constructor :param config: dict of options for the reward component's constructor
:type config: Dict :type config: Dict
:param game: Reference to the PrimAITE Game object
:type game: PrimaiteGame
:return: The reward component. :return: The reward component.
:rtype: DatabaseFileIntegrity :rtype: DatabaseFileIntegrity
""" """
@@ -179,13 +172,11 @@ class WebServer404Penalty(AbstractReward):
return 0.0 return 0.0
@classmethod @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. """Create a reward function component from a config dictionary.
:param config: dict of options for the reward component's constructor :param config: dict of options for the reward component's constructor
:type config: Dict :type config: Dict
:param game: Reference to the PrimAITE Game object
:type game: PrimaiteGame
:return: The reward component. :return: The reward component.
:rtype: WebServer404Penalty :rtype: WebServer404Penalty
""" """
@@ -202,6 +193,55 @@ class WebServer404Penalty(AbstractReward):
return cls(node_hostname=node_hostname, service_name=service_name) 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.
:param config: Configuration dictionary.
:type config: Dict
"""
node_hostname = config.get("node_hostname")
return cls(node_hostname=node_hostname)
class RewardFunction: class RewardFunction:
"""Manages the reward function for the agent.""" """Manages the reward function for the agent."""
@@ -209,6 +249,7 @@ class RewardFunction:
"DUMMY": DummyReward, "DUMMY": DummyReward,
"DATABASE_FILE_INTEGRITY": DatabaseFileIntegrity, "DATABASE_FILE_INTEGRITY": DatabaseFileIntegrity,
"WEB_SERVER_404_PENALTY": WebServer404Penalty, "WEB_SERVER_404_PENALTY": WebServer404Penalty,
"WEBPAGE_UNAVAILABLE_PENALTY": WebpageUnavailablePenalty,
} }
def __init__(self): def __init__(self):
@@ -243,13 +284,11 @@ class RewardFunction:
return self.current_reward return self.current_reward
@classmethod @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. """Create a reward function from a config dictionary.
:param config: dict of options for the reward manager's constructor :param config: dict of options for the reward manager's constructor
:type config: Dict :type config: Dict
:param game: Reference to the PrimAITE Game object
:type game: PrimaiteGame
:return: The reward manager. :return: The reward manager.
:rtype: RewardFunction :rtype: RewardFunction
""" """
@@ -259,6 +298,6 @@ class RewardFunction:
rew_type = rew_component_cfg["type"] rew_type = rew_component_cfg["type"]
weight = rew_component_cfg.get("weight", 1.0) weight = rew_component_cfg.get("weight", 1.0)
rew_class = cls.__rew_class_identifiers[rew_type] 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) new.register_component(component=rew_instance, weight=weight)
return new return new

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

@@ -345,7 +345,7 @@ class PrimaiteGame:
action_space = ActionManager.from_config(game, action_space_cfg) action_space = ActionManager.from_config(game, action_space_cfg)
# CREATE REWARD FUNCTION # CREATE REWARD FUNCTION
rew_function = RewardFunction.from_config(reward_function_cfg, game=game) reward_function = RewardFunction.from_config(reward_function_cfg)
# OTHER AGENT SETTINGS # OTHER AGENT SETTINGS
agent_settings = AgentSettings.from_config(agent_cfg.get("agent_settings")) agent_settings = AgentSettings.from_config(agent_cfg.get("agent_settings"))
@@ -357,7 +357,7 @@ class PrimaiteGame:
agent_name=agent_cfg["ref"], agent_name=agent_cfg["ref"],
action_space=action_space, action_space=action_space,
observation_space=obs_space, observation_space=obs_space,
reward_function=rew_function, reward_function=reward_function,
agent_settings=agent_settings, agent_settings=agent_settings,
) )
game.agents.append(new_agent) game.agents.append(new_agent)
@@ -366,7 +366,7 @@ class PrimaiteGame:
agent_name=agent_cfg["ref"], agent_name=agent_cfg["ref"],
action_space=action_space, action_space=action_space,
observation_space=obs_space, observation_space=obs_space,
reward_function=rew_function, reward_function=reward_function,
agent_settings=agent_settings, agent_settings=agent_settings,
) )
game.agents.append(new_agent) game.agents.append(new_agent)
@@ -376,7 +376,7 @@ class PrimaiteGame:
agent_name=agent_cfg["ref"], agent_name=agent_cfg["ref"],
action_space=action_space, action_space=action_space,
observation_space=obs_space, observation_space=obs_space,
reward_function=rew_function, reward_function=reward_function,
agent_settings=agent_settings, agent_settings=agent_settings,
) )
game.agents.append(new_agent) game.agents.append(new_agent)

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

@@ -46,7 +46,7 @@
"source": [ "source": [
"## Green agent\n", "## Green agent\n",
"\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 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": [ "source": [
"## Blue agent\n", "## Blue agent\n",
"\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 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", "The red agent does not use information about the state of the network to decide its action.\n",
"\n", "\n",
"### Green\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", "\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", "- `28-37`: Remove ACL rules 1-10\n",
"- `42`: Disconnect client 1 from the network\n", "- `42`: Disconnect client 1 from the network\n",
"\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,9 @@
"\n", "\n",
"The blue agent's reward is calculated using two measures:\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", "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", "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",
"These two components are averaged to get the final reward.\n" "\n",
"The file status reward and the two green-agent-related rewards are averaged to get a total step reward.\n"
] ]
}, },
{ {
@@ -326,7 +327,7 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": 1, "execution_count": null,
"metadata": {}, "metadata": {},
"outputs": [], "outputs": [],
"source": [ "source": [
@@ -336,20 +337,9 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": 2, "execution_count": null,
"metadata": {}, "metadata": {},
"outputs": [ "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"
]
}
],
"source": [ "source": [
"# Imports\n", "# Imports\n",
"from primaite.config.load import example_config_path\n", "from primaite.config.load import example_config_path\n",
@@ -370,134 +360,9 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": 3, "execution_count": null,
"metadata": {}, "metadata": {},
"outputs": [ "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"
]
}
],
"source": [ "source": [
"# create the env\n", "# create the env\n",
"with open(example_config_path(), 'r') as f:\n", "with open(example_config_path(), 'r') as f:\n",
@@ -523,48 +388,9 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": 4, "execution_count": null,
"metadata": {}, "metadata": {},
"outputs": [ "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"
]
}
],
"source": [ "source": [
"for step in range(32):\n", "for step in range(32):\n",
" obs, reward, terminated, truncated, info = env.step(0)\n", " obs, reward, terminated, truncated, info = env.step(0)\n",
@@ -580,44 +406,9 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": 5, "execution_count": null,
"metadata": {}, "metadata": {},
"outputs": [ "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"
]
}
],
"source": [ "source": [
"pprint(obs['NODES'])" "pprint(obs['NODES'])"
] ]
@@ -631,44 +422,9 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": 6, "execution_count": null,
"metadata": {}, "metadata": {},
"outputs": [ "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"
]
}
],
"source": [ "source": [
"obs, reward, terminated, truncated, info = env.step(9) # scan database file\n", "obs, reward, terminated, truncated, info = env.step(9) # scan database file\n",
"obs, reward, terminated, truncated, info = env.step(1) # scan webapp service\n", "obs, reward, terminated, truncated, info = env.step(1) # scan webapp service\n",
@@ -692,24 +448,14 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": 7, "execution_count": null,
"metadata": {}, "metadata": {},
"outputs": [ "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"
]
}
],
"source": [ "source": [
"obs, reward, terminated, truncated, info = env.step(13) # patch the database\n", "obs, reward, terminated, truncated, info = env.step(13) # patch the database\n",
"print(f\"step: {env.game.step_counter}\")\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\"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\"Green action: {info['agent_actions']['client_2_green_user'][0]}\" )\n",
"print(f\"Blue reward:{reward}\" )" "print(f\"Blue reward:{reward}\" )"
] ]
@@ -727,25 +473,15 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": 8, "execution_count": null,
"metadata": {}, "metadata": {},
"outputs": [ "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"
]
}
],
"source": [ "source": [
"obs, reward, terminated, truncated, info = env.step(0) # patch the database\n", "obs, reward, terminated, truncated, info = env.step(0) # patch the database\n",
"print(f\"step: {env.game.step_counter}\")\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\"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_2_green_user'][0]}\" )\n",
"print(f\"Green action: {info['agent_actions']['client_1_green_user'][0]}\" )\n",
"print(f\"Blue reward:{reward}\" )" "print(f\"Blue reward:{reward}\" )"
] ]
}, },
@@ -758,48 +494,9 @@
}, },
{ {
"cell_type": "code", "cell_type": "code",
"execution_count": 9, "execution_count": null,
"metadata": {}, "metadata": {},
"outputs": [ "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"
]
}
],
"source": [ "source": [
"env.step(13) # Patch the database\n", "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", "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 +523,14 @@
"Let's also have a look at the ACL observation to verify our new ACL rule at position 5." "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", "cell_type": "code",
"execution_count": null, "execution_count": null,
"metadata": {}, "metadata": {},
"outputs": [], "outputs": [],
"source": [] "source": [
"obs['ACL']"
]
} }
], ],
"metadata": { "metadata": {

48
src/primaite/simulator/file_system/file_system.py
View File

@@ -94,12 +94,12 @@ class FileSystem(SimComponent):
self._restore_manager.add_request( self._restore_manager.add_request(
name="file", name="file",
request_type=RequestType( request_type=RequestType(
func=lambda request, context: self.restore_file(folder_uuid=request[0], file_uuid=request[1]) func=lambda request, context: self.restore_file(folder_name=request[0], file_name=request[1])
), ),
) )
self._restore_manager.add_request( self._restore_manager.add_request(
name="folder", name="folder",
request_type=RequestType(func=lambda request, context: self.restore_folder(folder_uuid=request[0])), request_type=RequestType(func=lambda request, context: self.restore_folder(folder_name=request[0])),
) )
rm.add_request( rm.add_request(
name="restore", name="restore",
@@ -209,7 +209,7 @@ class FileSystem(SimComponent):
folder = self.get_folder_by_id(folder_uuid=folder_uuid) folder = self.get_folder_by_id(folder_uuid=folder_uuid)
self.delete_folder(folder_name=folder.name) self.delete_folder(folder_name=folder.name)
def get_folder(self, folder_name: str) -> Optional[Folder]: def get_folder(self, folder_name: str, include_deleted: bool = False) -> Optional[Folder]:
""" """
Get a folder by its name if it exists. Get a folder by its name if it exists.
@@ -219,9 +219,13 @@ class FileSystem(SimComponent):
for folder in self.folders.values(): for folder in self.folders.values():
if folder.name == folder_name: if folder.name == folder_name:
return folder return folder
if include_deleted:
for folder in self.deleted_folders.values():
if folder.name == folder_name:
return folder
return None return None
def get_folder_by_id(self, folder_uuid: str, include_deleted: bool = False) -> Optional[Folder]: def get_folder_by_id(self, folder_uuid: str, include_deleted: Optional[bool] = False) -> Optional[Folder]:
""" """
Get a folder by its uuid if it exists. Get a folder by its uuid if it exists.
@@ -283,7 +287,7 @@ class FileSystem(SimComponent):
self._file_request_manager.add_request(name=file.name, request_type=RequestType(func=file._request_manager)) self._file_request_manager.add_request(name=file.name, request_type=RequestType(func=file._request_manager))
return file return file
def get_file(self, folder_name: str, file_name: str) -> Optional[File]: def get_file(self, folder_name: str, file_name: str, include_deleted: Optional[bool] = False) -> Optional[File]:
""" """
Retrieve a file by its name from a specific folder. Retrieve a file by its name from a specific folder.
@@ -291,9 +295,9 @@ class FileSystem(SimComponent):
:param file_name: The name of the file to be retrieved, including its extension. :param file_name: The name of the file to be retrieved, including its extension.
:return: An instance of File if it exists, otherwise `None`. :return: An instance of File if it exists, otherwise `None`.
""" """
folder = self.get_folder(folder_name) folder = self.get_folder(folder_name, include_deleted=include_deleted)
if folder: if folder:
return folder.get_file(file_name) return folder.get_file(file_name, include_deleted=include_deleted)
self.sys_log.info(f"File not found /{folder_name}/{file_name}") self.sys_log.info(f"File not found /{folder_name}/{file_name}")
def get_file_by_id( def get_file_by_id(
@@ -455,46 +459,44 @@ class FileSystem(SimComponent):
for folder_id in self.folders: for folder_id in self.folders:
self.folders[folder_id].reveal_to_red(instant_scan=instant_scan) self.folders[folder_id].reveal_to_red(instant_scan=instant_scan)
def restore_folder(self, folder_uuid: str): def restore_folder(self, folder_name: str):
""" """
Restore a folder. Restore a folder.
Checks the current folder's status and applies the correct fix for the folder. Checks the current folder's status and applies the correct fix for the folder.
:param: folder_uuid: id of the folder to restore :param: folder_name: name of the folder to restore
:type: folder_uuid: str :type: folder_uuid: str
""" """
folder = self.get_folder_by_id(folder_uuid=folder_uuid, include_deleted=True) folder = self.get_folder(folder_name=folder_name, include_deleted=True)
if folder is None: if folder is None:
self.sys_log.error(f"Unable to restore folder with uuid {folder_uuid}. Folder does not exist.") self.sys_log.error(f"Unable to restore folder {folder_name}. Folder is not in deleted folder list.")
return return
self.deleted_folders.pop(folder.uuid, None)
folder.restore() folder.restore()
self.folders[folder.uuid] = folder self.folders[folder.uuid] = folder
if folder.deleted: def restore_file(self, folder_name: str, file_name: str):
self.deleted_folders.pop(folder.uuid)
def restore_file(self, folder_uuid: str, file_uuid: str):
""" """
Restore a file. Restore a file.
Checks the current file's status and applies the correct fix for the file. Checks the current file's status and applies the correct fix for the file.
:param: folder_uuid: id of the folder where the file is stored :param: folder_name: name of the folder where the file is stored
:type: folder_uuid: str :type: folder_name: str
:param: file_uuid: id of the file to restore :param: file_name: name of the file to restore
:type: file_uuid: str :type: file_name: str
""" """
folder = self.get_folder_by_id(folder_uuid=folder_uuid, include_deleted=True) folder = self.get_folder(folder_name=folder_name)
if folder: if folder:
file = folder.get_file_by_id(file_uuid=file_uuid, include_deleted=True) file = folder.get_file(file_name=file_name, include_deleted=True)
if file is None: if file is None:
self.sys_log.error(f"Unable to restore file with uuid {file_uuid}. File does not exist.") self.sys_log.error(f"Unable to restore file {file_name}. File does not exist.")
return return
folder.restore_file(file_uuid=file_uuid) folder.restore_file(file_name=file_name)

24
src/primaite/simulator/file_system/folder.py
View File

@@ -193,19 +193,19 @@ class Folder(FileSystemItemABC):
if self.restore_countdown == 0: if self.restore_countdown == 0:
# repair all files # repair all files
for file_id in self.files: for file_id, file in self.files.items():
self.restore_file(file_uuid=file_id) self.restore_file(file_name=file.name)
deleted_files = self.deleted_files.copy() deleted_files = self.deleted_files.copy()
for file_id in deleted_files: for file_id, file in deleted_files.items():
self.restore_file(file_uuid=file_id) self.restore_file(file_name=file.name)
if self.deleted: if self.deleted:
self.deleted = False self.deleted = False
elif self.health_status in [FileSystemItemHealthStatus.CORRUPT, FileSystemItemHealthStatus.RESTORING]: elif self.health_status in [FileSystemItemHealthStatus.CORRUPT, FileSystemItemHealthStatus.RESTORING]:
self.health_status = FileSystemItemHealthStatus.GOOD self.health_status = FileSystemItemHealthStatus.GOOD
def get_file(self, file_name: str) -> Optional[File]: def get_file(self, file_name: str, include_deleted: Optional[bool] = False) -> Optional[File]:
""" """
Get a file by its name. Get a file by its name.
@@ -218,6 +218,10 @@ class Folder(FileSystemItemABC):
for file in self.files.values(): for file in self.files.values():
if file.name == file_name: if file.name == file_name:
return file return file
if include_deleted:
for file in self.deleted_files.values():
if file.name == file_name:
return file
return None return None
def get_file_by_id(self, file_uuid: str, include_deleted: Optional[bool] = False) -> File: def get_file_by_id(self, file_uuid: str, include_deleted: Optional[bool] = False) -> File:
@@ -297,23 +301,23 @@ class Folder(FileSystemItemABC):
self.files = {} self.files = {}
def restore_file(self, file_uuid: str): def restore_file(self, file_name: str):
""" """
Restores a file. Restores a file.
:param file_uuid: The id of the file to restore :param file_name: The name of the file to restore
""" """
# if the file was not deleted, run a repair # if the file was not deleted, run a repair
file = self.get_file_by_id(file_uuid=file_uuid, include_deleted=True) file = self.get_file(file_name=file_name, include_deleted=True)
if not file: if not file:
self.sys_log.error(f"Unable to restore file with uuid {file_uuid}. File does not exist.") self.sys_log.error(f"Unable to restore file {file_name}. File does not exist.")
return return
file.restore() file.restore()
self.files[file.uuid] = file self.files[file.uuid] = file
if file.deleted: if file.deleted:
self.deleted_files.pop(file_uuid) self.deleted_files.pop(file.uuid)
def quarantine(self): def quarantine(self):
"""Quarantines the File System Folder.""" """Quarantines the File System Folder."""

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

@@ -1,7 +1,10 @@
from enum import Enum
from ipaddress import IPv4Address from ipaddress import IPv4Address
from typing import Dict, Optional from typing import Dict, List, Optional
from urllib.parse import urlparse from urllib.parse import urlparse
from pydantic import BaseModel, ConfigDict
from primaite import getLogger from primaite import getLogger
from primaite.simulator.core import RequestManager, RequestType from primaite.simulator.core import RequestManager, RequestType
from primaite.simulator.network.protocols.http import ( from primaite.simulator.network.protocols.http import (
@@ -33,6 +36,9 @@ class WebBrowser(Application):
latest_response: Optional[HttpResponsePacket] = None latest_response: Optional[HttpResponsePacket] = None
"""Keeps track of the latest HTTP response.""" """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): def __init__(self, **kwargs):
kwargs["name"] = "WebBrowser" kwargs["name"] = "WebBrowser"
kwargs["protocol"] = IPProtocol.TCP kwargs["protocol"] = IPProtocol.TCP
@@ -71,7 +77,7 @@ class WebBrowser(Application):
:return: A dictionary capturing the current state of the WebBrowser and its child objects. :return: A dictionary capturing the current state of the WebBrowser and its child objects.
""" """
state = super().describe_state() 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 return state
def reset_component_for_episode(self, episode: int): def reset_component_for_episode(self, episode: int):
@@ -119,7 +125,8 @@ class WebBrowser(Application):
# create HTTPRequest payload # create HTTPRequest payload
payload = HttpRequestPacket(request_method=HttpRequestMethod.GET, request_url=url) 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( if self.send(
payload=payload, payload=payload,
dest_ip_address=self.domain_name_ip_address, dest_ip_address=self.domain_name_ip_address,
@@ -129,9 +136,21 @@ class WebBrowser(Application):
f"{self.name}: Received HTTP {payload.request_method.name} " f"{self.name}: Received HTTP {payload.request_method.name} "
f"Response {payload.request_url} - {self.latest_response.status_code.value}" f"Response {payload.request_url} - {self.latest_response.status_code.value}"
) )
self.history.append(
WebBrowser.BrowserHistoryItem(
url=url,
status=self.BrowserHistoryItem._HistoryItemStatus.LOADED,
response_code=self.latest_response.status_code,
)
)
return self.latest_response.status_code is HttpStatusCode.OK return self.latest_response.status_code is HttpStatusCode.OK
else: else:
self.sys_log.error(f"Error sending Http Packet {str(payload)}") self.sys_log.error(f"Error sending Http Packet {str(payload)}")
self.history.append(
WebBrowser.BrowserHistoryItem(
url=url, status=self.BrowserHistoryItem._HistoryItemStatus.SERVER_UNREACHABLE
)
)
return False return False
def send( def send(
@@ -172,3 +191,31 @@ class WebBrowser(Application):
self.sys_log.info(f"{self.name}: Received HTTP {payload.status_code.value}") self.sys_log.info(f"{self.name}: Received HTTP {payload.status_code.value}")
self.latest_response = payload self.latest_response = payload
return True 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"""
class _HistoryItemStatus(Enum):
NOT_SENT = "NOT_SENT"
PENDING = "PENDING"
SERVER_UNREACHABLE = "SERVER_UNREACHABLE"
LOADED = "LOADED"
status: _HistoryItemStatus = _HistoryItemStatus.PENDING
response_code: Optional[HttpStatusCode] = None
"""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 self.status == self._HistoryItemStatus.LOADED:
outcome = self.response_code
else:
outcome = self.status.value
return {"url": self.url, "outcome": outcome}

4
src/primaite/simulator/system/services/database/database_service.py
View File

@@ -203,6 +203,10 @@ class DatabaseService(Service):
""" """
self.sys_log.info(f"{self.name}: Running {query}") self.sys_log.info(f"{self.name}: Running {query}")
if not self.db_file:
self.sys_log.info(f"{self.name}: Failed to run {query} because the database file is missing.")
return {"status_code": 404, "data": False}
if query == "SELECT": if query == "SELECT":
if self.db_file.health_status == FileSystemItemHealthStatus.GOOD: if self.db_file.health_status == FileSystemItemHealthStatus.GOOD:
return { return {

233
tests/conftest.py
View File

@@ -6,6 +6,10 @@ import pytest
import yaml import yaml
from primaite import getLogger from primaite import getLogger
from primaite.game.agent.actions import ActionManager
from primaite.game.agent.interface import AbstractAgent
from primaite.game.agent.observations import ICSObservation, ObservationManager
from primaite.game.agent.rewards import RewardFunction
from primaite.game.game import PrimaiteGame from primaite.game.game import PrimaiteGame
from primaite.session.session import PrimaiteSession from primaite.session.session import PrimaiteSession
@@ -20,9 +24,14 @@ from primaite.simulator.network.hardware.nodes.switch import Switch
from primaite.simulator.network.networks import arcd_uc2_network from primaite.simulator.network.networks import arcd_uc2_network
from primaite.simulator.network.transmission.network_layer import IPProtocol from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.sim_container import Simulation
from primaite.simulator.system.applications.application import Application from primaite.simulator.system.applications.application import Application
from primaite.simulator.system.applications.web_browser import WebBrowser
from primaite.simulator.system.core.sys_log import SysLog from primaite.simulator.system.core.sys_log import SysLog
from primaite.simulator.system.services.dns.dns_client import DNSClient
from primaite.simulator.system.services.dns.dns_server import DNSServer
from primaite.simulator.system.services.service import Service from primaite.simulator.system.services.service import Service
from primaite.simulator.system.services.web_server.web_server import WebServer
from tests.mock_and_patch.get_session_path_mock import temp_user_sessions_path from tests.mock_and_patch.get_session_path_mock import temp_user_sessions_path
ACTION_SPACE_NODE_VALUES = 1 ACTION_SPACE_NODE_VALUES = 1
@@ -237,3 +246,227 @@ def example_network() -> Network:
router_1.acl.add_rule(action=ACLAction.PERMIT, protocol=IPProtocol.ICMP, position=23) router_1.acl.add_rule(action=ACLAction.PERMIT, protocol=IPProtocol.ICMP, position=23)
return network return network
class ControlledAgent(AbstractAgent):
"""Agent that can be controlled by the tests."""
def __init__(
self,
agent_name: str,
action_space: ActionManager,
observation_space: ObservationManager,
reward_function: RewardFunction,
) -> None:
super().__init__(
agent_name=agent_name,
action_space=action_space,
observation_space=observation_space,
reward_function=reward_function,
)
self.most_recent_action: Tuple[str, Dict]
def get_action(self, obs: None, reward: float = 0.0) -> Tuple[str, Dict]:
"""Return the agent's most recent action, formatted in CAOS format."""
return self.most_recent_action
def store_action(self, action: Tuple[str, Dict]):
"""Store the most recent action."""
self.most_recent_action = action
def install_stuff_to_sim(sim: Simulation):
"""Create a simulation with a computer, two servers, two switches, and a router."""
# 0: Pull out the network
network = sim.network
# 1: Set up network hardware
# 1.1: Configure the router
router = Router(hostname="router", num_ports=3, operating_state=NodeOperatingState.ON)
router.power_on()
router.configure_port(port=1, ip_address="10.0.1.1", subnet_mask="255.255.255.0")
router.configure_port(port=2, ip_address="10.0.2.1", subnet_mask="255.255.255.0")
# 1.2: Create and connect switches
switch_1 = Switch(hostname="switch_1", num_ports=6, operating_state=NodeOperatingState.ON)
switch_1.power_on()
network.connect(endpoint_a=router.ethernet_ports[1], endpoint_b=switch_1.switch_ports[6])
router.enable_port(1)
switch_2 = Switch(hostname="switch_2", num_ports=6, operating_state=NodeOperatingState.ON)
switch_2.power_on()
network.connect(endpoint_a=router.ethernet_ports[2], endpoint_b=switch_2.switch_ports[6])
router.enable_port(2)
# 1.3: Create and connect computer
client_1 = Computer(
hostname="client_1",
ip_address="10.0.1.2",
subnet_mask="255.255.255.0",
default_gateway="10.0.1.1",
operating_state=NodeOperatingState.ON,
)
client_1.power_on()
network.connect(
endpoint_a=client_1.ethernet_port[1],
endpoint_b=switch_1.switch_ports[1],
)
# 1.4: Create and connect servers
server_1 = Server(
hostname="server_1",
ip_address="10.0.2.2",
subnet_mask="255.255.255.0",
default_gateway="10.0.2.1",
operating_state=NodeOperatingState.ON,
)
server_1.power_on()
network.connect(endpoint_a=server_1.ethernet_port[1], endpoint_b=switch_2.switch_ports[1])
server_2 = Server(
hostname="server_2",
ip_address="10.0.2.3",
subnet_mask="255.255.255.0",
default_gateway="10.0.2.1",
operating_state=NodeOperatingState.ON,
)
server_2.power_on()
network.connect(endpoint_a=server_2.ethernet_port[1], endpoint_b=switch_2.switch_ports[2])
# 2: Configure base ACL
router.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.ARP, dst_port=Port.ARP, position=22)
router.acl.add_rule(action=ACLAction.PERMIT, protocol=IPProtocol.ICMP, position=23)
router.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.DNS, dst_port=Port.DNS, position=1)
router.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.HTTP, dst_port=Port.HTTP, position=3)
# 3: Install server software
server_1.software_manager.install(DNSServer)
dns_service: DNSServer = server_1.software_manager.software.get("DNSServer") # noqa
dns_service.dns_register("www.example.com", server_2.ip_address)
server_2.software_manager.install(WebServer)
# 3.1: Ensure that the dns clients are configured correctly
client_1.software_manager.software.get("DNSClient").dns_server = server_1.ethernet_port[1].ip_address
server_2.software_manager.software.get("DNSClient").dns_server = server_1.ethernet_port[1].ip_address
# 4: Check that client came pre-installed with web browser and dns client
assert isinstance(client_1.software_manager.software.get("WebBrowser"), WebBrowser)
assert isinstance(client_1.software_manager.software.get("DNSClient"), DNSClient)
# 4.1: Create a file on the computer
client_1.file_system.create_file("cat.png", 300, folder_name="downloads")
# 5: Assert that the simulation starts off in the state that we expect
assert len(sim.network.nodes) == 6
assert len(sim.network.links) == 5
# 5.1: Assert the router is correctly configured
r = sim.network.routers[0]
for i, acl_rule in enumerate(r.acl.acl):
if i == 1:
assert acl_rule.src_port == acl_rule.dst_port == Port.DNS
elif i == 3:
assert acl_rule.src_port == acl_rule.dst_port == Port.HTTP
elif i == 22:
assert acl_rule.src_port == acl_rule.dst_port == Port.ARP
elif i == 23:
assert acl_rule.protocol == IPProtocol.ICMP
elif i == 24:
...
else:
assert acl_rule is None
# 5.2: Assert the client is correctly configured
c: Computer = [node for node in sim.network.nodes.values() if node.hostname == "client_1"][0]
assert c.software_manager.software.get("WebBrowser") is not None
assert c.software_manager.software.get("DNSClient") is not None
assert str(c.ethernet_port[1].ip_address) == "10.0.1.2"
# 5.3: Assert that server_1 is correctly configured
s1: Server = [node for node in sim.network.nodes.values() if node.hostname == "server_1"][0]
assert str(s1.ethernet_port[1].ip_address) == "10.0.2.2"
assert s1.software_manager.software.get("DNSServer") is not None
# 5.4: Assert that server_2 is correctly configured
s2: Server = [node for node in sim.network.nodes.values() if node.hostname == "server_2"][0]
assert str(s2.ethernet_port[1].ip_address) == "10.0.2.3"
assert s2.software_manager.software.get("WebServer") is not None
# 6: Return the simulation
return sim
@pytest.fixture
def game_and_agent():
"""Create a game with a simple agent that can be controlled by the tests."""
game = PrimaiteGame()
sim = game.simulation
install_stuff_to_sim(sim)
actions = [
{"type": "DONOTHING"},
{"type": "NODE_SERVICE_SCAN"},
{"type": "NODE_SERVICE_STOP"},
{"type": "NODE_SERVICE_START"},
{"type": "NODE_SERVICE_PAUSE"},
{"type": "NODE_SERVICE_RESUME"},
{"type": "NODE_SERVICE_RESTART"},
{"type": "NODE_SERVICE_DISABLE"},
{"type": "NODE_SERVICE_ENABLE"},
{"type": "NODE_SERVICE_PATCH"},
{"type": "NODE_APPLICATION_EXECUTE"},
{"type": "NODE_FILE_SCAN"},
{"type": "NODE_FILE_CHECKHASH"},
{"type": "NODE_FILE_DELETE"},
{"type": "NODE_FILE_REPAIR"},
{"type": "NODE_FILE_RESTORE"},
{"type": "NODE_FILE_CORRUPT"},
{"type": "NODE_FOLDER_SCAN"},
{"type": "NODE_FOLDER_CHECKHASH"},
{"type": "NODE_FOLDER_REPAIR"},
{"type": "NODE_FOLDER_RESTORE"},
{"type": "NODE_OS_SCAN"},
{"type": "NODE_SHUTDOWN"},
{"type": "NODE_STARTUP"},
{"type": "NODE_RESET"},
{"type": "NETWORK_ACL_ADDRULE", "options": {"target_router_hostname": "router"}},
{"type": "NETWORK_ACL_REMOVERULE", "options": {"target_router_hostname": "router"}},
{"type": "NETWORK_NIC_ENABLE"},
{"type": "NETWORK_NIC_DISABLE"},
]
action_space = ActionManager(
game=game,
actions=actions, # ALL POSSIBLE ACTIONS
nodes=[
{
"node_name": "client_1",
"applications": [{"application_name": "WebBrowser"}],
"folders": [{"folder_name": "downloads", "files": [{"file_name": "cat.png"}]}],
},
{"node_name": "server_1", "services": [{"service_name": "DNSServer"}]},
{"node_name": "server_2", "services": [{"service_name": "WebServer"}]},
],
max_folders_per_node=2,
max_files_per_folder=2,
max_services_per_node=2,
max_applications_per_node=2,
max_nics_per_node=2,
max_acl_rules=10,
protocols=["TCP", "UDP", "ICMP"],
ports=["HTTP", "DNS", "ARP"],
ip_address_list=["10.0.1.1", "10.0.1.2", "10.0.2.1", "10.0.2.2", "10.0.2.3"],
act_map={},
)
observation_space = ObservationManager(ICSObservation())
reward_function = RewardFunction()
test_agent = ControlledAgent(
agent_name="test_agent",
action_space=action_space,
observation_space=observation_space,
reward_function=reward_function,
)
game.agents.append(test_agent)
return (game, test_agent)

228
tests/integration_tests/game_layer/test_actions.py
View File

@@ -35,234 +35,6 @@ from primaite.simulator.system.services.web_server.web_server import WebServer
from primaite.simulator.system.software import SoftwareHealthState from primaite.simulator.system.software import SoftwareHealthState
class ControlledAgent(AbstractAgent):
"""Agent that can be controlled by the tests."""
def __init__(
self,
agent_name: str,
action_space: ActionManager,
observation_space: ObservationManager,
reward_function: RewardFunction,
) -> None:
super().__init__(
agent_name=agent_name,
action_space=action_space,
observation_space=observation_space,
reward_function=reward_function,
)
self.most_recent_action: Tuple[str, Dict]
def get_action(self, obs: None, reward: float = 0.0) -> Tuple[str, Dict]:
"""Return the agent's most recent action, formatted in CAOS format."""
return self.most_recent_action
def store_action(self, action: Tuple[str, Dict]):
"""Store the most recent action."""
self.most_recent_action = action
def install_stuff_to_sim(sim: Simulation):
"""Create a simulation with a computer, two servers, two switches, and a router."""
# 0: Pull out the network
network = sim.network
# 1: Set up network hardware
# 1.1: Configure the router
router = Router(hostname="router", num_ports=3, operating_state=NodeOperatingState.ON)
router.power_on()
router.configure_port(port=1, ip_address="10.0.1.1", subnet_mask="255.255.255.0")
router.configure_port(port=2, ip_address="10.0.2.1", subnet_mask="255.255.255.0")
# 1.2: Create and connect switches
switch_1 = Switch(hostname="switch_1", num_ports=6, operating_state=NodeOperatingState.ON)
switch_1.power_on()
network.connect(endpoint_a=router.ethernet_ports[1], endpoint_b=switch_1.switch_ports[6])
router.enable_port(1)
switch_2 = Switch(hostname="switch_2", num_ports=6, operating_state=NodeOperatingState.ON)
switch_2.power_on()
network.connect(endpoint_a=router.ethernet_ports[2], endpoint_b=switch_2.switch_ports[6])
router.enable_port(2)
# 1.3: Create and connect computer
client_1 = Computer(
hostname="client_1",
ip_address="10.0.1.2",
subnet_mask="255.255.255.0",
default_gateway="10.0.1.1",
operating_state=NodeOperatingState.ON,
)
client_1.power_on()
network.connect(
endpoint_a=client_1.ethernet_port[1],
endpoint_b=switch_1.switch_ports[1],
)
# 1.4: Create and connect servers
server_1 = Server(
hostname="server_1",
ip_address="10.0.2.2",
subnet_mask="255.255.255.0",
default_gateway="10.0.2.1",
operating_state=NodeOperatingState.ON,
)
server_1.power_on()
network.connect(endpoint_a=server_1.ethernet_port[1], endpoint_b=switch_2.switch_ports[1])
server_2 = Server(
hostname="server_2",
ip_address="10.0.2.3",
subnet_mask="255.255.255.0",
default_gateway="10.0.2.1",
operating_state=NodeOperatingState.ON,
)
server_2.power_on()
network.connect(endpoint_a=server_2.ethernet_port[1], endpoint_b=switch_2.switch_ports[2])
# 2: Configure base ACL
router.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.ARP, dst_port=Port.ARP, position=22)
router.acl.add_rule(action=ACLAction.PERMIT, protocol=IPProtocol.ICMP, position=23)
router.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.DNS, dst_port=Port.DNS, position=1)
router.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.HTTP, dst_port=Port.HTTP, position=3)
# 3: Install server software
server_1.software_manager.install(DNSServer)
dns_service: DNSServer = server_1.software_manager.software.get("DNSServer") # noqa
dns_service.dns_register("www.example.com", server_2.ip_address)
server_2.software_manager.install(WebServer)
# 3.1: Ensure that the dns clients are configured correctly
client_1.software_manager.software.get("DNSClient").dns_server = server_1.ethernet_port[1].ip_address
server_2.software_manager.software.get("DNSClient").dns_server = server_1.ethernet_port[1].ip_address
# 4: Check that client came pre-installed with web browser and dns client
assert isinstance(client_1.software_manager.software.get("WebBrowser"), WebBrowser)
assert isinstance(client_1.software_manager.software.get("DNSClient"), DNSClient)
# 4.1: Create a file on the computer
client_1.file_system.create_file("cat.png", 300, folder_name="downloads")
# 5: Assert that the simulation starts off in the state that we expect
assert len(sim.network.nodes) == 6
assert len(sim.network.links) == 5
# 5.1: Assert the router is correctly configured
r = sim.network.routers[0]
for i, acl_rule in enumerate(r.acl.acl):
if i == 1:
assert acl_rule.src_port == acl_rule.dst_port == Port.DNS
elif i == 3:
assert acl_rule.src_port == acl_rule.dst_port == Port.HTTP
elif i == 22:
assert acl_rule.src_port == acl_rule.dst_port == Port.ARP
elif i == 23:
assert acl_rule.protocol == IPProtocol.ICMP
elif i == 24:
...
else:
assert acl_rule is None
# 5.2: Assert the client is correctly configured
c: Computer = [node for node in sim.network.nodes.values() if node.hostname == "client_1"][0]
assert c.software_manager.software.get("WebBrowser") is not None
assert c.software_manager.software.get("DNSClient") is not None
assert str(c.ethernet_port[1].ip_address) == "10.0.1.2"
# 5.3: Assert that server_1 is correctly configured
s1: Server = [node for node in sim.network.nodes.values() if node.hostname == "server_1"][0]
assert str(s1.ethernet_port[1].ip_address) == "10.0.2.2"
assert s1.software_manager.software.get("DNSServer") is not None
# 5.4: Assert that server_2 is correctly configured
s2: Server = [node for node in sim.network.nodes.values() if node.hostname == "server_2"][0]
assert str(s2.ethernet_port[1].ip_address) == "10.0.2.3"
assert s2.software_manager.software.get("WebServer") is not None
# 6: Return the simulation
return sim
@pytest.fixture
def game_and_agent():
"""Create a game with a simple agent that can be controlled by the tests."""
game = PrimaiteGame()
sim = game.simulation
install_stuff_to_sim(sim)
actions = [
{"type": "DONOTHING"},
{"type": "NODE_SERVICE_SCAN"},
{"type": "NODE_SERVICE_STOP"},
{"type": "NODE_SERVICE_START"},
{"type": "NODE_SERVICE_PAUSE"},
{"type": "NODE_SERVICE_RESUME"},
{"type": "NODE_SERVICE_RESTART"},
{"type": "NODE_SERVICE_DISABLE"},
{"type": "NODE_SERVICE_ENABLE"},
{"type": "NODE_SERVICE_PATCH"},
{"type": "NODE_APPLICATION_EXECUTE"},
{"type": "NODE_FILE_SCAN"},
{"type": "NODE_FILE_CHECKHASH"},
{"type": "NODE_FILE_DELETE"},
{"type": "NODE_FILE_REPAIR"},
{"type": "NODE_FILE_RESTORE"},
{"type": "NODE_FILE_CORRUPT"},
{"type": "NODE_FOLDER_SCAN"},
{"type": "NODE_FOLDER_CHECKHASH"},
{"type": "NODE_FOLDER_REPAIR"},
{"type": "NODE_FOLDER_RESTORE"},
{"type": "NODE_OS_SCAN"},
{"type": "NODE_SHUTDOWN"},
{"type": "NODE_STARTUP"},
{"type": "NODE_RESET"},
{"type": "NETWORK_ACL_ADDRULE", "options": {"target_router_hostname": "router"}},
{"type": "NETWORK_ACL_REMOVERULE", "options": {"target_router_hostname": "router"}},
{"type": "NETWORK_NIC_ENABLE"},
{"type": "NETWORK_NIC_DISABLE"},
]
action_space = ActionManager(
game=game,
actions=actions, # ALL POSSIBLE ACTIONS
nodes=[
{
"node_name": "client_1",
"applications": [{"application_name": "WebBrowser"}],
"folders": [{"folder_name": "downloads", "files": [{"file_name": "cat.png"}]}],
},
{"node_name": "server_1", "services": [{"service_name": "DNSServer"}]},
{"node_name": "server_2", "services": [{"service_name": "WebServer"}]},
],
max_folders_per_node=2,
max_files_per_folder=2,
max_services_per_node=2,
max_applications_per_node=2,
max_nics_per_node=2,
max_acl_rules=10,
protocols=["TCP", "UDP", "ICMP"],
ports=["HTTP", "DNS", "ARP"],
ip_address_list=["10.0.1.1", "10.0.1.2", "10.0.2.1", "10.0.2.2", "10.0.2.3"],
act_map={},
)
observation_space = ObservationManager(ICSObservation())
reward_function = RewardFunction()
test_agent = ControlledAgent(
agent_name="test_agent",
action_space=action_space,
observation_space=observation_space,
reward_function=reward_function,
)
game.agents.append(test_agent)
return (game, test_agent)
# def test_test(game_and_agent:Tuple[PrimaiteGame, ProxyAgent]):
# game, agent = game_and_agent
def test_do_nothing_integration(game_and_agent: Tuple[PrimaiteGame, ProxyAgent]): def test_do_nothing_integration(game_and_agent: Tuple[PrimaiteGame, ProxyAgent]):
"""Test that the DoNothingAction can form a request and that it is accepted by the simulation.""" """Test that the DoNothingAction can form a request and that it is accepted by the simulation."""
game, agent = game_and_agent game, agent = game_and_agent

37
tests/integration_tests/game_layer/test_rewards.py Normal file
View File

@@ -0,0 +1,37 @@
from primaite.game.agent.rewards import RewardFunction, WebpageUnavailablePenalty
from primaite.simulator.network.hardware.nodes.router import ACLAction, Router
from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port
from tests.conftest import ControlledAgent
def test_WebpageUnavailablePenalty(game_and_agent):
"""Test that we get the right reward for failing to fetch a website."""
game, agent = game_and_agent
agent: ControlledAgent
comp = WebpageUnavailablePenalty(node_hostname="client_1")
agent.reward_function.register_component(comp, 0.7)
action = ("DONOTHING", {})
agent.store_action(action)
game.step()
# client 1 has not attempted to fetch webpage yet!
assert agent.reward_function.current_reward == 0.0
client_1 = game.simulation.network.get_node_by_hostname("client_1")
browser = client_1.software_manager.software.get("WebBrowser")
browser.run()
browser.target_url = "http://www.example.com"
assert browser.get_webpage()
action = ("DONOTHING", {})
agent.store_action(action)
game.step()
assert agent.reward_function.current_reward == 0.7
router: Router = game.simulation.network.get_node_by_hostname("router")
router.acl.add_rule(action=ACLAction.DENY, protocol=IPProtocol.TCP, src_port=Port.HTTP, dst_port=Port.HTTP)
assert not browser.get_webpage()
agent.store_action(action)
game.step()
assert agent.reward_function.current_reward == -0.7

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

@@ -3,6 +3,7 @@ from typing import Tuple
import pytest import pytest
from primaite.simulator.network.container import Network
from primaite.simulator.network.hardware.base import Link from primaite.simulator.network.hardware.base import Link
from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
from primaite.simulator.network.hardware.nodes.computer import Computer from primaite.simulator.network.hardware.nodes.computer import Computer
@@ -18,7 +19,7 @@ from primaite.simulator.system.services.web_server.web_server import WebServer
@pytest.fixture(scope="function") @pytest.fixture(scope="function")
def web_client_web_server_database(example_network) -> Tuple[Computer, Server, Server]: def web_client_web_server_database(example_network) -> Tuple[Network, Computer, Server, Server]:
# add rules to network router # add rules to network router
router_1: Router = example_network.get_node_by_hostname("router_1") router_1: Router = example_network.get_node_by_hostname("router_1")
router_1.acl.add_rule( router_1.acl.add_rule(
@@ -97,12 +98,52 @@ def web_client_web_server_database(example_network) -> Tuple[Computer, Server, S
assert dns_client.check_domain_exists("arcd.com") assert dns_client.check_domain_exists("arcd.com")
assert db_client.connect() 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): 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") web_browser: WebBrowser = computer.software_manager.software.get("WebBrowser")
assert web_browser.get_webpage() 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].status == WebBrowser.BrowserHistoryItem._HistoryItemStatus.LOADED
assert web_browser.history[-1].response_code == 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
# with current NIC behaviour, even if you block communication, you won't get SERVER_UNREACHABLE because
# application.send always returns true, even if communication fails. we should change what is returned from NICs
assert web_browser.history[-1].status == WebBrowser.BrowserHistoryItem._HistoryItemStatus.LOADED
assert web_browser.history[-1].response_code == 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

8
tests/unit_tests/_primaite/_simulator/_file_system/test_file_actions.py
View File

@@ -61,18 +61,18 @@ def test_file_restore_request(populated_file_system):
fs, folder, file = populated_file_system fs, folder, file = populated_file_system
assert fs.get_file_by_id(folder_uuid=folder.uuid, file_uuid=file.uuid) is not None assert fs.get_file_by_id(folder_uuid=folder.uuid, file_uuid=file.uuid) is not None
fs.apply_request(request=["delete", "file", folder.uuid, file.uuid]) fs.apply_request(request=["delete", "file", folder.name, file.name])
assert fs.get_file(folder_name=folder.name, file_name=file.name) is None assert fs.get_file(folder_name=folder.name, file_name=file.name) is None
assert fs.get_file_by_id(folder_uuid=folder.uuid, file_uuid=file.uuid, include_deleted=True).deleted is True assert fs.get_file_by_id(folder_uuid=folder.uuid, file_uuid=file.uuid, include_deleted=True).deleted is True
fs.apply_request(request=["restore", "file", folder.uuid, file.uuid]) fs.apply_request(request=["restore", "file", folder.name, file.name])
assert fs.get_file(folder_name=folder.name, file_name=file.name) is not None assert fs.get_file(folder_name=folder.name, file_name=file.name) is not None
assert fs.get_file(folder_name=folder.name, file_name=file.name).deleted is False assert fs.get_file(folder_name=folder.name, file_name=file.name).deleted is False
fs.apply_request(request=["file", file.name, "corrupt"]) fs.apply_request(request=["file", file.name, "corrupt"])
assert fs.get_file(folder_name=folder.name, file_name=file.name).health_status == FileSystemItemHealthStatus.CORRUPT assert fs.get_file(folder_name=folder.name, file_name=file.name).health_status == FileSystemItemHealthStatus.CORRUPT
fs.apply_request(request=["restore", "file", folder.uuid, file.uuid]) fs.apply_request(request=["restore", "file", folder.name, file.name])
assert fs.get_file(folder_name=folder.name, file_name=file.name).health_status == FileSystemItemHealthStatus.GOOD assert fs.get_file(folder_name=folder.name, file_name=file.name).health_status == FileSystemItemHealthStatus.GOOD
@@ -95,7 +95,7 @@ def test_deleted_file_cannot_be_interacted_with(populated_file_system):
== FileSystemItemHealthStatus.GOOD == FileSystemItemHealthStatus.GOOD
) )
fs.apply_request(request=["delete", "file", folder.uuid, file.uuid]) fs.apply_request(request=["delete", "file", folder.name, file.name])
assert fs.get_file(folder_name=folder.name, file_name=file.name) is None assert fs.get_file(folder_name=folder.name, file_name=file.name) is None
fs.apply_request(request=["file", file.name, "repair"]) fs.apply_request(request=["file", file.name, "repair"])

4
tests/unit_tests/_primaite/_simulator/_file_system/test_file_system_actions.py
View File

@@ -21,7 +21,7 @@ def test_file_delete_request(populated_file_system):
fs, folder, file = populated_file_system fs, folder, file = populated_file_system
assert fs.get_file(folder_name=folder.name, file_name=file.name) is not None assert fs.get_file(folder_name=folder.name, file_name=file.name) is not None
fs.apply_request(request=["delete", "file", folder.uuid, file.uuid]) fs.apply_request(request=["delete", "file", folder.name, file.name])
assert fs.get_file(folder_name=folder.name, file_name=file.name) is None assert fs.get_file(folder_name=folder.name, file_name=file.name) is None
fs.show(full=True) fs.show(full=True)
@@ -33,7 +33,7 @@ def test_folder_delete_request(populated_file_system):
assert folder.get_file_by_id(file_uuid=file.uuid) is not None assert folder.get_file_by_id(file_uuid=file.uuid) is not None
assert fs.get_folder_by_id(folder_uuid=folder.uuid) is not None assert fs.get_folder_by_id(folder_uuid=folder.uuid) is not None
fs.apply_request(request=["delete", "folder", folder.uuid]) fs.apply_request(request=["delete", "folder", folder.name])
assert fs.get_folder_by_id(folder_uuid=folder.uuid) is None assert fs.get_folder_by_id(folder_uuid=folder.uuid) is None
assert fs.get_file_by_id(folder_uuid=folder.uuid, file_uuid=file.uuid) is None assert fs.get_file_by_id(folder_uuid=folder.uuid, file_uuid=file.uuid) is None

8
tests/unit_tests/_primaite/_simulator/_file_system/test_folder_actions.py
View File

@@ -82,7 +82,7 @@ def test_folder_restore_request(populated_file_system):
assert fs.get_file_by_id(folder_uuid=folder.uuid, file_uuid=file.uuid) is not None assert fs.get_file_by_id(folder_uuid=folder.uuid, file_uuid=file.uuid) is not None
# delete folder # delete folder
fs.apply_request(request=["delete", "folder", folder.uuid]) fs.apply_request(request=["delete", "folder", folder.name])
assert fs.get_folder(folder_name=folder.name) is None assert fs.get_folder(folder_name=folder.name) is None
assert fs.get_folder_by_id(folder_uuid=folder.uuid, include_deleted=True).deleted is True assert fs.get_folder_by_id(folder_uuid=folder.uuid, include_deleted=True).deleted is True
@@ -90,7 +90,7 @@ def test_folder_restore_request(populated_file_system):
assert fs.get_file_by_id(folder_uuid=folder.uuid, file_uuid=file.uuid, include_deleted=True).deleted is True assert fs.get_file_by_id(folder_uuid=folder.uuid, file_uuid=file.uuid, include_deleted=True).deleted is True
# restore folder # restore folder
fs.apply_request(request=["restore", "folder", folder.uuid]) fs.apply_request(request=["restore", "folder", folder.name])
fs.apply_timestep(timestep=0) fs.apply_timestep(timestep=0)
assert fs.get_folder(folder_name=folder.name) is not None assert fs.get_folder(folder_name=folder.name) is not None
assert ( assert (
@@ -121,7 +121,7 @@ def test_folder_restore_request(populated_file_system):
assert fs.get_file(folder_name=folder.name, file_name=file.name).health_status == FileSystemItemHealthStatus.CORRUPT assert fs.get_file(folder_name=folder.name, file_name=file.name).health_status == FileSystemItemHealthStatus.CORRUPT
# restore folder # restore folder
fs.apply_request(request=["restore", "folder", folder.uuid]) fs.apply_request(request=["restore", "folder", folder.name])
fs.apply_timestep(timestep=0) fs.apply_timestep(timestep=0)
assert fs.get_folder(folder_name=folder.name).health_status == FileSystemItemHealthStatus.RESTORING assert fs.get_folder(folder_name=folder.name).health_status == FileSystemItemHealthStatus.RESTORING
assert fs.get_file(folder_name=folder.name, file_name=file.name).health_status == FileSystemItemHealthStatus.CORRUPT assert fs.get_file(folder_name=folder.name, file_name=file.name).health_status == FileSystemItemHealthStatus.CORRUPT
@@ -156,7 +156,7 @@ def test_deleted_folder_and_its_files_cannot_be_interacted_with(populated_file_s
fs.apply_request(request=["file", file.name, "corrupt"]) fs.apply_request(request=["file", file.name, "corrupt"])
assert fs.get_file(folder_name=folder.name, file_name=file.name).health_status == FileSystemItemHealthStatus.CORRUPT assert fs.get_file(folder_name=folder.name, file_name=file.name).health_status == FileSystemItemHealthStatus.CORRUPT
fs.apply_request(request=["delete", "folder", folder.uuid]) fs.apply_request(request=["delete", "folder", folder.name])
assert fs.get_file(folder_name=folder.name, file_name=file.name) is None assert fs.get_file(folder_name=folder.name, file_name=file.name) is None
fs.apply_request(request=["file", file.name, "repair"]) fs.apply_request(request=["file", file.name, "repair"])