From 94ca13c0f932d118d4282e8dc8b6568718bad6c9 Mon Sep 17 00:00:00 2001 From: Czar Echavez Date: Thu, 18 Apr 2024 15:14:40 +0100 Subject: [PATCH] #2470: add log level via config + test --- src/primaite/session/environment.py | 3 - src/primaite/session/io.py | 9 +- .../configs/basic_switched_network.yaml | 97 +++++++++++++++++++ .../test_io_settings.py | 36 +++++++ 4 files changed, 141 insertions(+), 4 deletions(-) create mode 100644 tests/integration_tests/configuration_file_parsing/test_io_settings.py diff --git a/src/primaite/session/environment.py b/src/primaite/session/environment.py index 9311e1f7..678a444a 100644 --- a/src/primaite/session/environment.py +++ b/src/primaite/session/environment.py @@ -26,9 +26,6 @@ class PrimaiteGymEnv(gymnasium.Env): def __init__(self, game_config: Dict): """Initialise the environment.""" super().__init__() - self.io = PrimaiteIO.from_config(game_config.get("io_settings", {})) - """Handles IO for the environment. This produces sys logs, agent logs, etc.""" - self.game_config: Dict = game_config """PrimaiteGame definition. This can be changed between episodes to enable curriculum learning.""" self.io = PrimaiteIO.from_config(game_config.get("io_settings", {})) diff --git a/src/primaite/session/io.py b/src/primaite/session/io.py index 0b22f784..ffc07e4e 100644 --- a/src/primaite/session/io.py +++ b/src/primaite/session/io.py @@ -6,7 +6,7 @@ from typing import Dict, List, Optional from pydantic import BaseModel, ConfigDict from primaite import getLogger, PRIMAITE_PATHS -from primaite.simulator import SIM_OUTPUT +from primaite.simulator import LogLevel, SIM_OUTPUT _LOGGER = getLogger(__name__) @@ -35,6 +35,8 @@ class PrimaiteIO: """Whether to save system logs.""" write_sys_log_to_terminal: bool = False """Whether to write the sys log to the terminal.""" + log_level: LogLevel = LogLevel.INFO + """The level of log that should be included in the logfiles/logged into terminal.""" def __init__(self, settings: Optional[Settings] = None) -> None: """ @@ -50,6 +52,7 @@ class PrimaiteIO: SIM_OUTPUT.save_pcap_logs = self.settings.save_pcap_logs SIM_OUTPUT.save_sys_logs = self.settings.save_sys_logs SIM_OUTPUT.write_sys_log_to_terminal = self.settings.write_sys_log_to_terminal + SIM_OUTPUT.log_level = self.settings.log_level def generate_session_path(self, timestamp: Optional[datetime] = None) -> Path: """Create a folder for the session and return the path to it.""" @@ -96,6 +99,10 @@ class PrimaiteIO: def from_config(cls, config: Dict) -> "PrimaiteIO": """Create an instance of PrimaiteIO based on a configuration dict.""" config = config or {} + + if config.get("log_level"): + config["log_level"] = LogLevel[config["log_level"].upper()] # convert to enum + new = cls(settings=cls.Settings(**config)) return new diff --git a/tests/assets/configs/basic_switched_network.yaml b/tests/assets/configs/basic_switched_network.yaml index 15dd377e..4e45b008 100644 --- a/tests/assets/configs/basic_switched_network.yaml +++ b/tests/assets/configs/basic_switched_network.yaml @@ -8,6 +8,7 @@ io_settings: save_step_metadata: false save_pcap_logs: true save_sys_logs: true + log_level: WARNING game: @@ -60,6 +61,102 @@ agents: frequency: 4 variance: 3 + + + - ref: defender + team: BLUE + type: ProxyAgent + + observation_space: + type: CUSTOM + options: + components: + - type: NODES + label: NODES + options: + hosts: + - hostname: client_1 + - hostname: client_2 + - hostname: client_3 + num_services: 1 + num_applications: 0 + num_folders: 1 + num_files: 1 + num_nics: 2 + include_num_access: false + include_nmne: true + routers: + - hostname: router_1 + num_ports: 0 + ip_list: + - 192.168.10.21 + - 192.168.10.22 + - 192.168.10.23 + wildcard_list: + - 0.0.0.1 + port_list: + - 80 + - 5432 + protocol_list: + - ICMP + - TCP + - UDP + num_rules: 10 + + - type: LINKS + label: LINKS + options: + link_references: + - switch_1:eth-1<->client_1:eth-1 + - switch_1:eth-2<->client_2:eth-1 + - type: "NONE" + label: ICS + options: {} + + action_space: + action_list: + - type: DONOTHING + + action_map: + 0: + action: DONOTHING + options: {} + options: + nodes: + - node_name: switch + - node_name: client_1 + - node_name: client_2 + - node_name: client_3 + max_folders_per_node: 2 + max_files_per_folder: 2 + max_services_per_node: 2 + max_nics_per_node: 8 + max_acl_rules: 10 + ip_list: + - 192.168.10.21 + - 192.168.10.22 + - 192.168.10.23 + + reward_function: + reward_components: + - type: DATABASE_FILE_INTEGRITY + weight: 0.5 + options: + node_hostname: database_server + folder_name: database + file_name: database.db + + + - type: WEB_SERVER_404_PENALTY + weight: 0.5 + options: + node_hostname: web_server + service_name: web_server_web_service + + + agent_settings: + flatten_obs: true + simulation: network: nodes: diff --git a/tests/integration_tests/configuration_file_parsing/test_io_settings.py b/tests/integration_tests/configuration_file_parsing/test_io_settings.py new file mode 100644 index 00000000..83df31ff --- /dev/null +++ b/tests/integration_tests/configuration_file_parsing/test_io_settings.py @@ -0,0 +1,36 @@ +from pathlib import Path +from typing import Union + +import yaml + +from primaite.config.load import data_manipulation_config_path +from primaite.game.game import PrimaiteGame +from primaite.session.environment import PrimaiteGymEnv +from primaite.simulator import LogLevel +from tests import TEST_ASSETS_ROOT + +BASIC_CONFIG = TEST_ASSETS_ROOT / "configs/basic_switched_network.yaml" + + +def load_config(config_path: Union[str, Path]) -> PrimaiteGame: + """Returns a PrimaiteGame object which loads the contents of a given yaml path.""" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + return PrimaiteGame.from_config(cfg) + + +def test_io_settings(): + """Test that the io_settings are loaded correctly.""" + with open(BASIC_CONFIG, "r") as f: + cfg = yaml.safe_load(f) + env = PrimaiteGymEnv(game_config=cfg) + + assert env.io.settings is not None + + assert env.io.settings.log_level is LogLevel.WARNING + assert env.io.settings.save_pcap_logs + assert env.io.settings.save_sys_logs + assert env.io.settings.save_step_metadata is False + + assert env.io.settings.write_sys_log_to_terminal is False # false by default