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

Merge branch 'dev' into bugfix/2455-notebook_updates

Browse Source
This commit is contained in:
Nick Todd
2024-04-19 11:52:09 +01:00
parent 9334b1e79b 56d03aaf7b
commit 873b70c2c9
40 changed files with 351 additions and 1839 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
src/primaite/cli.py
View File

@@ -114,23 +114,3 @@ def setup(overwrite_existing: bool = True) -> None:
reset_example_configs.run(overwrite_existing=True)
_LOGGER.info("PrimAITE setup complete!")
@app.command()
def session(
config: Optional[str] = None,
agent_load_file: Optional[str] = None,
) -> None:
"""
Run a PrimAITE session.
:param config: The path to the config file. Optional, if None, the example config will be used.
:type config: Optional[str]
"""
from primaite.config.load import data_manipulation_config_path
from primaite.main import run
if not config:
config = data_manipulation_config_path()
print(config)
run(config_path=config, agent_load_path=agent_load_file)

4
src/primaite/game/game.py
View File

@@ -210,8 +210,8 @@ class PrimaiteGame:
"""Create a PrimaiteGame object from a config dictionary.
The config dictionary should have the following top-level keys:
1. training_config: options for training the RL agent.
2. game_config: options for the game itself. Used by PrimaiteGame.
1. io_settings: options for logging data during training
2. game_config: options for the game itself, such as agents.
3. simulation: defines the network topology and the initial state of the simulation.
The specification for each of the three major areas is described in a separate documentation page.

47
src/primaite/main.py
View File

@@ -1,47 +0,0 @@
# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
"""The main PrimAITE session runner module."""
import argparse
from pathlib import Path
from typing import Optional, Union
from primaite import getLogger
from primaite.config.load import data_manipulation_config_path, load
from primaite.session.session import PrimaiteSession
# from primaite.primaite_session import PrimaiteSession
_LOGGER = getLogger(__name__)
def run(
config_path: Optional[Union[str, Path]] = "",
agent_load_path: Optional[Union[str, Path]] = None,
) -> None:
"""
Run the PrimAITE Session.
:param training_config_path: YAML file containing configurable items defined in
`primaite.config.training_config.TrainingConfig`
:type training_config_path: Union[path, str]
:param lay_down_config_path: YAML file containing configurable items for generating network laydown.
:type lay_down_config_path: Union[path, str]
:param session_path: directory path of the session to load
:param legacy_training_config: True if the training config file is a legacy file from PrimAITE < 2.0,
otherwise False.
:param legacy_lay_down_config: True if the lay_down config file is a legacy file from PrimAITE < 2.0,
otherwise False.
"""
cfg = load(config_path)
sess = PrimaiteSession.from_config(cfg=cfg, agent_load_path=agent_load_path)
sess.start_session()
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--config")
args = parser.parse_args()
if not args.config:
args.config = data_manipulation_config_path()
run(args.config)

76
src/primaite/notebooks/Training-an-SB3-Agent.ipynb
View File

@@ -1,5 +1,21 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Training an SB3 Agent\n",
"\n",
"This notebook will demonstrate how to use primaite to create and train a PPO agent, using a pre-defined configuration file."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### First, we import the inital packages and read in our configuration file."
]
},
{
"cell_type": "code",
"execution_count": null,
@@ -27,7 +43,14 @@
"outputs": [],
"source": [
"with open(data_manipulation_config_path(), 'r') as f:\n",
" cfg = yaml.safe_load(f)\n"
" cfg = yaml.safe_load(f)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Using the given configuration, we generate the environment our agent will train in."
]
},
{
@@ -39,6 +62,13 @@
"gym = PrimaiteGymEnv(game_config=cfg)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Lets define training parameters for the agent."
]
},
{
"cell_type": "code",
"execution_count": null,
@@ -48,9 +78,9 @@
"from stable_baselines3 import PPO\n",
"\n",
"EPISODE_LEN = 128\n",
"NO_STEPS = EPISODE_LEN * 10\n",
"BATCH_SIZE = EPISODE_LEN * 10\n",
"TOTAL_TIMESTEPS = 5e3 * EPISODE_LEN\n",
"NUM_EPISODES = 10\n",
"NO_STEPS = EPISODE_LEN * NUM_EPISODES\n",
"BATCH_SIZE = 32\n",
"LEARNING_RATE = 3e-4"
]
},
@@ -60,7 +90,14 @@
"metadata": {},
"outputs": [],
"source": [
"model = PPO('MlpPolicy', gym, learning_rate=LEARNING_RATE, n_steps=NO_STEPS, batch_size=BATCH_SIZE, verbose=0, tensorboard_log=\"./PPO_UC2/\")\n"
"model = PPO('MlpPolicy', gym, learning_rate=LEARNING_RATE, n_steps=NO_STEPS, batch_size=BATCH_SIZE, verbose=0, tensorboard_log=\"./PPO_UC2/\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"With the agent configured, let's train for our defined number of episodes."
]
},
{
@@ -69,7 +106,14 @@
"metadata": {},
"outputs": [],
"source": [
"model.learn(total_timesteps=TOTAL_TIMESTEPS)\n"
"model.learn(total_timesteps=NO_STEPS)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Next, let's save the agent to a zip file that can be used in future evaluation."
]
},
{
@@ -78,7 +122,14 @@
"metadata": {},
"outputs": [],
"source": [
"model.save(\"PrimAITE-v3.0.0b7-PPO\")"
"model.save(\"PrimAITE-PPO-example-agent\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now, we load the saved agent and run it in evaluation mode."
]
},
{
@@ -88,7 +139,14 @@
"outputs": [],
"source": [
"eval_model = PPO(\"MlpPolicy\", gym)\n",
"eval_model = PPO.load(\"PrimAITE-v3.0.0b7-PPO\", gym)"
"eval_model = PPO.load(\"PrimAITE-PPO-example-agent\", gym)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Finally, evaluate the agent."
]
},
{
@@ -119,7 +177,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.9.18"
"version": "3.10.11"
}
},
"nbformat": 4,

2
src/primaite/session/environment.py
View File

@@ -48,7 +48,7 @@ class PrimaiteGymEnv(gymnasium.Env):
def step(self, action: ActType) -> Tuple[ObsType, SupportsFloat, bool, bool, Dict[str, Any]]:
"""Perform a step in the environment."""
# make ProxyAgent store the action chosen my the RL policy
# make ProxyAgent store the action chosen by the RL policy
step = self.game.step_counter
self.agent.store_action(action)
# apply_agent_actions accesses the action we just stored

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

@@ -95,6 +95,7 @@ class PrimaiteIO:
@classmethod
def from_config(cls, config: Dict) -> "PrimaiteIO":
"""Create an instance of PrimaiteIO based on a configuration dict."""
config = config or {}
new = cls(settings=cls.Settings(**config))
return new

4
src/primaite/session/policy/__init__.py
View File

@@ -1,4 +0,0 @@
from primaite.session.policy.rllib import RaySingleAgentPolicy
from primaite.session.policy.sb3 import SB3Policy
__all__ = ["SB3Policy", "RaySingleAgentPolicy"]

82
src/primaite/session/policy/policy.py
View File

@@ -1,82 +0,0 @@
"""Base class and common logic for RL policies."""
from abc import ABC, abstractmethod
from pathlib import Path
from typing import Any, Dict, Type, TYPE_CHECKING
if TYPE_CHECKING:
from primaite.session.session import PrimaiteSession, TrainingOptions
class PolicyABC(ABC):
"""Base class for reinforcement learning agents."""
_registry: Dict[str, Type["PolicyABC"]] = {}
"""
Registry of policy types, keyed by name.
Automatically populated when PolicyABC subclasses are defined. Used for defining from_config.
"""
def __init_subclass__(cls, identifier: str, **kwargs: Any) -> None:
"""
Register a policy subclass.
:param name: Identifier used by from_config to create an instance of the policy.
:type name: str
:raises ValueError: When attempting to create a policy with a duplicate name.
"""
super().__init_subclass__(**kwargs)
if identifier in cls._registry:
raise ValueError(f"Duplicate policy name {identifier}")
cls._registry[identifier] = cls
return
@abstractmethod
def __init__(self, session: "PrimaiteSession") -> None:
"""
Initialize a reinforcement learning policy.
:param session: The session context.
:type session: PrimaiteSession
:param agents: The agents to train.
:type agents: List[RLAgent]
"""
self.session: "PrimaiteSession" = session
"""Reference to the session."""
@abstractmethod
def learn(self, n_episodes: int, timesteps_per_episode: int) -> None:
"""Train the agent."""
pass
@abstractmethod
def eval(self, n_episodes: int, timesteps_per_episode: int, deterministic: bool) -> None:
"""Evaluate the agent."""
pass
@abstractmethod
def save(self, save_path: Path) -> None:
"""Save the agent."""
pass
@abstractmethod
def load(self) -> None:
"""Load agent from a file."""
pass
def close(self) -> None:
"""Close the agent."""
pass
@classmethod
def from_config(cls, config: "TrainingOptions", session: "PrimaiteSession") -> "PolicyABC":
"""
Create an RL policy from a config by calling the relevant subclass's from_config method.
Subclasses should not call super().from_config(), they should just handle creation form config.
"""
# Assume that basically the contents of training_config are passed into here.
# I should really define a config schema class using pydantic.
PolicyType = cls._registry[config.rl_framework]
return PolicyType.from_config(config=config, session=session)

111
src/primaite/session/policy/rllib.py
View File

@@ -1,111 +0,0 @@
from pathlib import Path
from typing import Literal, Optional, TYPE_CHECKING
from primaite.session.environment import PrimaiteRayEnv, PrimaiteRayMARLEnv
from primaite.session.policy.policy import PolicyABC
if TYPE_CHECKING:
from primaite.session.session import PrimaiteSession, TrainingOptions
import ray
from ray import air, tune
from ray.rllib.algorithms import ppo
from ray.rllib.algorithms.ppo import PPOConfig
from primaite import getLogger
_LOGGER = getLogger(__name__)
class RaySingleAgentPolicy(PolicyABC, identifier="RLLIB_single_agent"):
"""Single agent RL policy using Ray RLLib."""
def __init__(self, session: "PrimaiteSession", algorithm: Literal["PPO", "A2C"], seed: Optional[int] = None):
super().__init__(session=session)
self.config = {
"env": PrimaiteRayEnv,
"env_config": {"game": session.game},
"disable_env_checking": True,
"num_rollout_workers": 0,
}
ray.shutdown()
ray.init()
def learn(self, n_episodes: int, timesteps_per_episode: int) -> None:
"""Train the agent."""
self.config["training_iterations"] = n_episodes * timesteps_per_episode
self.config["train_batch_size"] = 128
self._algo = ppo.PPO(config=self.config)
_LOGGER.info("Starting RLLIB training session")
self._algo.train()
def eval(self, n_episodes: int, deterministic: bool) -> None:
"""Evaluate the agent."""
for ep in range(n_episodes):
obs, info = self.session.env.reset()
for step in range(self.session.game.options.max_episode_length):
action = self._algo.compute_single_action(observation=obs, explore=False)
obs, rew, term, trunc, info = self.session.env.step(action)
def save(self, save_path: Path) -> None:
"""Save the policy to a file."""
self._algo.save(save_path)
def load(self, model_path: Path) -> None:
"""Load policy parameters from a file."""
raise NotImplementedError
@classmethod
def from_config(cls, config: "TrainingOptions", session: "PrimaiteSession") -> "RaySingleAgentPolicy":
"""Create a policy from a config."""
return cls(session=session, algorithm=config.rl_algorithm, seed=config.seed)
class RayMultiAgentPolicy(PolicyABC, identifier="RLLIB_multi_agent"):
"""Mutli agent RL policy using Ray RLLib."""
def __init__(self, session: "PrimaiteSession", algorithm: Literal["PPO"], seed: Optional[int] = None):
"""Initialise multi agent policy wrapper."""
super().__init__(session=session)
self.config = (
PPOConfig()
.environment(env=PrimaiteRayMARLEnv, env_config={"game": session.game})
.rollouts(num_rollout_workers=0)
.multi_agent(
policies={agent.agent_name for agent in session.game.rl_agents},
policy_mapping_fn=lambda agent_id, episode, worker, **kw: agent_id,
)
.training(train_batch_size=128)
)
def learn(self, n_episodes: int, timesteps_per_episode: int) -> None:
"""Train the agent."""
checkpoint_freq = self.session.io_manager.settings.checkpoint_interval
tune.Tuner(
"PPO",
run_config=air.RunConfig(
stop={"training_iteration": n_episodes * timesteps_per_episode},
checkpoint_config=air.CheckpointConfig(checkpoint_frequency=checkpoint_freq),
),
param_space=self.config,
).fit()
def load(self, model_path: Path) -> None:
"""Load policy parameters from a file."""
return NotImplemented
def eval(self, n_episodes: int, deterministic: bool) -> None:
"""Evaluate trained policy."""
return NotImplemented
def save(self, save_path: Path) -> None:
"""Save policy parameters to a file."""
return NotImplemented
@classmethod
def from_config(cls, config: "TrainingOptions", session: "PrimaiteSession") -> "RayMultiAgentPolicy":
"""Create policy from config."""
return cls(session=session, algorithm=config.rl_algorithm, seed=config.seed)

79
src/primaite/session/policy/sb3.py
View File

@@ -1,79 +0,0 @@
"""Stable baselines 3 policy."""
from pathlib import Path
from typing import Literal, Optional, Type, TYPE_CHECKING, Union
from stable_baselines3 import A2C, PPO
from stable_baselines3.a2c import MlpPolicy as A2C_MLP
from stable_baselines3.common.callbacks import CheckpointCallback
from stable_baselines3.common.evaluation import evaluate_policy
from stable_baselines3.ppo import MlpPolicy as PPO_MLP
from primaite.session.policy.policy import PolicyABC
if TYPE_CHECKING:
from primaite.session.session import PrimaiteSession, TrainingOptions
class SB3Policy(PolicyABC, identifier="SB3"):
"""Single agent RL policy using stable baselines 3."""
def __init__(self, session: "PrimaiteSession", algorithm: Literal["PPO", "A2C"], seed: Optional[int] = None):
"""Initialize a stable baselines 3 policy."""
super().__init__(session=session)
self._agent_class: Type[Union[PPO, A2C]]
if algorithm == "PPO":
self._agent_class = PPO
policy = PPO_MLP
elif algorithm == "A2C":
self._agent_class = A2C
policy = A2C_MLP
else:
raise ValueError(f"Unknown algorithm `{algorithm}` for stable_baselines3 policy")
self._agent = self._agent_class(
policy=policy,
env=self.session.env,
n_steps=128, # this is not the number of steps in an episode, but the number of steps in a batch
seed=seed,
)
def learn(self, n_episodes: int, timesteps_per_episode: int) -> None:
"""Train the agent."""
if self.session.save_checkpoints:
checkpoint_callback = CheckpointCallback(
save_freq=timesteps_per_episode * self.session.checkpoint_interval,
save_path=self.session.io_manager.generate_model_save_path("sb3"),
name_prefix="sb3_model",
)
else:
checkpoint_callback = None
self._agent.learn(total_timesteps=n_episodes * timesteps_per_episode, callback=checkpoint_callback)
def eval(self, n_episodes: int, deterministic: bool) -> None:
"""Evaluate the agent."""
_ = evaluate_policy(
self._agent,
self.session.env,
n_eval_episodes=n_episodes,
deterministic=deterministic,
return_episode_rewards=True,
)
def save(self, save_path: Path) -> None:
"""
Save the current policy parameters.
Warning: The recommended way to save model checkpoints is to use a callback within the `learn()` method. Please
refer to https://stable-baselines3.readthedocs.io/en/master/guide/callbacks.html for more information.
Therefore, this method is only used to save the final model.
"""
self._agent.save(save_path)
def load(self, model_path: Path) -> None:
"""Load agent from a checkpoint."""
self._agent = self._agent_class.load(model_path, env=self.session.env)
@classmethod
def from_config(cls, config: "TrainingOptions", session: "PrimaiteSession") -> "SB3Policy":
"""Create an agent from config file."""
return cls(session=session, algorithm=config.rl_algorithm, seed=config.seed)

119
src/primaite/session/session.py
View File

@@ -1,119 +0,0 @@
# raise DeprecationWarning("This module is deprecated")
from enum import Enum
from pathlib import Path
from typing import Dict, List, Literal, Optional, Union
from pydantic import BaseModel, ConfigDict
from primaite.session.environment import PrimaiteGymEnv, PrimaiteRayEnv, PrimaiteRayMARLEnv
from primaite.session.io import PrimaiteIO
# from primaite.game.game import PrimaiteGame
from primaite.session.policy.policy import PolicyABC
class TrainingOptions(BaseModel):
"""Options for training the RL agent."""
model_config = ConfigDict(extra="forbid")
rl_framework: Literal["SB3", "RLLIB_single_agent", "RLLIB_multi_agent"]
rl_algorithm: Literal["PPO", "A2C"]
n_learn_episodes: int
n_eval_episodes: Optional[int] = None
max_steps_per_episode: int
# checkpoint_freq: Optional[int] = None
deterministic_eval: bool
seed: Optional[int]
n_agents: int
agent_references: List[str]
class SessionMode(Enum):
"""Helper to keep track of the current session mode."""
TRAIN = "train"
EVAL = "eval"
MANUAL = "manual"
class PrimaiteSession:
"""The main entrypoint for PrimAITE sessions, this manages a simulation, policy training, and environments."""
def __init__(self, game_cfg: Dict):
"""Initialise PrimaiteSession object."""
self.training_options: TrainingOptions
"""Options specific to agent training."""
self.mode: SessionMode = SessionMode.MANUAL
"""Current session mode."""
self.env: Union[PrimaiteGymEnv, PrimaiteRayEnv, PrimaiteRayMARLEnv]
"""The environment that the RL algorithm can consume."""
self.policy: PolicyABC
"""The reinforcement learning policy."""
self.io_manager: Optional["PrimaiteIO"] = None
"""IO manager for the session."""
self.game_cfg: Dict = game_cfg
"""Primaite Game object for managing main simulation loop and agents."""
self.save_checkpoints: bool = False
"""Whether to save checkpoints."""
self.checkpoint_interval: int = 10
"""If save_checkpoints is true, checkpoints will be saved every checkpoint_interval episodes."""
def start_session(self) -> None:
"""Commence the training/eval session."""
print("Starting Primaite Session")
self.mode = SessionMode.TRAIN
n_learn_episodes = self.training_options.n_learn_episodes
n_eval_episodes = self.training_options.n_eval_episodes
max_steps_per_episode = self.training_options.max_steps_per_episode
deterministic_eval = self.training_options.deterministic_eval
self.policy.learn(
n_episodes=n_learn_episodes,
timesteps_per_episode=max_steps_per_episode,
)
self.save_models()
self.mode = SessionMode.EVAL
if n_eval_episodes > 0:
self.policy.eval(n_episodes=n_eval_episodes, deterministic=deterministic_eval)
self.mode = SessionMode.MANUAL
def save_models(self) -> None:
"""Save the RL models."""
save_path = self.io_manager.generate_model_save_path("temp_model_name")
self.policy.save(save_path)
@classmethod
def from_config(cls, cfg: Dict, agent_load_path: Optional[str] = None) -> "PrimaiteSession":
"""Create a PrimaiteSession object from a config dictionary."""
# READ IO SETTINGS (this sets the global session path as well) # TODO: GLOBAL SIDE EFFECTS...
io_manager = PrimaiteIO.from_config(cfg.get("io_settings", {}))
sess = cls(game_cfg=cfg)
sess.io_manager = io_manager
sess.training_options = TrainingOptions(**cfg["training_config"])
sess.save_checkpoints = cfg.get("io_settings", {}).get("save_checkpoints")
sess.checkpoint_interval = cfg.get("io_settings", {}).get("checkpoint_interval")
# CREATE ENVIRONMENT
if sess.training_options.rl_framework == "RLLIB_single_agent":
sess.env = PrimaiteRayEnv(env_config=cfg)
elif sess.training_options.rl_framework == "RLLIB_multi_agent":
sess.env = PrimaiteRayMARLEnv(env_config=cfg)
elif sess.training_options.rl_framework == "SB3":
sess.env = PrimaiteGymEnv(game_config=cfg)
sess.policy = PolicyABC.from_config(sess.training_options, session=sess)
if agent_load_path:
sess.policy.load(Path(agent_load_path))
return sess

34
src/primaite/simulator/_package_data/create-simulation_demo.ipynb
View File

@@ -6,7 +6,7 @@
"source": [
"# Build a simulation using the Python API\n",
"\n",
"Currently, this notbook manipulates the simulation by directly placing objects inside of the attributes of the network and domain. It should be refactored when proper methods exist for adding these objects.\n"
"Currently, this notebook manipulates the simulation by directly placing objects inside of the attributes of the network and domain. It should be refactored when proper methods exist for adding these objects.\n"
]
},
{
@@ -58,7 +58,8 @@
"metadata": {},
"outputs": [],
"source": [
"from primaite.simulator.network.hardware.base import Node\n"
"from primaite.simulator.network.hardware.nodes.host.computer import Computer\n",
"from primaite.simulator.network.hardware.nodes.host.server import Server"
]
},
{
@@ -67,9 +68,9 @@
"metadata": {},
"outputs": [],
"source": [
"my_pc = Node(hostname=\"primaite_pc\",)\n",
"my_pc = Computer(hostname=\"Computer\", ip_address=\"192.168.1.10\", subnet_mask=\"255.255.255.0\")\n",
"net.add_node(my_pc)\n",
"my_server = Node(hostname=\"google_server\")\n",
"my_server = Server(hostname=\"Server\", ip_address=\"192.168.1.11\", subnet_mask=\"255.255.255.0\")\n",
"net.add_node(my_server)\n"
]
},
@@ -86,7 +87,8 @@
"metadata": {},
"outputs": [],
"source": [
"from primaite.simulator.network.hardware.base import NIC, Link, Switch\n"
"from primaite.simulator.network.hardware.nodes.host.host_node import NIC\n",
"from primaite.simulator.network.hardware.nodes.network.switch import Switch\n"
]
},
{
@@ -95,19 +97,17 @@
"metadata": {},
"outputs": [],
"source": [
"my_swtich = Switch(hostname=\"switch1\", num_ports=12)\n",
"net.add_node(my_swtich)\n",
"my_switch = Switch(hostname=\"switch1\", num_ports=12)\n",
"net.add_node(my_switch)\n",
"\n",
"pc_nic = NIC(ip_address=\"130.1.1.1\", gateway=\"130.1.1.255\", subnet_mask=\"255.255.255.0\")\n",
"my_pc.connect_nic(pc_nic)\n",
"\n",
"\n",
"server_nic = NIC(ip_address=\"130.1.1.2\", gateway=\"130.1.1.255\", subnet_mask=\"255.255.255.0\")\n",
"my_server.connect_nic(server_nic)\n",
"\n",
"\n",
"net.connect(pc_nic, my_swtich.switch_ports[1])\n",
"net.connect(server_nic, my_swtich.switch_ports[2])\n"
"net.connect(pc_nic, my_switch.network_interface[1])\n",
"net.connect(server_nic, my_switch.network_interface[2])\n"
]
},
{
@@ -124,7 +124,8 @@
"outputs": [],
"source": [
"from primaite.simulator.file_system.file_type import FileType\n",
"from primaite.simulator.file_system.file_system import File"
"from primaite.simulator.file_system.file_system import File\n",
"from primaite.simulator.system.core.sys_log import SysLog"
]
},
{
@@ -134,7 +135,7 @@
"outputs": [],
"source": [
"my_pc_downloads_folder = my_pc.file_system.create_folder(\"downloads\")\n",
"my_pc_downloads_folder.add_file(File(name=\"firefox_installer.zip\",file_type=FileType.ZIP))"
"my_pc_downloads_folder.add_file(File(name=\"firefox_installer.zip\",folder_id=\"Test\", folder_name=\"downloads\" ,file_type=FileType.ZIP, sys_log=SysLog(hostname=\"Test\")))"
]
},
{
@@ -160,9 +161,12 @@
"metadata": {},
"outputs": [],
"source": [
"from pathlib import Path\n",
"from primaite.simulator.system.applications.application import Application, ApplicationOperatingState\n",
"from primaite.simulator.system.software import SoftwareHealthState, SoftwareCriticality\n",
"from primaite.simulator.network.transmission.transport_layer import Port\n",
"from primaite.simulator.network.transmission.network_layer import IPProtocol\n",
"from primaite.simulator.file_system.file_system import FileSystem\n",
"\n",
"# no applications exist yet so we will create our own.\n",
"class MSPaint(Application):\n",
@@ -176,7 +180,7 @@
"metadata": {},
"outputs": [],
"source": [
"mspaint = MSPaint(name = \"mspaint\", health_state_actual=SoftwareHealthState.GOOD, health_state_visible=SoftwareHealthState.GOOD, criticality=SoftwareCriticality.MEDIUM, ports={Port.HTTP}, operating_state=ApplicationOperatingState.RUNNING,execution_control_status='manual')"
"mspaint = MSPaint(name = \"mspaint\", health_state_actual=SoftwareHealthState.GOOD, health_state_visible=SoftwareHealthState.GOOD, criticality=SoftwareCriticality.MEDIUM, port=Port.HTTP, protocol = IPProtocol.NONE,operating_state=ApplicationOperatingState.RUNNING,execution_control_status='manual', file_system=FileSystem(sys_log=SysLog(hostname=\"Test\"), sim_root=Path(__name__).parent),)"
]
},
{
@@ -257,7 +261,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.12"
"version": "3.10.11"
}
},
"nbformat": 4,

208
src/primaite/simulator/_package_data/network_simulator_demo.ipynb
View File

@@ -238,7 +238,7 @@
"id": "20",
"metadata": {},
"source": [
"Calling `switch.show()` displays the Switch orts on the Switch."
"Calling `switch.show()` displays the Switch ports on the Switch."
]
},
{
@@ -256,11 +256,9 @@
{
"cell_type": "markdown",
"id": "22",
"metadata": {
"tags": []
},
"metadata": {},
"source": [
"Calling `switch.arp.show()` displays the Switch ARP Cache."
"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=<number of log entries>`."
]
},
{
@@ -271,33 +269,13 @@
"tags": []
},
"outputs": [],
"source": [
"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=<number of log entries>`."
]
},
{
"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",
"id": "24",
"metadata": {},
"source": [
"### Computer/Server Nodes\n",
@@ -307,7 +285,7 @@
},
{
"cell_type": "markdown",
"id": "27",
"id": "25",
"metadata": {
"tags": []
},
@@ -315,6 +293,26 @@
"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,
@@ -324,7 +322,7 @@
},
"outputs": [],
"source": [
"network.get_node_by_hostname(\"security_suite\").show()"
"network.get_node_by_hostname(\"security_suite\").arp.show()"
]
},
{
@@ -332,7 +330,7 @@
"id": "29",
"metadata": {},
"source": [
"Calling `computer.arp.show()` displays the Computer/Server ARP Cache."
"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=<number of log entries>`."
]
},
{
@@ -344,7 +342,7 @@
},
"outputs": [],
"source": [
"network.get_node_by_hostname(\"security_suite\").arp.show()"
"network.get_node_by_hostname(\"security_suite\").sys_log.show()"
]
},
{
@@ -352,7 +350,9 @@
"id": "31",
"metadata": {},
"source": [
"Calling `switch.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=<number of log entries>`."
"## Basic Network Comms Check\n",
"\n",
"We can perform a good old ping to check that Nodes are able to communicate with each other."
]
},
{
@@ -364,7 +364,7 @@
},
"outputs": [],
"source": [
"network.get_node_by_hostname(\"security_suite\").sys_log.show()"
"network.show(nodes=False, links=False)"
]
},
{
@@ -372,9 +372,7 @@
"id": "33",
"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."
"We'll first ping client_1's default gateway."
]
},
{
@@ -386,27 +384,27 @@
},
"outputs": [],
"source": [
"network.show(nodes=False, links=False)"
]
},
{
"cell_type": "markdown",
"id": "35",
"metadata": {},
"source": [
"We'll first ping client_1's default gateway."
"network.get_node_by_hostname(\"client_1\").ping(\"192.168.10.1\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "36",
"id": "35",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"network.get_node_by_hostname(\"client_1\").ping(\"192.168.10.1\")"
"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)."
]
},
{
@@ -418,7 +416,7 @@
},
"outputs": [],
"source": [
"network.get_node_by_hostname(\"client_1\").sys_log.show(15)"
"network.get_node_by_hostname(\"client_1\").ping(\"192.168.1.1\")"
]
},
{
@@ -426,7 +424,7 @@
"id": "38",
"metadata": {},
"source": [
"Next, we'll ping the interface of the 192.168.1.0/24 Network on the Router (port 1)."
"And finally, we'll ping the web server."
]
},
{
@@ -438,7 +436,7 @@
},
"outputs": [],
"source": [
"network.get_node_by_hostname(\"client_1\").ping(\"192.168.1.1\")"
"network.get_node_by_hostname(\"client_1\").ping(\"192.168.1.12\")"
]
},
{
@@ -446,7 +444,7 @@
"id": "40",
"metadata": {},
"source": [
"And finally, we'll ping the web server."
"To confirm that the ping was received and processed by the web_server, we can view the sys log"
]
},
{
@@ -458,7 +456,7 @@
},
"outputs": [],
"source": [
"network.get_node_by_hostname(\"client_1\").ping(\"192.168.1.12\")"
"network.get_node_by_hostname(\"web_server\").sys_log.show()"
]
},
{
@@ -466,29 +464,29 @@
"id": "42",
"metadata": {},
"source": [
"To confirm that the ping was received and processed by the web_server, we can view the sys log"
"## 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": "43",
"id": "44",
"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 advaced things."
"network.get_node_by_hostname(\"client_2\").ping(\"192.168.1.12\")"
]
},
{
@@ -496,7 +494,7 @@
"id": "45",
"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..."
"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:"
]
},
{
@@ -508,7 +506,7 @@
},
"outputs": [],
"source": [
"network.get_node_by_hostname(\"client_2\").ping(\"192.168.1.12\")"
"network.get_node_by_hostname(\"client_2\").sys_log.show()"
]
},
{
@@ -516,7 +514,7 @@
"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:"
"Now we'll add an ACL to block ICMP from 192.168.10.22"
]
},
{
@@ -527,30 +525,10 @@
"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"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "50",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"from primaite.simulator.network.transmission.network_layer import IPProtocol\n",
"from primaite.simulator.network.transmission.transport_layer import Port\n",
"from primaite.simulator.network.hardware.nodes.router import ACLAction\n",
"from primaite.simulator.network.hardware.nodes.network.router import ACLAction\n",
"network.get_node_by_hostname(\"router_1\").acl.add_rule(\n",
" action=ACLAction.DENY,\n",
" protocol=IPProtocol.ICMP,\n",
@@ -562,7 +540,7 @@
{
"cell_type": "code",
"execution_count": null,
"id": "51",
"id": "49",
"metadata": {
"tags": []
},
@@ -573,12 +551,32 @@
},
{
"cell_type": "markdown",
"id": "52",
"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"
]
},
{
"cell_type": "code",
"execution_count": null,
@@ -588,7 +586,7 @@
},
"outputs": [],
"source": [
"network.get_node_by_hostname(\"client_2\").ping(\"192.168.1.12\")"
"network.get_node_by_hostname(\"client_2\").sys_log.show()"
]
},
{
@@ -596,7 +594,7 @@
"id": "54",
"metadata": {},
"source": [
"We can check that the ping was actually sent by client_2 by viewing the sys log"
"We can check the router sys log to see why the traffic was blocked"
]
},
{
@@ -608,7 +606,7 @@
},
"outputs": [],
"source": [
"network.get_node_by_hostname(\"client_2\").sys_log.show()"
"network.get_node_by_hostname(\"router_1\").sys_log.show()"
]
},
{
@@ -616,7 +614,7 @@
"id": "56",
"metadata": {},
"source": [
"We can check the router sys log to see why the traffic was blocked"
"Now a final check to ensure that client_1 can still ping the web_server."
]
},
{
@@ -627,26 +625,6 @@
"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\")"
]
@@ -654,7 +632,7 @@
{
"cell_type": "code",
"execution_count": null,
"id": "60",
"id": "58",
"metadata": {
"tags": []
},

2
src/primaite/simulator/network/airspace.py
View File

@@ -271,7 +271,7 @@ class IPWirelessNetworkInterface(WirelessNetworkInterface, Layer3Interface, ABC)
# Update the state with information from Layer3Interface
state.update(Layer3Interface.describe_state(self))
state["frequency"] = self.frequency
state["frequency"] = self.frequency.value
return state

10
src/primaite/simulator/network/hardware/nodes/host/host_node.py
View File

@@ -316,6 +316,16 @@ class HostNode(Node):
super().__init__(**kwargs)
self.connect_nic(NIC(ip_address=ip_address, subnet_mask=subnet_mask))
@property
def arp(self) -> Optional[ARP]:
"""
Return the ARP Cache of the HostNode.
:return: ARP Cache for given HostNode
:rtype: Optional[ARP]
"""
return self.software_manager.software.get("ARP")
def _install_system_software(self):
"""
Installs the system software and network services typically found on an operating system.

12
src/primaite/simulator/network/hardware/nodes/network/network_node.py
View File

@@ -1,7 +1,9 @@
from abc import abstractmethod
from typing import Optional
from primaite.simulator.network.hardware.base import NetworkInterface, Node
from primaite.simulator.network.transmission.data_link_layer import Frame
from primaite.simulator.system.services.arp.arp import ARP
class NetworkNode(Node):
@@ -28,3 +30,13 @@ class NetworkNode(Node):
:type from_network_interface: NetworkInterface
"""
pass
@property
def arp(self) -> Optional[ARP]:
"""
Return the ARP Cache of the NetworkNode.
:return: ARP Cache for given NetworkNode
:rtype: Optional[ARP]
"""
return self.software_manager.software.get("ARP")

3
src/primaite/simulator/network/hardware/nodes/network/router.py
View File

@@ -1237,8 +1237,7 @@ class Router(NetworkNode):
icmp: RouterICMP = self.software_manager.icmp # noqa
icmp.router = self
self.software_manager.install(RouterARP)
arp: RouterARP = self.software_manager.arp # noqa
arp.router = self
self.arp.router = self
def _set_default_acl(self):
"""

4
src/primaite/utils/session_metadata_parser.py
View File

@@ -1,3 +1,7 @@
# flake8: noqa
raise DeprecationWarning(
"Benchmarking depends on deprecated functionality and it has not been updated to primaite v3 yet."
)
# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
import json
from pathlib import Path

4
src/primaite/utils/session_output_reader.py
View File

@@ -1,3 +1,7 @@
# flake8: noqa
raise DeprecationWarning(
"Benchmarking depends on deprecated functionality and it has not been updated to primaite v3 yet."
)
# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
from pathlib import Path
from typing import Any, Dict, Tuple, Union

4
src/primaite/utils/session_output_writer.py
View File

@@ -1,3 +1,7 @@
# flake8: noqa
raise DeprecationWarning(
"Benchmarking depends on deprecated functionality and it has not been updated to primaite v3 yet."
)
# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
import csv
from logging import Logger