oxbrbtx/PrimAITE
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

#2716 Initial Implementation + Initial Tests and updated changelog and sphinx documentation.

Browse Source
This commit is contained in:
Archer.Bowen
2024-07-09 16:46:31 +01:00
parent 1fdd05a0d1
commit 48645d2e72
11 changed files with 387 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

168
src/primaite/game/agent/agent_log.py Normal file
View File

@@ -0,0 +1,168 @@
# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
import logging
from pathlib import Path
from prettytable import MARKDOWN, PrettyTable
from primaite.simulator import LogLevel, SIM_OUTPUT
class _NotJSONFilter(logging.Filter):
def filter(self, record: logging.LogRecord) -> bool:
"""
Determines if a log message does not start and end with '{' and '}' (i.e., it is not a JSON-like message).
:param record: LogRecord object containing all the information pertinent to the event being logged.
:return: True if log message is not JSON-like, False otherwise.
"""
return not record.getMessage().startswith("{") and not record.getMessage().endswith("}")
class AgentLog:
"""
A Agent Log class is a simple logger dedicated to managing and writing logging updates and information for an agent.
Each log message is written to a file located at: <simulation output directory>/agent/agent_name.log
"""
def __init__(self, agent_name: str):
"""
Constructs a Agent Log instance for a given hostname.
:param hostname: The hostname associated with the system logs being recorded.
"""
self.agent_name = agent_name
self.current_episode: int = 1
self.setup_logger()
def setup_logger(self):
"""
Configures the logger for this Agent Log instance.
The logger is set to the DEBUG level, and is equipped with a handler that writes to a file and filters out
JSON-like messages.
"""
if not SIM_OUTPUT.save_agent_logs:
return
log_path = self._get_log_path()
file_handler = logging.FileHandler(filename=log_path)
file_handler.setLevel(logging.DEBUG)
log_format = "%(timestep)s::%(levelname)s::%(message)s"
file_handler.setFormatter(logging.Formatter(log_format))
self.logger = logging.getLogger(f"{self.agent_name}_log")
for handler in self.logger.handlers:
self.logger.removeHandler(handler)
self.logger.setLevel(logging.DEBUG)
self.logger.addHandler(file_handler)
def _get_log_path(self) -> Path:
"""
Constructs the path for the log file based on the hostname.
:return: Path object representing the location of the log file.
"""
root = SIM_OUTPUT.path / f"episode_{self.current_episode}" / "agent_logs" / self.agent_name
root.mkdir(exist_ok=True, parents=True)
return root / f"{self.agent_name}.log"
def _write_to_terminal(self, msg: str, timestep: int, level: str, to_terminal: bool = False):
if to_terminal or SIM_OUTPUT.write_agent_log_to_terminal:
print(f"{self.agent_name}: ({timestep}) ({level}) {msg}")
def debug(self, msg: str, time_step: int, to_terminal: bool = False):
"""
Logs a message with the DEBUG level.
:param msg: The message to be logged.
:param to_terminal: If True, prints to the terminal too.
"""
if SIM_OUTPUT.agent_log_level > LogLevel.DEBUG:
return
if SIM_OUTPUT.save_agent_logs:
self.logger.debug(msg, extra={"timestep": time_step})
self._write_to_terminal(msg, "DEBUG", to_terminal)
def info(self, msg: str, time_step: int, to_terminal: bool = False):
"""
Logs a message with the INFO level.
:param msg: The message to be logged.
:param timestep: The current timestep.
:param to_terminal: If True, prints to the terminal too.
"""
if SIM_OUTPUT.agent_log_level > LogLevel.INFO:
return
if SIM_OUTPUT.save_agent_logs:
self.logger.info(msg, extra={"timestep": time_step})
self._write_to_terminal(msg, "INFO", to_terminal)
def warning(self, msg: str, time_step: int, to_terminal: bool = False):
"""
Logs a message with the WARNING level.
:param msg: The message to be logged.
:param timestep: The current timestep.
:param to_terminal: If True, prints to the terminal too.
"""
if SIM_OUTPUT.agent_log_level > LogLevel.WARNING:
return
if SIM_OUTPUT.save_agent_logs:
self.logger.warning(msg, extra={"timestep": time_step})
self._write_to_terminal(msg, "WARNING", to_terminal)
def error(self, msg: str, time_step: int, to_terminal: bool = False):
"""
Logs a message with the ERROR level.
:param msg: The message to be logged.
:param timestep: The current timestep.
:param to_terminal: If True, prints to the terminal too.
"""
if SIM_OUTPUT.agent_log_level > LogLevel.ERROR:
return
if SIM_OUTPUT.save_agent_logs:
self.logger.error(msg, extra={"timestep": time_step})
self._write_to_terminal(msg, "ERROR", to_terminal)
def critical(self, msg: str, time_step: int, to_terminal: bool = False):
"""
Logs a message with the CRITICAL level.
:param msg: The message to be logged.
:param timestep: The current timestep.
:param to_terminal: If True, prints to the terminal too.
"""
if LogLevel.CRITICAL < SIM_OUTPUT.agent_log_level:
return
if SIM_OUTPUT.save_agent_logs:
self.logger.critical(msg, extra={"timestep": time_step})
self._write_to_terminal(msg, "CRITICAL", to_terminal)
def show(self, last_n: int = 10, markdown: bool = False):
"""
Print an Agents Log as a table.
Generate and print PrettyTable instance that shows the agents behaviour log, with columns Time step,
Level and Message.
:param markdown: Use Markdown style in table output. Defaults to False.
"""
table = PrettyTable(["Time Step", "Level", "Message"])
if markdown:
table.set_style(MARKDOWN)
table.align = "l"
table.title = f"{self.agent_name} Behaviour Log"
if self._get_log_path().exists():
with open(self._get_log_path()) as file:
lines = file.readlines()
for line in lines[-last_n:]:
table.add_row(line.strip().split("::"))
print(table)

2
src/primaite/game/agent/interface.py
View File

@@ -7,6 +7,7 @@ from gymnasium.core import ActType, ObsType
from pydantic import BaseModel, model_validator
from primaite.game.agent.actions import ActionManager
from primaite.game.agent.agent_log import AgentLog
from primaite.game.agent.observations.observation_manager import ObservationManager
from primaite.game.agent.rewards import RewardFunction
from primaite.interface.request import RequestFormat, RequestResponse
@@ -116,6 +117,7 @@ class AbstractAgent(ABC):
self.reward_function: Optional[RewardFunction] = reward_function
self.agent_settings = agent_settings or AgentSettings()
self.history: List[AgentHistoryItem] = []
self.logger = AgentLog(agent_name)
def update_observation(self, state: Dict) -> ObsType:
"""

4
src/primaite/game/agent/scripted_agents/data_manipulation_bot.py
View File

@@ -38,10 +38,11 @@ class DataManipulationAgent(AbstractScriptedAgent):
:rtype: Tuple[str, Dict]
"""
if timestep < self.next_execution_timestep:
self.logger.debug(msg="Performing do NOTHING", time_step=timestep)
return "DONOTHING", {}
self._set_next_execution_timestep(timestep + self.agent_settings.start_settings.frequency)
self.logger.info(msg="Performing a data manipulation attack!", time_step=timestep)
return "NODE_APPLICATION_EXECUTE", {"node_id": self.starting_node_idx, "application_id": 0}
def setup_agent(self) -> None:
@@ -54,3 +55,4 @@ class DataManipulationAgent(AbstractScriptedAgent):
# we are assuming that every node in the node manager has a data manipulation application at idx 0
num_nodes = len(self.action_manager.node_names)
self.starting_node_idx = random.randint(0, num_nodes - 1)
self.logger.debug(msg=f"Select Start Node ID: {self.starting_node_idx}", time_step=0)

1
src/primaite/game/agent/scripted_agents/probabilistic_agent.py
View File

@@ -85,4 +85,5 @@ class ProbabilisticAgent(AbstractScriptedAgent):
:rtype: Tuple[str, Dict]
"""
choice = self.rng.choice(len(self.action_manager.action_map), p=self.probabilities)
self.logger.info(f"Performing Action: {choice}", time_step=timestep)
return self.action_manager.get_action(choice)

14
src/primaite/session/io.py
View File

@@ -35,10 +35,16 @@ class PrimaiteIO:
"""Whether to save PCAP logs."""
save_sys_logs: bool = True
"""Whether to save system logs."""
save_agent_logs: bool = True
"""Whether to save agent logs."""
write_sys_log_to_terminal: bool = False
"""Whether to write the sys log to the terminal."""
write_agent_log_to_terminal: bool = False
"""Whether to write the agent log to the terminal."""
sys_log_level: LogLevel = LogLevel.INFO
"""The level of log that should be included in the logfiles/logged into terminal."""
"""The level of sys logs that should be included in the logfiles/logged into terminal."""
agent_log_level: LogLevel = LogLevel.INFO
"""The level of agent logs that should be included in the logfiles/logged into terminal."""
def __init__(self, settings: Optional[Settings] = None) -> None:
"""
@@ -53,8 +59,11 @@ class PrimaiteIO:
SIM_OUTPUT.path = self.session_path / "simulation_output"
SIM_OUTPUT.save_pcap_logs = self.settings.save_pcap_logs
SIM_OUTPUT.save_sys_logs = self.settings.save_sys_logs
SIM_OUTPUT.save_agent_logs = self.settings.save_agent_logs
SIM_OUTPUT.write_agent_log_to_terminal = self.settings.write_agent_log_to_terminal
SIM_OUTPUT.write_sys_log_to_terminal = self.settings.write_sys_log_to_terminal
SIM_OUTPUT.sys_log_level = self.settings.sys_log_level
SIM_OUTPUT.agent_log_level = self.settings.agent_log_level
def generate_session_path(self, timestamp: Optional[datetime] = None) -> Path:
"""Create a folder for the session and return the path to it."""
@@ -115,6 +124,9 @@ class PrimaiteIO:
if config.get("sys_log_level"):
config["sys_log_level"] = LogLevel[config["sys_log_level"].upper()] # convert to enum
if config.get("agent_log_level"):
config["agent_log_level"] = LogLevel[config["agent_log_level"].upper()] # convert to enum
new = cls(settings=cls.Settings(**config))
return new

34
src/primaite/simulator/__init__.py
View File

@@ -36,8 +36,11 @@ class _SimOutput:
self._path = path
self._save_pcap_logs: bool = False
self._save_sys_logs: bool = False
self._save_agent_logs: bool = False
self._write_sys_log_to_terminal: bool = False
self._write_agent_log_to_terminal: bool = False
self._sys_log_level: LogLevel = LogLevel.WARNING # default log level is at WARNING
self._agent_log_level: LogLevel = LogLevel.WARNING
@property
def path(self) -> Path:
@@ -81,6 +84,16 @@ class _SimOutput:
def save_sys_logs(self, save_sys_logs: bool) -> None:
self._save_sys_logs = save_sys_logs
@property
def save_agent_logs(self) -> bool:
if is_dev_mode():
return PRIMAITE_CONFIG.get("developer_mode").get("output_agent_logs")
return self._save_agent_logs
@save_agent_logs.setter
def save_agent_logs(self, save_agent_logs: bool) -> None:
self._save_agent_logs = save_agent_logs
@property
def write_sys_log_to_terminal(self) -> bool:
if is_dev_mode():
@@ -91,6 +104,17 @@ class _SimOutput:
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
# Should this be separate from sys_log?
@property
def write_agent_log_to_terminal(self) -> bool:
if is_dev_mode():
return PRIMAITE_CONFIG.get("developer_mode").get("output_to_terminal")
return self._write_agent_log_to_terminal
@write_agent_log_to_terminal.setter
def write_agent_log_to_terminal(self, write_agent_log_to_terminal: bool) -> None:
self._write_agent_log_to_terminal = write_agent_log_to_terminal
@property
def sys_log_level(self) -> LogLevel:
if is_dev_mode():
@@ -101,5 +125,15 @@ class _SimOutput:
def sys_log_level(self, sys_log_level: LogLevel) -> None:
self._sys_log_level = sys_log_level
@property
def agent_log_level(self) -> LogLevel:
if is_dev_mode():
return LogLevel[PRIMAITE_CONFIG.get("developer_mode").get("agent_log_level")]
return self._agent_log_level
@agent_log_level.setter
def agent_log_level(self, agent_log_level: LogLevel) -> None:
self._agent_log_level = agent_log_level
SIM_OUTPUT = _SimOutput()