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

#2248 - synced wth dev

Browse Source
This commit is contained in:
Chris McCarthy
2024-02-08 16:15:57 +00:00
parent 114fb8c077 c3ea06aeec
commit 5eed467271
20 changed files with 607 additions and 804 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
.azure/azure-ci-build-pipeline.yaml
View File

@@ -113,6 +113,7 @@ stages:
testRunner: JUnit
testResultsFiles: 'junit/**.xml'
testRunTitle: 'Publish test results'
failTaskOnFailedTests: true
- publish: $(System.DefaultWorkingDirectory)/htmlcov/
# publish the html report - so we can debug the coverage if needed

1
CHANGELOG.md
View File

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

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

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

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

@@ -26,16 +26,13 @@ the structure:
```
"""
from abc import abstractmethod
from typing import Dict, List, Tuple, Type, TYPE_CHECKING
from typing import Dict, List, Tuple, Type
from primaite import getLogger
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
_LOGGER = getLogger(__name__)
if TYPE_CHECKING:
from primaite.game.game import PrimaiteGame
class AbstractReward:
"""Base class for reward function components."""
@@ -47,13 +44,11 @@ class AbstractReward:
@classmethod
@abstractmethod
def from_config(cls, config: dict, game: "PrimaiteGame") -> "AbstractReward":
def from_config(cls, config: dict) -> "AbstractReward":
"""Create a reward function component from a config dictionary.
:param config: dict of options for the reward component's constructor
:type config: dict
:param game: Reference to the PrimAITE Game object
:type game: PrimaiteGame
:return: The reward component.
:rtype: AbstractReward
"""
@@ -68,13 +63,13 @@ 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: The reward component.
:rtype: DummyReward
"""
return cls()
@@ -126,13 +121,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 +172,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 +193,55 @@ 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.
:param config: Configuration dictionary.
:type config: Dict
"""
node_hostname = config.get("node_hostname")
return cls(node_hostname=node_hostname)
class RewardFunction:
"""Manages the reward function for the agent."""
@@ -209,6 +249,7 @@ class RewardFunction:
"DUMMY": DummyReward,
"DATABASE_FILE_INTEGRITY": DatabaseFileIntegrity,
"WEB_SERVER_404_PENALTY": WebServer404Penalty,
"WEBPAGE_UNAVAILABLE_PENALTY": WebpageUnavailablePenalty,
}
def __init__(self):
@@ -243,13 +284,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 +298,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

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

@@ -346,7 +346,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)
reward_function = RewardFunction.from_config(reward_function_cfg)
# OTHER AGENT SETTINGS
agent_settings = AgentSettings.from_config(agent_cfg.get("agent_settings"))
@@ -358,7 +358,7 @@ class PrimaiteGame:
agent_name=agent_cfg["ref"],
action_space=action_space,
observation_space=obs_space,
reward_function=rew_function,
reward_function=reward_function,
agent_settings=agent_settings,
)
game.agents.append(new_agent)
@@ -367,7 +367,7 @@ class PrimaiteGame:
agent_name=agent_cfg["ref"],
action_space=action_space,
observation_space=obs_space,
reward_function=rew_function,
reward_function=reward_function,
agent_settings=agent_settings,
)
game.agents.append(new_agent)
@@ -377,7 +377,7 @@ class PrimaiteGame:
agent_name=agent_cfg["ref"],
action_space=action_space,
observation_space=obs_space,
reward_function=rew_function,
reward_function=reward_function,
agent_settings=agent_settings,
)
game.agents.append(new_agent)

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

@@ -46,7 +46,7 @@
"source": [
"## Green agent\n",
"\n",
"The green agent is logged onto client 2. It sometimes uses the web browser on client 2 to navigate to `http://arcd.com/users`. The web server replies with a status code 200 if the data is available on the database or 404 if not available."
"There are green agents 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 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,9 @@
"\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",
"\n",
"The file status reward and the two green-agent-related rewards are averaged to get a total step reward.\n"
]
},
{
@@ -326,8 +327,10 @@
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"execution_count": null,
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"%load_ext autoreload\n",
@@ -336,20 +339,11 @@
},
{
"cell_type": "code",
"execution_count": 2,
"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"
]
}
],
"execution_count": null,
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"# Imports\n",
"from primaite.config.load import example_config_path\n",
@@ -370,134 +364,11 @@
},
{
"cell_type": "code",
"execution_count": 3,
"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",
" 'NETWORK_INTERFACES': {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",
" 'NETWORK_INTERFACES': {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",
" 'NETWORK_INTERFACES': {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",
" 'NETWORK_INTERFACES': {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",
" 'NETWORK_INTERFACES': {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",
" 'NETWORK_INTERFACES': {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",
" 'NETWORK_INTERFACES': {1: {'nic_status': 1}, 2: {'nic_status': 0}},\n",
" 'SERVICES': {1: {'health_status': 0, 'operating_status': 0}},\n",
" 'operating_status': 1}}}\n"
]
}
],
"execution_count": null,
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"# create the env\n",
"with open(example_config_path(), 'r') as f:\n",
@@ -523,48 +394,11 @@
},
{
"cell_type": "code",
"execution_count": 4,
"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"
]
}
],
"execution_count": null,
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"for step in range(32):\n",
" obs, reward, terminated, truncated, info = env.step(0)\n",
@@ -580,44 +414,11 @@
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"{1: {'FOLDERS': {1: {'FILES': {1: {'health_status': 0}}, 'health_status': 0}},\n",
" 'NETWORK_INTERFACES': {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",
" 'NETWORK_INTERFACES': {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",
" 'NETWORK_INTERFACES': {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",
" 'NETWORK_INTERFACES': {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",
" 'NETWORK_INTERFACES': {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",
" 'NETWORK_INTERFACES': {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",
" 'NETWORK_INTERFACES': {1: {'nic_status': 1}, 2: {'nic_status': 0}},\n",
" 'SERVICES': {1: {'health_status': 0, 'operating_status': 0}},\n",
" 'operating_status': 1}}\n"
]
}
],
"execution_count": null,
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"pprint(obs['NODES'])"
]
@@ -631,44 +432,11 @@
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"{1: {'FOLDERS': {1: {'FILES': {1: {'health_status': 0}}, 'health_status': 0}},\n",
" 'NETWORK_INTERFACES': {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",
" 'NETWORK_INTERFACES': {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",
" 'NETWORK_INTERFACES': {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",
" 'NETWORK_INTERFACES': {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",
" 'NETWORK_INTERFACES': {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",
" 'NETWORK_INTERFACES': {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",
" 'NETWORK_INTERFACES': {1: {'nic_status': 1}, 2: {'nic_status': 0}},\n",
" 'SERVICES': {1: {'health_status': 0, 'operating_status': 0}},\n",
" 'operating_status': 1}}\n"
]
}
],
"execution_count": null,
"metadata": {
"tags": []
},
"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 +460,16 @@
},
{
"cell_type": "code",
"execution_count": 7,
"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"
]
}
],
"execution_count": null,
"metadata": {
"tags": []
},
"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 +487,17 @@
},
{
"cell_type": "code",
"execution_count": 8,
"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"
]
}
],
"execution_count": null,
"metadata": {
"tags": []
},
"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 +510,11 @@
},
{
"cell_type": "code",
"execution_count": 9,
"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"
]
}
],
"execution_count": null,
"metadata": {
"tags": []
},
"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,106 +541,19 @@
"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": {
"kernelspec": {
"display_name": "venv",
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
@@ -939,9 +567,9 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.12"
"version": "3.8.10"
}
},
"nbformat": 4,
"nbformat_minor": 2
"nbformat_minor": 4
}

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

@@ -94,12 +94,12 @@ class FileSystem(SimComponent):
self._restore_manager.add_request(
name="file",
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(
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(
name="restore",
@@ -209,7 +209,7 @@ class FileSystem(SimComponent):
folder = self.get_folder_by_id(folder_uuid=folder_uuid)
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.
@@ -219,9 +219,13 @@ class FileSystem(SimComponent):
for folder in self.folders.values():
if folder.name == folder_name:
return folder
if include_deleted:
for folder in self.deleted_folders.values():
if folder.name == folder_name:
return folder
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.
@@ -283,7 +287,7 @@ class FileSystem(SimComponent):
self._file_request_manager.add_request(name=file.name, request_type=RequestType(func=file._request_manager))
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.
@@ -291,9 +295,9 @@ class FileSystem(SimComponent):
:param file_name: The name of the file to be retrieved, including its extension.
: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:
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}")
def get_file_by_id(
@@ -455,46 +459,44 @@ class FileSystem(SimComponent):
for folder_id in self.folders:
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.
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
"""
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:
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
self.deleted_folders.pop(folder.uuid, None)
folder.restore()
self.folders[folder.uuid] = folder
if folder.deleted:
self.deleted_folders.pop(folder.uuid)
def restore_file(self, folder_uuid: str, file_uuid: str):
def restore_file(self, folder_name: str, file_name: str):
"""
Restore a 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
:type: folder_uuid: str
:param: folder_name: name of the folder where the file is stored
:type: folder_name: str
:param: file_uuid: id of the file to restore
:type: file_uuid: str
:param: file_name: name of the file to restore
: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:
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:
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
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:
# repair all files
for file_id in self.files:
self.restore_file(file_uuid=file_id)
for file_id, file in self.files.items():
self.restore_file(file_name=file.name)
deleted_files = self.deleted_files.copy()
for file_id in deleted_files:
self.restore_file(file_uuid=file_id)
for file_id, file in deleted_files.items():
self.restore_file(file_name=file.name)
if self.deleted:
self.deleted = False
elif self.health_status in [FileSystemItemHealthStatus.CORRUPT, FileSystemItemHealthStatus.RESTORING]:
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.
@@ -218,6 +218,10 @@ class Folder(FileSystemItemABC):
for file in self.files.values():
if file.name == file_name:
return file
if include_deleted:
for file in self.deleted_files.values():
if file.name == file_name:
return file
return None
def get_file_by_id(self, file_uuid: str, include_deleted: Optional[bool] = False) -> File:
@@ -297,23 +301,23 @@ class Folder(FileSystemItemABC):
self.files = {}
def restore_file(self, file_uuid: str):
def restore_file(self, file_name: str):
"""
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
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:
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
file.restore()
self.files[file.uuid] = file
if file.deleted:
self.deleted_files.pop(file_uuid)
self.deleted_files.pop(file.uuid)
def quarantine(self):
"""Quarantines the File System Folder."""

7
src/primaite/simulator/network/hardware/nodes/host/host_node.py
View File

@@ -153,13 +153,6 @@ class HostARP(ARP):
)
return
# Matched ARP request
# TODO: try taking this out
self.add_arp_cache_entry(
ip_address=arp_packet.sender_ip_address,
mac_address=arp_packet.sender_mac_addr,
network_interface=from_network_interface,
)
arp_packet = arp_packet.generate_reply(from_network_interface.mac_address)
self.send_arp_reply(arp_packet)

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

@@ -1,7 +1,10 @@
from enum import Enum
from ipaddress import IPv4Address
from typing import Dict, Optional
from typing import Dict, List, Optional
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 +36,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 +77,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 +125,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 +136,21 @@ 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,
status=self.BrowserHistoryItem._HistoryItemStatus.LOADED,
response_code=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, status=self.BrowserHistoryItem._HistoryItemStatus.SERVER_UNREACHABLE
)
)
return False
def send(
@@ -172,3 +191,31 @@ 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"""
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}")
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 self.db_file.health_status == FileSystemItemHealthStatus.GOOD:
return {

237
tests/conftest.py
View File

@@ -6,11 +6,13 @@ import pytest
import yaml
from primaite import getLogger, PRIMAITE_PATHS
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.session.session import PrimaiteSession
from primaite.simulator.file_system.file_system import FileSystem
# from primaite.environment.primaite_env import Primaite
# from primaite.primaite_session import PrimaiteSession
from primaite.simulator.network.container import Network
from primaite.simulator.network.hardware.nodes.host.computer import Computer
from primaite.simulator.network.hardware.nodes.host.server import Server
@@ -19,9 +21,14 @@ from primaite.simulator.network.hardware.nodes.network.switch import Switch
from primaite.simulator.network.networks import arcd_uc2_network
from primaite.simulator.network.transmission.network_layer import IPProtocol
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.web_browser import WebBrowser
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.web_server.web_server import WebServer
from tests.mock_and_patch.get_session_path_mock import temp_user_sessions_path
ACTION_SPACE_NODE_VALUES = 1
@@ -282,3 +289,227 @@ def example_network() -> Network:
assert all(link.is_up for link in network.links.values())
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, start_up_duration=0)
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, start_up_duration=0)
switch_1.power_on()
network.connect(endpoint_a=router.network_interface[1], endpoint_b=switch_1.network_interface[6])
router.enable_port(1)
switch_2 = Switch(hostname="switch_2", num_ports=6, start_up_duration=0)
switch_2.power_on()
network.connect(endpoint_a=router.network_interface[2], endpoint_b=switch_2.network_interface[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",
start_up_duration=0,
)
client_1.power_on()
network.connect(
endpoint_a=client_1.network_interface[1],
endpoint_b=switch_1.network_interface[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",
start_up_duration=0,
)
server_1.power_on()
network.connect(endpoint_a=server_1.network_interface[1], endpoint_b=switch_2.network_interface[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",
start_up_duration=0,
)
server_2.power_on()
network.connect(endpoint_a=server_2.network_interface[1], endpoint_b=switch_2.network_interface[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.network_interface[1].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.network_interface[1].ip_address
server_2.software_manager.software.get("DNSClient").dns_server = server_1.network_interface[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.network_interface[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.network_interface[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.network_interface[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)

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

@@ -35,238 +35,6 @@ from primaite.simulator.system.services.web_server.web_server import WebServer
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, start_up_duration=0)
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, start_up_duration=0)
switch_1.power_on()
network.connect(endpoint_a=router.network_interface[1], endpoint_b=switch_1.network_interface[6])
router.enable_port(1)
switch_2 = Switch(
hostname="switch_2",
num_ports=6,
start_up_duration=0,
)
switch_2.power_on()
network.connect(endpoint_a=router.network_interface[2], endpoint_b=switch_2.network_interface[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",
start_up_duration=0,
)
client_1.power_on()
network.connect(
endpoint_a=client_1.network_interface[1],
endpoint_b=switch_1.network_interface[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",
start_up_duration=0,
)
server_1.power_on()
network.connect(endpoint_a=server_1.network_interface[1], endpoint_b=switch_2.network_interface[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",
start_up_duration=0,
)
server_2.power_on()
network.connect(endpoint_a=server_2.network_interface[1], endpoint_b=switch_2.network_interface[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.network_interface[1].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.network_interface[1].ip_address
server_2.software_manager.software.get("DNSClient").dns_server = server_1.network_interface[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.network_interface[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.network_interface[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.network_interface[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]):
"""Test that the DoNothingAction can form a request and that it is accepted by the simulation."""
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 WebpageUnavailablePenalty
from primaite.simulator.network.hardware.nodes.network.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

23
tests/integration_tests/system/test_application_on_node.py
View File

@@ -15,6 +15,7 @@ def populated_node(application_class) -> Tuple[Application, Computer]:
subnet_mask="255.255.255.0",
default_gateway="192.168.1.1",
start_up_duration=0,
shut_down_duration=0,
)
computer.power_on()
computer.software_manager.install(application_class)
@@ -33,6 +34,7 @@ def test_application_on_offline_node(application_class):
subnet_mask="255.255.255.0",
default_gateway="192.168.1.1",
start_up_duration=0,
shut_down_duration=0,
)
computer.software_manager.install(application_class)
@@ -40,9 +42,6 @@ def test_application_on_offline_node(application_class):
computer.power_off()
for i in range(computer.shut_down_duration + 1):
computer.apply_timestep(timestep=i)
assert computer.operating_state is NodeOperatingState.OFF
assert app.operating_state is ApplicationOperatingState.CLOSED
@@ -59,9 +58,6 @@ def test_server_turns_off_application(populated_node):
computer.power_off()
for i in range(computer.shut_down_duration + 1):
computer.apply_timestep(timestep=i)
assert computer.operating_state is NodeOperatingState.OFF
assert app.operating_state is ApplicationOperatingState.CLOSED
@@ -75,9 +71,6 @@ def test_application_cannot_be_turned_on_when_computer_is_off(populated_node):
computer.power_off()
for i in range(computer.shut_down_duration + 1):
computer.apply_timestep(timestep=i)
assert computer.operating_state is NodeOperatingState.OFF
assert app.operating_state is ApplicationOperatingState.CLOSED
@@ -96,28 +89,20 @@ def test_computer_runs_applications(populated_node):
computer.power_off()
for i in range(computer.shut_down_duration + 1):
computer.apply_timestep(timestep=i)
assert computer.operating_state is NodeOperatingState.OFF
assert app.operating_state is ApplicationOperatingState.CLOSED
computer.power_on()
for i in range(computer.start_up_duration + 1):
computer.apply_timestep(timestep=i)
assert computer.operating_state is NodeOperatingState.ON
assert app.operating_state is ApplicationOperatingState.RUNNING
computer.power_off()
for i in range(computer.start_up_duration + 1):
computer.apply_timestep(timestep=i)
assert computer.operating_state is NodeOperatingState.OFF
assert app.operating_state is ApplicationOperatingState.CLOSED
computer.power_on()
for i in range(computer.start_up_duration + 1):
computer.apply_timestep(timestep=i)
assert computer.operating_state is NodeOperatingState.ON
assert app.operating_state is ApplicationOperatingState.RUNNING

30
tests/integration_tests/system/test_service_on_node.py
View File

@@ -12,7 +12,13 @@ from primaite.simulator.system.services.service import Service, ServiceOperating
def populated_node(
service_class,
) -> Tuple[Server, Service]:
server = Server(hostname="server", ip_address="192.168.0.1", subnet_mask="255.255.255.0", start_up_duration=0)
server = Server(
hostname="server",
ip_address="192.168.0.1",
subnet_mask="255.255.255.0",
start_up_duration=0,
shut_down_duration=0,
)
server.power_on()
server.software_manager.install(service_class)
@@ -30,6 +36,7 @@ def test_service_on_offline_node(service_class):
subnet_mask="255.255.255.0",
default_gateway="192.168.1.1",
start_up_duration=0,
shut_down_duration=0,
)
computer.power_on()
computer.software_manager.install(service_class)
@@ -38,9 +45,6 @@ def test_service_on_offline_node(service_class):
computer.power_off()
for i in range(computer.shut_down_duration + 1):
computer.apply_timestep(timestep=i)
assert computer.operating_state is NodeOperatingState.OFF
assert service.operating_state is ServiceOperatingState.STOPPED
@@ -66,9 +70,6 @@ def test_server_turns_off_service(populated_node):
server.power_off()
for i in range(server.shut_down_duration + 1):
server.apply_timestep(timestep=i)
assert server.operating_state is NodeOperatingState.OFF
assert service.operating_state is ServiceOperatingState.STOPPED
@@ -82,9 +83,6 @@ def test_service_cannot_be_turned_on_when_server_is_off(populated_node):
server.power_off()
for i in range(server.shut_down_duration + 1):
server.apply_timestep(timestep=i)
assert server.operating_state is NodeOperatingState.OFF
assert service.operating_state is ServiceOperatingState.STOPPED
@@ -103,28 +101,20 @@ def test_server_turns_on_service(populated_node):
server.power_off()
for i in range(server.shut_down_duration + 1):
server.apply_timestep(timestep=i)
assert server.operating_state is NodeOperatingState.OFF
assert service.operating_state is ServiceOperatingState.STOPPED
server.power_on()
for i in range(server.start_up_duration + 1):
server.apply_timestep(timestep=i)
assert server.operating_state is NodeOperatingState.ON
assert service.operating_state is ServiceOperatingState.RUNNING
server.power_off()
for i in range(server.start_up_duration + 1):
server.apply_timestep(timestep=i)
assert server.operating_state is NodeOperatingState.OFF
assert service.operating_state is ServiceOperatingState.STOPPED
server.power_on()
for i in range(server.start_up_duration + 1):
server.apply_timestep(timestep=i)
assert server.operating_state is NodeOperatingState.ON
assert service.operating_state is ServiceOperatingState.RUNNING

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

@@ -3,6 +3,7 @@ from typing import Tuple
import pytest
from primaite.simulator.network.container import Network
from primaite.simulator.network.hardware.base import Link
from primaite.simulator.network.hardware.nodes.host.computer import Computer
from primaite.simulator.network.hardware.nodes.host.server import Server
@@ -17,7 +18,7 @@ from primaite.simulator.system.services.web_server.web_server import WebServer
@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
router_1: Router = example_network.get_node_by_hostname("router_1")
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 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].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
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_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).deleted is False
fs.apply_request(request=["file", file.name, "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
@@ -95,7 +95,7 @@ def test_deleted_file_cannot_be_interacted_with(populated_file_system):
== 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
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
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
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 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_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
# 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_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
# restore folder
fs.apply_request(request=["restore", "folder", folder.uuid])
fs.apply_request(request=["restore", "folder", folder.name])
fs.apply_timestep(timestep=0)
assert fs.get_folder(folder_name=folder.name) is not None
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
# restore folder
fs.apply_request(request=["restore", "folder", folder.uuid])
fs.apply_request(request=["restore", "folder", folder.name])
fs.apply_timestep(timestep=0)
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
@@ -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"])
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
fs.apply_request(request=["file", file.name, "repair"])