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

Merge origin/dev into feature/1859

Browse Source
This commit is contained in:
Marek Wolan
2023-11-27 22:13:42 +00:00
parent 517f99b04b 95f6cf6691
commit 21b865d8a8
45 changed files with 877 additions and 142 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
CHANGELOG.md
View File

@@ -36,10 +36,11 @@ SessionManager.
- DNS Services: `DNSClient` and `DNSServer`
- FTP Services: `FTPClient` and `FTPServer`
- HTTP Services: `WebBrowser` to simulate a web client and `WebServer`
- Fixed an issue where the services were still able to run even though the node the service is installed on is turned off
### Removed
- Removed legacy simulation modules: `acl`, `common`, `environment`, `links`, `nodes`, `pol`
- Removed legacy training modules, they are replaced by the new ARCD GATE dependency
- Removed legacy training modules
- Removed tests for legacy code

2
README.md
View File

@@ -22,8 +22,6 @@ PrimAITE presents the following features:
- Routers with traffic routing and firewall capabilities
- Integration with ARCD GATE for agent training
- Support for multiple agents, each having their own customisable observation space, action space, and reward function definition, and either deterministic or RL-directed behaviour
## Getting Started with PrimAITE

11
docs/index.rst
View File

@@ -30,7 +30,7 @@ PrimAITE incorporates the following features:
- Uses the concept of Information Exchange Requirements (IERs) to model background pattern of life and adversarial behaviour;
- An Access Control List (ACL) function, mimicking the behaviour of a network firewall, is applied across the model, following standard ACL rule format (e.g. DENY/ALLOW, source IP address, destination IP address, protocol and port);
- Application of traffic to the links of the platform / system laydown adheres to the ACL ruleset;
- Presents both an OpenAI gym and Ray RLLib interface to the environment, allowing integration with any compliant defensive agents;
- Presents both a Gymnasium and Ray RLLib interface to the environment, allowing integration with any compliant defensive agents;
- Allows for the saving and loading of trained defensive agents;
- Stochastic adversarial agent behaviour;
- Full capture of discrete logs relating to agent training or evaluation (system state, agent actions taken, instantaneous and average reward for every step of every episode);
@@ -40,18 +40,18 @@ PrimAITE incorporates the following features:
Architecture
^^^^^^^^^^^^
PrimAITE is a Python application and is therefore Operating System agnostic. The OpenAI gym and Ray RLLib frameworks are employed to provide an interface and source for AI agents. Configuration of PrimAITE is achieved via included YAML files which support full control over the platform / system laydown being modelled, background pattern of life, adversarial (red agent) behaviour, and step and episode count. NetworkX based nodes and links host Python classes to present attributes and methods, and hence a more representative platform / system can be modelled within the simulation.
PrimAITE is a Python application and is therefore Operating System agnostic. The Gymnasium and Ray RLLib frameworks are employed to provide an interface and source for AI agents. Configuration of PrimAITE is achieved via included YAML files which support full control over the platform / system laydown being modelled, background pattern of life, adversarial (red agent) behaviour, and step and episode count. NetworkX based nodes and links host Python classes to present attributes and methods, and hence a more representative platform / system can be modelled within the simulation.
Training & Evaluation Capability
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
PrimAITE provides a training and evaluation capability to AI agents in the context of cyber-attack, via its OpenAI Gym and RLLib compliant interface. Scenarios can be constructed to reflect platform / system laydowns consisting of any configuration of nodes (e.g. PCs, servers, switches etc.) and network links between them. All nodes can be configured to model services (and their status) and the traffic loading between them over the network links. Traffic loading is broken down into a per service granularity, relating directly to a protocol (e.g. Service A would be configured as a TCP service, and TCP traffic then flows between instances of Service A under the direction of a tailored IER). Highlights of PrimAITE’s training and evaluation capability are:
PrimAITE provides a training and evaluation capability to AI agents in the context of cyber-attack, via its Gymnasium and RLLib compliant interface. Scenarios can be constructed to reflect platform / system laydowns consisting of any configuration of nodes (e.g. PCs, servers, switches etc.) and network links between them. All nodes can be configured to model services (and their status) and the traffic loading between them over the network links. Traffic loading is broken down into a per service granularity, relating directly to a protocol (e.g. Service A would be configured as a TCP service, and TCP traffic then flows between instances of Service A under the direction of a tailored IER). Highlights of PrimAITE’s training and evaluation capability are:
- The scenario is not bound to a representation of any platform, system, or technology;
- Fully configurable (network / system laydown, IERs, node pattern-of-life, ACL, number of episodes, steps per episode) and repeatable to suit the requirements of AI agents;
- Can integrate with any OpenAI Gym or RLLib compliant AI agent.
- Can integrate with any Gymnasium or RLLib compliant AI agent.
Use of PrimAITE default scenarios within ARCD is supported by a “Use Case Profile” tailored to the scenario.
@@ -75,7 +75,7 @@ Logs are available in CSV format and provide coverage of the above data for ever
What is PrimAITE built with
---------------------------
* `OpenAI's Gym <https://gym.openai.com/>`_ is used as the basis for AI blue agent interaction with the PrimAITE environment
* `Gymnasium <https://gymnasium.farama.org/>`_ is used as the basis for AI blue agent interaction with the PrimAITE environment
* `Networkx <https://github.com/networkx/networkx>`_ is used as the underlying data structure used for the PrimAITE environment
* `Stable Baselines 3 <https://github.com/DLR-RM/stable-baselines3>`_ is used as a default source of RL algorithms (although PrimAITE is not limited to SB3 agents)
* `Ray RLlib <https://github.com/ray-project/ray>`_ is used as an additional source of RL algorithms
@@ -107,7 +107,6 @@ Head over to the :ref:`getting-started` page to install and setup PrimAITE!
source/primaite_session
source/simulation
source/game_layer
source/custom_agent
source/config
.. toctree::

7
docs/source/about.rst
View File

@@ -18,7 +18,6 @@ PrimAITE provides the following features:
* Highly configurable network hosts, including definition of software, file system, and network interfaces,
* Realistic network traffic simulation, including address and sending packets via internet protocols like TCP, UDP, ICMP, etc.
* Routers with traffic routing and firewall capabilities
* Interfaces with ARCD GATE to allow training of agents
* Simulation of customisable deterministic agents
* Support for multiple agents, each having their own customisable observation space, action space, and reward function definition.
@@ -148,7 +147,7 @@ The game layer is built on top of the simulator and it consumes the simulation a
Observation Spaces
******************
The observation space provides the blue agent with information about the current status of nodes and links.
PrimAITE builds on top of Gym Spaces to create an observation space that is easily configurable for users. It's made up of components which are managed by the :py:class:`primaite.environment.observations.ObservationsHandler`. Each training scenario can define its own observation space, and the user can choose which information to inlude, and how it should be formatted.
PrimAITE builds on top of Gymnasium Spaces to create an observation space that is easily configurable for users. It's made up of components which are managed by the :py:class:`primaite.environment.observations.ObservationsHandler`. Each training scenario can define its own observation space, and the user can choose which information to inlude, and how it should be formatted.
NodeLinkTable component
-----------------------
For example, the :py:class:`primaite.environment.observations.NodeLinkTable` component represents the status of nodes and links as a ``gym.spaces.Box`` with an example format shown below:
@@ -279,7 +278,7 @@ The game layer is built on top of the simulator and it consumes the simulation a
3. Any (Agent can take both node-based and ACL-based actions)
The choice of action space used during a training session is determined in the config_[name].yaml file.
**Node-Based**
The agent is able to influence the status of nodes by switching them off, resetting, or patching operating systems and services. In this instance, the action space is an OpenAI Gym spaces.Discrete type, as follows:
The agent is able to influence the status of nodes by switching them off, resetting, or patching operating systems and services. In this instance, the action space is a Gymnasium spaces.Discrete type, as follows:
* Dictionary item {... ,1: [x1, x2, x3,x4] ...}
The placeholders inside the list under the key '1' mean the following:
* [0, num nodes] - Node ID (0 = nothing, node ID)
@@ -287,7 +286,7 @@ The game layer is built on top of the simulator and it consumes the simulation a
* [0, 3] - Action on property (0 = nothing, 1 = on / scan, 2 = off / repair, 3 = reset / patch / restore)
* [0, num services] - Resolves to service ID (0 = nothing, resolves to service)
**Access Control List**
The blue agent is able to influence the configuration of the Access Control List rule set (which implements a system-wide firewall). In this instance, the action space is an OpenAI spaces.Discrete type, as follows:
The blue agent is able to influence the configuration of the Access Control List rule set (which implements a system-wide firewall). In this instance, the action space is an Gymnasium spaces.Discrete type, as follows:
* Dictionary item {... ,1: [x1, x2, x3, x4, x5, x6] ...}
The placeholders inside the list under the key '1' mean the following:
* [0, 2] - Action (0 = do nothing, 1 = create rule, 2 = delete rule)

14
docs/source/custom_agent.rst
View File

@@ -1,14 +0,0 @@
.. only:: comment
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
Custom Agents
=============
Integrating a user defined blue agent
*************************************
.. note::
PrimAITE uses ARCD GATE for agent integration. In order to use a custom agent with PrimAITE, you must integrate it with ARCD GATE. Please look at the ARCD GATE documentation for more information.

9
docs/source/game_layer.rst
View File

@@ -4,9 +4,9 @@ PrimAITE Game layer
The Primaite codebase consists of two main modules:
* ``simulator``: The simulation logic including the network topology, the network state, and behaviour of various hardware and software classes.
* ``game``: The agent-training infrastructure which helps reinforcement learning agents interface with the simulation. This includes the observation, action, and rewards, for RL agents, but also scripted deterministic agents. The game layer orchestrates all the interactions between modules, including ARCD GATE.
* ``game``: The agent-training infrastructure which helps reinforcement learning agents interface with the simulation. This includes the observation, action, and rewards, for RL agents, but also scripted deterministic agents. The game layer orchestrates all the interactions between modules.
These two components have been decoupled to allow the agent training code in ARCD GATE to be reused with other simulators. The simulator and game layer communicate using the PrimAITE State API and the PrimAITE Request API. The game layer communicates with ARCD gate using the `Farama Gymnasium Spaces API <https://gymnasium.farama.org/api/spaces/>`_.
The simulator and game layer communicate using the PrimAITE State API and the PrimAITE Request API.
..
TODO: write up these APIs and link them here.
@@ -20,13 +20,14 @@ The game layer is responsible for managing agents and getting them to interface
PrimAITE Session
^^^^^^^^^^^^^^^
``PrimaiteSession`` is the main entry point into Primaite and it allows the simultaneous coordination of a simulation and agents that interact with it. It also sends messages to ARCD GATE to perform reinforcement learning. ``PrimaiteSession`` keeps track of multiple agents of different types.
``PrimaiteSession`` is the main entry point into Primaite and it allows the simultaneous coordination of a simulation and agents that interact with it. ``PrimaiteSession`` keeps track of multiple agents of different types.
Agents
^^^^^^
All agents inherit from the :py:class:`primaite.game.agent.interface.AbstractAgent` class, which mandates that they have an ObservationManager, ActionManager, and RewardManager. The agent behaviour depends on the type of agent, but there are two main types:
* RL agents action during each step is decided by an RL algorithm which lives inside of ARCD GATE. The agent within PrimAITE just acts to format and forward actions decided by an RL policy.
* RL agents action during each step is decided by an appropriate RL algorithm. The agent within PrimAITE just acts to format and forward actions decided by an RL policy.
* Deterministic agents perform all of their decision making within the PrimAITE game layer. They typically have a scripted policy which always performs the same action or a rule-based policy which performs actions based on the current state of the simulation. They can have a stochastic element, and their seed will be settable.
..

30
docs/source/getting_started.rst
View File

@@ -87,22 +87,7 @@ Install PrimAITE
pip install path\to\your\primaite.whl
5. Install ARCD GATE from wheel file
.. code-block:: bash
:caption: Unix
pip install path/to/your/arcd_gate-0.1.0-py3-none-any.whl
.. code-block:: powershell
:caption: Windows (Powershell)
pip install path\to\your\arcd_gate-0.1.0-py3-none-any.whl
6. Perform the PrimAITE setup
5. Perform the PrimAITE setup
.. code-block:: bash
:caption: Unix
@@ -153,17 +138,4 @@ of your choice:
pip install -e .[dev]
4. Install ARCD GATE from wheel file
.. code-block:: bash
:caption: Unix
pip install GATE/arcd_gate-0.1.0-py3-none-any.whl
.. code-block:: powershell
:caption: Windows (Powershell)
pip install GATE\arcd_gate-0.1.0-py3-none-any.whl
To view the complete list of packages installed during PrimAITE installation, go to the dependencies page (:ref:`Dependencies`).

6
docs/source/glossary.rst
View File

@@ -74,8 +74,8 @@ Glossary
Laydown
The laydown is a file which defines the training scenario. It contains the network topology, firewall rules, services, protocols, and details about green and red agent behaviours.
Gym
PrimAITE uses the Gym reinforcement learning framework API to create a training environment and interface with RL agents. Gym defines a common way of creating observations, actions, and rewards.
Gymnasium
PrimAITE uses the Gymnasium reinforcement learning framework API to create a training environment and interface with RL agents. Gymnasium defines a common way of creating observations, actions, and rewards.
User app home
PrimAITE supports upgrading software version while retaining user data. The user data directory is where configs, notebooks, and results are stored, this location is `~/primaite<version>` on linux/darwin and `C:\Users\<username>\primaite\<version>` on Windows.
PrimAITE supports upgrading software version while retaining user data. The user data directory is where configs, notebooks, and results are stored, this location is `~/primaite<version>` on linux/darwin and `C:\\Users\\<username>\\primaite\\<version>` on Windows.

6
docs/source/request_system.rst
View File

@@ -5,12 +5,12 @@
Request System
==============
``SimComponent`` in the simulation are decoupled from the agent training logic. However, they still need a managed means of accepting requests to perform actions. For this, they use ``RequestManager`` and ``RequestType``.
``SimComponent`` objects in the simulation are decoupled from the agent training logic. However, they still need a managed means of accepting requests to perform actions. For this, they use ``RequestManager`` and ``RequestType``.
Just like other aspects of SimComponent, the request typess are not managed centrally for the whole simulation, but instead they are dynamically created and updated based on the nodes, links, and other components that currently exist. This was achieved in the following way:
Just like other aspects of SimComponent, the request types are not managed centrally for the whole simulation, but instead they are dynamically created and updated based on the nodes, links, and other components that currently exist. This was achieved in the following way:
- API
An ``RequestType`` contains two elements:
A ``RequestType`` contains two elements:
1. ``request`` - selects which action you want to take on this ``SimComponent``. This is formatted as a list of strings such as `['network', 'node', '<node-uuid>', 'service', '<service-uuid>', 'restart']`.
2. ``context`` - optional extra information that can be used to decide how to process the request. This is formatted as a dictionary. For example, if the request requires authentication, the context can include information about the user that initiated the request to decide if their permissions are sufficient.

5
docs/source/simulation.rst
View File

@@ -23,8 +23,3 @@ Contents
simulation_components/network/network
simulation_components/system/internal_frame_processing
simulation_components/system/software
simulation_components/system/data_manipulation_bot
simulation_components/system/database_client_server
simulation_components/system/dns_client_server
simulation_components/system/ftp_client_server
simulation_components/system/web_browser_and_web_server_service

62
docs/source/simulation_components/network/base_hardware.rst
View File

@@ -109,6 +109,67 @@ e.g.
instant_start_node = Node(hostname="client", start_up_duration=0, shut_down_duration=0)
instant_start_node.power_on() # node will still need to be powered on
.. _Node Start up and Shut down:
---------------------------
Node Start up and Shut down
---------------------------
Nodes are powered on and off over multiple timesteps. By default, the node ``start_up_duration`` and ``shut_down_duration`` is 3 timesteps.
Example code where a node is turned on:
.. code-block:: python
from primaite.simulator.network.hardware.base import Node
from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
node = Node(hostname="pc_a")
assert node.operating_state is NodeOperatingState.OFF # By default, node is instantiated in an OFF state
node.power_on() # power on the node
assert node.operating_state is NodeOperatingState.BOOTING # node is booting up
for i in range(node.start_up_duration + 1):
# apply timestep until the node start up duration
node.apply_timestep(timestep=i)
assert node.operating_state is NodeOperatingState.ON # node is in ON state
If the node needs to be instantiated in an on state:
.. code-block:: python
from primaite.simulator.network.hardware.base import Node
from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
node = Node(hostname="pc_a", operating_state=NodeOperatingState.ON)
assert node.operating_state is NodeOperatingState.ON # node is in ON state
Setting ``start_up_duration`` and/or ``shut_down_duration`` to ``0`` will allow for the ``power_on`` and ``power_off`` methods to be completed instantly without applying timesteps:
.. code-block:: python
from primaite.simulator.network.hardware.base import Node
from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
node = Node(hostname="pc_a", start_up_duration=0, shut_down_duration=0)
assert node.operating_state is NodeOperatingState.OFF # node is in OFF state
node.power_on()
assert node.operating_state is NodeOperatingState.ON # node is in ON state
node.power_off()
assert node.operating_state is NodeOperatingState.OFF # node is in OFF state
------------------
Network Interfaces
------------------
@@ -357,7 +418,6 @@ Creating the four nodes results in:
2023-08-08 15:50:08,359 INFO: Connected NIC 84:20:7c:ec:a5:c6/192.168.0.13
---------------
Create Switches
---------------

10
docs/source/simulation_components/system/data_manipulation_bot.rst
View File

@@ -16,8 +16,7 @@ The bot is intended to simulate a malicious actor carrying out attacks like:
- Dropping tables
- Deleting records
- Modifying data
On a database server by abusing an application's trusted database connectivity.
on a database server by abusing an application's trusted database connectivity.
The bot performs attacks in the following stages to simulate the real pattern of an attack:
@@ -47,9 +46,12 @@ Example
.. code-block:: python
client_1 = Computer(
hostname="client_1", ip_address="192.168.10.21", subnet_mask="255.255.255.0", default_gateway="192.168.10.1"
hostname="client_1",
ip_address="192.168.10.21",
subnet_mask="255.255.255.0",
default_gateway="192.168.10.1"
operating_state=NodeOperatingState.ON # initialise the computer in an ON state
)
client_1.power_on()
network.connect(endpoint_b=client_1.ethernet_port[1], endpoint_a=switch_2.switch_ports[1])
client_1.software_manager.install(DataManipulationBot)
data_manipulation_bot: DataManipulationBot = client_1.software_manager.software["DataManipulationBot"]

6
docs/source/simulation_components/system/database_client_server.rst
View File

@@ -45,17 +45,14 @@ Key features
^^^^^^^^^^^^
- Connects to the ``DatabaseService`` via the ``SoftwareManager``.
- Handles connecting and disconnecting.
- Executes SQL queries and retrieves result sets.
- Handles connecting, querying, and disconnecting.
- Provides a simple ``query`` method for running SQL.
Usage
^^^^^
- Initialise with server IP address and optional password.
- Connect to the ``DatabaseService`` with ``connect``.
- Execute SQL queries via ``query``.
- Retrieve results in a dictionary.
- Disconnect when finished.
@@ -70,6 +67,5 @@ Implementation
- Leverages ``SoftwareManager`` for sending payloads over the network.
- Connect and disconnect methods manage sessions.
- Provides easy interface for applications to query database.
- Payloads serialised as dictionaries for transmission.
- Extends base Application class.

17
docs/source/simulation_components/system/ftp_client_server.rst
View File

@@ -77,6 +77,7 @@ Dependencies
from primaite.simulator.network.hardware.nodes.server import Server
from primaite.simulator.system.services.ftp.ftp_server import FTPServer
from primaite.simulator.system.services.ftp.ftp_client import FTPClient
from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
Example peer to peer network
^^^^^^^^^^^^^^^^^^^^^^^^^^^^
@@ -85,10 +86,18 @@ Example peer to peer network
net = Network()
pc1 = Computer(hostname="pc1", ip_address="120.10.10.10", subnet_mask="255.255.255.0")
srv = Server(hostname="srv", ip_address="120.10.10.20", subnet_mask="255.255.255.0")
pc1.power_on()
srv.power_on()
pc1 = Computer(
hostname="pc1",
ip_address="120.10.10.10",
subnet_mask="255.255.255.0",
operating_state=NodeOperatingState.ON # initialise the computer in an ON state
)
srv = Server(
hostname="srv",
ip_address="120.10.10.20",
subnet_mask="255.255.255.0",
operating_state=NodeOperatingState.ON # initialise the server in an ON state
)
net.connect(pc1.ethernet_port[1], srv.ethernet_port[1])
Install the FTP Server

39
docs/source/simulation_components/system/software.rst
View File

@@ -6,14 +6,45 @@
Software
========
-------------
Base Software
-------------
All software which inherits ``IOSoftware`` installed on a node will not work unless the node has been turned on.
See :ref:`Node Start up and Shut down`
.. code-block:: python
from primaite.simulator.network.hardware.base import Node
from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
from primaite.simulator.system.services.service import ServiceOperatingState
from primaite.simulator.system.services.web_server.web_server import WebServer
node = Node(hostname="pc_a", start_up_duration=0, shut_down_duration=0)
node.power_on()
assert node.operating_state is NodeOperatingState.ON
node.software_manager.install(WebServer)
web_server: WebServer = node.software_manager.software["WebServer"]
assert web_server.operating_state is ServiceOperatingState.RUNNING # service is immediately ran after install
node.power_off()
assert node.operating_state is NodeOperatingState.OFF
assert web_server.operating_state is ServiceOperatingState.STOPPED # service stops when node is powered off
node.power_on()
assert node.operating_state is NodeOperatingState.ON
assert web_server.operating_state is ServiceOperatingState.RUNNING # service turned back on when node is powered on
Contents
########
Services, Processes and Applications:
#####################################
.. toctree::
:maxdepth: 8
:maxdepth: 2
database_client_server
data_manipulation_bot

4
docs/source/state_system.rst
View File

@@ -5,9 +5,9 @@
Simulation State
==============
``SimComponent`` in the simulation have a method called ``describe_state`` which returns a dictionary of the state of the component. This is used to report pertinent data that could impact agent's actions or rewards. For instance, the name and health status of a node is reported, which can be used by a reward function to punish corrupted or compromised nodes and reward healthy nodes. Each ``SimComponent`` reports not only it's own attributes in the state but also that of its child components. I.e. a computer node will report the state of its ``FileSystem``, and the ``FileSystem`` will report the state of its files and folders. This happens by recursively calling childrens' own ``describe_state`` methods.
``SimComponent`` objects in the simulation have a method called ``describe_state`` which return a dictionary of the state of the component. This is used to report pertinent data that could impact an agent's actions or rewards. For instance, the name and health status of a node is reported, which can be used by a reward function to punish corrupted or compromised nodes and reward healthy nodes. Each ``SimComponent`` object reports not only its own attributes in the state but also those of its child components. I.e. a computer node will report the state of its ``FileSystem`` and the ``FileSystem`` will report the state of its files and folders. This happens by recursively calling the childrens' own ``describe_state`` methods.
The game layer calls ``describe_state`` on the trunk ``SimComponent`` (the top-level parent) and then pass the state to the agents once per simulation step. For this reason, all ``SimComponent`` must have a ``describe_state`` method, and they must all be linked to the trunk ``SimComponent``.
The game layer calls ``describe_state`` on the trunk ``SimComponent`` (the top-level parent) and then passes the state to the agents once per simulation step. For this reason, all ``SimComponent`` objetcs must have a ``describe_state`` method, and they must all be linked to the trunk ``SimComponent``.
This code snippet demonstrates how the state information is defined within the ``SimComponent`` class:

2
src/primaite/VERSION
View File

@@ -1 +1 @@
3.0.0a1
3.0.0b2dev

2
src/primaite/exceptions.py
View File

@@ -1,6 +1,6 @@
# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
class PrimaiteError(Exception):
"""The root PrimAITe Error."""
"""The root PrimAITE Error."""
pass

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

@@ -263,7 +263,7 @@ class FolderObservation(AbstractObservation):
self.files.append(FileObservation())
while len(self.files) > num_files_per_folder:
truncated_file = self.files.pop()
msg = f"Too many files in folde observation. Truncating file {truncated_file}"
msg = f"Too many files in folder observation. Truncating file {truncated_file}"
_LOGGER.warn(msg)
self.default_observation = {

11
src/primaite/simulator/network/container.py
View File

@@ -66,6 +66,17 @@ class Network(SimComponent):
)
return rm
def apply_timestep(self, timestep: int) -> None:
"""Apply a timestep evolution to this the network and its nodes and links."""
super().apply_timestep(timestep=timestep)
# apply timestep to nodes
for node_id in self.nodes:
self.nodes[node_id].apply_timestep(timestep=timestep)
# apply timestep to links
for link_id in self.links:
self.links[link_id].apply_timestep(timestep=timestep)
@property
def routers(self) -> List[Router]:
"""The Routers in the Network."""

52
src/primaite/simulator/network/hardware/base.py
View File

@@ -2,7 +2,6 @@ from __future__ import annotations
import re
import secrets
from enum import Enum
from ipaddress import IPv4Address, IPv4Network
from pathlib import Path
from typing import Any, Dict, Literal, Optional, Tuple, Union
@@ -15,6 +14,7 @@ from primaite.simulator import SIM_OUTPUT
from primaite.simulator.core import RequestManager, RequestType, SimComponent
from primaite.simulator.domain.account import Account
from primaite.simulator.file_system.file_system import FileSystem
from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
from primaite.simulator.network.protocols.arp import ARPEntry, ARPPacket
from primaite.simulator.network.transmission.data_link_layer import EthernetHeader, Frame
from primaite.simulator.network.transmission.network_layer import ICMPPacket, ICMPType, IPPacket, IPProtocol
@@ -886,19 +886,6 @@ class ICMP:
return sequence, icmp_packet.identifier
class NodeOperatingState(Enum):
"""Enumeration of Node Operating States."""
ON = 1
"The node is powered on."
OFF = 2
"The node is powered off."
BOOTING = 3
"The node is in the process of booting up."
SHUTTING_DOWN = 4
"The node is in the process of shutting down."
class Node(SimComponent):
"""
A basic Node class that represents a node on the network.
@@ -1176,18 +1163,21 @@ class Node(SimComponent):
else:
if self.operating_state == NodeOperatingState.BOOTING:
self.operating_state = NodeOperatingState.ON
self.sys_log.info("Turned on")
self.sys_log.info(f"{self.hostname}: Turned on")
for nic in self.nics.values():
if nic._connected_link:
nic.enable()
self._start_up_actions()
# count down to shut down
if self.shut_down_countdown > 0:
self.shut_down_countdown -= 1
else:
if self.operating_state == NodeOperatingState.SHUTTING_DOWN:
self.operating_state = NodeOperatingState.OFF
self.sys_log.info("Turned off")
self.sys_log.info(f"{self.hostname}: Turned off")
self._shut_down_actions()
# if resetting turn back on
if self.is_resetting:
@@ -1283,6 +1273,7 @@ class Node(SimComponent):
self.start_up_countdown = self.start_up_duration
if self.start_up_duration <= 0:
self._start_up_actions()
self.operating_state = NodeOperatingState.ON
self.sys_log.info("Turned on")
for nic in self.nics.values():
@@ -1298,6 +1289,7 @@ class Node(SimComponent):
self.shut_down_countdown = self.shut_down_duration
if self.shut_down_duration <= 0:
self._shut_down_actions()
self.operating_state = NodeOperatingState.OFF
self.sys_log.info("Turned off")
@@ -1504,6 +1496,34 @@ class Node(SimComponent):
_LOGGER.info(f"Removed application {application.uuid} from node {self.uuid}")
self._application_request_manager.remove_request(application.uuid)
def _shut_down_actions(self):
"""Actions to perform when the node is shut down."""
# Turn off all the services in the node
for service_id in self.services:
self.services[service_id].stop()
# Turn off all the applications in the node
for app_id in self.applications:
self.applications[app_id].close()
# Turn off all processes in the node
# for process_id in self.processes:
# self.processes[process_id]
def _start_up_actions(self):
"""Actions to perform when the node is starting up."""
# Turn on all the services in the node
for service_id in self.services:
self.services[service_id].start()
# Turn on all the applications in the node
for app_id in self.applications:
self.applications[app_id].run()
# Turn off all processes in the node
# for process_id in self.processes:
# self.processes[process_id]
def __contains__(self, item: Any) -> bool:
if isinstance(item, Service):
return item.uuid in self.services

14
src/primaite/simulator/network/hardware/node_operating_state.py Normal file
View File

@@ -0,0 +1,14 @@
from enum import Enum
class NodeOperatingState(Enum):
"""Enumeration of Node Operating States."""
ON = 1
"The node is powered on."
OFF = 2
"The node is powered off."
BOOTING = 3
"The node is in the process of booting up."
SHUTTING_DOWN = 4
"The node is in the process of shutting down."

3
src/primaite/simulator/network/protocols/ftp.py
View File

@@ -35,6 +35,9 @@ class FTPCommand(Enum):
class FTPStatusCode(Enum):
"""Status code of the current FTP request."""
NOT_FOUND = 14
"""Destination not found."""
OK = 200
"""Command successful."""

32
src/primaite/simulator/system/applications/application.py
View File

@@ -2,8 +2,11 @@ from abc import abstractmethod
from enum import Enum
from typing import Any, Dict, Set
from primaite import getLogger
from primaite.simulator.system.software import IOSoftware, SoftwareHealthState
_LOGGER = getLogger(__name__)
class ApplicationOperatingState(Enum):
"""Enumeration of Application Operating States."""
@@ -57,7 +60,7 @@ class Application(IOSoftware):
state = super().describe_state()
state.update(
{
"opearting_state": self.operating_state.value,
"operating_state": self.operating_state.value,
"execution_control_status": self.execution_control_status,
"num_executions": self.num_executions,
"groups": list(self.groups),
@@ -65,8 +68,30 @@ class Application(IOSoftware):
)
return state
def _can_perform_action(self) -> bool:
"""
Checks if the application can perform actions.
This is done by checking if the application is operating properly or the node it is installed
in is operational.
Returns true if the software can perform actions.
"""
if not super()._can_perform_action():
return False
if self.operating_state is not self.operating_state.RUNNING:
# service is not running
_LOGGER.error(f"Cannot perform action: {self.name} is {self.operating_state.name}")
return False
return True
def run(self) -> None:
"""Open the Application."""
if not super()._can_perform_action():
return
if self.operating_state == ApplicationOperatingState.CLOSED:
self.sys_log.info(f"Running Application {self.name}")
self.operating_state = ApplicationOperatingState.RUNNING
@@ -83,6 +108,9 @@ class Application(IOSoftware):
def install(self) -> None:
"""Install Application."""
if self._can_perform_action():
return
super().install()
if self.operating_state == ApplicationOperatingState.CLOSED:
self.sys_log.info(f"Installing Application {self.name}")
@@ -98,4 +126,4 @@ class Application(IOSoftware):
:param payload: The payload to receive.
:return: True if successful, False otherwise.
"""
pass
return super().receive(payload=payload, session_id=session_id, **kwargs)

24
src/primaite/simulator/system/applications/database_client.py
View File

@@ -61,7 +61,10 @@ class DatabaseClient(Application):
def connect(self) -> bool:
"""Connect to a Database Service."""
if not self.connected and self.operating_state == ApplicationOperatingState.RUNNING:
if not self._can_perform_action():
return False
if not self.connected:
return self._connect(self.server_ip_address, self.server_password)
return False
@@ -142,20 +145,31 @@ class DatabaseClient(Application):
if self.operating_state == ApplicationOperatingState.RUNNING:
self.connect()
def query(self, sql: str) -> bool:
def query(self, sql: str, is_reattempt: bool = False) -> bool:
"""
Send a query to the Database Service.
:param sql: The SQL query.
:param: sql: The SQL query.
:param: is_reattempt: If true, the action has been reattempted.
:return: True if the query was successful, otherwise False.
"""
if self.connected and self.operating_state == ApplicationOperatingState.RUNNING:
if not self._can_perform_action():
return False
if self.connected:
query_id = str(uuid4())
# Initialise the tracker of this ID to False
self._query_success_tracker[query_id] = False
return self._query(sql=sql, query_id=query_id)
return False
else:
if is_reattempt:
return False
if not self.connect():
return False
self.query(sql=sql, is_reattempt=True)
def receive(self, payload: Any, session_id: str, **kwargs) -> bool:
"""

24
src/primaite/simulator/system/applications/web_browser.py
View File

@@ -3,7 +3,12 @@ from typing import Dict, Optional
from urllib.parse import urlparse
from primaite.simulator.core import RequestManager, RequestType
from primaite.simulator.network.protocols.http import HttpRequestMethod, HttpRequestPacket, HttpResponsePacket
from primaite.simulator.network.protocols.http import (
HttpRequestMethod,
HttpRequestPacket,
HttpResponsePacket,
HttpStatusCode,
)
from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.system.applications.application import Application
@@ -72,8 +77,11 @@ class WebBrowser(Application):
:type: url: str
"""
url = self.target_url
if not self._can_perform_action():
return False
# reset latest response
self.latest_response = None
self.latest_response = HttpResponsePacket(status_code=HttpStatusCode.NOT_FOUND)
try:
parsed_url = urlparse(url)
@@ -102,11 +110,19 @@ class WebBrowser(Application):
payload = HttpRequestPacket(request_method=HttpRequestMethod.GET, request_url=url)
# send request
return self.send(
if self.send(
payload=payload,
dest_ip_address=self.domain_name_ip_address,
dest_port=parsed_url.port if parsed_url.port else Port.HTTP,
)
):
self.sys_log.info(
f"{self.name}: Received HTTP {payload.request_method.name} "
f"Response {payload.request_url} - {self.latest_response.status_code.value}"
)
return self.latest_response.status_code is HttpStatusCode.OK
else:
self.sys_log.error(f"Error sending Http Packet {str(payload)}")
return False
def send(
self,

11
src/primaite/simulator/system/services/database/database_service.py
View File

@@ -65,6 +65,10 @@ class DatabaseService(Service):
def backup_database(self) -> bool:
"""Create a backup of the database to the configured backup server."""
# check if this action can be performed
if not self._can_perform_action():
return False
# check if the backup server was configured
if self.backup_server is None:
self.sys_log.error(f"{self.name} - {self.sys_log.hostname}: not configured.")
@@ -90,6 +94,10 @@ class DatabaseService(Service):
def restore_backup(self) -> bool:
"""Restore a backup from backup server."""
# check if this action can be performed
if not self._can_perform_action():
return False
software_manager: SoftwareManager = self.software_manager
ftp_client_service: FTPClient = software_manager.software["FTPClient"]
@@ -190,6 +198,9 @@ class DatabaseService(Service):
:param session_id: The session identifier.
:return: True if the Status Code is 200, otherwise False.
"""
if not super().receive(payload=payload, session_id=session_id, **kwargs):
return False
result = {"status_code": 500, "data": []}
if isinstance(payload, dict) and payload.get("type"):
if payload["type"] == "connect_request":

18
src/primaite/simulator/system/services/dns/dns_client.py
View File

@@ -53,14 +53,27 @@ class DNSClient(Service):
state = super().describe_state()
return state
def add_domain_to_cache(self, domain_name: str, ip_address: IPv4Address):
def reset_component_for_episode(self, episode: int):
"""
Resets the Service component for a new episode.
This method ensures the Service is ready for a new episode, including resetting any
stateful properties or statistics, and clearing any message queues.
"""
pass
def add_domain_to_cache(self, domain_name: str, ip_address: IPv4Address) -> bool:
"""
Adds a domain name to the DNS Client cache.
:param: domain_name: The domain name to save to cache
:param: ip_address: The IP Address to attach the domain name to
"""
if not self._can_perform_action():
return False
self.dns_cache[domain_name] = ip_address
return True
def check_domain_exists(
self,
@@ -74,6 +87,9 @@ class DNSClient(Service):
:param: session_id: The Session ID the payload is to originate from. Optional.
:param: is_reattempt: Checks if the request has been reattempted. Default is False.
"""
if not self._can_perform_action():
return False
# check if DNS server is configured
if self.dns_server is None:
self.sys_log.error(f"{self.name}: DNS Server is not configured")

10
src/primaite/simulator/system/services/dns/dns_server.py
View File

@@ -64,6 +64,9 @@ class DNSServer(Service):
:param target_domain: The single domain name requested by a DNS client.
:return ip_address: The IP address of that domain name or None.
"""
if not self._can_perform_action():
return
return self.dns_table.get(target_domain)
def dns_register(self, domain_name: str, domain_ip_address: IPv4Address):
@@ -76,6 +79,9 @@ class DNSServer(Service):
:param: domain_ip_address: The IP address that the domain should route to
:type: domain_ip_address: IPv4Address
"""
if not self._can_perform_action():
return
self.dns_table[domain_name] = domain_ip_address
def receive(
@@ -95,10 +101,14 @@ class DNSServer(Service):
:return: True if DNS request returns a valid IP, otherwise, False
"""
if not super().receive(payload=payload, session_id=session_id, **kwargs):
return False
# The payload should be a DNS packet
if not isinstance(payload, DNSPacket):
_LOGGER.debug(f"{payload} is not a DNSPacket")
return False
# cast payload into a DNS packet
payload: DNSPacket = payload
if payload.dns_request is not None:

8
src/primaite/simulator/system/services/ftp/ftp_client.py
View File

@@ -72,10 +72,7 @@ class FTPClient(FTPServiceABC):
# normally FTP will choose a random port for the transfer, but using the FTP command port will do for now
# create FTP packet
payload: FTPPacket = FTPPacket(
ftp_command=FTPCommand.PORT,
ftp_command_args=Port.FTP,
)
payload: FTPPacket = FTPPacket(ftp_command=FTPCommand.PORT, ftp_command_args=Port.FTP)
if self.send(payload=payload, dest_ip_address=dest_ip_address, dest_port=dest_port, session_id=session_id):
if payload.status_code == FTPStatusCode.OK:
@@ -271,7 +268,10 @@ class FTPClient(FTPServiceABC):
the same node.
"""
if payload.status_code is None:
self.sys_log.error(f"FTP Server could not be found - Error Code: {payload.status_code.value}")
return False
self.sys_log.info(f"{self.name}: Received FTP Response {payload.ftp_command.name} {payload.status_code.value}")
self._process_ftp_command(payload=payload, session_id=session_id)
return True

3
src/primaite/simulator/system/services/ftp/ftp_server.py
View File

@@ -86,6 +86,9 @@ class FTPServer(FTPServiceABC):
prevents an FTP request loop - FTP client and servers can exist on
the same node.
"""
if not super().receive(payload=payload, session_id=session_id, **kwargs):
return False
if payload.status_code is not None:
return False

40
src/primaite/simulator/system/services/service.py
View File

@@ -1,5 +1,5 @@
from enum import Enum
from typing import Dict, Optional
from typing import Any, Dict, Optional
from primaite import getLogger
from primaite.simulator.core import RequestManager, RequestType
@@ -40,6 +40,40 @@ class Service(IOSoftware):
restart_countdown: Optional[int] = None
"If currently restarting, how many timesteps remain until the restart is finished."
def _can_perform_action(self) -> bool:
"""
Checks if the service can perform actions.
This is done by checking if the service is operating properly or the node it is installed
in is operational.
Returns true if the software can perform actions.
"""
if not super()._can_perform_action():
return False
if self.operating_state is not self.operating_state.RUNNING:
# service is not running
_LOGGER.error(f"Cannot perform action: {self.name} is {self.operating_state.name}")
return False
return True
def receive(self, payload: Any, session_id: str, **kwargs) -> bool:
"""
Receives a payload from the SessionManager.
The specifics of how the payload is processed and whether a response payload
is generated should be implemented in subclasses.
:param payload: The payload to receive.
:param session_id: The identifier of the session that the payload is associated with.
:param kwargs: Additional keyword arguments specific to the implementation.
:return: True if the payload was successfully received and processed, False otherwise.
"""
return super().receive(payload=payload, session_id=session_id, **kwargs)
def __init__(self, **kwargs):
super().__init__(**kwargs)
@@ -88,6 +122,10 @@ class Service(IOSoftware):
def start(self, **kwargs) -> None:
"""Start the service."""
# cant start the service if the node it is on is off
if not super()._can_perform_action():
return
if self.operating_state == ServiceOperatingState.STOPPED:
self.sys_log.info(f"Starting service {self.name}")
self.operating_state = ServiceOperatingState.RUNNING

3
src/primaite/simulator/system/services/web_server/web_server.py
View File

@@ -170,6 +170,9 @@ class WebServer(Service):
:param: payload: The payload to send.
:param: session_id: The id of the session. Optional.
"""
if not super().receive(payload=payload, session_id=session_id, **kwargs):
return False
# check if the payload is an HTTPPacket
if not isinstance(payload, HttpRequestPacket):
self.sys_log.error("Payload is not an HTTPPacket")

24
src/primaite/simulator/system/software.py
View File

@@ -3,8 +3,9 @@ from enum import Enum
from ipaddress import IPv4Address
from typing import Any, Dict, Optional
from primaite.simulator.core import RequestManager, RequestType, SimComponent
from primaite.simulator.core import _LOGGER, RequestManager, RequestType, SimComponent
from primaite.simulator.file_system.file_system import FileSystem, Folder
from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.system.core.session_manager import Session
from primaite.simulator.system.core.sys_log import SysLog
@@ -234,6 +235,21 @@ class IOSoftware(Software):
)
return state
@abstractmethod
def _can_perform_action(self) -> bool:
"""
Checks if the software can perform actions.
This is done by checking if the software is operating properly or the node it is installed
in is operational.
Returns true if the software can perform actions.
"""
if self.software_manager and self.software_manager.node.operating_state is NodeOperatingState.OFF:
_LOGGER.debug(f"{self.name} Error: {self.software_manager.node.hostname} is not online.")
return False
return True
def send(
self,
payload: Any,
@@ -252,6 +268,9 @@ class IOSoftware(Software):
:return: True if successful, False otherwise.
"""
if not self._can_perform_action():
return False
return self.software_manager.send_payload_to_session_manager(
payload=payload, dest_ip_address=dest_ip_address, dest_port=dest_port, session_id=session_id
)
@@ -270,4 +289,5 @@ class IOSoftware(Software):
:param kwargs: Additional keyword arguments specific to the implementation.
:return: True if the payload was successfully received and processed, False otherwise.
"""
pass
# return false if not allowed to perform actions
return self._can_perform_action()

23
tests/conftest.py
View File

@@ -13,6 +13,7 @@ from primaite.session.session import PrimaiteSession
# from primaite.primaite_session import PrimaiteSession
from primaite.simulator.network.container import Network
from primaite.simulator.network.networks import arcd_uc2_network
from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.system.applications.application import Application
from primaite.simulator.system.core.sys_log import SysLog
@@ -34,6 +35,12 @@ from primaite.simulator.network.hardware.base import Node
class TestService(Service):
"""Test Service class"""
def __init__(self, **kwargs):
kwargs["name"] = "TestService"
kwargs["port"] = Port.HTTP
kwargs["protocol"] = IPProtocol.TCP
super().__init__(**kwargs)
def receive(self, payload: Any, session_id: str, **kwargs) -> bool:
pass
@@ -41,6 +48,12 @@ class TestService(Service):
class TestApplication(Application):
"""Test Application class"""
def __init__(self, **kwargs):
kwargs["name"] = "TestApplication"
kwargs["port"] = Port.HTTP
kwargs["protocol"] = IPProtocol.TCP
super().__init__(**kwargs)
def describe_state(self) -> Dict:
pass
@@ -57,6 +70,11 @@ def service(file_system) -> TestService:
)
@pytest.fixture(scope="function")
def service_class():
return TestService
@pytest.fixture(scope="function")
def application(file_system) -> TestApplication:
return TestApplication(
@@ -64,6 +82,11 @@ def application(file_system) -> TestApplication:
)
@pytest.fixture(scope="function")
def application_class():
return TestApplication
@pytest.fixture(scope="function")
def file_system() -> FileSystem:
return Node(hostname="fs_node").file_system

1
tests/e2e_integration_tests/test_uc2_data_manipulation_scenario.py
View File

@@ -23,6 +23,7 @@ def test_data_manipulation(uc2_network):
# Now we run the DataManipulationBot
db_manipulation_bot.run()
db_manipulation_bot.attack()
# Now check that the DB client on the web_server cannot query the users table on the database
assert not db_client.query("SELECT")

121
tests/integration_tests/system/test_application_on_node.py Normal file
View File

@@ -0,0 +1,121 @@
from typing import Tuple
import pytest
from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
from primaite.simulator.network.hardware.nodes.computer import Computer
from primaite.simulator.system.applications.application import Application, ApplicationOperatingState
@pytest.fixture(scope="function")
def populated_node(application_class) -> Tuple[Application, Computer]:
computer: Computer = Computer(
hostname="test_computer",
ip_address="192.168.1.2",
subnet_mask="255.255.255.0",
default_gateway="192.168.1.1",
operating_state=NodeOperatingState.ON,
)
computer.software_manager.install(application_class)
app = computer.software_manager.software["TestApplication"]
app.run()
return app, computer
def test_service_on_offline_node(application_class):
"""Test to check that the service cannot be interacted with when node it is on is off."""
computer: Computer = Computer(
hostname="test_computer",
ip_address="192.168.1.2",
subnet_mask="255.255.255.0",
default_gateway="192.168.1.1",
operating_state=NodeOperatingState.ON,
)
computer.software_manager.install(application_class)
app: Application = computer.software_manager.software["TestApplication"]
computer.power_off()
for i in range(computer.shut_down_duration + 1):
computer.apply_timestep(timestep=i)
assert computer.operating_state is NodeOperatingState.OFF
assert app.operating_state is ApplicationOperatingState.CLOSED
app.run()
assert app.operating_state is ApplicationOperatingState.CLOSED
def test_server_turns_off_service(populated_node):
"""Check that the service is turned off when the server is turned off"""
app, computer = populated_node
assert computer.operating_state is NodeOperatingState.ON
assert app.operating_state is ApplicationOperatingState.RUNNING
computer.power_off()
for i in range(computer.shut_down_duration + 1):
computer.apply_timestep(timestep=i)
assert computer.operating_state is NodeOperatingState.OFF
assert app.operating_state is ApplicationOperatingState.CLOSED
def test_service_cannot_be_turned_on_when_server_is_off(populated_node):
"""Check that the service cannot be started when the server is off."""
app, computer = populated_node
assert computer.operating_state is NodeOperatingState.ON
assert app.operating_state is ApplicationOperatingState.RUNNING
computer.power_off()
for i in range(computer.shut_down_duration + 1):
computer.apply_timestep(timestep=i)
assert computer.operating_state is NodeOperatingState.OFF
assert app.operating_state is ApplicationOperatingState.CLOSED
app.run()
assert computer.operating_state is NodeOperatingState.OFF
assert app.operating_state is ApplicationOperatingState.CLOSED
def test_server_turns_on_service(populated_node):
"""Check that turning on the server turns on service."""
app, computer = populated_node
assert computer.operating_state is NodeOperatingState.ON
assert app.operating_state is ApplicationOperatingState.RUNNING
computer.power_off()
for i in range(computer.shut_down_duration + 1):
computer.apply_timestep(timestep=i)
assert computer.operating_state is NodeOperatingState.OFF
assert app.operating_state is ApplicationOperatingState.CLOSED
computer.power_on()
for i in range(computer.start_up_duration + 1):
computer.apply_timestep(timestep=i)
assert computer.operating_state is NodeOperatingState.ON
assert app.operating_state is ApplicationOperatingState.RUNNING
computer.start_up_duration = 0
computer.shut_down_duration = 0
computer.power_off()
assert computer.operating_state is NodeOperatingState.OFF
assert app.operating_state is ApplicationOperatingState.CLOSED
computer.power_on()
assert computer.operating_state is NodeOperatingState.ON
assert app.operating_state is ApplicationOperatingState.RUNNING

30
tests/integration_tests/system/test_database_on_node.py
View File

@@ -1,9 +1,11 @@
from ipaddress import IPv4Address
from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
from primaite.simulator.network.hardware.nodes.server import Server
from primaite.simulator.system.applications.database_client import DatabaseClient
from primaite.simulator.system.services.database.database_service import DatabaseService
from primaite.simulator.system.services.ftp.ftp_server import FTPServer
from primaite.simulator.system.services.service import ServiceOperatingState
def test_database_client_server_connection(uc2_network):
@@ -55,7 +57,8 @@ def test_database_client_query(uc2_network):
"""Tests DB query across the network returns HTTP status 200 and date."""
web_server: Server = uc2_network.get_node_by_hostname("web_server")
db_client: DatabaseClient = web_server.software_manager.software["DatabaseClient"]
db_client.connect()
assert db_client.connected
assert db_client.query("SELECT")
@@ -92,3 +95,28 @@ def test_restore_backup(uc2_network):
assert db_service.restore_backup() is True
assert db_service.file_system.get_file(folder_name="database", file_name="database.db") is not None
def test_database_client_cannot_query_offline_database_server(uc2_network):
"""Tests DB query across the network returns HTTP status 404 when db server is offline."""
db_server: Server = uc2_network.get_node_by_hostname("database_server")
db_service: DatabaseService = db_server.software_manager.software["DatabaseService"]
assert db_server.operating_state is NodeOperatingState.ON
assert db_service.operating_state is ServiceOperatingState.RUNNING
web_server: Server = uc2_network.get_node_by_hostname("web_server")
db_client: DatabaseClient = web_server.software_manager.software["DatabaseClient"]
assert db_client.connected
assert db_client.query("SELECT") is True
db_server.power_off()
for i in range(db_server.shut_down_duration + 1):
uc2_network.apply_timestep(timestep=i)
assert db_server.operating_state is NodeOperatingState.OFF
assert db_service.operating_state is ServiceOperatingState.STOPPED
assert db_client.query("SELECT") is False

37
tests/integration_tests/system/test_dns_client_server.py
View File

@@ -1,3 +1,4 @@
from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
from primaite.simulator.network.hardware.nodes.computer import Computer
from primaite.simulator.network.hardware.nodes.server import Server
from primaite.simulator.system.services.dns.dns_client import DNSClient
@@ -24,3 +25,39 @@ def test_dns_client_server(uc2_network):
# arcd.com is registered in dns server and should be saved to cache
assert dns_client.check_domain_exists(target_domain="arcd.com")
assert dns_client.dns_cache.get("arcd.com", None) is not None
assert len(dns_client.dns_cache) == 1
def test_dns_client_requests_offline_dns_server(uc2_network):
client_1: Computer = uc2_network.get_node_by_hostname("client_1")
domain_controller: Server = uc2_network.get_node_by_hostname("domain_controller")
dns_client: DNSClient = client_1.software_manager.software["DNSClient"]
dns_server: DNSServer = domain_controller.software_manager.software["DNSServer"]
assert dns_client.operating_state == ServiceOperatingState.RUNNING
assert dns_server.operating_state == ServiceOperatingState.RUNNING
dns_server.show()
# arcd.com is registered in dns server
assert dns_client.check_domain_exists(target_domain="arcd.com")
assert dns_client.dns_cache.get("arcd.com", None) is not None
assert len(dns_client.dns_cache) == 1
dns_client.dns_cache = {}
domain_controller.power_off()
for i in range(domain_controller.shut_down_duration + 1):
uc2_network.apply_timestep(timestep=i)
assert domain_controller.operating_state == NodeOperatingState.OFF
assert dns_server.operating_state == ServiceOperatingState.STOPPED
# this time it should not cache because dns server is not online
assert dns_client.check_domain_exists(target_domain="arcd.com") is False
assert dns_client.dns_cache.get("arcd.com", None) is None
assert len(dns_client.dns_cache) == 0

49
tests/integration_tests/system/test_ftp_client_server.py
View File

@@ -15,10 +15,10 @@ def test_ftp_client_store_file_in_server(uc2_network):
backup_server: Server = uc2_network.get_node_by_hostname("backup_server")
ftp_client: FTPClient = client_1.software_manager.software["FTPClient"]
ftp_server: FTPServer = backup_server.software_manager.software["FTPServer"]
ftp_server_service: FTPServer = backup_server.software_manager.software["FTPServer"]
assert ftp_client.operating_state == ServiceOperatingState.RUNNING
assert ftp_server.operating_state == ServiceOperatingState.RUNNING
assert ftp_server_service.operating_state == ServiceOperatingState.RUNNING
# create file on ftp client
ftp_client.file_system.create_file(file_name="test_file.txt")
@@ -31,7 +31,7 @@ def test_ftp_client_store_file_in_server(uc2_network):
dest_ip_address=backup_server.nics.get(next(iter(backup_server.nics))).ip_address,
)
assert ftp_server.file_system.get_file(folder_name="client_1_backup", file_name="test_file.txt")
assert ftp_server_service.file_system.get_file(folder_name="client_1_backup", file_name="test_file.txt")
def test_ftp_client_retrieve_file_from_server(uc2_network):
@@ -42,13 +42,13 @@ def test_ftp_client_retrieve_file_from_server(uc2_network):
backup_server: Server = uc2_network.get_node_by_hostname("backup_server")
ftp_client: FTPClient = client_1.software_manager.software["FTPClient"]
ftp_server: FTPServer = backup_server.software_manager.software["FTPServer"]
ftp_server_service: FTPServer = backup_server.software_manager.software["FTPServer"]
assert ftp_client.operating_state == ServiceOperatingState.RUNNING
assert ftp_server.operating_state == ServiceOperatingState.RUNNING
assert ftp_server_service.operating_state == ServiceOperatingState.RUNNING
# create file on ftp server
ftp_server.file_system.create_file(file_name="test_file.txt", folder_name="file_share")
ftp_server_service.file_system.create_file(file_name="test_file.txt", folder_name="file_share")
assert ftp_client.request_file(
src_folder_name="file_share",
@@ -60,3 +60,40 @@ def test_ftp_client_retrieve_file_from_server(uc2_network):
# client should have retrieved the file
assert ftp_client.file_system.get_file(folder_name="downloads", file_name="test_file.txt")
def test_ftp_client_tries_to_connect_to_offline_server(uc2_network):
"""Test checks to make sure that the client can't do anything when the server is offline."""
client_1: Computer = uc2_network.get_node_by_hostname("client_1")
backup_server: Server = uc2_network.get_node_by_hostname("backup_server")
ftp_client: FTPClient = client_1.software_manager.software["FTPClient"]
ftp_server_service: FTPServer = backup_server.software_manager.software["FTPServer"]
assert ftp_client.operating_state == ServiceOperatingState.RUNNING
assert ftp_server_service.operating_state == ServiceOperatingState.RUNNING
# create file on ftp server
ftp_server_service.file_system.create_file(file_name="test_file.txt", folder_name="file_share")
backup_server.power_off()
for i in range(backup_server.shut_down_duration + 1):
uc2_network.apply_timestep(timestep=i)
assert ftp_client.operating_state == ServiceOperatingState.RUNNING
assert ftp_server_service.operating_state == ServiceOperatingState.STOPPED
assert (
ftp_client.request_file(
src_folder_name="file_share",
src_file_name="test_file.txt",
dest_folder_name="downloads",
dest_file_name="test_file.txt",
dest_ip_address=backup_server.nics.get(next(iter(backup_server.nics))).ip_address,
)
is False
)
# client should have retrieved the file
assert ftp_client.file_system.get_file(folder_name="downloads", file_name="test_file.txt") is None

129
tests/integration_tests/system/test_service_on_node.py Normal file
View File

@@ -0,0 +1,129 @@
from typing import Tuple
import pytest
from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
from primaite.simulator.network.hardware.nodes.computer import Computer
from primaite.simulator.network.hardware.nodes.server import Server
from primaite.simulator.system.services.service import Service, ServiceOperatingState
@pytest.fixture(scope="function")
def populated_node(
service_class,
) -> Tuple[Server, Service]:
server = Server(
hostname="server", ip_address="192.168.0.1", subnet_mask="255.255.255.0", operating_state=NodeOperatingState.ON
)
server.software_manager.install(service_class)
service = server.software_manager.software["TestService"]
service.start()
return server, service
def test_service_on_offline_node(service_class):
"""Test to check that the service cannot be interacted with when node it is on is off."""
computer: Computer = Computer(
hostname="test_computer",
ip_address="192.168.1.2",
subnet_mask="255.255.255.0",
default_gateway="192.168.1.1",
operating_state=NodeOperatingState.ON,
)
computer.software_manager.install(service_class)
service: Service = computer.software_manager.software["TestService"]
computer.power_off()
for i in range(computer.shut_down_duration + 1):
computer.apply_timestep(timestep=i)
assert computer.operating_state is NodeOperatingState.OFF
assert service.operating_state is ServiceOperatingState.STOPPED
service.start()
assert service.operating_state is ServiceOperatingState.STOPPED
service.resume()
assert service.operating_state is ServiceOperatingState.STOPPED
service.restart()
assert service.operating_state is ServiceOperatingState.STOPPED
service.pause()
assert service.operating_state is ServiceOperatingState.STOPPED
def test_server_turns_off_service(populated_node):
"""Check that the service is turned off when the server is turned off"""
server, service = populated_node
assert server.operating_state is NodeOperatingState.ON
assert service.operating_state is ServiceOperatingState.RUNNING
server.power_off()
for i in range(server.shut_down_duration + 1):
server.apply_timestep(timestep=i)
assert server.operating_state is NodeOperatingState.OFF
assert service.operating_state is ServiceOperatingState.STOPPED
def test_service_cannot_be_turned_on_when_server_is_off(populated_node):
"""Check that the service cannot be started when the server is off."""
server, service = populated_node
assert server.operating_state is NodeOperatingState.ON
assert service.operating_state is ServiceOperatingState.RUNNING
server.power_off()
for i in range(server.shut_down_duration + 1):
server.apply_timestep(timestep=i)
assert server.operating_state is NodeOperatingState.OFF
assert service.operating_state is ServiceOperatingState.STOPPED
service.start()
assert server.operating_state is NodeOperatingState.OFF
assert service.operating_state is ServiceOperatingState.STOPPED
def test_server_turns_on_service(populated_node):
"""Check that turning on the server turns on service."""
server, service = populated_node
assert server.operating_state is NodeOperatingState.ON
assert service.operating_state is ServiceOperatingState.RUNNING
server.power_off()
for i in range(server.shut_down_duration + 1):
server.apply_timestep(timestep=i)
assert server.operating_state is NodeOperatingState.OFF
assert service.operating_state is ServiceOperatingState.STOPPED
server.power_on()
for i in range(server.start_up_duration + 1):
server.apply_timestep(timestep=i)
assert server.operating_state is NodeOperatingState.ON
assert service.operating_state is ServiceOperatingState.RUNNING
server.start_up_duration = 0
server.shut_down_duration = 0
server.power_off()
assert server.operating_state is NodeOperatingState.OFF
assert service.operating_state is ServiceOperatingState.STOPPED
server.power_on()
assert server.operating_state is NodeOperatingState.ON
assert service.operating_state is ServiceOperatingState.RUNNING

52
tests/integration_tests/system/test_web_client_server.py
View File

@@ -1,3 +1,4 @@
from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
from primaite.simulator.network.hardware.nodes.computer import Computer
from primaite.simulator.network.hardware.nodes.server import Server
from primaite.simulator.network.protocols.http import HttpStatusCode
@@ -48,6 +49,55 @@ def test_web_page_get_users_page_request_with_ip_address(uc2_network):
assert web_client.get_webpage() is True
# latest reponse should have status code 200
# latest response should have status code 200
assert web_client.latest_response is not None
assert web_client.latest_response.status_code == HttpStatusCode.OK
def test_web_page_request_from_shut_down_server(uc2_network):
"""Test to see that the web server does not respond when the server is off."""
client_1: Computer = uc2_network.get_node_by_hostname("client_1")
web_client: WebBrowser = client_1.software_manager.software["WebBrowser"]
web_client.run()
web_server: Server = uc2_network.get_node_by_hostname("web_server")
assert web_client.operating_state == ApplicationOperatingState.RUNNING
assert web_client.get_webpage("http://arcd.com/users/") is True
# latest response should have status code 200
assert web_client.latest_response.status_code == HttpStatusCode.OK
web_server.power_off()
for i in range(web_server.shut_down_duration + 1):
uc2_network.apply_timestep(timestep=i)
# node should be off
assert web_server.operating_state is NodeOperatingState.OFF
assert web_client.get_webpage("http://arcd.com/users/") is False
assert web_client.latest_response.status_code == HttpStatusCode.NOT_FOUND
def test_web_page_request_from_closed_web_browser(uc2_network):
client_1: Computer = uc2_network.get_node_by_hostname("client_1")
web_client: WebBrowser = client_1.software_manager.software["WebBrowser"]
web_client.run()
web_server: Server = uc2_network.get_node_by_hostname("web_server")
assert web_client.operating_state == ApplicationOperatingState.RUNNING
assert web_client.get_webpage("http://arcd.com/users/") is True
# latest response should have status code 200
assert web_client.latest_response.status_code == HttpStatusCode.OK
web_client.close()
# node should be off
assert web_client.operating_state is ApplicationOperatingState.CLOSED
assert web_client.get_webpage("http://arcd.com/users/") is False

49
tests/unit_tests/_primaite/_simulator/_system/_services/test_dns.py
View File

@@ -3,6 +3,7 @@ from ipaddress import IPv4Address
import pytest
from primaite.simulator.network.hardware.base import Node
from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
from primaite.simulator.network.hardware.nodes.computer import Computer
from primaite.simulator.network.hardware.nodes.server import Server
from primaite.simulator.network.protocols.dns import DNSPacket, DNSReply, DNSRequest
@@ -10,15 +11,19 @@ from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.system.services.dns.dns_client import DNSClient
from primaite.simulator.system.services.dns.dns_server import DNSServer
from primaite.simulator.system.services.service import ServiceOperatingState
@pytest.fixture(scope="function")
def dns_server() -> Node:
node = Server(
hostname="dns_server", ip_address="192.168.1.10", subnet_mask="255.255.255.0", default_gateway="192.168.1.1"
hostname="dns_server",
ip_address="192.168.1.10",
subnet_mask="255.255.255.0",
default_gateway="192.168.1.1",
operating_state=NodeOperatingState.ON,
)
node.software_manager.install(software_class=DNSServer)
node.software_manager.software["DNSServer"].start()
return node
@@ -50,6 +55,44 @@ def test_create_dns_client(dns_client):
assert dns_client_service.protocol is IPProtocol.TCP
def test_dns_client_add_domain_to_cache_when_not_running(dns_client):
dns_client_service: DNSClient = dns_client.software_manager.software["DNSClient"]
assert dns_client.operating_state is NodeOperatingState.OFF
assert dns_client_service.operating_state is ServiceOperatingState.STOPPED
assert (
dns_client_service.add_domain_to_cache(domain_name="test.com", ip_address=IPv4Address("192.168.1.100")) is False
)
assert dns_client_service.dns_cache.get("test.com") is None
def test_dns_client_check_domain_exists_when_not_running(dns_client):
dns_client.operating_state = NodeOperatingState.ON
dns_client_service: DNSClient = dns_client.software_manager.software["DNSClient"]
dns_client_service.start()
assert dns_client.operating_state is NodeOperatingState.ON
assert dns_client_service.operating_state is ServiceOperatingState.RUNNING
assert (
dns_client_service.add_domain_to_cache(domain_name="test.com", ip_address=IPv4Address("192.168.1.100"))
is not False
)
assert dns_client_service.check_domain_exists("test.com") is True
dns_client.power_off()
for i in range(dns_client.shut_down_duration + 1):
dns_client.apply_timestep(timestep=i)
assert dns_client.operating_state is NodeOperatingState.OFF
assert dns_client_service.operating_state is ServiceOperatingState.STOPPED
assert dns_client_service.check_domain_exists("test.com") is False
def test_dns_server_domain_name_registration(dns_server):
"""Test to check if the domain name registration works."""
dns_server_service: DNSServer = dns_server.software_manager.software["DNSServer"]
@@ -64,7 +107,9 @@ def test_dns_server_domain_name_registration(dns_server):
def test_dns_client_check_domain_in_cache(dns_client):
"""Test to make sure that the check_domain_in_cache returns the correct values."""
dns_client.operating_state = NodeOperatingState.ON
dns_client_service: DNSClient = dns_client.software_manager.software["DNSClient"]
dns_client_service.start()
# add a domain to the dns client cache
dns_client_service.add_domain_to_cache("real-domain.com", IPv4Address("192.168.1.12"))

14
tests/unit_tests/_primaite/_simulator/_system/_services/test_ftp.py
View File

@@ -3,6 +3,7 @@ from ipaddress import IPv4Address
import pytest
from primaite.simulator.network.hardware.base import Node
from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
from primaite.simulator.network.hardware.nodes.computer import Computer
from primaite.simulator.network.hardware.nodes.server import Server
from primaite.simulator.network.protocols.ftp import FTPCommand, FTPPacket, FTPStatusCode
@@ -15,17 +16,24 @@ from primaite.simulator.system.services.ftp.ftp_server import FTPServer
@pytest.fixture(scope="function")
def ftp_server() -> Node:
node = Server(
hostname="ftp_server", ip_address="192.168.1.10", subnet_mask="255.255.255.0", default_gateway="192.168.1.1"
hostname="ftp_server",
ip_address="192.168.1.10",
subnet_mask="255.255.255.0",
default_gateway="192.168.1.1",
operating_state=NodeOperatingState.ON,
)
node.software_manager.install(software_class=FTPServer)
node.software_manager.software["FTPServer"].start()
return node
@pytest.fixture(scope="function")
def ftp_client() -> Node:
node = Computer(
hostname="ftp_client", ip_address="192.168.1.11", subnet_mask="255.255.255.0", default_gateway="192.168.1.1"
hostname="ftp_client",
ip_address="192.168.1.11",
subnet_mask="255.255.255.0",
default_gateway="192.168.1.1",
operating_state=NodeOperatingState.ON,
)
return node

0
tests/unit_tests/_primaite/_simulator/test_sim_conatiner.py → tests/unit_tests/_primaite/_simulator/test_sim_container.py
View File