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"]