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