#2533: optimise so we are not reading from file all the time

This commit is contained in:
Czar Echavez
2024-04-30 19:36:22 +01:00
parent 729f9c5064
commit 95643d3255
4 changed files with 42 additions and 89 deletions

View File

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

View File

@@ -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)