From a7492ba39f33c4000ec9bc239c72fcb8787592f1 Mon Sep 17 00:00:00 2001 From: Czar Echavez Date: Mon, 29 Apr 2024 13:45:10 +0100 Subject: [PATCH 01/37] #2447: cli commands for dev mode --- .gitignore | 1 + docs/source/getting_started.rst | 4 +- src/primaite/cli.py | 60 ++----- src/primaite/session/io.py | 12 +- .../setup/_package_data/primaite_config.yaml | 7 +- src/primaite/simulator/__init__.py | 21 ++- src/primaite/utils/cli/__init__.py | 0 src/primaite/utils/cli/dev_cli.py | 151 ++++++++++++++++++ .../utils/cli/primaite_config_utils.py | 26 +++ src/primaite/utils/primaite_config_utils.py | 11 -- 10 files changed, 226 insertions(+), 67 deletions(-) create mode 100644 src/primaite/utils/cli/__init__.py create mode 100644 src/primaite/utils/cli/dev_cli.py create mode 100644 src/primaite/utils/cli/primaite_config_utils.py delete mode 100644 src/primaite/utils/primaite_config_utils.py diff --git a/.gitignore b/.gitignore index b3d9682a..5cd39f24 100644 --- a/.gitignore +++ b/.gitignore @@ -82,6 +82,7 @@ target/ # Jupyter Notebook .ipynb_checkpoints +PPO_UC2/ # IPython profile_default/ diff --git a/docs/source/getting_started.rst b/docs/source/getting_started.rst index d88be5c9..dccfd8aa 100644 --- a/docs/source/getting_started.rst +++ b/docs/source/getting_started.rst @@ -161,9 +161,9 @@ To set PrimAITE to run in development mode: .. code-block:: bash :caption: Unix - primaite mode --dev + primaite dev-mode enable .. code-block:: powershell :caption: Windows (Powershell) - primaite mode --dev + primaite dev-mode enable diff --git a/src/primaite/cli.py b/src/primaite/cli.py index b65a6c97..54f29067 100644 --- a/src/primaite/cli.py +++ b/src/primaite/cli.py @@ -9,9 +9,12 @@ import typer import yaml from typing_extensions import Annotated -from primaite import PRIMAITE_PATHS +from primaite import _PRIMAITE_ROOT, PRIMAITE_PATHS +from primaite.utils.cli import dev_cli +from primaite.utils.cli.primaite_config_utils import get_primaite_config_dict, update_primaite_config app = typer.Typer(no_args_is_help=True) +app.add_typer(dev_cli.dev, name="dev-mode") @app.command() @@ -113,48 +116,15 @@ def setup(overwrite_existing: bool = True) -> None: _LOGGER.info("Rebuilding the example notebooks...") reset_example_configs.run(overwrite_existing=True) + _LOGGER.info("Setting default simulation output") + config_dict = get_primaite_config_dict() + + if config_dict is None: + return + + config_dict["developer_mode"]["output_dir"] = str(_PRIMAITE_ROOT.parent.parent / "simulation_output") + print(f"PrimAITE dev-mode config updated output_dir {config_dict['developer_mode']['output_dir']}") + # update application config + update_primaite_config(config_dict) + _LOGGER.info("PrimAITE setup complete!") - - -@app.command() -def mode( - dev: Annotated[bool, typer.Option("--dev", help="Activates PrimAITE developer mode")] = None, - prod: Annotated[bool, typer.Option("--prod", help="Activates PrimAITE production mode")] = None, -) -> None: - """ - Switch PrimAITE between developer mode and production mode. - - By default, PrimAITE will be in production mode. - - To view the current mode, use: primaite mode - - To set to development mode, use: primaite mode --dev - - To return to production mode, use: primaite mode --prod - """ - if PRIMAITE_PATHS.app_config_file_path.exists(): - with open(PRIMAITE_PATHS.app_config_file_path, "r") as file: - primaite_config = yaml.safe_load(file) - if dev and prod: - print("Unable to activate developer and production modes concurrently.") - return - - if (dev is None) and (prod is None): - is_dev_mode = primaite_config["developer_mode"] - - if is_dev_mode: - print("PrimAITE is running in developer mode.") - else: - print("PrimAITE is running in production mode.") - if dev: - # activate dev mode - primaite_config["developer_mode"] = True - with open(PRIMAITE_PATHS.app_config_file_path, "w") as file: - yaml.dump(primaite_config, file) - print("PrimAITE is running in developer mode.") - if prod: - # activate prod mode - primaite_config["developer_mode"] = False - with open(PRIMAITE_PATHS.app_config_file_path, "w") as file: - yaml.dump(primaite_config, file) - print("PrimAITE is running in production mode.") diff --git a/src/primaite/session/io.py b/src/primaite/session/io.py index 22001fd2..65e900a8 100644 --- a/src/primaite/session/io.py +++ b/src/primaite/session/io.py @@ -5,9 +5,9 @@ from typing import Dict, List, Optional from pydantic import BaseModel, ConfigDict -from primaite import getLogger, PRIMAITE_PATHS +from primaite import _PRIMAITE_ROOT, getLogger, PRIMAITE_PATHS from primaite.simulator import LogLevel, SIM_OUTPUT -from src.primaite.utils.primaite_config_utils import is_dev_mode +from primaite.utils.cli.primaite_config_utils import get_primaite_config_dict, is_dev_mode _LOGGER = getLogger(__name__) @@ -64,8 +64,12 @@ class PrimaiteIO: # check if running in dev mode if is_dev_mode(): - # if dev mode, simulation output will be the current working directory - session_path = Path.cwd() / "simulation_output" / date_str / time_str + # if dev mode, simulation output will be the repository root or whichever path is configured + app_config = get_primaite_config_dict() + if app_config["developer_mode"]["output_dir"] is not None: + session_path = app_config["developer_mode"]["output_path"] + else: + session_path = _PRIMAITE_ROOT.parent.parent / "simulation_output" / date_str / time_str else: session_path = PRIMAITE_PATHS.user_sessions_path / date_str / time_str diff --git a/src/primaite/setup/_package_data/primaite_config.yaml b/src/primaite/setup/_package_data/primaite_config.yaml index f80f4d8a..84619f17 100644 --- a/src/primaite/setup/_package_data/primaite_config.yaml +++ b/src/primaite/setup/_package_data/primaite_config.yaml @@ -1,6 +1,11 @@ # The main PrimAITE application config file -developer_mode: False # false by default +developer_mode: + enabled: False # not enabled by default + output_sys_logs: False # system logs not output by default + output_pcap_logs: False # pcap logs not output by default + output_to_terminal: False # do not output to terminal by default + output_dir: null # none by default - none will print to repository root # Logging logging: diff --git a/src/primaite/simulator/__init__.py b/src/primaite/simulator/__init__.py index 3f371ee5..6758dab6 100644 --- a/src/primaite/simulator/__init__.py +++ b/src/primaite/simulator/__init__.py @@ -7,6 +7,8 @@ from primaite import _PRIMAITE_ROOT __all__ = ["SIM_OUTPUT"] +from primaite.utils.cli.primaite_config_utils import get_primaite_config_dict, is_dev_mode + class LogLevel(IntEnum): """Enum containing all the available log levels for PrimAITE simulation output.""" @@ -24,18 +26,29 @@ class LogLevel(IntEnum): class _SimOutput: + _default_path = _PRIMAITE_ROOT.parent.parent / "simulation_output" / datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + def __init__(self): - self._path: Path = ( - _PRIMAITE_ROOT.parent.parent / "simulation_output" / datetime.now().strftime("%Y-%m-%d_%H-%M-%S") - ) + self._path: Path = self._default_path self.save_pcap_logs: bool = False self.save_sys_logs: bool = False self.write_sys_log_to_terminal: bool = False self.sys_log_level: LogLevel = LogLevel.WARNING # default log level is at WARNING + if is_dev_mode(): + # if dev mode, override with the values configured via the primaite dev-mode command + dev_config = get_primaite_config_dict().get("developer_mode") + self.save_pcap_logs = dev_config["output_pcap_logs"] + self.save_sys_logs = dev_config["output_sys_logs"] + self.write_sys_log_to_terminal = dev_config["output_to_terminal"] + @property def path(self) -> Path: - return self._path + if not is_dev_mode(): + return self._path + if is_dev_mode(): + dev_config = get_primaite_config_dict().get("developer_mode") + return Path(dev_config["output_path"]) if dev_config["output_path"] else self._default_path @path.setter def path(self, new_path: Path) -> None: diff --git a/src/primaite/utils/cli/__init__.py b/src/primaite/utils/cli/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/primaite/utils/cli/dev_cli.py b/src/primaite/utils/cli/dev_cli.py new file mode 100644 index 00000000..a17efa46 --- /dev/null +++ b/src/primaite/utils/cli/dev_cli.py @@ -0,0 +1,151 @@ +import typer +from rich import print +from typing_extensions import Annotated + +from primaite import _PRIMAITE_ROOT +from primaite.utils.cli.primaite_config_utils import get_primaite_config_dict, is_dev_mode, update_primaite_config + +dev = typer.Typer() + +PRODUCTION_MODE_MESSAGE = ( + "\n[green]:monkey_face::monkey_face::monkey_face: " + " PrimAITE is running in Production mode " + " :monkey_face::monkey_face::monkey_face: [/green]\n" +) + +DEVELOPMENT_MODE_MESSAGE = ( + "\n[yellow] :construction::construction::construction: " + " PrimAITE is running in Development mode " + " :construction::construction::construction: [/yellow]\n" +) + + +def dev_mode(): + """ + CLI commands relevant to the dev-mode for PrimAITE. + + The dev-mode contains tools that help with the ease of developing or debugging PrimAITE. + + By default, PrimAITE will be in production mode. + + To enable development mode, use `primaite dev-mode enable` + """ + + +@dev.command() +def show(): + """Show if PrimAITE is in development mode or production mode.""" + # print if dev mode is enabled + print(DEVELOPMENT_MODE_MESSAGE if is_dev_mode() else PRODUCTION_MODE_MESSAGE) + print("\nTo see available options, use [medium_turquoise]`primaite dev-mode --help`[/medium_turquoise]\n") + + +@dev.command() +def enable(): + """Enable the development mode for PrimAITE.""" + config_dict = get_primaite_config_dict() + + if config_dict is None: + return + + # enable dev mode + config_dict["developer_mode"]["enabled"] = True + update_primaite_config(config_dict) + print(DEVELOPMENT_MODE_MESSAGE) + + +@dev.command() +def disable(): + """Disable the development mode for PrimAITE.""" + config_dict = get_primaite_config_dict() + + if config_dict is None: + return + + # disable dev mode + config_dict["developer_mode"]["enabled"] = False + update_primaite_config(config_dict) + print(PRODUCTION_MODE_MESSAGE) + + +def config_callback( + ctx: typer.Context, + output_sys_logs: Annotated[ + bool, typer.Option("--output-sys-logs/--no-sys-logs", "-sys/-nsys", help="Output system logs to file.") + ] = None, + output_pcap_logs: Annotated[ + bool, + typer.Option( + "--output-pcap-logs/--no-pcap-logs", "-pcap/-npcap", help="Output network packet capture logs to file." + ), + ] = None, + output_to_terminal: Annotated[ + bool, typer.Option("--output-to_terminal/--no-terminal", "-t/-nt", help="Output system logs to terminal.") + ] = None, +): + """Configure the development tools and environment.""" + config_dict = get_primaite_config_dict() + + if config_dict is None: + return + + if output_sys_logs is not None: + config_dict["developer_mode"]["output_sys_logs"] = output_sys_logs + print(f"PrimAITE dev-mode config updated {output_sys_logs=}") + + if output_pcap_logs is not None: + config_dict["developer_mode"]["output_pcap_logs"] = output_pcap_logs + print(f"PrimAITE dev-mode config updated {output_pcap_logs=}") + + if output_to_terminal is not None: + config_dict["developer_mode"]["output_to_terminal"] = output_to_terminal + print(f"PrimAITE dev-mode config updated {output_to_terminal=}") + + # update application config + update_primaite_config(config_dict) + + +config_typer = typer.Typer(callback=config_callback, name="config", invoke_without_command=True) +dev.add_typer(config_typer) + + +@config_typer.command() +def path( + directory: Annotated[ + str, + typer.Argument( + help="Directory where the system logs and PCAP logs will be output. By default, this will be where the" + "root of the PrimAITE repository is located.", + show_default=False, + ), + ] = None, + default: Annotated[ + bool, + typer.Option( + "--default", + "-root", + help="Set PrimAITE to output system logs and pcap logs to the PrimAITE repository root.", + ), + ] = None, +): + """Set the output directory for the PrimAITE system and PCAP logs.""" + config_dict = get_primaite_config_dict() + + if config_dict is None: + return + + if default: + config_dict["developer_mode"]["output_dir"] = None + print( + f"PrimAITE dev-mode config updated output directory will be in " + f"{str(_PRIMAITE_ROOT.parent.parent / 'simulation_output')}" + ) + # update application config + update_primaite_config(config_dict) + return + + if directory: + config_dict["developer_mode"]["output_dir"] = directory + print(f"PrimAITE dev-mode config updated output_dir={directory}") + # update application config + update_primaite_config(config_dict) diff --git a/src/primaite/utils/cli/primaite_config_utils.py b/src/primaite/utils/cli/primaite_config_utils.py new file mode 100644 index 00000000..e78a5f86 --- /dev/null +++ b/src/primaite/utils/cli/primaite_config_utils.py @@ -0,0 +1,26 @@ +from typing import Dict + +import yaml + +from primaite import PRIMAITE_PATHS + + +def get_primaite_config_dict() -> Dict: + """Returns a dict containing the PrimAITE application config.""" + if PRIMAITE_PATHS.app_config_file_path.exists(): + with open(PRIMAITE_PATHS.app_config_file_path, "r") as file: + return yaml.safe_load(file) + else: + print("PrimAITE application config was not found. Have you run [bold red]primaite setup[/bold red]?") + + +def is_dev_mode() -> bool: + """Returns True if PrimAITE is currently running in developer mode.""" + config = get_primaite_config_dict() + return config["developer_mode"]["enabled"] + + +def update_primaite_config(config: Dict) -> None: + """Update the PrimAITE application config file.""" + with open(PRIMAITE_PATHS.app_config_file_path, "w") as file: + yaml.dump(config, file) diff --git a/src/primaite/utils/primaite_config_utils.py b/src/primaite/utils/primaite_config_utils.py deleted file mode 100644 index 70a7e4ba..00000000 --- a/src/primaite/utils/primaite_config_utils.py +++ /dev/null @@ -1,11 +0,0 @@ -import yaml - -from primaite import PRIMAITE_PATHS - - -def is_dev_mode() -> bool: - """Returns True if PrimAITE is currently running in developer mode.""" - if PRIMAITE_PATHS.app_config_file_path.exists(): - with open(PRIMAITE_PATHS.app_config_file_path, "r") as file: - primaite_config = yaml.safe_load(file) - return primaite_config["developer_mode"] From 59990813f53bbbc6e2b8150909bebfac65cab217 Mon Sep 17 00:00:00 2001 From: Czar Echavez Date: Tue, 30 Apr 2024 15:43:25 +0100 Subject: [PATCH 02/37] #2533: added tests + log level for output --- src/primaite/__init__.py | 5 +- src/primaite/cli.py | 14 +- .../setup/_package_data/primaite_config.yaml | 1 + src/primaite/simulator/__init__.py | 2 +- src/primaite/utils/cli/dev_cli.py | 108 ++++++---- .../utils/cli/primaite_config_utils.py | 25 ++- tests/integration_tests/cli/__init__.py | 11 ++ tests/integration_tests/cli/test_dev_cli.py | 186 ++++++++++++++++++ 8 files changed, 294 insertions(+), 58 deletions(-) create mode 100644 tests/integration_tests/cli/__init__.py create mode 100644 tests/integration_tests/cli/test_dev_cli.py diff --git a/src/primaite/__init__.py b/src/primaite/__init__.py index c58f0103..59725ac7 100644 --- a/src/primaite/__init__.py +++ b/src/primaite/__init__.py @@ -123,9 +123,8 @@ PRIMAITE_PATHS: Final[_PrimaitePaths] = _PrimaitePaths() def _host_primaite_config() -> None: - if not PRIMAITE_PATHS.app_config_file_path.exists(): - pkg_config_path = Path(pkg_resources.resource_filename("primaite", "setup/_package_data/primaite_config.yaml")) - shutil.copy2(pkg_config_path, PRIMAITE_PATHS.app_config_file_path) + pkg_config_path = Path(pkg_resources.resource_filename("primaite", "setup/_package_data/primaite_config.yaml")) + shutil.copy(pkg_config_path, PRIMAITE_PATHS.app_config_file_path) _host_primaite_config() diff --git a/src/primaite/cli.py b/src/primaite/cli.py index 54f29067..6ad019dd 100644 --- a/src/primaite/cli.py +++ b/src/primaite/cli.py @@ -9,9 +9,8 @@ import typer import yaml from typing_extensions import Annotated -from primaite import _PRIMAITE_ROOT, PRIMAITE_PATHS +from primaite import PRIMAITE_PATHS from primaite.utils.cli import dev_cli -from primaite.utils.cli.primaite_config_utils import get_primaite_config_dict, update_primaite_config app = typer.Typer(no_args_is_help=True) app.add_typer(dev_cli.dev, name="dev-mode") @@ -116,15 +115,4 @@ def setup(overwrite_existing: bool = True) -> None: _LOGGER.info("Rebuilding the example notebooks...") reset_example_configs.run(overwrite_existing=True) - _LOGGER.info("Setting default simulation output") - config_dict = get_primaite_config_dict() - - if config_dict is None: - return - - config_dict["developer_mode"]["output_dir"] = str(_PRIMAITE_ROOT.parent.parent / "simulation_output") - print(f"PrimAITE dev-mode config updated output_dir {config_dict['developer_mode']['output_dir']}") - # update application config - update_primaite_config(config_dict) - _LOGGER.info("PrimAITE setup complete!") diff --git a/src/primaite/setup/_package_data/primaite_config.yaml b/src/primaite/setup/_package_data/primaite_config.yaml index 84619f17..c1caf1f4 100644 --- a/src/primaite/setup/_package_data/primaite_config.yaml +++ b/src/primaite/setup/_package_data/primaite_config.yaml @@ -2,6 +2,7 @@ developer_mode: enabled: False # not enabled by default + sys_log_level: DEBUG # level of output for system logs, DEBUG by default output_sys_logs: False # system logs not output by default output_pcap_logs: False # pcap logs not output by default output_to_terminal: False # do not output to terminal by default diff --git a/src/primaite/simulator/__init__.py b/src/primaite/simulator/__init__.py index 6758dab6..985bb663 100644 --- a/src/primaite/simulator/__init__.py +++ b/src/primaite/simulator/__init__.py @@ -48,7 +48,7 @@ class _SimOutput: return self._path if is_dev_mode(): dev_config = get_primaite_config_dict().get("developer_mode") - return Path(dev_config["output_path"]) if dev_config["output_path"] else self._default_path + return Path(dev_config["output_dir"]) if dev_config["output_dir"] else self._default_path @path.setter def path(self, new_path: Path) -> None: diff --git a/src/primaite/utils/cli/dev_cli.py b/src/primaite/utils/cli/dev_cli.py index a17efa46..886f1bf1 100644 --- a/src/primaite/utils/cli/dev_cli.py +++ b/src/primaite/utils/cli/dev_cli.py @@ -1,8 +1,12 @@ +from typing import List + +import click import typer from rich import print from typing_extensions import Annotated from primaite import _PRIMAITE_ROOT +from primaite.simulator import LogLevel from primaite.utils.cli.primaite_config_utils import get_primaite_config_dict, is_dev_mode, update_primaite_config dev = typer.Typer() @@ -69,19 +73,44 @@ def disable(): def config_callback( - ctx: typer.Context, - output_sys_logs: Annotated[ - bool, typer.Option("--output-sys-logs/--no-sys-logs", "-sys/-nsys", help="Output system logs to file.") - ] = None, - output_pcap_logs: Annotated[ - bool, - typer.Option( - "--output-pcap-logs/--no-pcap-logs", "-pcap/-npcap", help="Output network packet capture logs to file." - ), - ] = None, - output_to_terminal: Annotated[ - bool, typer.Option("--output-to_terminal/--no-terminal", "-t/-nt", help="Output system logs to terminal.") - ] = None, + ctx: typer.Context, + sys_log_level: Annotated[ + LogLevel, + typer.Option( + "--sys-log-level", + "-level", + click_type=click.Choice(LogLevel._member_names_, case_sensitive=False), + help="The level of system logs to output.", + show_default=False + ) + ] = None, + output_sys_logs: Annotated[ + bool, + typer.Option( + "--output-sys-logs/--no-sys-logs", + "-sys/-nsys", + help="Output system logs to file.", + show_default=False + ) + ] = None, + output_pcap_logs: Annotated[ + bool, + typer.Option( + "--output-pcap-logs/--no-pcap-logs", + "-pcap/-npcap", + help="Output network packet capture logs to file.", + show_default=False + ), + ] = None, + output_to_terminal: Annotated[ + bool, + typer.Option( + "--output-to-terminal/--no-terminal", + "-t/-nt", + help="Output system logs to terminal.", + show_default=False + ) + ] = None, ): """Configure the development tools and environment.""" config_dict = get_primaite_config_dict() @@ -89,6 +118,11 @@ def config_callback( if config_dict is None: return + if ctx.params.get("sys_log_level") is not None: + config_dict["developer_mode"]["sys_log_level"] = ctx.params.get("sys_log_level") + print(f"PrimAITE dev-mode config updated sys_log_level={ctx.params.get('sys_log_level')}") + + if output_sys_logs is not None: config_dict["developer_mode"]["output_sys_logs"] = output_sys_logs print(f"PrimAITE dev-mode config updated {output_sys_logs=}") @@ -105,28 +139,33 @@ def config_callback( update_primaite_config(config_dict) -config_typer = typer.Typer(callback=config_callback, name="config", invoke_without_command=True) +config_typer = typer.Typer( + callback=config_callback, + name="config", + no_args_is_help=True, + invoke_without_command=True, +) dev.add_typer(config_typer) @config_typer.command() def path( - directory: Annotated[ - str, - typer.Argument( - help="Directory where the system logs and PCAP logs will be output. By default, this will be where the" - "root of the PrimAITE repository is located.", - show_default=False, - ), - ] = None, - default: Annotated[ - bool, - typer.Option( - "--default", - "-root", - help="Set PrimAITE to output system logs and pcap logs to the PrimAITE repository root.", - ), - ] = None, + directory: Annotated[ + str, + typer.Argument( + help="Directory where the system logs and PCAP logs will be output. By default, this will be where the" + "root of the PrimAITE repository is located.", + show_default=False, + ), + ] = None, + default: Annotated[ + bool, + typer.Option( + "--default", + "-root", + help="Set PrimAITE to output system logs and pcap logs to the PrimAITE repository root.", + ), + ] = None, ): """Set the output directory for the PrimAITE system and PCAP logs.""" config_dict = get_primaite_config_dict() @@ -136,16 +175,15 @@ def path( if default: config_dict["developer_mode"]["output_dir"] = None - print( - f"PrimAITE dev-mode config updated output directory will be in " - f"{str(_PRIMAITE_ROOT.parent.parent / 'simulation_output')}" - ) # update application config update_primaite_config(config_dict) + print(f"PrimAITE dev-mode output_dir [medium_turquoise]" + f"{str(_PRIMAITE_ROOT.parent.parent / 'simulation_output')}" + f"[/medium_turquoise]") return if directory: config_dict["developer_mode"]["output_dir"] = directory - print(f"PrimAITE dev-mode config updated output_dir={directory}") # update application config update_primaite_config(config_dict) + print(f"PrimAITE dev-mode output_dir [medium_turquoise]{directory}[/medium_turquoise]") diff --git a/src/primaite/utils/cli/primaite_config_utils.py b/src/primaite/utils/cli/primaite_config_utils.py index e78a5f86..2eb590ad 100644 --- a/src/primaite/utils/cli/primaite_config_utils.py +++ b/src/primaite/utils/cli/primaite_config_utils.py @@ -1,17 +1,30 @@ -from typing import Dict +from pathlib import Path +from typing import Dict, Optional import yaml from primaite import PRIMAITE_PATHS -def get_primaite_config_dict() -> Dict: - """Returns a dict containing the PrimAITE application config.""" - if PRIMAITE_PATHS.app_config_file_path.exists(): - with open(PRIMAITE_PATHS.app_config_file_path, "r") as file: +def get_primaite_config_dict( + config_path: Optional[Path] = None +) -> Dict: + """ + Returns a dict containing the PrimAITE application config. + + :param: config_path: takes in a path object - leave empty to use the default app config path + """ + err_msg = "PrimAITE application config could not be loaded." + + if config_path is None: + config_path = PRIMAITE_PATHS.app_config_file_path + err_msg = "PrimAITE application config was not found. Have you run `primaite setup`?" + + if config_path.exists(): + with open(config_path, "r") as file: return yaml.safe_load(file) else: - print("PrimAITE application config was not found. Have you run [bold red]primaite setup[/bold red]?") + print(err_msg) def is_dev_mode() -> bool: diff --git a/tests/integration_tests/cli/__init__.py b/tests/integration_tests/cli/__init__.py new file mode 100644 index 00000000..07487650 --- /dev/null +++ b/tests/integration_tests/cli/__init__.py @@ -0,0 +1,11 @@ +from typing import List + +from typer.testing import CliRunner, Result + +from primaite.cli import app + + +def cli(args: List[str]) -> Result: + """Pass in a list of arguments and it will return the result.""" + runner = CliRunner() + return runner.invoke(app, args) diff --git a/tests/integration_tests/cli/test_dev_cli.py b/tests/integration_tests/cli/test_dev_cli.py new file mode 100644 index 00000000..cd697711 --- /dev/null +++ b/tests/integration_tests/cli/test_dev_cli.py @@ -0,0 +1,186 @@ +import os +import shutil +import tempfile +from pathlib import Path + +import pytest + +from primaite import PRIMAITE_PATHS, _PRIMAITE_ROOT +from primaite.utils.cli.primaite_config_utils import get_primaite_config_dict +from tests.integration_tests.cli import cli + + +@pytest.fixture(autouse=True) +def test_setup(): + """ + Setup this test by copying the + """ + original_config_path = PRIMAITE_PATHS.app_config_file_path # keep track of app config before test + + temp_dir = tempfile.gettempdir() + temp_config = Path(temp_dir) / "primaite_config.yaml" + shutil.copyfile( + _PRIMAITE_ROOT / "setup" / "_package_data" / "primaite_config.yaml", temp_config + ) # copy the default primaite config to temp directory + PRIMAITE_PATHS.app_config_file_path = temp_config # use the copy for the test + yield # run test + os.remove(temp_config) # clean up temp file + PRIMAITE_PATHS.app_config_file_path = original_config_path # restore app conf because other devs will yell at me + + +def test_dev_mode_enable_disable(): + """Test dev mode enable and disable.""" + # check defaults + config_dict = get_primaite_config_dict() + assert config_dict["developer_mode"]["enabled"] is False # not enabled by default + + result = cli(["dev-mode", "show"]) + assert "Production" in result.output # should print that it is in Production mode by default + + result = cli(["dev-mode", "enable"]) + + assert "Development" in result.output # should print that it is in Development mode + + config_dict = get_primaite_config_dict() + assert config_dict["developer_mode"]["enabled"] # config should reflect that dev mode is enabled + + result = cli(["dev-mode", "show"]) + assert "Development" in result.output # should print that it is in Development mode + + result = cli(["dev-mode", "disable"]) + + assert "Production" in result.output # should print that it is in Production mode + + config_dict = get_primaite_config_dict() + assert config_dict["developer_mode"]["enabled"] is False # config should reflect that dev mode is disabled + + result = cli(["dev-mode", "show"]) + assert "Production" in result.output # should print that it is in Production mode + + +def test_dev_mode_config_sys_log_level(): + """Check that the system log level can be changed via CLI.""" + # check defaults + config_dict = get_primaite_config_dict() + assert config_dict["developer_mode"]["sys_log_level"] == "DEBUG" # DEBUG by default + + result = cli(["dev-mode", "config", "-level", "WARNING"]) + + assert "sys_log_level=WARNING" in result.output # should print correct value + + config_dict = get_primaite_config_dict() + # config should reflect that log level is WARNING + assert config_dict["developer_mode"]["sys_log_level"] == "WARNING" + + result = cli(["dev-mode", "config", "--sys-log-level", "INFO"]) + + assert "sys_log_level=INFO" in result.output # should print correct value + + config_dict = get_primaite_config_dict() + # config should reflect that log level is WARNING + assert config_dict["developer_mode"]["sys_log_level"] == "INFO" + + +def test_dev_mode_config_sys_logs_enable_disable(): + """Test that the system logs output can be enabled or disabled.""" + # check defaults + config_dict = get_primaite_config_dict() + assert config_dict["developer_mode"]["output_sys_logs"] is False # False by default + + result = cli(["dev-mode", "config", "--output-sys-logs"]) + assert "output_sys_logs=True" in result.output # should print correct value + + config_dict = get_primaite_config_dict() + # config should reflect that output_sys_logs is True + assert config_dict["developer_mode"]["output_sys_logs"] + + result = cli(["dev-mode", "config", "--no-sys-logs"]) + assert "output_sys_logs=False" in result.output # should print correct value + + config_dict = get_primaite_config_dict() + # config should reflect that output_sys_logs is True + assert config_dict["developer_mode"]["output_sys_logs"] is False + + result = cli(["dev-mode", "config", "-sys"]) + assert "output_sys_logs=True" in result.output # should print correct value + + config_dict = get_primaite_config_dict() + # config should reflect that output_sys_logs is True + assert config_dict["developer_mode"]["output_sys_logs"] + + result = cli(["dev-mode", "config", "-nsys"]) + assert "output_sys_logs=False" in result.output # should print correct value + + config_dict = get_primaite_config_dict() + # config should reflect that output_sys_logs is True + assert config_dict["developer_mode"]["output_sys_logs"] is False + + +def test_dev_mode_config_pcap_logs_enable_disable(): + """Test that the pcap logs output can be enabled or disabled.""" + # check defaults + config_dict = get_primaite_config_dict() + assert config_dict["developer_mode"]["output_pcap_logs"] is False # False by default + + result = cli(["dev-mode", "config", "--output-pcap-logs"]) + assert "output_pcap_logs=True" in result.output # should print correct value + + config_dict = get_primaite_config_dict() + # config should reflect that output_pcap_logs is True + assert config_dict["developer_mode"]["output_pcap_logs"] + + result = cli(["dev-mode", "config", "--no-pcap-logs"]) + assert "output_pcap_logs=False" in result.output # should print correct value + + config_dict = get_primaite_config_dict() + # config should reflect that output_pcap_logs is True + assert config_dict["developer_mode"]["output_pcap_logs"] is False + + result = cli(["dev-mode", "config", "-pcap"]) + assert "output_pcap_logs=True" in result.output # should print correct value + + config_dict = get_primaite_config_dict() + # config should reflect that output_pcap_logs is True + assert config_dict["developer_mode"]["output_pcap_logs"] + + result = cli(["dev-mode", "config", "-npcap"]) + assert "output_pcap_logs=False" in result.output # should print correct value + + config_dict = get_primaite_config_dict() + # config should reflect that output_pcap_logs is True + assert config_dict["developer_mode"]["output_pcap_logs"] is False + + +def test_dev_mode_config_output_to_terminal_enable_disable(): + """Test that the output to terminal can be enabled or disabled.""" + # check defaults + config_dict = get_primaite_config_dict() + assert config_dict["developer_mode"]["output_to_terminal"] is False # False by default + + result = cli(["dev-mode", "config", "--output-to-terminal"]) + assert "output_to_terminal=True" in result.output # should print correct value + + config_dict = get_primaite_config_dict() + # config should reflect that output_to_terminal is True + assert config_dict["developer_mode"]["output_to_terminal"] + + result = cli(["dev-mode", "config", "--no-terminal"]) + assert "output_to_terminal=False" in result.output # should print correct value + + config_dict = get_primaite_config_dict() + # config should reflect that output_to_terminal is True + assert config_dict["developer_mode"]["output_to_terminal"] is False + + result = cli(["dev-mode", "config", "-t"]) + assert "output_to_terminal=True" in result.output # should print correct value + + config_dict = get_primaite_config_dict() + # config should reflect that output_to_terminal is True + assert config_dict["developer_mode"]["output_to_terminal"] + + result = cli(["dev-mode", "config", "-nt"]) + assert "output_to_terminal=False" in result.output # should print correct value + + config_dict = get_primaite_config_dict() + # config should reflect that output_to_terminal is True + assert config_dict["developer_mode"]["output_to_terminal"] is False From ab3d23785dc4676f38d89ceebc12806f68fd00cc Mon Sep 17 00:00:00 2001 From: Czar Echavez Date: Tue, 30 Apr 2024 15:50:53 +0100 Subject: [PATCH 03/37] #2533: precommit --- src/primaite/utils/cli/dev_cli.py | 113 ++++++++---------- .../utils/cli/primaite_config_utils.py | 4 +- tests/integration_tests/cli/test_dev_cli.py | 2 +- 3 files changed, 55 insertions(+), 64 deletions(-) diff --git a/src/primaite/utils/cli/dev_cli.py b/src/primaite/utils/cli/dev_cli.py index 886f1bf1..03567785 100644 --- a/src/primaite/utils/cli/dev_cli.py +++ b/src/primaite/utils/cli/dev_cli.py @@ -1,5 +1,3 @@ -from typing import List - import click import typer from rich import print @@ -73,44 +71,38 @@ def disable(): def config_callback( - ctx: typer.Context, - sys_log_level: Annotated[ - LogLevel, - typer.Option( - "--sys-log-level", - "-level", - click_type=click.Choice(LogLevel._member_names_, case_sensitive=False), - help="The level of system logs to output.", - show_default=False - ) - ] = None, - output_sys_logs: Annotated[ - bool, - typer.Option( - "--output-sys-logs/--no-sys-logs", - "-sys/-nsys", - help="Output system logs to file.", - show_default=False - ) - ] = None, - output_pcap_logs: Annotated[ - bool, - typer.Option( - "--output-pcap-logs/--no-pcap-logs", - "-pcap/-npcap", - help="Output network packet capture logs to file.", - show_default=False - ), - ] = None, - output_to_terminal: Annotated[ - bool, - typer.Option( - "--output-to-terminal/--no-terminal", - "-t/-nt", - help="Output system logs to terminal.", - show_default=False - ) - ] = None, + ctx: typer.Context, + sys_log_level: Annotated[ + LogLevel, + typer.Option( + "--sys-log-level", + "-level", + click_type=click.Choice(LogLevel._member_names_, case_sensitive=False), + help="The level of system logs to output.", + show_default=False, + ), + ] = None, + output_sys_logs: Annotated[ + bool, + typer.Option( + "--output-sys-logs/--no-sys-logs", "-sys/-nsys", help="Output system logs to file.", show_default=False + ), + ] = None, + output_pcap_logs: Annotated[ + bool, + typer.Option( + "--output-pcap-logs/--no-pcap-logs", + "-pcap/-npcap", + help="Output network packet capture logs to file.", + show_default=False, + ), + ] = None, + output_to_terminal: Annotated[ + bool, + typer.Option( + "--output-to-terminal/--no-terminal", "-t/-nt", help="Output system logs to terminal.", show_default=False + ), + ] = None, ): """Configure the development tools and environment.""" config_dict = get_primaite_config_dict() @@ -122,7 +114,6 @@ def config_callback( config_dict["developer_mode"]["sys_log_level"] = ctx.params.get("sys_log_level") print(f"PrimAITE dev-mode config updated sys_log_level={ctx.params.get('sys_log_level')}") - if output_sys_logs is not None: config_dict["developer_mode"]["output_sys_logs"] = output_sys_logs print(f"PrimAITE dev-mode config updated {output_sys_logs=}") @@ -150,22 +141,22 @@ dev.add_typer(config_typer) @config_typer.command() def path( - directory: Annotated[ - str, - typer.Argument( - help="Directory where the system logs and PCAP logs will be output. By default, this will be where the" - "root of the PrimAITE repository is located.", - show_default=False, - ), - ] = None, - default: Annotated[ - bool, - typer.Option( - "--default", - "-root", - help="Set PrimAITE to output system logs and pcap logs to the PrimAITE repository root.", - ), - ] = None, + directory: Annotated[ + str, + typer.Argument( + help="Directory where the system logs and PCAP logs will be output. By default, this will be where the" + "root of the PrimAITE repository is located.", + show_default=False, + ), + ] = None, + default: Annotated[ + bool, + typer.Option( + "--default", + "-root", + help="Set PrimAITE to output system logs and pcap logs to the PrimAITE repository root.", + ), + ] = None, ): """Set the output directory for the PrimAITE system and PCAP logs.""" config_dict = get_primaite_config_dict() @@ -177,9 +168,11 @@ def path( config_dict["developer_mode"]["output_dir"] = None # update application config update_primaite_config(config_dict) - print(f"PrimAITE dev-mode output_dir [medium_turquoise]" - f"{str(_PRIMAITE_ROOT.parent.parent / 'simulation_output')}" - f"[/medium_turquoise]") + print( + f"PrimAITE dev-mode output_dir [medium_turquoise]" + f"{str(_PRIMAITE_ROOT.parent.parent / 'simulation_output')}" + f"[/medium_turquoise]" + ) return if directory: diff --git a/src/primaite/utils/cli/primaite_config_utils.py b/src/primaite/utils/cli/primaite_config_utils.py index 2eb590ad..99e95bed 100644 --- a/src/primaite/utils/cli/primaite_config_utils.py +++ b/src/primaite/utils/cli/primaite_config_utils.py @@ -6,9 +6,7 @@ import yaml from primaite import PRIMAITE_PATHS -def get_primaite_config_dict( - config_path: Optional[Path] = None -) -> Dict: +def get_primaite_config_dict(config_path: Optional[Path] = None) -> Dict: """ Returns a dict containing the PrimAITE application config. diff --git a/tests/integration_tests/cli/test_dev_cli.py b/tests/integration_tests/cli/test_dev_cli.py index cd697711..fb051085 100644 --- a/tests/integration_tests/cli/test_dev_cli.py +++ b/tests/integration_tests/cli/test_dev_cli.py @@ -5,7 +5,7 @@ from pathlib import Path import pytest -from primaite import PRIMAITE_PATHS, _PRIMAITE_ROOT +from primaite import _PRIMAITE_ROOT, PRIMAITE_PATHS from primaite.utils.cli.primaite_config_utils import get_primaite_config_dict from tests.integration_tests.cli import cli From 729f9c5064617027aaeb9224ed31799956dd85f0 Mon Sep 17 00:00:00 2001 From: Czar Echavez Date: Tue, 30 Apr 2024 19:15:54 +0100 Subject: [PATCH 04/37] #2533: fix primaite config recreated when running setup --- src/primaite/__init__.py | 5 +- src/primaite/cli.py | 12 +++- src/primaite/session/io.py | 21 +++--- src/primaite/simulator/__init__.py | 67 +++++++++++++------ .../utils/cli/primaite_config_utils.py | 2 +- tests/integration_tests/cli/test_dev_cli.py | 6 +- 6 files changed, 74 insertions(+), 39 deletions(-) diff --git a/src/primaite/__init__.py b/src/primaite/__init__.py index 59725ac7..c58f0103 100644 --- a/src/primaite/__init__.py +++ b/src/primaite/__init__.py @@ -123,8 +123,9 @@ PRIMAITE_PATHS: Final[_PrimaitePaths] = _PrimaitePaths() def _host_primaite_config() -> None: - pkg_config_path = Path(pkg_resources.resource_filename("primaite", "setup/_package_data/primaite_config.yaml")) - shutil.copy(pkg_config_path, PRIMAITE_PATHS.app_config_file_path) + if not PRIMAITE_PATHS.app_config_file_path.exists(): + pkg_config_path = Path(pkg_resources.resource_filename("primaite", "setup/_package_data/primaite_config.yaml")) + shutil.copy2(pkg_config_path, PRIMAITE_PATHS.app_config_file_path) _host_primaite_config() diff --git a/src/primaite/cli.py b/src/primaite/cli.py index 6ad019dd..e2e5f8f6 100644 --- a/src/primaite/cli.py +++ b/src/primaite/cli.py @@ -2,9 +2,12 @@ """Provides a CLI using Typer as an entry point.""" import logging import os +import shutil from enum import Enum +from pathlib import Path from typing import Optional +import pkg_resources import typer import yaml from typing_extensions import Annotated @@ -91,7 +94,7 @@ def version() -> None: @app.command() -def setup(overwrite_existing: bool = True) -> None: +def setup(overwrite_existing: bool = False) -> None: """ Perform the PrimAITE first-time setup. @@ -104,11 +107,14 @@ def setup(overwrite_existing: bool = True) -> None: _LOGGER.info("Performing the PrimAITE first-time setup...") - _LOGGER.info("Building primaite_config.yaml...") - _LOGGER.info("Building the PrimAITE app directories...") PRIMAITE_PATHS.mkdirs() + _LOGGER.info("Building primaite_config.yaml...") + if overwrite_existing: + pkg_config_path = Path(pkg_resources.resource_filename("primaite", "setup/_package_data/primaite_config.yaml")) + shutil.copy(pkg_config_path, PRIMAITE_PATHS.app_config_file_path) + _LOGGER.info("Rebuilding the demo notebooks...") reset_demo_notebooks.run(overwrite_existing=True) diff --git a/src/primaite/session/io.py b/src/primaite/session/io.py index 65e900a8..0e7b90c9 100644 --- a/src/primaite/session/io.py +++ b/src/primaite/session/io.py @@ -5,9 +5,8 @@ from typing import Dict, List, Optional from pydantic import BaseModel, ConfigDict -from primaite import _PRIMAITE_ROOT, getLogger, PRIMAITE_PATHS +from primaite import getLogger, PRIMAITE_PATHS from primaite.simulator import LogLevel, SIM_OUTPUT -from primaite.utils.cli.primaite_config_utils import get_primaite_config_dict, is_dev_mode _LOGGER = getLogger(__name__) @@ -63,15 +62,15 @@ class PrimaiteIO: time_str = timestamp.strftime("%H-%M-%S") # check if running in dev mode - if is_dev_mode(): - # if dev mode, simulation output will be the repository root or whichever path is configured - app_config = get_primaite_config_dict() - if app_config["developer_mode"]["output_dir"] is not None: - session_path = app_config["developer_mode"]["output_path"] - else: - session_path = _PRIMAITE_ROOT.parent.parent / "simulation_output" / date_str / time_str - else: - session_path = PRIMAITE_PATHS.user_sessions_path / date_str / time_str + # if is_dev_mode(): + # # if dev mode, simulation output will be the repository root or whichever path is configured + # app_config = get_primaite_config_dict() + # if app_config["developer_mode"]["output_dir"] is not None: + # session_path = app_config["developer_mode"]["output_dir"] + # else: + # session_path = _PRIMAITE_ROOT.parent.parent / "simulation_output" / date_str / time_str + # else: + session_path = PRIMAITE_PATHS.user_sessions_path / date_str / time_str session_path.mkdir(exist_ok=True, parents=True) return session_path diff --git a/src/primaite/simulator/__init__.py b/src/primaite/simulator/__init__.py index 985bb663..cde7136c 100644 --- a/src/primaite/simulator/__init__.py +++ b/src/primaite/simulator/__init__.py @@ -26,34 +26,63 @@ class LogLevel(IntEnum): class _SimOutput: - _default_path = _PRIMAITE_ROOT.parent.parent / "simulation_output" / datetime.now().strftime("%Y-%m-%d_%H-%M-%S") - def __init__(self): - self._path: Path = self._default_path - self.save_pcap_logs: bool = False - self.save_sys_logs: bool = False - self.write_sys_log_to_terminal: bool = False - self.sys_log_level: LogLevel = LogLevel.WARNING # default log level is at WARNING - - if is_dev_mode(): - # if dev mode, override with the values configured via the primaite dev-mode command - dev_config = get_primaite_config_dict().get("developer_mode") - self.save_pcap_logs = dev_config["output_pcap_logs"] - self.save_sys_logs = dev_config["output_sys_logs"] - self.write_sys_log_to_terminal = dev_config["output_to_terminal"] + self._path: Path = ( + _PRIMAITE_ROOT.parent.parent / "simulation_output" / datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + ) + self._save_pcap_logs: bool = False + self._save_sys_logs: bool = False + self._write_sys_log_to_terminal: bool = False + self._sys_log_level: LogLevel = LogLevel.WARNING # default log level is at WARNING @property def path(self) -> Path: - if not is_dev_mode(): - return self._path - if is_dev_mode(): - dev_config = get_primaite_config_dict().get("developer_mode") - return Path(dev_config["output_dir"]) if dev_config["output_dir"] else self._default_path + return self._path @path.setter def path(self, new_path: Path) -> None: self._path = new_path self._path.mkdir(exist_ok=True, parents=True) + @property + def save_pcap_logs(self) -> bool: + if is_dev_mode(): + return get_primaite_config_dict().get("developer_mode").get("output_pcap_logs") + return self._save_pcap_logs + + @save_pcap_logs.setter + def save_pcap_logs(self, save_pcap_logs: bool) -> None: + self._save_pcap_logs = save_pcap_logs + + @property + def save_sys_logs(self) -> bool: + if is_dev_mode(): + return get_primaite_config_dict().get("developer_mode").get("output_sys_logs") + return self._save_sys_logs + + @save_sys_logs.setter + def save_sys_logs(self, save_sys_logs: bool) -> None: + self._save_sys_logs = save_sys_logs + + @property + def write_sys_log_to_terminal(self) -> bool: + if is_dev_mode(): + return get_primaite_config_dict().get("developer_mode").get("output_to_terminal") + return self._write_sys_log_to_terminal + + @write_sys_log_to_terminal.setter + def write_sys_log_to_terminal(self, write_sys_log_to_terminal: bool) -> None: + self._write_sys_log_to_terminal = write_sys_log_to_terminal + + @property + def sys_log_level(self) -> LogLevel: + if is_dev_mode(): + return LogLevel[get_primaite_config_dict().get("developer_mode").get("sys_log_level")] + return self._sys_log_level + + @sys_log_level.setter + def sys_log_level(self, sys_log_level: LogLevel) -> None: + self._sys_log_level = sys_log_level + SIM_OUTPUT = _SimOutput() diff --git a/src/primaite/utils/cli/primaite_config_utils.py b/src/primaite/utils/cli/primaite_config_utils.py index 99e95bed..a11c2ce3 100644 --- a/src/primaite/utils/cli/primaite_config_utils.py +++ b/src/primaite/utils/cli/primaite_config_utils.py @@ -28,7 +28,7 @@ def get_primaite_config_dict(config_path: Optional[Path] = None) -> Dict: def is_dev_mode() -> bool: """Returns True if PrimAITE is currently running in developer mode.""" config = get_primaite_config_dict() - return config["developer_mode"]["enabled"] + return config["developer_mode"]["enabled"] if config.get("developer_mode", {}).get("enabled") else False def update_primaite_config(config: Dict) -> None: diff --git a/tests/integration_tests/cli/test_dev_cli.py b/tests/integration_tests/cli/test_dev_cli.py index fb051085..d650bc90 100644 --- a/tests/integration_tests/cli/test_dev_cli.py +++ b/tests/integration_tests/cli/test_dev_cli.py @@ -3,6 +3,7 @@ import shutil import tempfile from pathlib import Path +import pkg_resources import pytest from primaite import _PRIMAITE_ROOT, PRIMAITE_PATHS @@ -19,9 +20,8 @@ def test_setup(): temp_dir = tempfile.gettempdir() temp_config = Path(temp_dir) / "primaite_config.yaml" - shutil.copyfile( - _PRIMAITE_ROOT / "setup" / "_package_data" / "primaite_config.yaml", temp_config - ) # copy the default primaite config to temp directory + pkg_config_path = Path(pkg_resources.resource_filename("primaite", "setup/_package_data/primaite_config.yaml")) + shutil.copyfile(pkg_config_path, temp_config) # copy the default primaite config to temp directory PRIMAITE_PATHS.app_config_file_path = temp_config # use the copy for the test yield # run test os.remove(temp_config) # clean up temp file From 95643d3255bdedf618cad0da36364e9e4a49729b Mon Sep 17 00:00:00 2001 From: Czar Echavez Date: Tue, 30 Apr 2024 19:36:22 +0100 Subject: [PATCH 05/37] #2533: optimise so we are not reading from file all the time --- src/primaite/__init__.py | 34 +++++-------- src/primaite/simulator/__init__.py | 12 ++--- src/primaite/utils/cli/dev_cli.py | 50 ++++++------------- .../utils/cli/primaite_config_utils.py | 35 +++---------- 4 files changed, 42 insertions(+), 89 deletions(-) diff --git a/src/primaite/__init__.py b/src/primaite/__init__.py index c58f0103..2cd44755 100644 --- a/src/primaite/__init__.py +++ b/src/primaite/__init__.py @@ -134,23 +134,15 @@ _host_primaite_config() def _get_primaite_config() -> Dict: config_path = PRIMAITE_PATHS.app_config_file_path if not config_path.exists(): + # load from package if config does not exist config_path = Path(pkg_resources.resource_filename("primaite", "setup/_package_data/primaite_config.yaml")) with open(config_path, "r") as file: + # load from config primaite_config = yaml.safe_load(file) - log_level_map = { - "NOTSET": logging.NOTSET, - "DEBUG": logging.DEBUG, - "INFO": logging.INFO, - "WARN": logging.WARN, - "WARNING": logging.WARN, - "ERROR": logging.ERROR, - "CRITICAL": logging.CRITICAL, - } - primaite_config["log_level"] = log_level_map[primaite_config["logging"]["log_level"]] - return primaite_config + return primaite_config -_PRIMAITE_CONFIG = _get_primaite_config() +PRIMAITE_CONFIG = _get_primaite_config() class _LevelFormatter(Formatter): @@ -177,11 +169,11 @@ class _LevelFormatter(Formatter): _LEVEL_FORMATTER: Final[_LevelFormatter] = _LevelFormatter( { - logging.DEBUG: _PRIMAITE_CONFIG["logging"]["logger_format"]["DEBUG"], - logging.INFO: _PRIMAITE_CONFIG["logging"]["logger_format"]["INFO"], - logging.WARNING: _PRIMAITE_CONFIG["logging"]["logger_format"]["WARNING"], - logging.ERROR: _PRIMAITE_CONFIG["logging"]["logger_format"]["ERROR"], - logging.CRITICAL: _PRIMAITE_CONFIG["logging"]["logger_format"]["CRITICAL"], + logging.DEBUG: PRIMAITE_CONFIG["logging"]["logger_format"]["DEBUG"], + logging.INFO: PRIMAITE_CONFIG["logging"]["logger_format"]["INFO"], + logging.WARNING: PRIMAITE_CONFIG["logging"]["logger_format"]["WARNING"], + logging.ERROR: PRIMAITE_CONFIG["logging"]["logger_format"]["ERROR"], + logging.CRITICAL: PRIMAITE_CONFIG["logging"]["logger_format"]["CRITICAL"], } ) @@ -193,10 +185,10 @@ _FILE_HANDLER: Final[RotatingFileHandler] = RotatingFileHandler( backupCount=9, # Max 100MB of logs encoding="utf8", ) -_STREAM_HANDLER.setLevel(_PRIMAITE_CONFIG["logging"]["log_level"]) -_FILE_HANDLER.setLevel(_PRIMAITE_CONFIG["logging"]["log_level"]) +_STREAM_HANDLER.setLevel(PRIMAITE_CONFIG["logging"]["log_level"]) +_FILE_HANDLER.setLevel(PRIMAITE_CONFIG["logging"]["log_level"]) -_LOG_FORMAT_STR: Final[str] = _PRIMAITE_CONFIG["logging"]["logger_format"] +_LOG_FORMAT_STR: Final[str] = PRIMAITE_CONFIG["logging"]["logger_format"] _STREAM_HANDLER.setFormatter(_LEVEL_FORMATTER) _FILE_HANDLER.setFormatter(_LEVEL_FORMATTER) @@ -215,6 +207,6 @@ def getLogger(name: str) -> Logger: # noqa logging config. """ logger = logging.getLogger(name) - logger.setLevel(_PRIMAITE_CONFIG["log_level"]) + logger.setLevel(PRIMAITE_CONFIG["logging"]["log_level"]) return logger diff --git a/src/primaite/simulator/__init__.py b/src/primaite/simulator/__init__.py index cde7136c..9f936249 100644 --- a/src/primaite/simulator/__init__.py +++ b/src/primaite/simulator/__init__.py @@ -3,11 +3,11 @@ from datetime import datetime from enum import IntEnum from pathlib import Path -from primaite import _PRIMAITE_ROOT +from primaite import _PRIMAITE_ROOT, PRIMAITE_CONFIG __all__ = ["SIM_OUTPUT"] -from primaite.utils.cli.primaite_config_utils import get_primaite_config_dict, is_dev_mode +from primaite.utils.cli.primaite_config_utils import is_dev_mode class LogLevel(IntEnum): @@ -47,7 +47,7 @@ class _SimOutput: @property def save_pcap_logs(self) -> bool: if is_dev_mode(): - return get_primaite_config_dict().get("developer_mode").get("output_pcap_logs") + return PRIMAITE_CONFIG.get("developer_mode").get("output_pcap_logs") return self._save_pcap_logs @save_pcap_logs.setter @@ -57,7 +57,7 @@ class _SimOutput: @property def save_sys_logs(self) -> bool: if is_dev_mode(): - return get_primaite_config_dict().get("developer_mode").get("output_sys_logs") + return PRIMAITE_CONFIG.get("developer_mode").get("output_sys_logs") return self._save_sys_logs @save_sys_logs.setter @@ -67,7 +67,7 @@ class _SimOutput: @property def write_sys_log_to_terminal(self) -> bool: if is_dev_mode(): - return get_primaite_config_dict().get("developer_mode").get("output_to_terminal") + return PRIMAITE_CONFIG.get("developer_mode").get("output_to_terminal") return self._write_sys_log_to_terminal @write_sys_log_to_terminal.setter @@ -77,7 +77,7 @@ class _SimOutput: @property def sys_log_level(self) -> LogLevel: if is_dev_mode(): - return LogLevel[get_primaite_config_dict().get("developer_mode").get("sys_log_level")] + return LogLevel[PRIMAITE_CONFIG.get("developer_mode").get("sys_log_level")] return self._sys_log_level @sys_log_level.setter diff --git a/src/primaite/utils/cli/dev_cli.py b/src/primaite/utils/cli/dev_cli.py index 03567785..8d426b2d 100644 --- a/src/primaite/utils/cli/dev_cli.py +++ b/src/primaite/utils/cli/dev_cli.py @@ -3,9 +3,9 @@ import typer from rich import print from typing_extensions import Annotated -from primaite import _PRIMAITE_ROOT +from primaite import _PRIMAITE_ROOT, PRIMAITE_CONFIG from primaite.simulator import LogLevel -from primaite.utils.cli.primaite_config_utils import get_primaite_config_dict, is_dev_mode, update_primaite_config +from primaite.utils.cli.primaite_config_utils import is_dev_mode, update_primaite_application_config dev = typer.Typer() @@ -45,28 +45,18 @@ def show(): @dev.command() def enable(): """Enable the development mode for PrimAITE.""" - config_dict = get_primaite_config_dict() - - if config_dict is None: - return - # enable dev mode - config_dict["developer_mode"]["enabled"] = True - update_primaite_config(config_dict) + PRIMAITE_CONFIG["developer_mode"]["enabled"] = True + update_primaite_application_config() print(DEVELOPMENT_MODE_MESSAGE) @dev.command() def disable(): """Disable the development mode for PrimAITE.""" - config_dict = get_primaite_config_dict() - - if config_dict is None: - return - # disable dev mode - config_dict["developer_mode"]["enabled"] = False - update_primaite_config(config_dict) + PRIMAITE_CONFIG["developer_mode"]["enabled"] = False + update_primaite_application_config() print(PRODUCTION_MODE_MESSAGE) @@ -105,29 +95,24 @@ def config_callback( ] = None, ): """Configure the development tools and environment.""" - config_dict = get_primaite_config_dict() - - if config_dict is None: - return - if ctx.params.get("sys_log_level") is not None: - config_dict["developer_mode"]["sys_log_level"] = ctx.params.get("sys_log_level") + PRIMAITE_CONFIG["developer_mode"]["sys_log_level"] = ctx.params.get("sys_log_level") print(f"PrimAITE dev-mode config updated sys_log_level={ctx.params.get('sys_log_level')}") if output_sys_logs is not None: - config_dict["developer_mode"]["output_sys_logs"] = output_sys_logs + PRIMAITE_CONFIG["developer_mode"]["output_sys_logs"] = output_sys_logs print(f"PrimAITE dev-mode config updated {output_sys_logs=}") if output_pcap_logs is not None: - config_dict["developer_mode"]["output_pcap_logs"] = output_pcap_logs + PRIMAITE_CONFIG["developer_mode"]["output_pcap_logs"] = output_pcap_logs print(f"PrimAITE dev-mode config updated {output_pcap_logs=}") if output_to_terminal is not None: - config_dict["developer_mode"]["output_to_terminal"] = output_to_terminal + PRIMAITE_CONFIG["developer_mode"]["output_to_terminal"] = output_to_terminal print(f"PrimAITE dev-mode config updated {output_to_terminal=}") # update application config - update_primaite_config(config_dict) + update_primaite_application_config() config_typer = typer.Typer( @@ -159,15 +144,10 @@ def path( ] = None, ): """Set the output directory for the PrimAITE system and PCAP logs.""" - config_dict = get_primaite_config_dict() - - if config_dict is None: - return - if default: - config_dict["developer_mode"]["output_dir"] = None + PRIMAITE_CONFIG["developer_mode"]["output_dir"] = None # update application config - update_primaite_config(config_dict) + update_primaite_application_config() print( f"PrimAITE dev-mode output_dir [medium_turquoise]" f"{str(_PRIMAITE_ROOT.parent.parent / 'simulation_output')}" @@ -176,7 +156,7 @@ def path( return if directory: - config_dict["developer_mode"]["output_dir"] = directory + PRIMAITE_CONFIG["developer_mode"]["output_dir"] = directory # update application config - update_primaite_config(config_dict) + update_primaite_application_config() print(f"PrimAITE dev-mode output_dir [medium_turquoise]{directory}[/medium_turquoise]") diff --git a/src/primaite/utils/cli/primaite_config_utils.py b/src/primaite/utils/cli/primaite_config_utils.py index a11c2ce3..e0f6fe56 100644 --- a/src/primaite/utils/cli/primaite_config_utils.py +++ b/src/primaite/utils/cli/primaite_config_utils.py @@ -1,37 +1,18 @@ -from pathlib import Path -from typing import Dict, Optional - import yaml -from primaite import PRIMAITE_PATHS - - -def get_primaite_config_dict(config_path: Optional[Path] = None) -> Dict: - """ - Returns a dict containing the PrimAITE application config. - - :param: config_path: takes in a path object - leave empty to use the default app config path - """ - err_msg = "PrimAITE application config could not be loaded." - - if config_path is None: - config_path = PRIMAITE_PATHS.app_config_file_path - err_msg = "PrimAITE application config was not found. Have you run `primaite setup`?" - - if config_path.exists(): - with open(config_path, "r") as file: - return yaml.safe_load(file) - else: - print(err_msg) +from primaite import PRIMAITE_CONFIG, PRIMAITE_PATHS def is_dev_mode() -> bool: """Returns True if PrimAITE is currently running in developer mode.""" - config = get_primaite_config_dict() - return config["developer_mode"]["enabled"] if config.get("developer_mode", {}).get("enabled") else False + return ( + PRIMAITE_CONFIG["developer_mode"]["enabled"] + if (PRIMAITE_CONFIG.get("developer_mode", {}).get("enabled")) + else False + ) -def update_primaite_config(config: Dict) -> None: +def update_primaite_application_config() -> None: """Update the PrimAITE application config file.""" with open(PRIMAITE_PATHS.app_config_file_path, "w") as file: - yaml.dump(config, file) + yaml.dump(PRIMAITE_CONFIG, file) From b7fa826d9512a109353b1bd2fd928b0b3e275bb2 Mon Sep 17 00:00:00 2001 From: Nick Todd Date: Wed, 1 May 2024 11:16:45 +0100 Subject: [PATCH 06/37] #2442: Initial commit of MP test script --- src/primaite/notebooks/mp.py | 42 ++++++++++++++++++++++++++++++++++++ 1 file changed, 42 insertions(+) create mode 100644 src/primaite/notebooks/mp.py diff --git a/src/primaite/notebooks/mp.py b/src/primaite/notebooks/mp.py new file mode 100644 index 00000000..ebed7122 --- /dev/null +++ b/src/primaite/notebooks/mp.py @@ -0,0 +1,42 @@ +import yaml +from stable_baselines3 import PPO +from stable_baselines3.common.utils import set_random_seed +from stable_baselines3.common.vec_env import SubprocVecEnv + +from primaite.session.environment import PrimaiteGymEnv + +EPISODE_LEN = 128 +NUM_EPISODES = 10 +NO_STEPS = EPISODE_LEN * NUM_EPISODES +BATCH_SIZE = 32 +LEARNING_RATE = 3e-4 + +with open("c:/projects/primaite/src/primaite/config/_package_data/data_manipulation.yaml", "r") as f: + cfg = yaml.safe_load(f) + + +def make_env(rank: int, seed: int = 0) -> callable: + """Wrapper script for _init function.""" + + def _init() -> PrimaiteGymEnv: + env = PrimaiteGymEnv(env_config=cfg) + env.reset(seed=seed + rank) + model = PPO( + "MlpPolicy", + env, + learning_rate=LEARNING_RATE, + n_steps=NO_STEPS, + batch_size=BATCH_SIZE, + verbose=0, + tensorboard_log="./PPO_UC2/", + ) + model.learn(total_timesteps=NO_STEPS) + return env + + set_random_seed(seed) + return _init + + +if __name__ == "__main__": + n_procs = 4 + train_env = SubprocVecEnv([make_env(i + n_procs) for i in range(n_procs)]) From 8d9ffab5148036cfafb8fb895dda2ae7aa2fbbde Mon Sep 17 00:00:00 2001 From: Charlie Crane Date: Wed, 1 May 2024 12:04:11 +0100 Subject: [PATCH 07/37] #2464 - Correcting Python versioning pre-req range on README.md as we do not currently support Python3.11 --- README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 2265538a..927aa17c 100644 --- a/README.md +++ b/README.md @@ -33,7 +33,7 @@ Currently, the PrimAITE wheel can only be installed from GitHub. This may change #### Windows (PowerShell) **Prerequisites:** -* Manual install of Python >= 3.8 < 3.12 +* Manual install of Python >= 3.8 < 3.11 **Install:** @@ -51,7 +51,7 @@ primaite setup #### Unix **Prerequisites:** -* Manual install of Python >= 3.8 < 3.12 +* Manual install of Python >= 3.8 < 3.11 ``` bash sudo add-apt-repository ppa:deadsnakes/ppa From a2fb04e6f6c88d2ee41e2d6b8c63dab5df00e576 Mon Sep 17 00:00:00 2001 From: Czar Echavez Date: Wed, 1 May 2024 14:02:40 +0100 Subject: [PATCH 08/37] #2533: set default dev output path + clean up --- .gitignore | 1 + src/primaite/__init__.py | 11 +-- src/primaite/session/io.py | 17 ++-- src/primaite/simulator/__init__.py | 16 ++- src/primaite/utils/cli/dev_cli.py | 6 +- .../utils/cli/primaite_config_utils.py | 6 +- tests/conftest.py | 15 --- tests/integration_tests/cli/test_dev_cli.py | 97 +++++++++---------- .../_simulator/_system/core/test_sys_log.py | 10 ++ 9 files changed, 84 insertions(+), 95 deletions(-) diff --git a/.gitignore b/.gitignore index 5cd39f24..b464566b 100644 --- a/.gitignore +++ b/.gitignore @@ -151,6 +151,7 @@ docs/source/primaite-dependencies.rst # outputs src/primaite/outputs/ simulation_output/ +sessions/ # benchmark session outputs benchmark/output diff --git a/src/primaite/__init__.py b/src/primaite/__init__.py index 2cd44755..98612040 100644 --- a/src/primaite/__init__.py +++ b/src/primaite/__init__.py @@ -122,20 +122,13 @@ class _PrimaitePaths: PRIMAITE_PATHS: Final[_PrimaitePaths] = _PrimaitePaths() -def _host_primaite_config() -> None: - if not PRIMAITE_PATHS.app_config_file_path.exists(): - pkg_config_path = Path(pkg_resources.resource_filename("primaite", "setup/_package_data/primaite_config.yaml")) - shutil.copy2(pkg_config_path, PRIMAITE_PATHS.app_config_file_path) - - -_host_primaite_config() - - def _get_primaite_config() -> Dict: config_path = PRIMAITE_PATHS.app_config_file_path if not config_path.exists(): # load from package if config does not exist config_path = Path(pkg_resources.resource_filename("primaite", "setup/_package_data/primaite_config.yaml")) + # generate app config + shutil.copy2(config_path, PRIMAITE_PATHS.app_config_file_path) with open(config_path, "r") as file: # load from config primaite_config = yaml.safe_load(file) diff --git a/src/primaite/session/io.py b/src/primaite/session/io.py index 0e7b90c9..9dcc529d 100644 --- a/src/primaite/session/io.py +++ b/src/primaite/session/io.py @@ -5,8 +5,9 @@ from typing import Dict, List, Optional from pydantic import BaseModel, ConfigDict -from primaite import getLogger, PRIMAITE_PATHS +from primaite import _PRIMAITE_ROOT, getLogger, PRIMAITE_PATHS from primaite.simulator import LogLevel, SIM_OUTPUT +from primaite.utils.cli.primaite_config_utils import is_dev_mode _LOGGER = getLogger(__name__) @@ -61,17 +62,13 @@ class PrimaiteIO: date_str = timestamp.strftime("%Y-%m-%d") time_str = timestamp.strftime("%H-%M-%S") - # check if running in dev mode - # if is_dev_mode(): - # # if dev mode, simulation output will be the repository root or whichever path is configured - # app_config = get_primaite_config_dict() - # if app_config["developer_mode"]["output_dir"] is not None: - # session_path = app_config["developer_mode"]["output_dir"] - # else: - # session_path = _PRIMAITE_ROOT.parent.parent / "simulation_output" / date_str / time_str - # else: session_path = PRIMAITE_PATHS.user_sessions_path / date_str / time_str + # check if running in dev mode + if is_dev_mode(): + # check if there is an output directory set in config + session_path = _PRIMAITE_ROOT.parent.parent / "sessions" / date_str / time_str + session_path.mkdir(exist_ok=True, parents=True) return session_path diff --git a/src/primaite/simulator/__init__.py b/src/primaite/simulator/__init__.py index 9f936249..ee165c9f 100644 --- a/src/primaite/simulator/__init__.py +++ b/src/primaite/simulator/__init__.py @@ -3,7 +3,7 @@ from datetime import datetime from enum import IntEnum from pathlib import Path -from primaite import _PRIMAITE_ROOT, PRIMAITE_CONFIG +from primaite import _PRIMAITE_ROOT, PRIMAITE_CONFIG, PRIMAITE_PATHS __all__ = ["SIM_OUTPUT"] @@ -27,9 +27,17 @@ class LogLevel(IntEnum): class _SimOutput: def __init__(self): - self._path: Path = ( - _PRIMAITE_ROOT.parent.parent / "simulation_output" / datetime.now().strftime("%Y-%m-%d_%H-%M-%S") - ) + date_str = datetime.now().strftime("%Y-%m-%d") + time_str = datetime.now().strftime("%H-%M-%S") + + path = PRIMAITE_PATHS.user_sessions_path / date_str / time_str + + if is_dev_mode(): + # if dev mode is enabled, if output dir is not set, print to primaite repo root + path: Path = _PRIMAITE_ROOT.parent.parent / "sessions" / date_str / time_str / "simulation_output" + # otherwise print to output dir + + self._path = path self._save_pcap_logs: bool = False self._save_sys_logs: bool = False self._write_sys_log_to_terminal: bool = False diff --git a/src/primaite/utils/cli/dev_cli.py b/src/primaite/utils/cli/dev_cli.py index 8d426b2d..09cafff0 100644 --- a/src/primaite/utils/cli/dev_cli.py +++ b/src/primaite/utils/cli/dev_cli.py @@ -15,7 +15,7 @@ PRODUCTION_MODE_MESSAGE = ( " :monkey_face::monkey_face::monkey_face: [/green]\n" ) -DEVELOPMENT_MODE_MESSAGE = ( +DEVELOPER_MODE_MESSAGE = ( "\n[yellow] :construction::construction::construction: " " PrimAITE is running in Development mode " " :construction::construction::construction: [/yellow]\n" @@ -38,7 +38,7 @@ def dev_mode(): def show(): """Show if PrimAITE is in development mode or production mode.""" # print if dev mode is enabled - print(DEVELOPMENT_MODE_MESSAGE if is_dev_mode() else PRODUCTION_MODE_MESSAGE) + print(DEVELOPER_MODE_MESSAGE if is_dev_mode() else PRODUCTION_MODE_MESSAGE) print("\nTo see available options, use [medium_turquoise]`primaite dev-mode --help`[/medium_turquoise]\n") @@ -48,7 +48,7 @@ def enable(): # enable dev mode PRIMAITE_CONFIG["developer_mode"]["enabled"] = True update_primaite_application_config() - print(DEVELOPMENT_MODE_MESSAGE) + print(DEVELOPER_MODE_MESSAGE) @dev.command() diff --git a/src/primaite/utils/cli/primaite_config_utils.py b/src/primaite/utils/cli/primaite_config_utils.py index e0f6fe56..2a94ece8 100644 --- a/src/primaite/utils/cli/primaite_config_utils.py +++ b/src/primaite/utils/cli/primaite_config_utils.py @@ -5,11 +5,7 @@ from primaite import PRIMAITE_CONFIG, PRIMAITE_PATHS def is_dev_mode() -> bool: """Returns True if PrimAITE is currently running in developer mode.""" - return ( - PRIMAITE_CONFIG["developer_mode"]["enabled"] - if (PRIMAITE_CONFIG.get("developer_mode", {}).get("enabled")) - else False - ) + return PRIMAITE_CONFIG["developer_mode"]["enabled"] def update_primaite_application_config() -> None: diff --git a/tests/conftest.py b/tests/conftest.py index 7de2bfde..c7b6ac04 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -40,21 +40,6 @@ ACTION_SPACE_NODE_ACTION_VALUES = 1 _LOGGER = getLogger(__name__) -@pytest.fixture(scope="function", autouse=True) -def set_syslog_output_to_true(): - """Will be run before each test.""" - monkeypatch = MonkeyPatch() - monkeypatch.setattr( - SIM_OUTPUT, - "path", - Path(TEST_ASSETS_ROOT.parent.parent / "simulation_output" / datetime.now().strftime("%Y-%m-%d_%H-%M-%S")), - ) - monkeypatch.setattr(SIM_OUTPUT, "save_pcap_logs", False) - monkeypatch.setattr(SIM_OUTPUT, "save_sys_logs", False) - - yield - - class TestService(Service): """Test Service class""" diff --git a/tests/integration_tests/cli/test_dev_cli.py b/tests/integration_tests/cli/test_dev_cli.py index d650bc90..acd086b9 100644 --- a/tests/integration_tests/cli/test_dev_cli.py +++ b/tests/integration_tests/cli/test_dev_cli.py @@ -5,34 +5,47 @@ from pathlib import Path import pkg_resources import pytest +import yaml +from _pytest.monkeypatch import MonkeyPatch -from primaite import _PRIMAITE_ROOT, PRIMAITE_PATHS -from primaite.utils.cli.primaite_config_utils import get_primaite_config_dict +import primaite from tests.integration_tests.cli import cli @pytest.fixture(autouse=True) def test_setup(): """ - Setup this test by copying the + Setup this test by using the default primaite app config in package """ - original_config_path = PRIMAITE_PATHS.app_config_file_path # keep track of app config before test + current_config = primaite.PRIMAITE_CONFIG # store the config before test + original_config_path = primaite.PRIMAITE_PATHS.app_config_file_path # keep track of app config before test temp_dir = tempfile.gettempdir() temp_config = Path(temp_dir) / "primaite_config.yaml" pkg_config_path = Path(pkg_resources.resource_filename("primaite", "setup/_package_data/primaite_config.yaml")) shutil.copyfile(pkg_config_path, temp_config) # copy the default primaite config to temp directory - PRIMAITE_PATHS.app_config_file_path = temp_config # use the copy for the test - yield # run test + primaite.PRIMAITE_PATHS.app_config_file_path = temp_config # use the copy for the test + + with open(pkg_config_path, "r") as file: + # load from config + config_dict = yaml.safe_load(file) + + primaite.PRIMAITE_CONFIG = config_dict + assert primaite.PRIMAITE_CONFIG == config_dict + + yield + os.remove(temp_config) # clean up temp file - PRIMAITE_PATHS.app_config_file_path = original_config_path # restore app conf because other devs will yell at me + primaite.PRIMAITE_CONFIG = current_config # restore config to prevent being yelled at + assert primaite.PRIMAITE_CONFIG == current_config + primaite.PRIMAITE_PATHS.app_config_file_path = original_config_path +@pytest.mark.skip(reason="borked") def test_dev_mode_enable_disable(): """Test dev mode enable and disable.""" # check defaults - config_dict = get_primaite_config_dict() - assert config_dict["developer_mode"]["enabled"] is False # not enabled by default + assert primaite.PRIMAITE_CONFIG["developer_mode"]["enabled"] is False # not enabled by default result = cli(["dev-mode", "show"]) assert "Production" in result.output # should print that it is in Production mode by default @@ -41,8 +54,7 @@ def test_dev_mode_enable_disable(): assert "Development" in result.output # should print that it is in Development mode - config_dict = get_primaite_config_dict() - assert config_dict["developer_mode"]["enabled"] # config should reflect that dev mode is enabled + assert primaite.PRIMAITE_CONFIG["developer_mode"]["enabled"] # config should reflect that dev mode is enabled result = cli(["dev-mode", "show"]) assert "Development" in result.output # should print that it is in Development mode @@ -51,136 +63,123 @@ def test_dev_mode_enable_disable(): assert "Production" in result.output # should print that it is in Production mode - config_dict = get_primaite_config_dict() - assert config_dict["developer_mode"]["enabled"] is False # config should reflect that dev mode is disabled + assert ( + primaite.PRIMAITE_CONFIG["developer_mode"]["enabled"] is False + ) # config should reflect that dev mode is disabled result = cli(["dev-mode", "show"]) assert "Production" in result.output # should print that it is in Production mode +@pytest.mark.skip(reason="borked") def test_dev_mode_config_sys_log_level(): """Check that the system log level can be changed via CLI.""" # check defaults - config_dict = get_primaite_config_dict() - assert config_dict["developer_mode"]["sys_log_level"] == "DEBUG" # DEBUG by default + assert primaite.PRIMAITE_CONFIG["developer_mode"]["sys_log_level"] == "DEBUG" # DEBUG by default result = cli(["dev-mode", "config", "-level", "WARNING"]) assert "sys_log_level=WARNING" in result.output # should print correct value - config_dict = get_primaite_config_dict() # config should reflect that log level is WARNING - assert config_dict["developer_mode"]["sys_log_level"] == "WARNING" + assert primaite.PRIMAITE_CONFIG["developer_mode"]["sys_log_level"] == "WARNING" result = cli(["dev-mode", "config", "--sys-log-level", "INFO"]) assert "sys_log_level=INFO" in result.output # should print correct value - config_dict = get_primaite_config_dict() # config should reflect that log level is WARNING - assert config_dict["developer_mode"]["sys_log_level"] == "INFO" + assert primaite.PRIMAITE_CONFIG["developer_mode"]["sys_log_level"] == "INFO" +@pytest.mark.skip(reason="borked") def test_dev_mode_config_sys_logs_enable_disable(): """Test that the system logs output can be enabled or disabled.""" # check defaults - config_dict = get_primaite_config_dict() - assert config_dict["developer_mode"]["output_sys_logs"] is False # False by default + assert primaite.PRIMAITE_CONFIG["developer_mode"]["output_sys_logs"] is False # False by default result = cli(["dev-mode", "config", "--output-sys-logs"]) assert "output_sys_logs=True" in result.output # should print correct value - config_dict = get_primaite_config_dict() # config should reflect that output_sys_logs is True - assert config_dict["developer_mode"]["output_sys_logs"] + assert primaite.PRIMAITE_CONFIG["developer_mode"]["output_sys_logs"] result = cli(["dev-mode", "config", "--no-sys-logs"]) assert "output_sys_logs=False" in result.output # should print correct value - config_dict = get_primaite_config_dict() # config should reflect that output_sys_logs is True - assert config_dict["developer_mode"]["output_sys_logs"] is False + assert primaite.PRIMAITE_CONFIG["developer_mode"]["output_sys_logs"] is False result = cli(["dev-mode", "config", "-sys"]) assert "output_sys_logs=True" in result.output # should print correct value - config_dict = get_primaite_config_dict() # config should reflect that output_sys_logs is True - assert config_dict["developer_mode"]["output_sys_logs"] + assert primaite.PRIMAITE_CONFIG["developer_mode"]["output_sys_logs"] result = cli(["dev-mode", "config", "-nsys"]) assert "output_sys_logs=False" in result.output # should print correct value - config_dict = get_primaite_config_dict() # config should reflect that output_sys_logs is True - assert config_dict["developer_mode"]["output_sys_logs"] is False + assert primaite.PRIMAITE_CONFIG["developer_mode"]["output_sys_logs"] is False +@pytest.mark.skip(reason="borked") def test_dev_mode_config_pcap_logs_enable_disable(): """Test that the pcap logs output can be enabled or disabled.""" # check defaults - config_dict = get_primaite_config_dict() - assert config_dict["developer_mode"]["output_pcap_logs"] is False # False by default + assert primaite.PRIMAITE_CONFIG["developer_mode"]["output_pcap_logs"] is False # False by default result = cli(["dev-mode", "config", "--output-pcap-logs"]) assert "output_pcap_logs=True" in result.output # should print correct value - config_dict = get_primaite_config_dict() # config should reflect that output_pcap_logs is True - assert config_dict["developer_mode"]["output_pcap_logs"] + assert primaite.PRIMAITE_CONFIG["developer_mode"]["output_pcap_logs"] result = cli(["dev-mode", "config", "--no-pcap-logs"]) assert "output_pcap_logs=False" in result.output # should print correct value - config_dict = get_primaite_config_dict() # config should reflect that output_pcap_logs is True - assert config_dict["developer_mode"]["output_pcap_logs"] is False + assert primaite.PRIMAITE_CONFIG["developer_mode"]["output_pcap_logs"] is False result = cli(["dev-mode", "config", "-pcap"]) assert "output_pcap_logs=True" in result.output # should print correct value - config_dict = get_primaite_config_dict() # config should reflect that output_pcap_logs is True - assert config_dict["developer_mode"]["output_pcap_logs"] + assert primaite.PRIMAITE_CONFIG["developer_mode"]["output_pcap_logs"] result = cli(["dev-mode", "config", "-npcap"]) assert "output_pcap_logs=False" in result.output # should print correct value - config_dict = get_primaite_config_dict() # config should reflect that output_pcap_logs is True - assert config_dict["developer_mode"]["output_pcap_logs"] is False + assert primaite.PRIMAITE_CONFIG["developer_mode"]["output_pcap_logs"] is False +@pytest.mark.skip(reason="borked") def test_dev_mode_config_output_to_terminal_enable_disable(): """Test that the output to terminal can be enabled or disabled.""" # check defaults - config_dict = get_primaite_config_dict() - assert config_dict["developer_mode"]["output_to_terminal"] is False # False by default + assert primaite.PRIMAITE_CONFIG["developer_mode"]["output_to_terminal"] is False # False by default result = cli(["dev-mode", "config", "--output-to-terminal"]) assert "output_to_terminal=True" in result.output # should print correct value - config_dict = get_primaite_config_dict() # config should reflect that output_to_terminal is True - assert config_dict["developer_mode"]["output_to_terminal"] + assert primaite.PRIMAITE_CONFIG["developer_mode"]["output_to_terminal"] result = cli(["dev-mode", "config", "--no-terminal"]) assert "output_to_terminal=False" in result.output # should print correct value - config_dict = get_primaite_config_dict() # config should reflect that output_to_terminal is True - assert config_dict["developer_mode"]["output_to_terminal"] is False + assert primaite.PRIMAITE_CONFIG["developer_mode"]["output_to_terminal"] is False result = cli(["dev-mode", "config", "-t"]) assert "output_to_terminal=True" in result.output # should print correct value - config_dict = get_primaite_config_dict() # config should reflect that output_to_terminal is True - assert config_dict["developer_mode"]["output_to_terminal"] + assert primaite.PRIMAITE_CONFIG["developer_mode"]["output_to_terminal"] result = cli(["dev-mode", "config", "-nt"]) assert "output_to_terminal=False" in result.output # should print correct value - config_dict = get_primaite_config_dict() # config should reflect that output_to_terminal is True - assert config_dict["developer_mode"]["output_to_terminal"] is False + assert primaite.PRIMAITE_CONFIG["developer_mode"]["output_to_terminal"] is False diff --git a/tests/unit_tests/_primaite/_simulator/_system/core/test_sys_log.py b/tests/unit_tests/_primaite/_simulator/_system/core/test_sys_log.py index 56b58d71..1009adc3 100644 --- a/tests/unit_tests/_primaite/_simulator/_system/core/test_sys_log.py +++ b/tests/unit_tests/_primaite/_simulator/_system/core/test_sys_log.py @@ -2,10 +2,20 @@ from uuid import uuid4 import pytest +from primaite import PRIMAITE_CONFIG from primaite.simulator import LogLevel, SIM_OUTPUT from primaite.simulator.system.core.sys_log import SysLog +@pytest.fixture(autouse=True) +def override_dev_mode_temporarily(): + """Temporarily turn off dev mode for this test.""" + primaite_dev_mode = PRIMAITE_CONFIG["developer_mode"]["enabled"] + PRIMAITE_CONFIG["developer_mode"]["enabled"] = False + yield # run tests + PRIMAITE_CONFIG["developer_mode"]["enabled"] = primaite_dev_mode + + @pytest.fixture(scope="function") def syslog() -> SysLog: return SysLog(hostname="test") From b6b5ce91c29b3069c29790ef91c5fd436f42e52b Mon Sep 17 00:00:00 2001 From: Czar Echavez Date: Wed, 1 May 2024 14:33:33 +0100 Subject: [PATCH 09/37] #2533: clean up temp items in tests + fixing the CLI tests --- .../utils/cli/primaite_config_utils.py | 14 +++- tests/conftest.py | 10 +-- .../test_rllib_single_agent_environment.py | 2 + .../environments/test_sb3_environment.py | 1 + tests/integration_tests/cli/test_dev_cli.py | 70 ++++++++----------- 5 files changed, 44 insertions(+), 53 deletions(-) diff --git a/src/primaite/utils/cli/primaite_config_utils.py b/src/primaite/utils/cli/primaite_config_utils.py index 2a94ece8..966376d9 100644 --- a/src/primaite/utils/cli/primaite_config_utils.py +++ b/src/primaite/utils/cli/primaite_config_utils.py @@ -1,3 +1,5 @@ +from typing import Dict, Optional + import yaml from primaite import PRIMAITE_CONFIG, PRIMAITE_PATHS @@ -8,7 +10,13 @@ def is_dev_mode() -> bool: return PRIMAITE_CONFIG["developer_mode"]["enabled"] -def update_primaite_application_config() -> None: - """Update the PrimAITE application config file.""" +def update_primaite_application_config(config: Optional[Dict] = None) -> None: + """ + Update the PrimAITE application config file. + + :params: config: Leave empty so that PRIMAITE_CONFIG is used - otherwise provide the Dict + """ with open(PRIMAITE_PATHS.app_config_file_path, "w") as file: - yaml.dump(PRIMAITE_CONFIG, file) + if not config: + config = PRIMAITE_CONFIG + yaml.dump(config, file) diff --git a/tests/conftest.py b/tests/conftest.py index c7b6ac04..37bc9581 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,19 +1,14 @@ # © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK -from datetime import datetime -from pathlib import Path -from typing import Any, Dict, Optional, Tuple, Union +from typing import Any, Dict, Tuple import pytest -import yaml -from _pytest.monkeypatch import MonkeyPatch -from primaite import getLogger, PRIMAITE_PATHS +from primaite import getLogger from primaite.game.agent.actions import ActionManager from primaite.game.agent.interface import AbstractAgent from primaite.game.agent.observations.observation_manager import NestedObservation, ObservationManager from primaite.game.agent.rewards import RewardFunction from primaite.game.game import PrimaiteGame -from primaite.simulator import SIM_OUTPUT from primaite.simulator.file_system.file_system import FileSystem from primaite.simulator.network.container import Network from primaite.simulator.network.hardware.nodes.host.computer import Computer @@ -32,7 +27,6 @@ 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 import TEST_ASSETS_ROOT -from tests.mock_and_patch.get_session_path_mock import temp_user_sessions_path ACTION_SPACE_NODE_VALUES = 1 ACTION_SPACE_NODE_ACTION_VALUES = 1 diff --git a/tests/e2e_integration_tests/environments/test_rllib_single_agent_environment.py b/tests/e2e_integration_tests/environments/test_rllib_single_agent_environment.py index 4c4b8d8d..d9057fef 100644 --- a/tests/e2e_integration_tests/environments/test_rllib_single_agent_environment.py +++ b/tests/e2e_integration_tests/environments/test_rllib_single_agent_environment.py @@ -38,3 +38,5 @@ def test_rllib_single_agent_compatibility(): save_file = Path(tempfile.gettempdir()) / "ray/" algo.save(save_file) assert save_file.exists() + + save_file.unlink() # clean up diff --git a/tests/e2e_integration_tests/environments/test_sb3_environment.py b/tests/e2e_integration_tests/environments/test_sb3_environment.py index f6ff595f..f654234b 100644 --- a/tests/e2e_integration_tests/environments/test_sb3_environment.py +++ b/tests/e2e_integration_tests/environments/test_sb3_environment.py @@ -25,3 +25,4 @@ def test_sb3_compatibility(): model.save(save_path) assert (save_path).exists() + save_path.unlink() # clean up diff --git a/tests/integration_tests/cli/test_dev_cli.py b/tests/integration_tests/cli/test_dev_cli.py index acd086b9..8f1bdec6 100644 --- a/tests/integration_tests/cli/test_dev_cli.py +++ b/tests/integration_tests/cli/test_dev_cli.py @@ -6,9 +6,9 @@ from pathlib import Path import pkg_resources import pytest import yaml -from _pytest.monkeypatch import MonkeyPatch -import primaite +from primaite import PRIMAITE_CONFIG +from primaite.utils.cli.primaite_config_utils import update_primaite_application_config from tests.integration_tests.cli import cli @@ -17,35 +17,27 @@ def test_setup(): """ Setup this test by using the default primaite app config in package """ - current_config = primaite.PRIMAITE_CONFIG # store the config before test - original_config_path = primaite.PRIMAITE_PATHS.app_config_file_path # keep track of app config before test + global PRIMAITE_CONFIG + current_config = PRIMAITE_CONFIG.copy() # store the config before test - temp_dir = tempfile.gettempdir() - temp_config = Path(temp_dir) / "primaite_config.yaml" pkg_config_path = Path(pkg_resources.resource_filename("primaite", "setup/_package_data/primaite_config.yaml")) - shutil.copyfile(pkg_config_path, temp_config) # copy the default primaite config to temp directory - primaite.PRIMAITE_PATHS.app_config_file_path = temp_config # use the copy for the test with open(pkg_config_path, "r") as file: # load from config config_dict = yaml.safe_load(file) - primaite.PRIMAITE_CONFIG = config_dict - assert primaite.PRIMAITE_CONFIG == config_dict + PRIMAITE_CONFIG["developer_mode"] = config_dict["developer_mode"] yield - os.remove(temp_config) # clean up temp file - primaite.PRIMAITE_CONFIG = current_config # restore config to prevent being yelled at - assert primaite.PRIMAITE_CONFIG == current_config - primaite.PRIMAITE_PATHS.app_config_file_path = original_config_path + PRIMAITE_CONFIG["developer_mode"] = current_config["developer_mode"] # restore config to prevent being yelled at + update_primaite_application_config(config=PRIMAITE_CONFIG) -@pytest.mark.skip(reason="borked") def test_dev_mode_enable_disable(): """Test dev mode enable and disable.""" # check defaults - assert primaite.PRIMAITE_CONFIG["developer_mode"]["enabled"] is False # not enabled by default + assert PRIMAITE_CONFIG["developer_mode"]["enabled"] is False # not enabled by default result = cli(["dev-mode", "show"]) assert "Production" in result.output # should print that it is in Production mode by default @@ -54,7 +46,7 @@ def test_dev_mode_enable_disable(): assert "Development" in result.output # should print that it is in Development mode - assert primaite.PRIMAITE_CONFIG["developer_mode"]["enabled"] # config should reflect that dev mode is enabled + assert PRIMAITE_CONFIG["developer_mode"]["enabled"] # config should reflect that dev mode is enabled result = cli(["dev-mode", "show"]) assert "Development" in result.output # should print that it is in Development mode @@ -63,123 +55,117 @@ def test_dev_mode_enable_disable(): assert "Production" in result.output # should print that it is in Production mode - assert ( - primaite.PRIMAITE_CONFIG["developer_mode"]["enabled"] is False - ) # config should reflect that dev mode is disabled + assert PRIMAITE_CONFIG["developer_mode"]["enabled"] is False # config should reflect that dev mode is disabled result = cli(["dev-mode", "show"]) assert "Production" in result.output # should print that it is in Production mode -@pytest.mark.skip(reason="borked") def test_dev_mode_config_sys_log_level(): """Check that the system log level can be changed via CLI.""" # check defaults - assert primaite.PRIMAITE_CONFIG["developer_mode"]["sys_log_level"] == "DEBUG" # DEBUG by default + assert PRIMAITE_CONFIG["developer_mode"]["sys_log_level"] == "DEBUG" # DEBUG by default result = cli(["dev-mode", "config", "-level", "WARNING"]) assert "sys_log_level=WARNING" in result.output # should print correct value # config should reflect that log level is WARNING - assert primaite.PRIMAITE_CONFIG["developer_mode"]["sys_log_level"] == "WARNING" + assert PRIMAITE_CONFIG["developer_mode"]["sys_log_level"] == "WARNING" result = cli(["dev-mode", "config", "--sys-log-level", "INFO"]) assert "sys_log_level=INFO" in result.output # should print correct value # config should reflect that log level is WARNING - assert primaite.PRIMAITE_CONFIG["developer_mode"]["sys_log_level"] == "INFO" + assert PRIMAITE_CONFIG["developer_mode"]["sys_log_level"] == "INFO" -@pytest.mark.skip(reason="borked") def test_dev_mode_config_sys_logs_enable_disable(): """Test that the system logs output can be enabled or disabled.""" # check defaults - assert primaite.PRIMAITE_CONFIG["developer_mode"]["output_sys_logs"] is False # False by default + assert PRIMAITE_CONFIG["developer_mode"]["output_sys_logs"] is False # False by default result = cli(["dev-mode", "config", "--output-sys-logs"]) assert "output_sys_logs=True" in result.output # should print correct value # config should reflect that output_sys_logs is True - assert primaite.PRIMAITE_CONFIG["developer_mode"]["output_sys_logs"] + assert PRIMAITE_CONFIG["developer_mode"]["output_sys_logs"] result = cli(["dev-mode", "config", "--no-sys-logs"]) assert "output_sys_logs=False" in result.output # should print correct value # config should reflect that output_sys_logs is True - assert primaite.PRIMAITE_CONFIG["developer_mode"]["output_sys_logs"] is False + assert PRIMAITE_CONFIG["developer_mode"]["output_sys_logs"] is False result = cli(["dev-mode", "config", "-sys"]) assert "output_sys_logs=True" in result.output # should print correct value # config should reflect that output_sys_logs is True - assert primaite.PRIMAITE_CONFIG["developer_mode"]["output_sys_logs"] + assert PRIMAITE_CONFIG["developer_mode"]["output_sys_logs"] result = cli(["dev-mode", "config", "-nsys"]) assert "output_sys_logs=False" in result.output # should print correct value # config should reflect that output_sys_logs is True - assert primaite.PRIMAITE_CONFIG["developer_mode"]["output_sys_logs"] is False + assert PRIMAITE_CONFIG["developer_mode"]["output_sys_logs"] is False -@pytest.mark.skip(reason="borked") def test_dev_mode_config_pcap_logs_enable_disable(): """Test that the pcap logs output can be enabled or disabled.""" # check defaults - assert primaite.PRIMAITE_CONFIG["developer_mode"]["output_pcap_logs"] is False # False by default + assert PRIMAITE_CONFIG["developer_mode"]["output_pcap_logs"] is False # False by default result = cli(["dev-mode", "config", "--output-pcap-logs"]) assert "output_pcap_logs=True" in result.output # should print correct value # config should reflect that output_pcap_logs is True - assert primaite.PRIMAITE_CONFIG["developer_mode"]["output_pcap_logs"] + assert PRIMAITE_CONFIG["developer_mode"]["output_pcap_logs"] result = cli(["dev-mode", "config", "--no-pcap-logs"]) assert "output_pcap_logs=False" in result.output # should print correct value # config should reflect that output_pcap_logs is True - assert primaite.PRIMAITE_CONFIG["developer_mode"]["output_pcap_logs"] is False + assert PRIMAITE_CONFIG["developer_mode"]["output_pcap_logs"] is False result = cli(["dev-mode", "config", "-pcap"]) assert "output_pcap_logs=True" in result.output # should print correct value # config should reflect that output_pcap_logs is True - assert primaite.PRIMAITE_CONFIG["developer_mode"]["output_pcap_logs"] + assert PRIMAITE_CONFIG["developer_mode"]["output_pcap_logs"] result = cli(["dev-mode", "config", "-npcap"]) assert "output_pcap_logs=False" in result.output # should print correct value # config should reflect that output_pcap_logs is True - assert primaite.PRIMAITE_CONFIG["developer_mode"]["output_pcap_logs"] is False + assert PRIMAITE_CONFIG["developer_mode"]["output_pcap_logs"] is False -@pytest.mark.skip(reason="borked") def test_dev_mode_config_output_to_terminal_enable_disable(): """Test that the output to terminal can be enabled or disabled.""" # check defaults - assert primaite.PRIMAITE_CONFIG["developer_mode"]["output_to_terminal"] is False # False by default + assert PRIMAITE_CONFIG["developer_mode"]["output_to_terminal"] is False # False by default result = cli(["dev-mode", "config", "--output-to-terminal"]) assert "output_to_terminal=True" in result.output # should print correct value # config should reflect that output_to_terminal is True - assert primaite.PRIMAITE_CONFIG["developer_mode"]["output_to_terminal"] + assert PRIMAITE_CONFIG["developer_mode"]["output_to_terminal"] result = cli(["dev-mode", "config", "--no-terminal"]) assert "output_to_terminal=False" in result.output # should print correct value # config should reflect that output_to_terminal is True - assert primaite.PRIMAITE_CONFIG["developer_mode"]["output_to_terminal"] is False + assert PRIMAITE_CONFIG["developer_mode"]["output_to_terminal"] is False result = cli(["dev-mode", "config", "-t"]) assert "output_to_terminal=True" in result.output # should print correct value # config should reflect that output_to_terminal is True - assert primaite.PRIMAITE_CONFIG["developer_mode"]["output_to_terminal"] + assert PRIMAITE_CONFIG["developer_mode"]["output_to_terminal"] result = cli(["dev-mode", "config", "-nt"]) assert "output_to_terminal=False" in result.output # should print correct value # config should reflect that output_to_terminal is True - assert primaite.PRIMAITE_CONFIG["developer_mode"]["output_to_terminal"] is False + assert PRIMAITE_CONFIG["developer_mode"]["output_to_terminal"] is False From efdc9794958fdead42b0e9b46f0fdf9efaa23c0f Mon Sep 17 00:00:00 2001 From: Charlie Crane Date: Wed, 1 May 2024 14:33:38 +0100 Subject: [PATCH 10/37] #2464 Changing pyproject.toml to allow the use of python 3.11 and tweaking the pipeline to use python 3.11. Seeing what happens... --- .azure/azure-ci-build-pipeline.yaml | 47 +++++++++++++++++++---------- pyproject.toml | 2 +- 2 files changed, 32 insertions(+), 17 deletions(-) diff --git a/.azure/azure-ci-build-pipeline.yaml b/.azure/azure-ci-build-pipeline.yaml index e9139d5b..41e27931 100644 --- a/.azure/azure-ci-build-pipeline.yaml +++ b/.azure/azure-ci-build-pipeline.yaml @@ -19,28 +19,43 @@ parameters: # img: 'ubuntu-latest' # every_time: false # publish_coverage: false - - job_name: 'UbuntuPython310' - py: '3.10' + # - job_name: 'UbuntuPython310' + # py: '3.10' + # img: 'ubuntu-latest' + # every_time: true + # publish_coverage: true + - job_name: 'UbuntuPython311 TEST' + py: '3.11' img: 'ubuntu-latest' every_time: true publish_coverage: true - - job_name: 'WindowsPython38' - py: '3.8' + # - job_name: 'WindowsPython38' + # py: '3.8' + # img: 'windows-latest' + # every_time: false + # publish_coverage: false + # - job_name: 'WindowsPython310' + # py: '3.10' + # img: 'windows-latest' + # every_time: false + # publish_coverage: false + - job_name: 'WindowsPython311 TEST' + py: '3.11' img: 'windows-latest' every_time: false publish_coverage: false - - job_name: 'WindowsPython310' - py: '3.10' - img: 'windows-latest' - every_time: false - publish_coverage: false - - job_name: 'MacOSPython38' - py: '3.8' - img: 'macOS-latest' - every_time: false - publish_coverage: false - - job_name: 'MacOSPython310' - py: '3.10' + # - job_name: 'MacOSPython38' + # py: '3.8' + # img: 'macOS-latest' + # every_time: false + # publish_coverage: false + # - job_name: 'MacOSPython310' + # py: '3.10' + # img: 'macOS-latest' + # every_time: false + # publish_coverage: false + - job_name: 'MacOSPython311 TEST' + py: '3.11' img: 'macOS-latest' every_time: false publish_coverage: false diff --git a/pyproject.toml b/pyproject.toml index 333132bc..fd75252c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ name = "primaite" description = "PrimAITE (Primary-level AI Training Environment) is a simulation environment for training AI under the ARCD programme." authors = [{name="Defence Science and Technology Laboratory UK", email="oss@dstl.gov.uk"}] license = {file = "LICENSE"} -requires-python = ">=3.8, <3.11" +requires-python = ">=3.8, <3.12" dynamic = ["version", "readme"] classifiers = [ "License :: OSI Approved :: MIT License", From 392b83fc64241d28cc25b125442cdedd8bac6483 Mon Sep 17 00:00:00 2001 From: Charlie Crane Date: Wed, 1 May 2024 14:34:38 +0100 Subject: [PATCH 11/37] #2464 - Correcting use of spaces in pipeline job names --- .azure/azure-ci-build-pipeline.yaml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.azure/azure-ci-build-pipeline.yaml b/.azure/azure-ci-build-pipeline.yaml index 41e27931..e190471d 100644 --- a/.azure/azure-ci-build-pipeline.yaml +++ b/.azure/azure-ci-build-pipeline.yaml @@ -24,7 +24,7 @@ parameters: # img: 'ubuntu-latest' # every_time: true # publish_coverage: true - - job_name: 'UbuntuPython311 TEST' + - job_name: 'UbuntuPython311_TEST' py: '3.11' img: 'ubuntu-latest' every_time: true @@ -39,7 +39,7 @@ parameters: # img: 'windows-latest' # every_time: false # publish_coverage: false - - job_name: 'WindowsPython311 TEST' + - job_name: 'WindowsPython311_TEST' py: '3.11' img: 'windows-latest' every_time: false @@ -54,7 +54,7 @@ parameters: # img: 'macOS-latest' # every_time: false # publish_coverage: false - - job_name: 'MacOSPython311 TEST' + - job_name: 'MacOSPython311_TEST' py: '3.11' img: 'macOS-latest' every_time: false From 5516fbc6fc709e3727fbcdee4b3400054928e88a Mon Sep 17 00:00:00 2001 From: Czar Echavez Date: Wed, 1 May 2024 14:48:10 +0100 Subject: [PATCH 12/37] #2533: output directory for sessions --- src/primaite/simulator/__init__.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/src/primaite/simulator/__init__.py b/src/primaite/simulator/__init__.py index ee165c9f..e9d44ec3 100644 --- a/src/primaite/simulator/__init__.py +++ b/src/primaite/simulator/__init__.py @@ -36,6 +36,14 @@ class _SimOutput: # if dev mode is enabled, if output dir is not set, print to primaite repo root path: Path = _PRIMAITE_ROOT.parent.parent / "sessions" / date_str / time_str / "simulation_output" # otherwise print to output dir + if PRIMAITE_CONFIG["developer_mode"]["output_dir"]: + path: Path = ( + Path(PRIMAITE_CONFIG["developer_mode"]["output_dir"]) + / "sessions" + / date_str + / time_str + / "simulation_output" + ) self._path = path self._save_pcap_logs: bool = False From 57cff8d4e3094008d02ff5d7805de2e400d1528f Mon Sep 17 00:00:00 2001 From: Charlie Crane Date: Wed, 1 May 2024 16:01:08 +0100 Subject: [PATCH 13/37] #2464 - Seeing if python 3.12 will also run... --- .azure/azure-ci-build-pipeline.yaml | 12 ++++++------ pyproject.toml | 2 +- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/.azure/azure-ci-build-pipeline.yaml b/.azure/azure-ci-build-pipeline.yaml index e190471d..1a7b4870 100644 --- a/.azure/azure-ci-build-pipeline.yaml +++ b/.azure/azure-ci-build-pipeline.yaml @@ -24,8 +24,8 @@ parameters: # img: 'ubuntu-latest' # every_time: true # publish_coverage: true - - job_name: 'UbuntuPython311_TEST' - py: '3.11' + - job_name: 'UbuntuPython312_TEST' + py: '3.12' img: 'ubuntu-latest' every_time: true publish_coverage: true @@ -39,8 +39,8 @@ parameters: # img: 'windows-latest' # every_time: false # publish_coverage: false - - job_name: 'WindowsPython311_TEST' - py: '3.11' + - job_name: 'WindowsPython312_TEST' + py: '3.12' img: 'windows-latest' every_time: false publish_coverage: false @@ -54,8 +54,8 @@ parameters: # img: 'macOS-latest' # every_time: false # publish_coverage: false - - job_name: 'MacOSPython311_TEST' - py: '3.11' + - job_name: 'MacOSPython312_TEST' + py: '3.12' img: 'macOS-latest' every_time: false publish_coverage: false diff --git a/pyproject.toml b/pyproject.toml index fd75252c..c71b55a6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ name = "primaite" description = "PrimAITE (Primary-level AI Training Environment) is a simulation environment for training AI under the ARCD programme." authors = [{name="Defence Science and Technology Laboratory UK", email="oss@dstl.gov.uk"}] license = {file = "LICENSE"} -requires-python = ">=3.8, <3.12" +requires-python = ">=3.8, <3.13" dynamic = ["version", "readme"] classifiers = [ "License :: OSI Approved :: MIT License", From 61e7a4e4394deeb583be936ff5f17da39be9401f Mon Sep 17 00:00:00 2001 From: Czar Echavez Date: Wed, 1 May 2024 16:04:18 +0100 Subject: [PATCH 14/37] #2533: documentation on use of dev-mode --- docs/index.rst | 1 + docs/source/developer_tools.rst | 210 ++++++++++++++++++++++++++++++ docs/source/getting_started.rst | 2 + src/primaite/utils/cli/dev_cli.py | 1 + 4 files changed, 214 insertions(+) create mode 100644 docs/source/developer_tools.rst diff --git a/docs/index.rst b/docs/index.rst index 4cc81b13..52bc8b57 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -116,6 +116,7 @@ Head over to the :ref:`getting-started` page to install and setup PrimAITE! :caption: Developer information: :hidden: + source/developer_tools source/state_system source/request_system PrimAITE API diff --git a/docs/source/developer_tools.rst b/docs/source/developer_tools.rst new file mode 100644 index 00000000..fcb52443 --- /dev/null +++ b/docs/source/developer_tools.rst @@ -0,0 +1,210 @@ +.. only:: comment + + © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK + +.. _Developer Tools: + +Developer Tools +*************** + +PrimAITE includes developer CLI tools that are intended to be used by developers. + +dev-mode +======== + +The dev-mode contains configuration which override any of the config files during runtime. + +This is intended to make debugging easier by removing the need to find the relevant configuration file/settings. + +Enabling dev-mode +----------------- + +The PrimAITE dev-mode can be enabled via the use of + +.. code-block:: + + primaite dev-mode enable + +Disabling dev-mode +------------------ + +The PrimAITE dev-mode can be disabled via the use of + +.. code-block:: + + primaite dev-mode disable + +Show current mode +----------------- + +To show if the dev-mode is enabled or not, use +The PrimAITE dev-mode can be disabled via the use of + +.. code-block:: + + primaite dev-mode show + +dev-mode configuration +====================== + +The following configures some specific items that the dev-mode overrides, if enabled. + +`--sys-log-level` or `-level` +---------------------------- + +The level of system logs can be overridden by dev-mode. + +By default, this is set to DEBUG + +The available options are [DEBUG|INFO|WARNING|ERROR|CRITICAL] + +.. code-block:: + + primaite dev-mode config -level INFO + +or + +.. code-block:: + + primaite dev-mode config --sys-log-level INFO + +`--output-sys-logs` or `-sys` +----------------------------- + +The outputting of system logs can be overridden by dev-mode. + +By default, this is set to False + +Enabling system logs +"""""""""""""""""""" + +To enable outputting of system logs + +.. code-block:: + + primaite dev-mode config --output-sys-logs + +or + +.. code-block:: + + primaite dev-mode config -sys + +Disabling system logs +""""""""""""""""""""" + +To disable outputting of system logs + +.. code-block:: + + primaite dev-mode config --no-sys-logs + +or + +.. code-block:: + + primaite dev-mode config -nsys + +`--output-pcap-logs` or `-pcap` +------------------------------- + +The outputting of packet capture logs can be overridden by dev-mode. + +By default, this is set to False + +Enabling PCAP logs +"""""""""""""""""" + +To enable outputting of packet capture logs + +.. code-block:: + + primaite dev-mode config --output-pcap-logs + +or + +.. code-block:: + + primaite dev-mode config -pcap + +Disabling PCAP logs +""""""""""""""""""" + +To disable outputting of packet capture logs + +.. code-block:: + + primaite dev-mode config --no-pcap-logs + +or + +.. code-block:: + + primaite dev-mode config -npcap + +`--output-to-terminal` or `-t` +------------------------------ + +The outputting of system logs to the terminal can be overridden by dev-mode. + +By default, this is set to False + +Enabling system log output to terminal +"""""""""""""""""""""""""""""""""""""" + +To enable outputting of system logs to terminal + +.. code-block:: + + primaite dev-mode config --output-to-terminal + +or + +.. code-block:: + + primaite dev-mode config -t + +Disabling system log output to terminal +""""""""""""""""""""""""""""""""""""""" + +To disable outputting of system logs to terminal + +.. code-block:: + + primaite dev-mode config --no-terminal + +or + +.. code-block:: + + primaite dev-mode config -nt + +path +---- + +PrimAITE dev-mode can override where sessions are output. + +By default, PrimAITE will output the sessions in USER_HOME/primaite/sessions + +With dev-mode enabled, by default, this will be changed to PRIMAITE_REPOSITORY_ROOT/sessions + +However, providing a path will let dev-mode output sessions to the given path e.g. + +.. code-block:: bash + :caption: Unix + + primaite dev-mode config path ~/output/path + +.. code-block:: powershell + :caption: Windows (Powershell) + + primaite dev-mode config path ~\output\path + +default path +"""""""""""" + +To reset the path to use the PRIMAITE_REPOSITORY_ROOT/sessions, run the command + +.. code-block:: + + primaite dev-mode config path --default diff --git a/docs/source/getting_started.rst b/docs/source/getting_started.rst index dccfd8aa..7c91498c 100644 --- a/docs/source/getting_started.rst +++ b/docs/source/getting_started.rst @@ -167,3 +167,5 @@ To set PrimAITE to run in development mode: :caption: Windows (Powershell) primaite dev-mode enable + +More information about :ref:`Developer Tools` diff --git a/src/primaite/utils/cli/dev_cli.py b/src/primaite/utils/cli/dev_cli.py index 09cafff0..7437d2fc 100644 --- a/src/primaite/utils/cli/dev_cli.py +++ b/src/primaite/utils/cli/dev_cli.py @@ -39,6 +39,7 @@ def show(): """Show if PrimAITE is in development mode or production mode.""" # print if dev mode is enabled print(DEVELOPER_MODE_MESSAGE if is_dev_mode() else PRODUCTION_MODE_MESSAGE) + print(f"Current Settings: {PRIMAITE_CONFIG['developer_mode']}") print("\nTo see available options, use [medium_turquoise]`primaite dev-mode --help`[/medium_turquoise]\n") From d7572dc18b9ed43c59cac9d43588c068d944531c Mon Sep 17 00:00:00 2001 From: Czar Echavez Date: Wed, 1 May 2024 16:10:27 +0100 Subject: [PATCH 15/37] #2533: changelog --- CHANGELOG.md | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f79a1fd3..8426f5dc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,13 +14,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Added ability to define scenarios that change depending on the episode number. - Standardised Environment API by renaming the config parameter of `PrimaiteGymEnv` from `game_config` to `env_config` - Database Connection ID's are now created/issued by DatabaseService and not DatabaseClient -- added ability to set PrimAITE between development and production modes via PrimAITE CLI ``mode`` command - Updated DatabaseClient so that it can now have a single native DatabaseClientConnection along with a collection of DatabaseClientConnection's. - Implemented the uninstall functionality for DatabaseClient so that all connections are terminated at the DatabaseService. - Added the ability for a DatabaseService to terminate a connection. - Added active_connection to DatabaseClientConnection so that if the connection is terminated active_connection is set to False and the object can no longer be used. - Added additional show functions to enable connection inspection. - Updates to agent logging, to include the reward both per step and per episode. +- Introduced Developer CLI tools to assist with developing/debugging PrimAITE + - Can be enabled via `primaite dev-mode enable` + - Activating dev-mode will change the location where the sessions will be output - by default will output where the PrimAITE repository is located ## [Unreleased] From c2cc5228ae981650b210744883441cde311ffd65 Mon Sep 17 00:00:00 2001 From: Charlie Crane Date: Wed, 1 May 2024 16:20:01 +0100 Subject: [PATCH 16/37] #2464 - Linting changes to see if the python 3.12 pipeline will pass --- src/primaite/game/agent/rewards.py | 2 +- src/primaite/simulator/network/container.py | 2 +- src/primaite/simulator/network/hardware/base.py | 6 +++--- src/primaite/simulator/network/protocols/icmp.py | 4 ++-- src/primaite/simulator/system/services/dns/dns_client.py | 2 +- src/primaite/simulator/system/services/ftp/ftp_service.py | 2 +- src/primaite/simulator/system/services/icmp/icmp.py | 2 +- 7 files changed, 10 insertions(+), 10 deletions(-) diff --git a/src/primaite/game/agent/rewards.py b/src/primaite/game/agent/rewards.py index 0222bfcc..36a984dd 100644 --- a/src/primaite/game/agent/rewards.py +++ b/src/primaite/game/agent/rewards.py @@ -231,7 +231,7 @@ class WebpageUnavailablePenalty(AbstractReward): # If the last request did actually go through, then check if the webpage also loaded 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: + if web_browser_state is NOT_PRESENT_IN_STATE or "history" not in web_browser_state: # noqa _LOGGER.debug( "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", diff --git a/src/primaite/simulator/network/container.py b/src/primaite/simulator/network/container.py index 17308c97..717d04b3 100644 --- a/src/primaite/simulator/network/container.py +++ b/src/primaite/simulator/network/container.py @@ -248,7 +248,7 @@ class Network(SimComponent): hostname_b = node_b.hostname if node_b else None port_a = link.endpoint_a.port_num port_b = link.endpoint_b.port_num - link_key = f"{hostname_a}:eth-{port_a}<->{hostname_b}:eth-{port_b}" + link_key = f"{hostname_a}: eth-{port_a}<->{hostname_b}: eth-{port_b}" state["links"][link_key] = link.describe_state() state["links"][link_key]["hostname_a"] = hostname_a state["links"][link_key]["hostname_b"] = hostname_b diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py index 78a0bf2d..27a31e84 100644 --- a/src/primaite/simulator/network/hardware/base.py +++ b/src/primaite/simulator/network/hardware/base.py @@ -56,7 +56,7 @@ def generate_mac_address(oui: Optional[str] = None) -> str: if oui: oui_pattern = re.compile(r"^([0-9A-Fa-f]{2}[:-]){2}[0-9A-Fa-f]{2}$") if not oui_pattern.match(oui): - msg = f"Invalid oui. The oui should be in the format xx:xx:xx, where x is a hexadecimal digit, got '{oui}'" + msg = f"Invalid oui. The oui should be in the format xx:xx:xx, where x is a hexadecimal digit, got '{oui}'" # noqa _LOGGER.error(msg) raise ValueError(msg) oui_bytes = [int(chunk, 16) for chunk in oui.split(":")] @@ -64,7 +64,7 @@ def generate_mac_address(oui: Optional[str] = None) -> str: else: mac = random_bytes - return ":".join(f"{b:02x}" for b in mac) + return ":".join(f"{b:02x}" for b in mac) # noqa class NetworkInterface(SimComponent, ABC): @@ -599,7 +599,7 @@ class Link(SimComponent): @property def current_load_percent(self) -> str: """Get the current load formatted as a percentage string.""" - return f"{self.current_load / self.bandwidth:.5f}%" + return f"{self.current_load / self.bandwidth:.5f}%" # noqa def endpoint_up(self): """Let the Link know and endpoint has been brought up.""" diff --git a/src/primaite/simulator/network/protocols/icmp.py b/src/primaite/simulator/network/protocols/icmp.py index 35b0a05d..b891fe5a 100644 --- a/src/primaite/simulator/network/protocols/icmp.py +++ b/src/primaite/simulator/network/protocols/icmp.py @@ -100,7 +100,7 @@ class ICMPPacket(BaseModel): icmp_type = info.data["icmp_type"] if get_icmp_type_code_description(icmp_type, v): return v - msg = f"No Matching ICMP code for type:{icmp_type.name}, code:{v}" + msg = f"No Matching ICMP code for type: {icmp_type.name}, code: {v}" _LOGGER.error(msg) raise ValueError(msg) @@ -109,6 +109,6 @@ class ICMPPacket(BaseModel): description = get_icmp_type_code_description(self.icmp_type, self.icmp_code) if description: return description - msg = f"No Matching ICMP code for type:{self.icmp_type.name}, code:{self.icmp_code}" + msg = f"No Matching ICMP code for type: {self.icmp_type.name}, code: {self.icmp_code}" _LOGGER.error(msg) raise ValueError(msg) diff --git a/src/primaite/simulator/system/services/dns/dns_client.py b/src/primaite/simulator/system/services/dns/dns_client.py index 063ff74f..1baef868 100644 --- a/src/primaite/simulator/system/services/dns/dns_client.py +++ b/src/primaite/simulator/system/services/dns/dns_client.py @@ -81,7 +81,7 @@ class DNSClient(Service): # check if the domain is already in the DNS cache if target_domain in self.dns_cache: self.sys_log.info( - f"{self.name}: Domain lookup for {target_domain} successful," + f"{self.name}: Domain lookup for {target_domain} successful, " f"resolves to {self.dns_cache[target_domain]}" ) return True diff --git a/src/primaite/simulator/system/services/ftp/ftp_service.py b/src/primaite/simulator/system/services/ftp/ftp_service.py index 70ba74d7..75c8daab 100644 --- a/src/primaite/simulator/system/services/ftp/ftp_service.py +++ b/src/primaite/simulator/system/services/ftp/ftp_service.py @@ -147,7 +147,7 @@ class FTPServiceABC(Service, ABC): retrieved_file: File = self.file_system.get_file(folder_name=folder_name, file_name=file_name) # if file does not exist, return an error - if not retrieved_file: + if not retrieved_file: # noqa self.sys_log.error( f"File {payload.ftp_command_args['dest_folder_name']}/" f"{payload.ftp_command_args['dest_file_name']} does not exist in {self.sys_log.hostname}" diff --git a/src/primaite/simulator/system/services/icmp/icmp.py b/src/primaite/simulator/system/services/icmp/icmp.py index c4b4173f..14aa47e4 100644 --- a/src/primaite/simulator/system/services/icmp/icmp.py +++ b/src/primaite/simulator/system/services/icmp/icmp.py @@ -60,7 +60,7 @@ class ICMP(Service): if target_ip_address.is_loopback: self.sys_log.info("Pinging loopback address") return any(network_interface.enabled for network_interface in self.network_interfaces.values()) - self.sys_log.info(f"Pinging {target_ip_address}:", to_terminal=True) + self.sys_log.info(f"Pinging {target_ip_address}: ", to_terminal=True) sequence, identifier = 0, None while sequence < pings: sequence, identifier = self._send_icmp_echo_request(target_ip_address, sequence, identifier, pings) From a2482dc7446712e658cba58674f7867bc4058af9 Mon Sep 17 00:00:00 2001 From: Charlie Crane Date: Wed, 1 May 2024 16:22:26 +0100 Subject: [PATCH 17/37] #2464 - Forgot to save and add a file that would cause flake8 failure --- src/primaite/simulator/network/hardware/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py index 27a31e84..1d305598 100644 --- a/src/primaite/simulator/network/hardware/base.py +++ b/src/primaite/simulator/network/hardware/base.py @@ -652,7 +652,7 @@ class Link(SimComponent): # Load the frame size on the link self.current_load += frame_size _LOGGER.debug( - f"Added {frame_size:.3f} Mbits to {self}, current load {self.current_load:.3f} Mbits " + f"Added {frame_size:.3f} Mbits to {self}, current load {self.current_load:.3f} Mbits " # noqa f"({self.current_load_percent})" ) return True From 952f6ee225a11bf843f355ae262b0fcaf678c41c Mon Sep 17 00:00:00 2001 From: Czar Echavez Date: Wed, 1 May 2024 16:23:44 +0100 Subject: [PATCH 18/37] #2533: fix is_dev_mode for missing or outdated configs --- src/primaite/utils/cli/primaite_config_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/primaite/utils/cli/primaite_config_utils.py b/src/primaite/utils/cli/primaite_config_utils.py index 966376d9..fa9b68ff 100644 --- a/src/primaite/utils/cli/primaite_config_utils.py +++ b/src/primaite/utils/cli/primaite_config_utils.py @@ -7,7 +7,7 @@ from primaite import PRIMAITE_CONFIG, PRIMAITE_PATHS def is_dev_mode() -> bool: """Returns True if PrimAITE is currently running in developer mode.""" - return PRIMAITE_CONFIG["developer_mode"]["enabled"] + return PRIMAITE_CONFIG.get("developer_mode", {}).get("enabled", False) def update_primaite_application_config(config: Optional[Dict] = None) -> None: From 0fd85722aeae3d9bff062f33c18a5bf2a5aa92af Mon Sep 17 00:00:00 2001 From: Charlie Crane Date: Wed, 1 May 2024 16:37:16 +0100 Subject: [PATCH 19/37] #2464 I missed a linting error --- .flake8 | 1 + docs/conf.py | 3 +-- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.flake8 b/.flake8 index c2d9e4bb..ba2f3864 100644 --- a/.flake8 +++ b/.flake8 @@ -7,6 +7,7 @@ extend-ignore = D104 E203 E712 + E713 D401 F811 ANN002 diff --git a/docs/conf.py b/docs/conf.py index a666e460..93e8a447 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -30,8 +30,7 @@ release = version # set global variables rst_prolog = f""" -.. |VERSION| replace:: {release} -""" +.. |VERSION| replace:: {release} """ # noqa html_title = f"{project} v{release} docs" From ec6d2bf6401eb9d5e3a43fd707a701573df2fb40 Mon Sep 17 00:00:00 2001 From: Czar Echavez Date: Wed, 1 May 2024 16:53:37 +0100 Subject: [PATCH 20/37] #2533: fix missing implementation of output_dir override in io.py + moving the dev mode override in SIM_OUTPUT to prevent the override from being overridden --- src/primaite/session/io.py | 7 +++++-- src/primaite/simulator/__init__.py | 20 +++++++++++--------- 2 files changed, 16 insertions(+), 11 deletions(-) diff --git a/src/primaite/session/io.py b/src/primaite/session/io.py index 9dcc529d..8bbc1b07 100644 --- a/src/primaite/session/io.py +++ b/src/primaite/session/io.py @@ -5,7 +5,7 @@ from typing import Dict, List, Optional from pydantic import BaseModel, ConfigDict -from primaite import _PRIMAITE_ROOT, getLogger, PRIMAITE_PATHS +from primaite import _PRIMAITE_ROOT, getLogger, PRIMAITE_CONFIG, PRIMAITE_PATHS from primaite.simulator import LogLevel, SIM_OUTPUT from primaite.utils.cli.primaite_config_utils import is_dev_mode @@ -66,9 +66,12 @@ class PrimaiteIO: # check if running in dev mode if is_dev_mode(): - # check if there is an output directory set in config session_path = _PRIMAITE_ROOT.parent.parent / "sessions" / date_str / time_str + # check if there is an output directory set in config + if PRIMAITE_CONFIG["developer_mode"]["output_dir"]: + session_path = Path(PRIMAITE_CONFIG["developer_mode"]["output_dir"]) / "sessions" / date_str / time_str + session_path.mkdir(exist_ok=True, parents=True) return session_path diff --git a/src/primaite/simulator/__init__.py b/src/primaite/simulator/__init__.py index e9d44ec3..bbcf9af4 100644 --- a/src/primaite/simulator/__init__.py +++ b/src/primaite/simulator/__init__.py @@ -32,7 +32,17 @@ class _SimOutput: path = PRIMAITE_PATHS.user_sessions_path / date_str / time_str + self._path = path + self._save_pcap_logs: bool = False + self._save_sys_logs: bool = False + self._write_sys_log_to_terminal: bool = False + self._sys_log_level: LogLevel = LogLevel.WARNING # default log level is at WARNING + + @property + def path(self) -> Path: if is_dev_mode(): + date_str = datetime.now().strftime("%Y-%m-%d") + time_str = datetime.now().strftime("%H-%M-%S") # if dev mode is enabled, if output dir is not set, print to primaite repo root path: Path = _PRIMAITE_ROOT.parent.parent / "sessions" / date_str / time_str / "simulation_output" # otherwise print to output dir @@ -44,15 +54,7 @@ class _SimOutput: / time_str / "simulation_output" ) - - self._path = path - self._save_pcap_logs: bool = False - self._save_sys_logs: bool = False - self._write_sys_log_to_terminal: bool = False - self._sys_log_level: LogLevel = LogLevel.WARNING # default log level is at WARNING - - @property - def path(self) -> Path: + self._path = path return self._path @path.setter From d25a8aa0efab14ef295a28b78ddceffce49034fa Mon Sep 17 00:00:00 2001 From: Charlie Crane Date: Wed, 1 May 2024 17:06:58 +0100 Subject: [PATCH 21/37] #2464 - Reverting changes from Python3.12 investigations, updating README and prproject.toml so that python 3.11 is allowed --- .azure/azure-ci-build-pipeline.yaml | 57 +++++++------------ README.md | 2 +- docs/conf.py | 3 +- pyproject.toml | 3 +- src/primaite/game/agent/rewards.py | 2 +- src/primaite/simulator/network/container.py | 2 +- .../simulator/network/hardware/base.py | 8 +-- .../simulator/network/protocols/icmp.py | 4 +- .../system/services/dns/dns_client.py | 2 +- .../system/services/ftp/ftp_service.py | 2 +- .../simulator/system/services/icmp/icmp.py | 2 +- 11 files changed, 37 insertions(+), 50 deletions(-) diff --git a/.azure/azure-ci-build-pipeline.yaml b/.azure/azure-ci-build-pipeline.yaml index 1a7b4870..f0a1793e 100644 --- a/.azure/azure-ci-build-pipeline.yaml +++ b/.azure/azure-ci-build-pipeline.yaml @@ -14,48 +14,33 @@ parameters: - name: matrix type: object default: - # - job_name: 'UbuntuPython38' - # py: '3.8' - # img: 'ubuntu-latest' - # every_time: false - # publish_coverage: false - # - job_name: 'UbuntuPython310' - # py: '3.10' - # img: 'ubuntu-latest' - # every_time: true - # publish_coverage: true - - job_name: 'UbuntuPython312_TEST' - py: '3.12' + - job_name: 'UbuntuPython38' + py: '3.8' + img: 'ubuntu-latest' + every_time: false + publish_coverage: false + - job_name: 'UbuntuPython311' + py: '3.11' img: 'ubuntu-latest' every_time: true publish_coverage: true - # - job_name: 'WindowsPython38' - # py: '3.8' - # img: 'windows-latest' - # every_time: false - # publish_coverage: false - # - job_name: 'WindowsPython310' - # py: '3.10' - # img: 'windows-latest' - # every_time: false - # publish_coverage: false - - job_name: 'WindowsPython312_TEST' - py: '3.12' + - job_name: 'WindowsPython38' + py: '3.8' img: 'windows-latest' every_time: false publish_coverage: false - # - job_name: 'MacOSPython38' - # py: '3.8' - # img: 'macOS-latest' - # every_time: false - # publish_coverage: false - # - job_name: 'MacOSPython310' - # py: '3.10' - # img: 'macOS-latest' - # every_time: false - # publish_coverage: false - - job_name: 'MacOSPython312_TEST' - py: '3.12' + - job_name: 'WindowsPython311' + py: '3.11' + img: 'windows-latest' + every_time: false + publish_coverage: false + - job_name: 'MacOSPython38' + py: '3.8' + img: 'macOS-latest' + every_time: false + publish_coverage: false + - job_name: 'MacOSPython311' + py: '3.11' img: 'macOS-latest' every_time: false publish_coverage: false diff --git a/README.md b/README.md index 927aa17c..c57d9c6b 100644 --- a/README.md +++ b/README.md @@ -51,7 +51,7 @@ primaite setup #### Unix **Prerequisites:** -* Manual install of Python >= 3.8 < 3.11 +* Manual install of Python >= 3.8 < 3.12 ``` bash sudo add-apt-repository ppa:deadsnakes/ppa diff --git a/docs/conf.py b/docs/conf.py index 93e8a447..a666e460 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -30,7 +30,8 @@ release = version # set global variables rst_prolog = f""" -.. |VERSION| replace:: {release} """ # noqa +.. |VERSION| replace:: {release} +""" html_title = f"{project} v{release} docs" diff --git a/pyproject.toml b/pyproject.toml index c71b55a6..2619da90 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ name = "primaite" description = "PrimAITE (Primary-level AI Training Environment) is a simulation environment for training AI under the ARCD programme." authors = [{name="Defence Science and Technology Laboratory UK", email="oss@dstl.gov.uk"}] license = {file = "LICENSE"} -requires-python = ">=3.8, <3.13" +requires-python = ">=3.8, <3.12" dynamic = ["version", "readme"] classifiers = [ "License :: OSI Approved :: MIT License", @@ -20,6 +20,7 @@ classifiers = [ "Programming Language :: Python :: 3.8", "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3 :: Only", ] diff --git a/src/primaite/game/agent/rewards.py b/src/primaite/game/agent/rewards.py index 36a984dd..0222bfcc 100644 --- a/src/primaite/game/agent/rewards.py +++ b/src/primaite/game/agent/rewards.py @@ -231,7 +231,7 @@ class WebpageUnavailablePenalty(AbstractReward): # If the last request did actually go through, then check if the webpage also loaded 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: # noqa + if web_browser_state is NOT_PRESENT_IN_STATE or "history" not in web_browser_state: _LOGGER.debug( "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", diff --git a/src/primaite/simulator/network/container.py b/src/primaite/simulator/network/container.py index 717d04b3..17308c97 100644 --- a/src/primaite/simulator/network/container.py +++ b/src/primaite/simulator/network/container.py @@ -248,7 +248,7 @@ class Network(SimComponent): hostname_b = node_b.hostname if node_b else None port_a = link.endpoint_a.port_num port_b = link.endpoint_b.port_num - link_key = f"{hostname_a}: eth-{port_a}<->{hostname_b}: eth-{port_b}" + link_key = f"{hostname_a}:eth-{port_a}<->{hostname_b}:eth-{port_b}" state["links"][link_key] = link.describe_state() state["links"][link_key]["hostname_a"] = hostname_a state["links"][link_key]["hostname_b"] = hostname_b diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py index 1d305598..78a0bf2d 100644 --- a/src/primaite/simulator/network/hardware/base.py +++ b/src/primaite/simulator/network/hardware/base.py @@ -56,7 +56,7 @@ def generate_mac_address(oui: Optional[str] = None) -> str: if oui: oui_pattern = re.compile(r"^([0-9A-Fa-f]{2}[:-]){2}[0-9A-Fa-f]{2}$") if not oui_pattern.match(oui): - msg = f"Invalid oui. The oui should be in the format xx:xx:xx, where x is a hexadecimal digit, got '{oui}'" # noqa + msg = f"Invalid oui. The oui should be in the format xx:xx:xx, where x is a hexadecimal digit, got '{oui}'" _LOGGER.error(msg) raise ValueError(msg) oui_bytes = [int(chunk, 16) for chunk in oui.split(":")] @@ -64,7 +64,7 @@ def generate_mac_address(oui: Optional[str] = None) -> str: else: mac = random_bytes - return ":".join(f"{b:02x}" for b in mac) # noqa + return ":".join(f"{b:02x}" for b in mac) class NetworkInterface(SimComponent, ABC): @@ -599,7 +599,7 @@ class Link(SimComponent): @property def current_load_percent(self) -> str: """Get the current load formatted as a percentage string.""" - return f"{self.current_load / self.bandwidth:.5f}%" # noqa + return f"{self.current_load / self.bandwidth:.5f}%" def endpoint_up(self): """Let the Link know and endpoint has been brought up.""" @@ -652,7 +652,7 @@ class Link(SimComponent): # Load the frame size on the link self.current_load += frame_size _LOGGER.debug( - f"Added {frame_size:.3f} Mbits to {self}, current load {self.current_load:.3f} Mbits " # noqa + f"Added {frame_size:.3f} Mbits to {self}, current load {self.current_load:.3f} Mbits " f"({self.current_load_percent})" ) return True diff --git a/src/primaite/simulator/network/protocols/icmp.py b/src/primaite/simulator/network/protocols/icmp.py index b891fe5a..35b0a05d 100644 --- a/src/primaite/simulator/network/protocols/icmp.py +++ b/src/primaite/simulator/network/protocols/icmp.py @@ -100,7 +100,7 @@ class ICMPPacket(BaseModel): icmp_type = info.data["icmp_type"] if get_icmp_type_code_description(icmp_type, v): return v - msg = f"No Matching ICMP code for type: {icmp_type.name}, code: {v}" + msg = f"No Matching ICMP code for type:{icmp_type.name}, code:{v}" _LOGGER.error(msg) raise ValueError(msg) @@ -109,6 +109,6 @@ class ICMPPacket(BaseModel): description = get_icmp_type_code_description(self.icmp_type, self.icmp_code) if description: return description - msg = f"No Matching ICMP code for type: {self.icmp_type.name}, code: {self.icmp_code}" + msg = f"No Matching ICMP code for type:{self.icmp_type.name}, code:{self.icmp_code}" _LOGGER.error(msg) raise ValueError(msg) diff --git a/src/primaite/simulator/system/services/dns/dns_client.py b/src/primaite/simulator/system/services/dns/dns_client.py index 1baef868..063ff74f 100644 --- a/src/primaite/simulator/system/services/dns/dns_client.py +++ b/src/primaite/simulator/system/services/dns/dns_client.py @@ -81,7 +81,7 @@ class DNSClient(Service): # check if the domain is already in the DNS cache if target_domain in self.dns_cache: self.sys_log.info( - f"{self.name}: Domain lookup for {target_domain} successful, " + f"{self.name}: Domain lookup for {target_domain} successful," f"resolves to {self.dns_cache[target_domain]}" ) return True diff --git a/src/primaite/simulator/system/services/ftp/ftp_service.py b/src/primaite/simulator/system/services/ftp/ftp_service.py index 75c8daab..70ba74d7 100644 --- a/src/primaite/simulator/system/services/ftp/ftp_service.py +++ b/src/primaite/simulator/system/services/ftp/ftp_service.py @@ -147,7 +147,7 @@ class FTPServiceABC(Service, ABC): retrieved_file: File = self.file_system.get_file(folder_name=folder_name, file_name=file_name) # if file does not exist, return an error - if not retrieved_file: # noqa + if not retrieved_file: self.sys_log.error( f"File {payload.ftp_command_args['dest_folder_name']}/" f"{payload.ftp_command_args['dest_file_name']} does not exist in {self.sys_log.hostname}" diff --git a/src/primaite/simulator/system/services/icmp/icmp.py b/src/primaite/simulator/system/services/icmp/icmp.py index 14aa47e4..c4b4173f 100644 --- a/src/primaite/simulator/system/services/icmp/icmp.py +++ b/src/primaite/simulator/system/services/icmp/icmp.py @@ -60,7 +60,7 @@ class ICMP(Service): if target_ip_address.is_loopback: self.sys_log.info("Pinging loopback address") return any(network_interface.enabled for network_interface in self.network_interfaces.values()) - self.sys_log.info(f"Pinging {target_ip_address}: ", to_terminal=True) + self.sys_log.info(f"Pinging {target_ip_address}:", to_terminal=True) sequence, identifier = 0, None while sequence < pings: sequence, identifier = self._send_icmp_echo_request(target_ip_address, sequence, identifier, pings) From ec582c0002c18010f524cd83d0180a8ff27a7922 Mon Sep 17 00:00:00 2001 From: Charlie Crane Date: Wed, 1 May 2024 17:08:08 +0100 Subject: [PATCH 22/37] #2464 - Reverting flake8 change --- .flake8 | 1 - 1 file changed, 1 deletion(-) diff --git a/.flake8 b/.flake8 index ba2f3864..c2d9e4bb 100644 --- a/.flake8 +++ b/.flake8 @@ -7,7 +7,6 @@ extend-ignore = D104 E203 E712 - E713 D401 F811 ANN002 From 5cfa9c7fa261ac5c9caa5f8c8c46aff41dce54ca Mon Sep 17 00:00:00 2001 From: Charlie Crane Date: Thu, 2 May 2024 08:54:37 +0100 Subject: [PATCH 23/37] #2464 - Version typo in README.md --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index c57d9c6b..2265538a 100644 --- a/README.md +++ b/README.md @@ -33,7 +33,7 @@ Currently, the PrimAITE wheel can only be installed from GitHub. This may change #### Windows (PowerShell) **Prerequisites:** -* Manual install of Python >= 3.8 < 3.11 +* Manual install of Python >= 3.8 < 3.12 **Install:** From e599e03c10bf491e5faffed52ca62e3d4550bfe5 Mon Sep 17 00:00:00 2001 From: Czar Echavez Date: Thu, 2 May 2024 12:21:24 +0100 Subject: [PATCH 24/37] #2533: cleaner output text for dev-mode show + consistent colors --- src/primaite/utils/cli/dev_cli.py | 22 +++++++++++++++------- 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/src/primaite/utils/cli/dev_cli.py b/src/primaite/utils/cli/dev_cli.py index 7437d2fc..d2c8e370 100644 --- a/src/primaite/utils/cli/dev_cli.py +++ b/src/primaite/utils/cli/dev_cli.py @@ -1,6 +1,7 @@ import click import typer from rich import print +from rich.table import Table from typing_extensions import Annotated from primaite import _PRIMAITE_ROOT, PRIMAITE_CONFIG @@ -10,9 +11,9 @@ from primaite.utils.cli.primaite_config_utils import is_dev_mode, update_primait dev = typer.Typer() PRODUCTION_MODE_MESSAGE = ( - "\n[green]:monkey_face::monkey_face::monkey_face: " + "\n[green]:rocket::rocket::rocket: " " PrimAITE is running in Production mode " - " :monkey_face::monkey_face::monkey_face: [/green]\n" + " :rocket::rocket::rocket: [/green]\n" ) DEVELOPER_MODE_MESSAGE = ( @@ -39,8 +40,15 @@ def show(): """Show if PrimAITE is in development mode or production mode.""" # print if dev mode is enabled print(DEVELOPER_MODE_MESSAGE if is_dev_mode() else PRODUCTION_MODE_MESSAGE) - print(f"Current Settings: {PRIMAITE_CONFIG['developer_mode']}") - print("\nTo see available options, use [medium_turquoise]`primaite dev-mode --help`[/medium_turquoise]\n") + + table = Table(title="Current Dev-Mode Settings") + table.add_column("Setting", style="cyan") + table.add_column("Value", style="default") + for setting, value in PRIMAITE_CONFIG["developer_mode"].items(): + table.add_row(setting, str(value)) + + print(table) + print("\nTo see available options, use [cyan]`primaite dev-mode --help`[/cyan]\n") @dev.command() @@ -150,9 +158,9 @@ def path( # update application config update_primaite_application_config() print( - f"PrimAITE dev-mode output_dir [medium_turquoise]" + f"PrimAITE dev-mode output_dir [cyan]" f"{str(_PRIMAITE_ROOT.parent.parent / 'simulation_output')}" - f"[/medium_turquoise]" + f"[/cyan]" ) return @@ -160,4 +168,4 @@ def path( PRIMAITE_CONFIG["developer_mode"]["output_dir"] = directory # update application config update_primaite_application_config() - print(f"PrimAITE dev-mode output_dir [medium_turquoise]{directory}[/medium_turquoise]") + print(f"PrimAITE dev-mode output_dir [cyan]{directory}[/cyan]") From db29d336290f75e27fa65937b9481a60cf9b192b Mon Sep 17 00:00:00 2001 From: Nick Todd Date: Thu, 2 May 2024 15:44:18 +0100 Subject: [PATCH 25/37] #2442: Add multi-processing notebook example. --- src/primaite/notebooks/multi-processing.ipynb | 148 ++++++++++++++++++ 1 file changed, 148 insertions(+) create mode 100644 src/primaite/notebooks/multi-processing.ipynb diff --git a/src/primaite/notebooks/multi-processing.ipynb b/src/primaite/notebooks/multi-processing.ipynb new file mode 100644 index 00000000..83366b54 --- /dev/null +++ b/src/primaite/notebooks/multi-processing.ipynb @@ -0,0 +1,148 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Simple multi-processing demo using SubprocVecEnv from SB3\n", + "Based on a code example provided by Rachael Proctor." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "#!primaite setup" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Import packages and read config file." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Set up training data." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import multiprocessing as mp\n", + "mp.get_all_start_methods()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import yaml\n", + "from stable_baselines3 import PPO\n", + "from stable_baselines3.common.utils import set_random_seed\n", + "from stable_baselines3.common.vec_env import SubprocVecEnv\n", + "\n", + "from primaite.session.environment import PrimaiteGymEnv\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "EPISODE_LEN = 128\n", + "NUM_EPISODES = 10\n", + "NO_STEPS = EPISODE_LEN * NUM_EPISODES\n", + "BATCH_SIZE = 32\n", + "LEARNING_RATE = 3e-4\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "with open(\"c:/projects/primaite/src/primaite/config/_package_data/data_manipulation.yaml\", \"r\") as f:\n", + " cfg = yaml.safe_load(f)\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "\n", + "def make_env(rank: int, seed: int = 0) -> callable:\n", + " \"\"\"Wrapper script for _init function.\"\"\"\n", + "\n", + " def _init() -> PrimaiteGymEnv:\n", + " env = PrimaiteGymEnv(env_config=cfg)\n", + " env.reset(seed=seed + rank)\n", + " model = PPO(\n", + " \"MlpPolicy\",\n", + " env,\n", + " learning_rate=LEARNING_RATE,\n", + " n_steps=NO_STEPS,\n", + " batch_size=BATCH_SIZE,\n", + " verbose=0,\n", + " tensorboard_log=\"./PPO_UC2/\",\n", + " )\n", + " model.learn(total_timesteps=NO_STEPS)\n", + " return env\n", + "\n", + " set_random_seed(seed)\n", + " return _init\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "\n", + "n_procs = 4\n", + "train_env = SubprocVecEnv([make_env(i + n_procs) for i in range(n_procs)])\n", + "print(train_env)\n" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": ".venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.11" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} From b9b922e776a3e0e07099e04761bbaad730f8a0c2 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Thu, 2 May 2024 15:51:06 +0100 Subject: [PATCH 26/37] Make ftp client only default on computers, not servers --- .../simulator/network/hardware/nodes/host/computer.py | 5 +++++ .../simulator/network/hardware/nodes/host/host_node.py | 2 -- tests/conftest.py | 6 +++++- .../network/test_multi_lan_internet_example_network.py | 2 +- tests/integration_tests/system/test_database_on_node.py | 3 ++- 5 files changed, 13 insertions(+), 5 deletions(-) diff --git a/src/primaite/simulator/network/hardware/nodes/host/computer.py b/src/primaite/simulator/network/hardware/nodes/host/computer.py index 0b13163e..7ce64867 100644 --- a/src/primaite/simulator/network/hardware/nodes/host/computer.py +++ b/src/primaite/simulator/network/hardware/nodes/host/computer.py @@ -1,4 +1,7 @@ +from typing import ClassVar, Dict + from primaite.simulator.network.hardware.nodes.host.host_node import HostNode +from primaite.simulator.system.services.ftp.ftp_client import FTPClient class Computer(HostNode): @@ -29,4 +32,6 @@ class Computer(HostNode): * Web Browser """ + SYSTEM_SOFTWARE: ClassVar[Dict] = {**HostNode.SYSTEM_SOFTWARE, "FTPClient": FTPClient} + pass diff --git a/src/primaite/simulator/network/hardware/nodes/host/host_node.py b/src/primaite/simulator/network/hardware/nodes/host/host_node.py index df254b1e..caea2dd7 100644 --- a/src/primaite/simulator/network/hardware/nodes/host/host_node.py +++ b/src/primaite/simulator/network/hardware/nodes/host/host_node.py @@ -10,7 +10,6 @@ from primaite.simulator.network.transmission.data_link_layer import Frame from primaite.simulator.system.applications.web_browser import WebBrowser from primaite.simulator.system.services.arp.arp import ARP, ARPPacket from primaite.simulator.system.services.dns.dns_client import DNSClient -from primaite.simulator.system.services.ftp.ftp_client import FTPClient from primaite.simulator.system.services.icmp.icmp import ICMP from primaite.simulator.system.services.ntp.ntp_client import NTPClient from primaite.utils.validators import IPV4Address @@ -301,7 +300,6 @@ class HostNode(Node): "HostARP": HostARP, "ICMP": ICMP, "DNSClient": DNSClient, - "FTPClient": FTPClient, "NTPClient": NTPClient, "WebBrowser": WebBrowser, } diff --git a/tests/conftest.py b/tests/conftest.py index 7de2bfde..a9ed42c6 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -86,7 +86,11 @@ class TestApplication(Application): @pytest.fixture(scope="function") def uc2_network() -> Network: - return arcd_uc2_network() + with open(PRIMAITE_PATHS.user_config_path / "example_config" / "data_manipulation.yaml") as f: + cfg = yaml.safe_load(f) + game = PrimaiteGame.from_config(cfg) + return game.simulation.network + # return arcd_uc2_network() @pytest.fixture(scope="function") diff --git a/tests/integration_tests/network/test_multi_lan_internet_example_network.py b/tests/integration_tests/network/test_multi_lan_internet_example_network.py index f56fcbf2..f6d702d8 100644 --- a/tests/integration_tests/network/test_multi_lan_internet_example_network.py +++ b/tests/integration_tests/network/test_multi_lan_internet_example_network.py @@ -83,7 +83,7 @@ def test_sometech_webserver_cannot_access_ftp_on_sometech_storage_server(): some_tech_storage_srv.file_system.create_file(file_name="test.png") web_server: Server = network.get_node_by_hostname("some_tech_web_srv") - + web_server.software_manager.install(FTPClient) web_ftp_client: FTPClient = web_server.software_manager.software["FTPClient"] assert not web_ftp_client.request_file( diff --git a/tests/integration_tests/system/test_database_on_node.py b/tests/integration_tests/system/test_database_on_node.py index 9a396fae..8a2a9793 100644 --- a/tests/integration_tests/system/test_database_on_node.py +++ b/tests/integration_tests/system/test_database_on_node.py @@ -101,7 +101,7 @@ def test_database_client_native_connection_query(uc2_network): """Tests DB query across the network returns HTTP status 200 and date.""" web_server: Server = uc2_network.get_node_by_hostname("web_server") db_client: DatabaseClient = web_server.software_manager.software["DatabaseClient"] - + db_client.connect() assert db_client.query(sql="SELECT") assert db_client.query(sql="INSERT") @@ -222,6 +222,7 @@ def test_database_client_cannot_query_offline_database_server(uc2_network): web_server: Server = uc2_network.get_node_by_hostname("web_server") db_client: DatabaseClient = web_server.software_manager.software.get("DatabaseClient") + db_client.connect() assert len(db_client.client_connections) # Establish a new connection to the DatabaseService From a37ed4e487d2f5f7ba716ac2eeaca0c8a05ccf53 Mon Sep 17 00:00:00 2001 From: Nick Todd Date: Thu, 2 May 2024 16:16:36 +0100 Subject: [PATCH 27/37] #2442: Fix config file path to be OS independent. --- src/primaite/notebooks/multi-processing.ipynb | 47 +++++++++---------- 1 file changed, 21 insertions(+), 26 deletions(-) diff --git a/src/primaite/notebooks/multi-processing.ipynb b/src/primaite/notebooks/multi-processing.ipynb index 83366b54..fc3ef088 100644 --- a/src/primaite/notebooks/multi-processing.ipynb +++ b/src/primaite/notebooks/multi-processing.ipynb @@ -21,24 +21,14 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "Import packages and read config file." + "Set up training data." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ - "Set up training data." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import multiprocessing as mp\n", - "mp.get_all_start_methods()" + "Import packages and read config file." ] }, { @@ -61,12 +51,17 @@ "metadata": {}, "outputs": [], "source": [ - "\n", - "EPISODE_LEN = 128\n", - "NUM_EPISODES = 10\n", - "NO_STEPS = EPISODE_LEN * NUM_EPISODES\n", - "BATCH_SIZE = 32\n", - "LEARNING_RATE = 3e-4\n" + "from primaite.config.load import data_manipulation_config_path" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "with open(data_manipulation_config_path(), 'r') as f:\n", + " cfg = yaml.safe_load(f)" ] }, { @@ -76,8 +71,11 @@ "outputs": [], "source": [ "\n", - "with open(\"c:/projects/primaite/src/primaite/config/_package_data/data_manipulation.yaml\", \"r\") as f:\n", - " cfg = yaml.safe_load(f)\n" + "EPISODE_LEN = 128\n", + "NUM_EPISODES = 10\n", + "NO_STEPS = EPISODE_LEN * NUM_EPISODES\n", + "BATCH_SIZE = 32\n", + "LEARNING_RATE = 3e-4\n" ] }, { @@ -116,11 +114,8 @@ "metadata": {}, "outputs": [], "source": [ - "\n", - "\n", - "n_procs = 4\n", - "train_env = SubprocVecEnv([make_env(i + n_procs) for i in range(n_procs)])\n", - "print(train_env)\n" + "n_procs = 2\n", + "train_env = SubprocVecEnv([make_env(i + n_procs) for i in range(n_procs)])\n" ] } ], @@ -140,7 +135,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.11" + "version": "3.10.12" } }, "nbformat": 4, From 7c689a0d35c7d477a0b064bcb4e4a6f001c9ffbb Mon Sep 17 00:00:00 2001 From: Nick Todd Date: Thu, 2 May 2024 16:32:54 +0100 Subject: [PATCH 28/37] #2442: Markdown changes. --- src/primaite/notebooks/multi-processing.ipynb | 30 ++++++++++++++----- 1 file changed, 22 insertions(+), 8 deletions(-) diff --git a/src/primaite/notebooks/multi-processing.ipynb b/src/primaite/notebooks/multi-processing.ipynb index fc3ef088..d5b56cbb 100644 --- a/src/primaite/notebooks/multi-processing.ipynb +++ b/src/primaite/notebooks/multi-processing.ipynb @@ -17,13 +17,6 @@ "#!primaite setup" ] }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Set up training data." - ] - }, { "cell_type": "markdown", "metadata": {}, @@ -64,6 +57,13 @@ " cfg = yaml.safe_load(f)" ] }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Set up training data." + ] + }, { "cell_type": "code", "execution_count": null, @@ -78,6 +78,13 @@ "LEARNING_RATE = 3e-4\n" ] }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Define an environment function." + ] + }, { "cell_type": "code", "execution_count": null, @@ -108,6 +115,13 @@ " return _init\n" ] }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Run experiment." + ] + }, { "cell_type": "code", "execution_count": null, @@ -135,7 +149,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.12" + "version": "3.10.11" } }, "nbformat": 4, From 5a3b6ade2f46f4c0ae7bf00fb9734412ad36f547 Mon Sep 17 00:00:00 2001 From: Nick Todd Date: Thu, 2 May 2024 16:52:35 +0100 Subject: [PATCH 29/37] #2442: Updated changelog. --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7a96b27c..d7d88e5f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -20,7 +20,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Added the ability for a DatabaseService to terminate a connection. - Added active_connection to DatabaseClientConnection so that if the connection is terminated active_connection is set to False and the object can no longer be used. - Added additional show functions to enable connection inspection. - +- Added notebook to demonstrate use of SubprocVecEnv from SB3 to vectorise environments to speed up training. ## [Unreleased] - Made requests fail to reach their target if the node is off From 4a02d1d8e5a7adebf6a41270c9067aa51a5470a9 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Thu, 2 May 2024 17:00:29 +0100 Subject: [PATCH 30/37] Readd import statement --- tests/conftest.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/conftest.py b/tests/conftest.py index 51740d7b..c806dfc2 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -2,8 +2,9 @@ from typing import Any, Dict, Tuple import pytest +import yaml -from primaite import getLogger +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.observation_manager import NestedObservation, ObservationManager From 689400e95cfb8ee469fec12c1dd246a3cf2c6829 Mon Sep 17 00:00:00 2001 From: Nick Todd Date: Thu, 2 May 2024 17:04:43 +0100 Subject: [PATCH 31/37] Merge branch 'dev' into bugfix/2442-add_SubprocVecEnv_support --- src/primaite/notebooks/mp.py | 42 ------------------------------------ 1 file changed, 42 deletions(-) delete mode 100644 src/primaite/notebooks/mp.py diff --git a/src/primaite/notebooks/mp.py b/src/primaite/notebooks/mp.py deleted file mode 100644 index ebed7122..00000000 --- a/src/primaite/notebooks/mp.py +++ /dev/null @@ -1,42 +0,0 @@ -import yaml -from stable_baselines3 import PPO -from stable_baselines3.common.utils import set_random_seed -from stable_baselines3.common.vec_env import SubprocVecEnv - -from primaite.session.environment import PrimaiteGymEnv - -EPISODE_LEN = 128 -NUM_EPISODES = 10 -NO_STEPS = EPISODE_LEN * NUM_EPISODES -BATCH_SIZE = 32 -LEARNING_RATE = 3e-4 - -with open("c:/projects/primaite/src/primaite/config/_package_data/data_manipulation.yaml", "r") as f: - cfg = yaml.safe_load(f) - - -def make_env(rank: int, seed: int = 0) -> callable: - """Wrapper script for _init function.""" - - def _init() -> PrimaiteGymEnv: - env = PrimaiteGymEnv(env_config=cfg) - env.reset(seed=seed + rank) - model = PPO( - "MlpPolicy", - env, - learning_rate=LEARNING_RATE, - n_steps=NO_STEPS, - batch_size=BATCH_SIZE, - verbose=0, - tensorboard_log="./PPO_UC2/", - ) - model.learn(total_timesteps=NO_STEPS) - return env - - set_random_seed(seed) - return _init - - -if __name__ == "__main__": - n_procs = 4 - train_env = SubprocVecEnv([make_env(i + n_procs) for i in range(n_procs)]) From 4f3818495906a4b1ec2310bcae490cfbf63cf366 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Fri, 3 May 2024 09:08:26 +0100 Subject: [PATCH 32/37] remove redundant comment --- tests/conftest.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/conftest.py b/tests/conftest.py index c806dfc2..83778748 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -70,7 +70,6 @@ def uc2_network() -> Network: cfg = yaml.safe_load(f) game = PrimaiteGame.from_config(cfg) return game.simulation.network - # return arcd_uc2_network() @pytest.fixture(scope="function") From 535e064af56438c4b3226086984f8d08e07aae8c Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Fri, 3 May 2024 11:05:33 +0100 Subject: [PATCH 33/37] unflake test --- .../test_uc2_data_manipulation_scenario.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/e2e_integration_tests/test_uc2_data_manipulation_scenario.py b/tests/e2e_integration_tests/test_uc2_data_manipulation_scenario.py index e598dd19..4e203669 100644 --- a/tests/e2e_integration_tests/test_uc2_data_manipulation_scenario.py +++ b/tests/e2e_integration_tests/test_uc2_data_manipulation_scenario.py @@ -26,6 +26,9 @@ def test_data_manipulation(uc2_network): # First check that the DB client on the web_server can successfully query the users table on the database assert db_connection.query("SELECT") + db_manipulation_bot.data_manipulation_p_of_success = 1.0 + db_manipulation_bot.port_scan_p_of_success = 1.0 + # Now we run the DataManipulationBot db_manipulation_bot.attack() From e510c3ceffaa0edeb6f81b15d0d53259fc332f84 Mon Sep 17 00:00:00 2001 From: Nick Todd Date: Fri, 3 May 2024 15:10:24 +0100 Subject: [PATCH 34/37] #2442: Removed 'primaite setup' call --- src/primaite/notebooks/multi-processing.ipynb | 9 --------- 1 file changed, 9 deletions(-) diff --git a/src/primaite/notebooks/multi-processing.ipynb b/src/primaite/notebooks/multi-processing.ipynb index d5b56cbb..71addce6 100644 --- a/src/primaite/notebooks/multi-processing.ipynb +++ b/src/primaite/notebooks/multi-processing.ipynb @@ -8,15 +8,6 @@ "Based on a code example provided by Rachael Proctor." ] }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "#!primaite setup" - ] - }, { "cell_type": "markdown", "metadata": {}, From 4d3c85bc141464eef5ab592c7d7994705ac4b07b Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Tue, 7 May 2024 15:35:02 +0100 Subject: [PATCH 35/37] #2550 Backport changes into core --- .../notebooks/Training-an-RLLib-Agent.ipynb | 14 +- .../create-simulation_demo.ipynb | 2 +- .../network_simulator_demo.ipynb | 218 ++++++++++-------- src/primaite/simulator/network/airspace.py | 9 +- .../simulator/network/hardware/base.py | 2 +- .../system/applications/application.py | 13 +- .../applications/red_applications/dos_bot.py | 1 + .../red_applications/ransomware_script.py | 8 - .../test_uc2_data_manipulation_scenario.py | 5 + 9 files changed, 151 insertions(+), 121 deletions(-) diff --git a/src/primaite/notebooks/Training-an-RLLib-Agent.ipynb b/src/primaite/notebooks/Training-an-RLLib-Agent.ipynb index 2fe84655..60737ee5 100644 --- a/src/primaite/notebooks/Training-an-RLLib-Agent.ipynb +++ b/src/primaite/notebooks/Training-an-RLLib-Agent.ipynb @@ -45,15 +45,9 @@ "metadata": {}, "outputs": [], "source": [ - "print(cfg['agents'][2]['agent_settings'])" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ + "for agent in cfg['agents']:\n", + " if agent[\"ref\"] == \"defender\":\n", + " agent['agent_settings']['flatten_obs'] = True\n", "env_config = cfg\n", "\n", "config = (\n", @@ -80,7 +74,7 @@ "tune.Tuner(\n", " \"PPO\",\n", " run_config=air.RunConfig(\n", - " stop={\"timesteps_total\": 512}\n", + " stop={\"timesteps_total\": 1e3 * 128}\n", " ),\n", " param_space=config\n", ").fit()\n" diff --git a/src/primaite/simulator/_package_data/create-simulation_demo.ipynb b/src/primaite/simulator/_package_data/create-simulation_demo.ipynb index 06ecd4be..31173022 100644 --- a/src/primaite/simulator/_package_data/create-simulation_demo.ipynb +++ b/src/primaite/simulator/_package_data/create-simulation_demo.ipynb @@ -261,7 +261,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.11" + "version": "3.10.12" } }, "nbformat": 4, diff --git a/src/primaite/simulator/_package_data/network_simulator_demo.ipynb b/src/primaite/simulator/_package_data/network_simulator_demo.ipynb index 7f4cf3b1..107a2565 100644 --- a/src/primaite/simulator/_package_data/network_simulator_demo.ipynb +++ b/src/primaite/simulator/_package_data/network_simulator_demo.ipynb @@ -256,9 +256,11 @@ { "cell_type": "markdown", "id": "22", - "metadata": {}, + "metadata": { + "tags": [] + }, "source": [ - "Calling `switch.sys_log.show()` displays the Switch system log. By default, only the last 10 log entries are displayed, this can be changed by passing `last_n=`." + "Calling `switch.arp.show()` displays the Switch ARP Cache." ] }, { @@ -270,13 +272,33 @@ }, "outputs": [], "source": [ - "network.get_node_by_hostname(\"switch_1\").sys_log.show()" + "network.get_node_by_hostname(\"switch_1\").arp.show()" ] }, { "cell_type": "markdown", "id": "24", "metadata": {}, + "source": [ + "Calling `switch.sys_log.show()` displays the Switch system log. By default, only the last 10 log entries are displayed, this can be changed by passing `last_n=`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "25", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "network.get_node_by_hostname(\"switch_1\").sys_log.show()" + ] + }, + { + "cell_type": "markdown", + "id": "26", + "metadata": {}, "source": [ "### Computer/Server Nodes\n", "\n", @@ -285,7 +307,7 @@ }, { "cell_type": "markdown", - "id": "25", + "id": "27", "metadata": { "tags": [] }, @@ -293,26 +315,6 @@ "Calling `computer.show()` displays the NICs on the Computer/Server." ] }, - { - "cell_type": "code", - "execution_count": null, - "id": "26", - "metadata": { - "tags": [] - }, - "outputs": [], - "source": [ - "network.get_node_by_hostname(\"security_suite\").show()" - ] - }, - { - "cell_type": "markdown", - "id": "27", - "metadata": {}, - "source": [ - "Calling `computer.arp.show()` displays the Computer/Server ARP Cache." - ] - }, { "cell_type": "code", "execution_count": null, @@ -322,7 +324,7 @@ }, "outputs": [], "source": [ - "network.get_node_by_hostname(\"security_suite\").arp.show()" + "network.get_node_by_hostname(\"security_suite\").show()" ] }, { @@ -330,7 +332,7 @@ "id": "29", "metadata": {}, "source": [ - "Calling `computer.sys_log.show()` displays the Computer/Server system log. By default, only the last 10 log entries are displayed, this can be changed by passing `last_n=`." + "Calling `computer.arp.show()` displays the Computer/Server ARP Cache." ] }, { @@ -342,7 +344,7 @@ }, "outputs": [], "source": [ - "network.get_node_by_hostname(\"security_suite\").sys_log.show()" + "network.get_node_by_hostname(\"security_suite\").arp.show()" ] }, { @@ -350,9 +352,7 @@ "id": "31", "metadata": {}, "source": [ - "## Basic Network Comms Check\n", - "\n", - "We can perform a good old ping to check that Nodes are able to communicate with each other." + "Calling `computer.sys_log.show()` displays the Computer/Server system log. By default, only the last 10 log entries are displayed, this can be changed by passing `last_n=`." ] }, { @@ -364,7 +364,7 @@ }, "outputs": [], "source": [ - "network.show(nodes=False, links=False)" + "network.get_node_by_hostname(\"security_suite\").sys_log.show()" ] }, { @@ -372,7 +372,9 @@ "id": "33", "metadata": {}, "source": [ - "We'll first ping client_1's default gateway." + "## Basic Network Comms Check\n", + "\n", + "We can perform a good old ping to check that Nodes are able to communicate with each other." ] }, { @@ -384,27 +386,27 @@ }, "outputs": [], "source": [ - "network.get_node_by_hostname(\"client_1\").ping(\"192.168.10.1\")" + "network.show(nodes=False, links=False)" + ] + }, + { + "cell_type": "markdown", + "id": "35", + "metadata": {}, + "source": [ + "We'll first ping client_1's default gateway." ] }, { "cell_type": "code", "execution_count": null, - "id": "35", + "id": "36", "metadata": { "tags": [] }, "outputs": [], "source": [ - "network.get_node_by_hostname(\"client_1\").sys_log.show(15)" - ] - }, - { - "cell_type": "markdown", - "id": "36", - "metadata": {}, - "source": [ - "Next, we'll ping the interface of the 192.168.1.0/24 Network on the Router (port 1)." + "network.get_node_by_hostname(\"client_1\").ping(\"192.168.10.1\")" ] }, { @@ -416,7 +418,7 @@ }, "outputs": [], "source": [ - "network.get_node_by_hostname(\"client_1\").ping(\"192.168.1.1\")" + "network.get_node_by_hostname(\"client_1\").sys_log.show(15)" ] }, { @@ -424,7 +426,7 @@ "id": "38", "metadata": {}, "source": [ - "And finally, we'll ping the web server." + "Next, we'll ping the interface of the 192.168.1.0/24 Network on the Router (port 1)." ] }, { @@ -436,7 +438,7 @@ }, "outputs": [], "source": [ - "network.get_node_by_hostname(\"client_1\").ping(\"192.168.1.12\")" + "network.get_node_by_hostname(\"client_1\").ping(\"192.168.1.1\")" ] }, { @@ -444,7 +446,7 @@ "id": "40", "metadata": {}, "source": [ - "To confirm that the ping was received and processed by the web_server, we can view the sys log" + "And finally, we'll ping the web server." ] }, { @@ -456,45 +458,45 @@ }, "outputs": [], "source": [ - "network.get_node_by_hostname(\"web_server\").sys_log.show()" + "network.get_node_by_hostname(\"client_1\").ping(\"192.168.1.12\")" ] }, { "cell_type": "markdown", "id": "42", "metadata": {}, + "source": [ + "To confirm that the ping was received and processed by the web_server, we can view the sys log" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "43", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "network.get_node_by_hostname(\"web_server\").sys_log.show()" + ] + }, + { + "cell_type": "markdown", + "id": "44", + "metadata": {}, "source": [ "## Advanced Network Usage\n", "\n", "We can now use the Network to perform some more advanced things." ] }, - { - "cell_type": "markdown", - "id": "43", - "metadata": {}, - "source": [ - "Let's attempt to prevent client_2 from being able to ping the web server. First, we'll confirm that it can ping the server first..." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "44", - "metadata": { - "tags": [] - }, - "outputs": [], - "source": [ - "network.get_node_by_hostname(\"client_2\").ping(\"192.168.1.12\")" - ] - }, { "cell_type": "markdown", "id": "45", "metadata": {}, "source": [ - "If we look at the client_2 sys log we can see that the four ICMP echo requests were sent and four ICMP each replies were received:" + "Let's attempt to prevent client_2 from being able to ping the web server. First, we'll confirm that it can ping the server first..." ] }, { @@ -506,13 +508,33 @@ }, "outputs": [], "source": [ - "network.get_node_by_hostname(\"client_2\").sys_log.show()" + "network.get_node_by_hostname(\"client_2\").ping(\"192.168.1.12\")" ] }, { "cell_type": "markdown", "id": "47", "metadata": {}, + "source": [ + "If we look at the client_2 sys log we can see that the four ICMP echo requests were sent and four ICMP each replies were received:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "48", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "network.get_node_by_hostname(\"client_2\").sys_log.show()" + ] + }, + { + "cell_type": "markdown", + "id": "49", + "metadata": {}, "source": [ "Now we'll add an ACL to block ICMP from 192.168.10.22" ] @@ -520,7 +542,7 @@ { "cell_type": "code", "execution_count": null, - "id": "48", + "id": "50", "metadata": { "tags": [] }, @@ -540,7 +562,7 @@ { "cell_type": "code", "execution_count": null, - "id": "49", + "id": "51", "metadata": { "tags": [] }, @@ -549,32 +571,12 @@ "network.get_node_by_hostname(\"router_1\").acl.show()" ] }, - { - "cell_type": "markdown", - "id": "50", - "metadata": {}, - "source": [ - "Now we attempt (and fail) to ping the web server" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "51", - "metadata": { - "tags": [] - }, - "outputs": [], - "source": [ - "network.get_node_by_hostname(\"client_2\").ping(\"192.168.1.12\")" - ] - }, { "cell_type": "markdown", "id": "52", "metadata": {}, "source": [ - "We can check that the ping was actually sent by client_2 by viewing the sys log" + "Now we attempt (and fail) to ping the web server" ] }, { @@ -586,7 +588,7 @@ }, "outputs": [], "source": [ - "network.get_node_by_hostname(\"client_2\").sys_log.show()" + "network.get_node_by_hostname(\"client_2\").ping(\"192.168.1.12\")" ] }, { @@ -594,7 +596,7 @@ "id": "54", "metadata": {}, "source": [ - "We can check the router sys log to see why the traffic was blocked" + "We can check that the ping was actually sent by client_2 by viewing the sys log" ] }, { @@ -606,7 +608,7 @@ }, "outputs": [], "source": [ - "network.get_node_by_hostname(\"router_1\").sys_log.show()" + "network.get_node_by_hostname(\"client_2\").sys_log.show()" ] }, { @@ -614,7 +616,7 @@ "id": "56", "metadata": {}, "source": [ - "Now a final check to ensure that client_1 can still ping the web_server." + "We can check the router sys log to see why the traffic was blocked" ] }, { @@ -625,6 +627,26 @@ "tags": [] }, "outputs": [], + "source": [ + "network.get_node_by_hostname(\"router_1\").sys_log.show()" + ] + }, + { + "cell_type": "markdown", + "id": "58", + "metadata": {}, + "source": [ + "Now a final check to ensure that client_1 can still ping the web_server." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "59", + "metadata": { + "tags": [] + }, + "outputs": [], "source": [ "network.get_node_by_hostname(\"client_1\").ping(\"192.168.1.12\")" ] @@ -632,7 +654,7 @@ { "cell_type": "code", "execution_count": null, - "id": "58", + "id": "60", "metadata": { "tags": [] }, diff --git a/src/primaite/simulator/network/airspace.py b/src/primaite/simulator/network/airspace.py index 8a00a4a4..907ab233 100644 --- a/src/primaite/simulator/network/airspace.py +++ b/src/primaite/simulator/network/airspace.py @@ -192,7 +192,6 @@ class WirelessNetworkInterface(NetworkInterface, ABC): # Cannot send Frame as the network interface is not enabled return False - @abstractmethod def receive_frame(self, frame: Frame) -> bool: """ Receives a network frame on the network interface. @@ -200,7 +199,13 @@ class WirelessNetworkInterface(NetworkInterface, ABC): :param frame: The network frame being received. :return: A boolean indicating whether the frame was successfully received. """ - pass + if self.enabled: + frame.set_sent_timestamp() + self.pcap.capture_inbound(frame) + self._connected_node.receive_frame(frame, self) + return True + # Cannot receive Frame as the network interface is not enabled + return False class IPWirelessNetworkInterface(WirelessNetworkInterface, Layer3Interface, ABC): diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py index 78a0bf2d..b4d32dc4 100644 --- a/src/primaite/simulator/network/hardware/base.py +++ b/src/primaite/simulator/network/hardware/base.py @@ -1378,7 +1378,7 @@ class Node(SimComponent): application_instance.configure(server_ip_address=IPv4Address(ip_address)) else: pass - + application_instance.install() if application_instance.name in self.software_manager.software: return True else: diff --git a/src/primaite/simulator/system/applications/application.py b/src/primaite/simulator/system/applications/application.py index 26603b43..294de27b 100644 --- a/src/primaite/simulator/system/applications/application.py +++ b/src/primaite/simulator/system/applications/application.py @@ -1,6 +1,6 @@ from abc import abstractmethod from enum import Enum -from typing import Any, Dict, Set +from typing import Any, Dict, Optional, Set from primaite.interface.request import RequestResponse from primaite.simulator.core import RequestManager, RequestType @@ -33,6 +33,10 @@ class Application(IOSoftware): "The number of times the application has been executed. Default is 0." groups: Set[str] = set() "The set of groups to which the application belongs." + install_duration: int = 2 + "How long it takes to install the application." + install_countdown: Optional[int] = None + "The countdown to the end of the installation process. None if not currently installing" def __init__(self, **kwargs): super().__init__(**kwargs) @@ -76,6 +80,12 @@ class Application(IOSoftware): :param timestep: The current timestep of the simulation. """ super().apply_timestep(timestep=timestep) + if self.operating_state is ApplicationOperatingState.INSTALLING: + self.install_countdown -= 1 + if self.install_countdown <= 0: + self.operating_state = ApplicationOperatingState.RUNNING + self.health_state_actual = SoftwareHealthState.GOOD + self.install_countdown = None def pre_timestep(self, timestep: int) -> None: """Apply pre-timestep logic.""" @@ -129,6 +139,7 @@ class Application(IOSoftware): super().install() if self.operating_state == ApplicationOperatingState.CLOSED: self.operating_state = ApplicationOperatingState.INSTALLING + self.install_countdown = self.install_duration def receive(self, payload: Any, session_id: str, **kwargs) -> bool: """ diff --git a/src/primaite/simulator/system/applications/red_applications/dos_bot.py b/src/primaite/simulator/system/applications/red_applications/dos_bot.py index 53fc9740..0e45aad9 100644 --- a/src/primaite/simulator/system/applications/red_applications/dos_bot.py +++ b/src/primaite/simulator/system/applications/red_applications/dos_bot.py @@ -177,4 +177,5 @@ class DoSBot(DatabaseClient): :param timestep: The timestep value to update the bot's state. """ + super().apply_timestep(timestep=timestep) self._application_loop() diff --git a/src/primaite/simulator/system/applications/red_applications/ransomware_script.py b/src/primaite/simulator/system/applications/red_applications/ransomware_script.py index 4c2d7927..8acc07b4 100644 --- a/src/primaite/simulator/system/applications/red_applications/ransomware_script.py +++ b/src/primaite/simulator/system/applications/red_applications/ransomware_script.py @@ -118,14 +118,6 @@ class RansomwareScript(Application): self.sys_log.info(f"{self.name}: Activated!") self.attack_stage = RansomwareAttackStage.ACTIVATE - def apply_timestep(self, timestep: int) -> None: - """ - Apply a timestep to the bot, triggering the application loop. - - :param timestep: The timestep value to update the bot's state. - """ - pass - def run(self) -> bool: """Calls the parent classes execute method before starting the application loop.""" super().run() diff --git a/tests/e2e_integration_tests/test_uc2_data_manipulation_scenario.py b/tests/e2e_integration_tests/test_uc2_data_manipulation_scenario.py index 4e203669..5df8f964 100644 --- a/tests/e2e_integration_tests/test_uc2_data_manipulation_scenario.py +++ b/tests/e2e_integration_tests/test_uc2_data_manipulation_scenario.py @@ -62,6 +62,11 @@ def test_application_install_uninstall_on_uc2(): _, _, _, _, info = env.step(78) assert "DoSBot" in domcon.software_manager.software + # installing takes 3 steps so let's wait for 3 steps + env.step(0) + env.step(0) + env.step(0) + # Test we can now execute the DoSBot app _, _, _, _, info = env.step(81) assert info["agent_actions"]["defender"].response.status == "success" From c6b6753aeb81f9c9c180738763c1daf22a6148fd Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Wed, 8 May 2024 11:13:32 +0000 Subject: [PATCH 36/37] Updated VERSION --- src/primaite/VERSION | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/primaite/VERSION b/src/primaite/VERSION index f08efbdc..9e0b71d0 100644 --- a/src/primaite/VERSION +++ b/src/primaite/VERSION @@ -1 +1 @@ -3.0.0b9dev +3.0.0b9 From ebbdea517e91ad6075e689b49d17a79c265b2eab Mon Sep 17 00:00:00 2001 From: Chris McCarthy Date: Mon, 13 May 2024 07:30:33 +0100 Subject: [PATCH 37/37] 2536 - Removed 'real files' from FileSystem and FTP --- src/primaite/simulator/file_system/file.py | 40 +++---------------- .../simulator/file_system/file_system.py | 13 ------ .../system/services/ftp/ftp_client.py | 1 - .../system/services/ftp/ftp_service.py | 10 ++--- .../_simulator/_file_system/test_file.py | 14 ------- .../_file_system/test_file_system.py | 2 +- .../_simulator/_file_system/test_folder.py | 18 --------- 7 files changed, 10 insertions(+), 88 deletions(-) diff --git a/src/primaite/simulator/file_system/file.py b/src/primaite/simulator/file_system/file.py index a4d54e79..328b052b 100644 --- a/src/primaite/simulator/file_system/file.py +++ b/src/primaite/simulator/file_system/file.py @@ -2,9 +2,7 @@ from __future__ import annotations import hashlib import json -import os.path import warnings -from pathlib import Path from typing import Dict, Optional from primaite import getLogger @@ -21,8 +19,6 @@ class File(FileSystemItemABC): :ivar Folder folder: The folder in which the file resides. :ivar FileType file_type: The type of the file. :ivar Optional[int] sim_size: The simulated file size. - :ivar bool real: Indicates if the file is actually a real file in the Node sim fs output. - :ivar Optional[Path] sim_path: The path if the file is real. """ folder_id: str @@ -33,12 +29,6 @@ class File(FileSystemItemABC): "The type of File." sim_size: Optional[int] = None "The simulated file size." - real: bool = False - "Indicates whether the File is actually a real file in the Node sim fs output." - sim_path: Optional[Path] = None - "The Path if real is True." - sim_root: Optional[Path] = None - "Root path of the simulation." num_access: int = 0 "Number of times the file was accessed in the current step." @@ -67,13 +57,6 @@ class File(FileSystemItemABC): if not kwargs.get("sim_size"): kwargs["sim_size"] = kwargs["file_type"].default_size super().__init__(**kwargs) - if self.real: - self.sim_path = self.sim_root / self.path - if not self.sim_path.exists(): - self.sim_path.parent.mkdir(exist_ok=True, parents=True) - with open(self.sim_path, mode="a"): - pass - self.sys_log.info(f"Created file /{self.path} (id: {self.uuid})") @property @@ -92,8 +75,6 @@ class File(FileSystemItemABC): :return: The size of the file in bytes. """ - if self.real: - return os.path.getsize(self.sim_path) return self.sim_size def apply_timestep(self, timestep: int) -> None: @@ -127,7 +108,7 @@ class File(FileSystemItemABC): self.num_access += 1 # file was accessed path = self.folder.name + "/" + self.name - self.sys_log.info(f"Scanning file {self.sim_path if self.sim_path else path}") + self.sys_log.info(f"Scanning file {path}") self.visible_health_status = self.health_status return True @@ -155,17 +136,8 @@ class File(FileSystemItemABC): return False current_hash = None - # if file is real, read the file contents - if self.real: - with open(self.sim_path, "rb") as f: - file_hash = hashlib.blake2b() - while chunk := f.read(8192): - file_hash.update(chunk) - - current_hash = file_hash.hexdigest() - else: - # otherwise get describe_state dict and hash that - current_hash = hashlib.blake2b(json.dumps(self.describe_state(), sort_keys=True).encode()).hexdigest() + # otherwise get describe_state dict and hash that + current_hash = hashlib.blake2b(json.dumps(self.describe_state(), sort_keys=True).encode()).hexdigest() # if the previous hash is None, set the current hash to previous if self.previous_hash is None: @@ -188,7 +160,7 @@ class File(FileSystemItemABC): self.num_access += 1 # file was accessed path = self.folder.name + "/" + self.name - self.sys_log.info(f"Repaired file {self.sim_path if self.sim_path else path}") + self.sys_log.info(f"Repaired file {path}") return True def corrupt(self) -> bool: @@ -203,7 +175,7 @@ class File(FileSystemItemABC): self.num_access += 1 # file was accessed path = self.folder.name + "/" + self.name - self.sys_log.info(f"Corrupted file {self.sim_path if self.sim_path else path}") + self.sys_log.info(f"Corrupted file {path}") return True def restore(self) -> bool: @@ -217,7 +189,7 @@ class File(FileSystemItemABC): self.num_access += 1 # file was accessed path = self.folder.name + "/" + self.name - self.sys_log.info(f"Restored file {self.sim_path if self.sim_path else path}") + self.sys_log.info(f"Restored file {path}") return True def delete(self) -> bool: diff --git a/src/primaite/simulator/file_system/file_system.py b/src/primaite/simulator/file_system/file_system.py index 0eae6009..cf4a4721 100644 --- a/src/primaite/simulator/file_system/file_system.py +++ b/src/primaite/simulator/file_system/file_system.py @@ -1,6 +1,5 @@ from __future__ import annotations -import shutil from pathlib import Path from typing import Dict, Optional @@ -230,7 +229,6 @@ class FileSystem(SimComponent): size: Optional[int] = None, file_type: Optional[FileType] = None, folder_name: Optional[str] = None, - real: bool = False, ) -> File: """ Creates a File and adds it to the list of files. @@ -239,7 +237,6 @@ class FileSystem(SimComponent): :param size: The size the file takes on disk in bytes. :param file_type: The type of the file. :param folder_name: The folder to add the file to. - :param real: "Indicates whether the File is actually a real file in the Node sim fs output." """ if folder_name: # check if file with name already exists @@ -258,8 +255,6 @@ class FileSystem(SimComponent): file_type=file_type, folder_id=folder.uuid, folder_name=folder.name, - real=real, - sim_path=self.sim_root if real else None, sim_root=self.sim_root, sys_log=self.sys_log, ) @@ -368,11 +363,6 @@ class FileSystem(SimComponent): # add file to dst dst_folder.add_file(file) self.num_file_creations += 1 - if file.real: - old_sim_path = file.sim_path - file.sim_path = file.sim_root / file.path - file.sim_path.parent.mkdir(exist_ok=True) - shutil.move(old_sim_path, file.sim_path) def copy_file(self, src_folder_name: str, src_file_name: str, dst_folder_name: str): """ @@ -401,9 +391,6 @@ class FileSystem(SimComponent): dst_folder.add_file(file_copy, force=True) - if file.real: - file_copy.sim_path.parent.mkdir(exist_ok=True) - shutil.copy2(file.sim_path, file_copy.sim_path) else: self.sys_log.error(f"Unable to copy file. {src_file_name} does not exist.") diff --git a/src/primaite/simulator/system/services/ftp/ftp_client.py b/src/primaite/simulator/system/services/ftp/ftp_client.py index 22a583e4..4eb18f6a 100644 --- a/src/primaite/simulator/system/services/ftp/ftp_client.py +++ b/src/primaite/simulator/system/services/ftp/ftp_client.py @@ -126,7 +126,6 @@ class FTPClient(FTPServiceABC): dest_file_name: str, dest_port: Optional[Port] = Port.FTP, session_id: Optional[str] = None, - real_file_path: Optional[str] = None, ) -> bool: """ Send a file to a target IP address. diff --git a/src/primaite/simulator/system/services/ftp/ftp_service.py b/src/primaite/simulator/system/services/ftp/ftp_service.py index 70ba74d7..b89ae1a2 100644 --- a/src/primaite/simulator/system/services/ftp/ftp_service.py +++ b/src/primaite/simulator/system/services/ftp/ftp_service.py @@ -1,4 +1,3 @@ -import shutil from abc import ABC from ipaddress import IPv4Address from typing import Dict, Optional @@ -55,19 +54,17 @@ class FTPServiceABC(Service, ABC): file_name = payload.ftp_command_args["dest_file_name"] folder_name = payload.ftp_command_args["dest_folder_name"] file_size = payload.ftp_command_args["file_size"] - real_file_path = payload.ftp_command_args.get("real_file_path") health_status = payload.ftp_command_args["health_status"] - is_real = real_file_path is not None file = self.file_system.create_file( - file_name=file_name, folder_name=folder_name, size=file_size, real=is_real + file_name=file_name, + folder_name=folder_name, + size=file_size, ) file.health_status = health_status self.sys_log.info( f"{self.name}: Created item in {self.sys_log.hostname}: {payload.ftp_command_args['dest_folder_name']}/" f"{payload.ftp_command_args['dest_file_name']}" ) - if is_real: - shutil.copy(real_file_path, file.sim_path) # file should exist return self.file_system.get_file(file_name=file_name, folder_name=folder_name) is not None except Exception as e: @@ -115,7 +112,6 @@ class FTPServiceABC(Service, ABC): "dest_folder_name": dest_folder_name, "dest_file_name": dest_file_name, "file_size": file.sim_size, - "real_file_path": file.sim_path if file.real else None, "health_status": file.health_status, }, packet_payload_size=file.sim_size, diff --git a/tests/unit_tests/_primaite/_simulator/_file_system/test_file.py b/tests/unit_tests/_primaite/_simulator/_file_system/test_file.py index 9d9228d8..72a5889e 100644 --- a/tests/unit_tests/_primaite/_simulator/_file_system/test_file.py +++ b/tests/unit_tests/_primaite/_simulator/_file_system/test_file.py @@ -57,20 +57,6 @@ def test_simulated_file_check_hash(file_system): assert file.health_status == FileSystemItemHealthStatus.CORRUPT -@pytest.mark.skip(reason="NODE_FILE_CHECKHASH not implemented") -def test_real_file_check_hash(file_system): - file: File = file_system.create_file(file_name="test_file.txt", real=True) - - file.check_hash() - assert file.health_status == FileSystemItemHealthStatus.GOOD - # change file content - with open(file.sim_path, "a") as f: - f.write("get hacked scrub lol xD\n") - - file.check_hash() - assert file.health_status == FileSystemItemHealthStatus.CORRUPT - - def test_file_corrupt_repair_restore(file_system): """Test the ability to corrupt and repair files.""" file: File = file_system.create_file(file_name="test_file.txt", folder_name="test_folder") diff --git a/tests/unit_tests/_primaite/_simulator/_file_system/test_file_system.py b/tests/unit_tests/_primaite/_simulator/_file_system/test_file_system.py index 9b2ecf45..0cb7dce7 100644 --- a/tests/unit_tests/_primaite/_simulator/_file_system/test_file_system.py +++ b/tests/unit_tests/_primaite/_simulator/_file_system/test_file_system.py @@ -191,7 +191,7 @@ def test_copy_file(file_system): file_system.create_folder(folder_name="src_folder") file_system.create_folder(folder_name="dst_folder") - file = file_system.create_file(file_name="test_file.txt", size=10, folder_name="src_folder", real=True) + file = file_system.create_file(file_name="test_file.txt", size=10, folder_name="src_folder") assert file_system.num_file_creations == 1 original_uuid = file.uuid diff --git a/tests/unit_tests/_primaite/_simulator/_file_system/test_folder.py b/tests/unit_tests/_primaite/_simulator/_file_system/test_folder.py index e00c002d..1d9f2d9c 100644 --- a/tests/unit_tests/_primaite/_simulator/_file_system/test_folder.py +++ b/tests/unit_tests/_primaite/_simulator/_file_system/test_folder.py @@ -132,21 +132,3 @@ def test_simulated_folder_check_hash(file_system): file.sim_size = 0 folder.check_hash() assert folder.health_status == FileSystemItemHealthStatus.CORRUPT - - -@pytest.mark.skip(reason="NODE_FILE_CHECKHASH not implemented") -def test_real_folder_check_hash(file_system): - folder: Folder = file_system.create_folder(folder_name="test_folder") - file_system.create_file(file_name="test_file.txt", folder_name="test_folder", real=True) - - folder.check_hash() - assert folder.health_status == FileSystemItemHealthStatus.GOOD - # change simulated file size - file = folder.get_file(file_name="test_file.txt") - - # change file content - with open(file.sim_path, "a") as f: - f.write("get hacked scrub lol xD\n") - - folder.check_hash() - assert folder.health_status == FileSystemItemHealthStatus.CORRUPT