Merge remote-tracking branch 'devops/dev' into downstream_github_sync
This commit is contained in:
Chris McCarthy
@@ -86,5 +86,5 @@ stages:
|
||||
displayName: 'Perform PrimAITE Setup'
|
||||
|
||||
- script: |
|
||||
pytest -n 4
|
||||
pytest -n auto
|
||||
displayName: 'Run tests'
|
||||
|
||||
@@ -9,5 +9,6 @@
|
||||
- [ ] I have performed **self-review** of the code
|
||||
- [ ] I have written **tests** for any new functionality added with this PR
|
||||
- [ ] I have updated the **documentation** if this PR changes or adds functionality
|
||||
- [ ] I have written/updated **design docs** if this PR implements new functionality.
|
||||
- [ ] I have written/updated **design docs** if this PR implements new functionality
|
||||
- [ ] I have update the **change log**
|
||||
- [ ] I have run **pre-commit** checks for code style
|
||||
|
||||
1
.gitignore
vendored
1
.gitignore
vendored
@@ -150,3 +150,4 @@ src/primaite/outputs/
|
||||
|
||||
# benchmark session outputs
|
||||
benchmark/output
|
||||
src/primaite/notebooks/scratch.ipynb
|
||||
|
||||
13
CHANGELOG.md
13
CHANGELOG.md
@@ -7,6 +7,19 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
||||
|
||||
## [Unreleased]
|
||||
|
||||
### Added
|
||||
- Network Hardware - Added base hardware module with NIC, SwitchPort, Node, Switch, and Link. Nodes and Switches have
|
||||
fundamental services like ARP, ICMP, and PCAP running them by default.
|
||||
- Network Transmission - Modelled OSI Model layers 1 through to 5 with various classes for creating network frames and
|
||||
transmitting them from a Service/Application, down through the layers, over the wire, and back up through the layers to
|
||||
a Service/Application another machine.
|
||||
- system - Added the core structure of Application, Services, and Components. Also added a SoftwareManager and
|
||||
SessionManager.
|
||||
- Permission System - each action can define criteria that will be used to permit or deny agent actions.
|
||||
- File System - ability to emulate a node's file system during a simulation
|
||||
- Example notebooks - There is currently 1 jupyter notebook which walks through using PrimAITE
|
||||
1. Creating a simulation - this notebook explains how to build up a simulation using the Python package. (WIP)
|
||||
|
||||
## [2.0.0] - 2023-07-26
|
||||
|
||||
### Added
|
||||
|
||||
BIN
docs/_static/component_relationship.png
vendored
Normal file
BIN
docs/_static/component_relationship.png
vendored
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 80 KiB |
BIN
docs/_static/four_node_two_switch_network.png
vendored
Normal file
BIN
docs/_static/four_node_two_switch_network.png
vendored
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 88 KiB |
@@ -6,7 +6,7 @@
|
||||
Simulation
|
||||
==========
|
||||
|
||||
.. TODO:: Add spiel here about what the simulation is.
|
||||
|
||||
|
||||
|
||||
Contents
|
||||
@@ -16,5 +16,5 @@ Contents
|
||||
:maxdepth: 8
|
||||
|
||||
simulation_structure
|
||||
simulation_components/network/physical_layer
|
||||
simulation_components/network/base_hardware
|
||||
simulation_components/network/transport_to_data_link_layer
|
||||
|
||||
626
docs/source/simulation_components/network/base_hardware.rst
Normal file
626
docs/source/simulation_components/network/base_hardware.rst
Normal file
@@ -0,0 +1,626 @@
|
||||
.. only:: comment
|
||||
|
||||
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
|
||||
|
||||
Base Hardware
|
||||
=============
|
||||
|
||||
The physical layer components are models of a NIC (Network Interface Card), SwitchPort, Node, Switch, and a Link.
|
||||
These components allow modelling of layer 1 (physical layer) in the OSI model and the nodes that connect to and
|
||||
transmit across layer 1.
|
||||
|
||||
NIC
|
||||
###
|
||||
The NIC class provides a realistic model of a Network Interface Card. The NIC acts as the interface between a Node and
|
||||
a Link, handling IP and MAC addressing, status, and sending/receiving frames.
|
||||
|
||||
Addressing
|
||||
**********
|
||||
|
||||
A NIC has both an IPv4 address and MAC address assigned:
|
||||
|
||||
- **ip_address** - The IPv4 address assigned to the NIC for communication on an IP network.
|
||||
- **subnet_mask** - The subnet mask that defines the network subnet.
|
||||
- **gateway** - The default gateway IP address for routing traffic beyond the local network.
|
||||
- **mac_address** - A unique MAC address assigned to the NIC by the manufacturer.
|
||||
|
||||
Status
|
||||
******
|
||||
|
||||
The status of the NIC is represented by:
|
||||
|
||||
- **enabled** - Indicates if the NIC is active/enabled or disabled/down. It must be enabled to send/receive frames.
|
||||
- **connected_node** - The Node instance the NIC is attached to.
|
||||
- **connected_link** - The Link instance the NIC is wired to.
|
||||
|
||||
Packet Capture
|
||||
**************
|
||||
|
||||
- **pcap** - A PacketCapture instance attached to the NIC for capturing all frames sent and received. This allows packet
|
||||
capture and analysis.
|
||||
|
||||
Sending/Receiving Frames
|
||||
************************
|
||||
|
||||
The NIC can send and receive Frames to/from the connected Link:
|
||||
|
||||
- **send_frame()** - Sends a Frame through the NIC onto the attached Link.
|
||||
- **receive_frame()** - Receives a Frame from the attached Link and processes it.
|
||||
|
||||
This allows a NIC to handle sending, receiving, and forwarding of network traffic at layer 2 of the OSI model.
|
||||
The Frames contain network data encapsulated with various protocol headers.
|
||||
|
||||
Basic Usage
|
||||
***********
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
nic1 = NIC(
|
||||
ip_address="192.168.0.100",
|
||||
subnet_mask="255.255.255.0",
|
||||
gateway="192.168.0.1"
|
||||
)
|
||||
nic1.enable()
|
||||
frame = Frame(...)
|
||||
nic1.send_frame(frame)
|
||||
|
||||
SwitchPort
|
||||
##########
|
||||
|
||||
The SwitchPort models a port on a network switch. It has similar attributes and methods to NIC for addressing, status,
|
||||
packet capture, sending/receiving frames, etc.
|
||||
|
||||
Key attributes:
|
||||
|
||||
- **port_num**: The port number on the switch.
|
||||
- **connected_switch**: The switch to which this port belongs.
|
||||
|
||||
Node
|
||||
####
|
||||
|
||||
The Node class represents a base node that communicates on the Network.
|
||||
|
||||
Network Interfaces
|
||||
******************
|
||||
|
||||
A Node will typically have one or more NICs attached to it for network connectivity:
|
||||
|
||||
- **nics** - A dictionary containing the NIC instances attached to the Node. NICs can be added/removed.
|
||||
|
||||
Configuration
|
||||
*************
|
||||
|
||||
- **hostname** - Configured hostname of the Node.
|
||||
- **operating_state** - Current operating state like ON or OFF. The NICs will be enabled/disabled based on this.
|
||||
|
||||
Network Services
|
||||
****************
|
||||
|
||||
A Node runs various network services and components for handling traffic:
|
||||
|
||||
- **session_manager** - Handles establishing sessions to/from the Node.
|
||||
- **software_manager** - Manages software and applications on the Node.
|
||||
- **arp** - ARP cache for resolving IP addresses to MAC addresses.
|
||||
- **icmp** - ICMP service for responding to pings and echo requests.
|
||||
- **sys_log** - System log service for logging internal events and messages.
|
||||
|
||||
The SysLog provides a logging mechanism for the Node:
|
||||
|
||||
The SysLog records informational, warning, and error events that occur on the Node during simulation. This allows
|
||||
debugging and tracing program execution and network activity for each simulated Node. Other Node services like ARP and
|
||||
ICMP, along with custom Applications, services, and Processes will log to the SysLog.
|
||||
|
||||
Sending/Receiving
|
||||
*****************
|
||||
|
||||
The Node handles sending and receiving Frames via its attached NICs:
|
||||
|
||||
- **send_frame()** - Sends a Frame to the network through one of the Node's NICs.
|
||||
- **receive_frame()** - Receives a Frame from the network through a NIC. The Node then processes it appropriately based
|
||||
on the protocols and payload.
|
||||
|
||||
Basic Usage
|
||||
***********
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
node1 = Node(hostname='server1')
|
||||
node1.operating_state = NodeOperatingState.ON
|
||||
|
||||
nic1 = NIC()
|
||||
node1.connect_nic(nic1)
|
||||
|
||||
Send a frame
|
||||
frame = Frame(...)
|
||||
node1.send_frame(frame)
|
||||
|
||||
The Node class brings together the NICs, configuration, and services to model a full network node that can send,
|
||||
receive, process, and forward traffic on a simulated network.
|
||||
|
||||
|
||||
Switch
|
||||
######
|
||||
|
||||
The Switch subclass models a network switch. It inherits from Node and acts at layer 2 of the OSI model to forward
|
||||
frames based on MAC addresses.
|
||||
|
||||
Inherits Node Capabilities
|
||||
**************************
|
||||
|
||||
Since Switch subclasses Node, it inherits all capabilities from Node like:
|
||||
|
||||
- **Managing NICs**
|
||||
- **Running network services like ARP, ICMP**
|
||||
- **Sending and receiving frames**
|
||||
- **Maintaining system logs**
|
||||
|
||||
Ports
|
||||
*****
|
||||
|
||||
A Switch has multiple ports implemented using SwitchPort instances:
|
||||
|
||||
- **switch_ports** - A dictionary mapping port numbers to SwitchPort instances.
|
||||
- **num_ports** - The number of ports the Switch has.
|
||||
|
||||
Forwarding
|
||||
**********
|
||||
|
||||
A Switch forwards frames between ports based on the destination MAC:
|
||||
|
||||
- **dst_mac_table** - MAC address table that maps MACs to SwitchPorts.
|
||||
- **forward_frame()** - Forwards a frame out the port associated with the destination MAC.
|
||||
|
||||
When a frame is received on a SwitchPort:
|
||||
|
||||
1. The source MAC address is extracted from the frame.
|
||||
2. An entry is added to dst_mac_table that maps this source MAC to the SwitchPort it was received on.
|
||||
3. When a frame with that destination MAC is received in the future, it will be forwarded out this SwitchPort.
|
||||
|
||||
This allows the Switch to dynamically build up a mapping table between MAC addresses and SwitchPorts based on traffic
|
||||
received. If no entry exists for a destination MAC, it floods the frame out all ports.
|
||||
|
||||
Link
|
||||
####
|
||||
|
||||
The Link class represents a physical link or connection between two network endpoints like NICs or SwitchPorts.
|
||||
|
||||
Endpoints
|
||||
*********
|
||||
|
||||
A Link connects two endpoints:
|
||||
|
||||
- **endpoint_a** - The first endpoint, a NIC or SwitchPort.
|
||||
- **endpoint_b** - The second endpoint, a NIC or SwitchPort.
|
||||
|
||||
Transmission
|
||||
************
|
||||
|
||||
Links transmit Frames between the endpoints:
|
||||
|
||||
- **transmit_frame()** - Sends a Frame from one endpoint to the other.
|
||||
|
||||
Uses bandwidth/load properties to determine if transmission is possible.
|
||||
|
||||
Bandwidth & Load
|
||||
****************
|
||||
|
||||
- **bandwidth** - The total capacity of the Link in Mbps.
|
||||
- **current_load** - The current bandwidth utilization of the Link in Mbps.
|
||||
|
||||
As Frames are sent over the Link, the load increases. The Link tracks if there is enough unused capacity to transmit a
|
||||
Frame based on its size and the current load.
|
||||
|
||||
Status
|
||||
******
|
||||
|
||||
- **up** - Boolean indicating if the Link is currently up/active based on the endpoint status.
|
||||
- **endpoint_up()/down()** - Notifies the Link when an endpoint goes up or down.
|
||||
|
||||
This allows the Link to realistically model the connection and transmission characteristics between two endpoints.
|
||||
|
||||
Putting it all Together
|
||||
#######################
|
||||
|
||||
We'll now demonstrate how the nodes, NICs, switches, and links connect in a network, including full code examples and
|
||||
syslog extracts to illustrate the step-by-step process.
|
||||
|
||||
To demonstrate successful network communication between nodes and switches, we'll model a standard network with four
|
||||
PC's and two switches.
|
||||
|
||||
|
||||
.. image:: ../../../_static/four_node_two_switch_network.png
|
||||
|
||||
Create Nodes & NICs
|
||||
*******************
|
||||
|
||||
First, we'll create the four nodes, each with a single NIC.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
pc_a = Node(hostname="pc_a")
|
||||
nic_a = NIC(ip_address="192.168.0.10", subnet_mask="255.255.255.0", gateway="192.168.0.1")
|
||||
pc_a.connect_nic(nic_a)
|
||||
pc_a.power_on()
|
||||
|
||||
pc_b = Node(hostname="pc_b")
|
||||
nic_b = NIC(ip_address="192.168.0.11", subnet_mask="255.255.255.0", gateway="192.168.0.1")
|
||||
pc_b.connect_nic(nic_b)
|
||||
pc_b.power_on()
|
||||
|
||||
pc_c = Node(hostname="pc_c")
|
||||
nic_c = NIC(ip_address="192.168.0.12", subnet_mask="255.255.255.0", gateway="192.168.0.1")
|
||||
pc_c.connect_nic(nic_c)
|
||||
pc_c.power_on()
|
||||
|
||||
pc_d = Node(hostname="pc_d")
|
||||
nic_d = NIC(ip_address="192.168.0.13", subnet_mask="255.255.255.0", gateway="192.168.0.1")
|
||||
pc_d.connect_nic(nic_d)
|
||||
pc_d.power_on()
|
||||
|
||||
|
||||
This produces:
|
||||
|
||||
**node_a NIC table**
|
||||
|
||||
+-------------------+--------------+---------------+-----------------+--------------+----------+
|
||||
| MAC Address | IP Address | Subnet Mask | Default Gateway | Speed (Mbps) | Status |
|
||||
+===================+==============+===============+=================+==============+==========+
|
||||
| 80:af:f2:f6:58:b7 | 102.169.0.10 | 255.255.255.0 | 192.168.0.1 | 100 | Disabled |
|
||||
+-------------------+--------------+---------------+-----------------+--------------+----------+
|
||||
|
||||
**node_a sys log**
|
||||
|
||||
.. code-block::
|
||||
|
||||
2023-08-08 15:50:08,355 INFO: Connected NIC 80:af:f2:f6:58:b7/192.168.0.10
|
||||
2023-08-08 15:50:08,355 INFO: Turned on
|
||||
|
||||
**node_b NIC table**
|
||||
|
||||
+-------------------+--------------+---------------+-----------------+--------------+----------+
|
||||
| MAC Address | IP Address | Subnet Mask | Default Gateway | Speed (Mbps) | Status |
|
||||
+===================+==============+===============+=================+==============+==========+
|
||||
| 98:ad:eb:7c:dc:cb | 102.169.0.11 | 255.255.255.0 | 192.168.0.1 | 100 | Disabled |
|
||||
+-------------------+--------------+---------------+-----------------+--------------+----------+
|
||||
|
||||
**node_b sys log**
|
||||
|
||||
.. code-block::
|
||||
|
||||
2023-08-08 15:50:08,357 INFO: Connected NIC 98:ad:eb:7c:dc:cb/192.168.0.11
|
||||
2023-08-08 15:50:08,357 INFO: Turned on
|
||||
|
||||
**node_c NIC table**
|
||||
|
||||
+-------------------+--------------+---------------+-----------------+--------------+----------+
|
||||
| MAC Address | IP Address | Subnet Mask | Default Gateway | Speed (Mbps) | Status |
|
||||
+===================+==============+===============+=================+==============+==========+
|
||||
| bc:72:82:5d:82:a4 | 102.169.0.12 | 255.255.255.0 | 192.168.0.1 | 100 | Disabled |
|
||||
+-------------------+--------------+---------------+-----------------+--------------+----------+
|
||||
|
||||
**node_c sys log**
|
||||
|
||||
.. code-block::
|
||||
|
||||
2023-08-08 15:50:08,358 INFO: Connected NIC bc:72:82:5d:82:a4/192.168.0.12
|
||||
2023-08-08 15:50:08,358 INFO: Turned on
|
||||
|
||||
**node_d NIC table**
|
||||
|
||||
+-------------------+--------------+---------------+-----------------+--------------+----------+
|
||||
| MAC Address | IP Address | Subnet Mask | Default Gateway | Speed (Mbps) | Status |
|
||||
+===================+==============+===============+=================+==============+==========+
|
||||
| 84:20:7c:ec:a5:c6 | 102.169.0.13 | 255.255.255.0 | 192.168.0.1 | 100 | Disabled |
|
||||
+-------------------+--------------+---------------+-----------------+--------------+----------+
|
||||
|
||||
**node_d sys log**
|
||||
|
||||
.. code-block::
|
||||
|
||||
2023-08-08 15:50:08,359 INFO: Connected NIC 84:20:7c:ec:a5:c6/192.168.0.13
|
||||
2023-08-08 15:50:08,360 INFO: Turned on
|
||||
|
||||
|
||||
Create Switches
|
||||
***************
|
||||
|
||||
Next, we'll create two six-port switches:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
switch_1 = Switch(hostname="switch_1", num_ports=6)
|
||||
switch_1.power_on()
|
||||
|
||||
switch_2 = Switch(hostname="switch_2", num_ports=6)
|
||||
switch_2.power_on()
|
||||
|
||||
This produces:
|
||||
|
||||
**switch_1 MAC table**
|
||||
|
||||
+------+-------------------+--------------+----------+
|
||||
| Port | MAC Address | Speed (Mbps) | Status |
|
||||
+======+===================+==============+==========+
|
||||
| 1 | 9d:ac:59:a0:05:13 | 100 | Disabled |
|
||||
+------+-------------------+--------------+----------+
|
||||
| 2 | 45:f5:8e:b6:f5:d3 | 100 | Disabled |
|
||||
+------+-------------------+--------------+----------+
|
||||
| 3 | ef:f5:b9:28:cb:ae | 100 | Disabled |
|
||||
+------+-------------------+--------------+----------+
|
||||
| 4 | 88:76:0a:72:fc:14 | 100 | Disabled |
|
||||
+------+-------------------+--------------+----------+
|
||||
| 5 | 79:de:da:bd:e2:ba | 100 | Disabled |
|
||||
+------+-------------------+--------------+----------+
|
||||
| 6 | 91:d5:83:a0:02:f2 | 100 | Disabled |
|
||||
+------+-------------------+--------------+----------+
|
||||
|
||||
**switch_1 sys log**
|
||||
|
||||
.. code-block::
|
||||
|
||||
2023-08-08 15:50:08,373 INFO: Turned on
|
||||
|
||||
**switch_2 MAC table**
|
||||
|
||||
+------+-------------------+--------------+----------+
|
||||
| Port | MAC Address | Speed (Mbps) | Status |
|
||||
+======+===================+==============+==========+
|
||||
| 1 | aa:58:fa:66:d7:be | 100 | Disabled |
|
||||
+------+-------------------+--------------+----------+
|
||||
| 2 | 72:d2:1e:88:e9:45 | 100 | Disabled |
|
||||
+------+-------------------+--------------+----------+
|
||||
| 3 | 8a:fc:2a:56:d5:c5 | 100 | Disabled |
|
||||
+------+-------------------+--------------+----------+
|
||||
| 4 | fb:b5:9a:04:4a:49 | 100 | Disabled |
|
||||
+------+-------------------+--------------+----------+
|
||||
| 5 | 88:aa:48:d0:21:9e | 100 | Disabled |
|
||||
+------+-------------------+--------------+----------+
|
||||
| 6 | 96:77:39:d1:de:44 | 100 | Disabled |
|
||||
+------+-------------------+--------------+----------+
|
||||
|
||||
**switch_2 sys log**
|
||||
|
||||
.. code-block::
|
||||
|
||||
2023-08-08 15:50:08,374 INFO: Turned on
|
||||
|
||||
Create Links
|
||||
************
|
||||
|
||||
Finally, we'll create the five links that connect the nodes and the switches:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
link_nic_a_switch_1 = Link(endpoint_a=nic_a, endpoint_b=switch_1.switch_ports[1])
|
||||
link_nic_b_switch_1 = Link(endpoint_a=nic_b, endpoint_b=switch_1.switch_ports[2])
|
||||
link_nic_c_switch_2 = Link(endpoint_a=nic_c, endpoint_b=switch_2.switch_ports[1])
|
||||
link_nic_d_switch_2 = Link(endpoint_a=nic_d, endpoint_b=switch_2.switch_ports[2])
|
||||
link_switch_1_switch_2 = Link(
|
||||
endpoint_a=switch_1.switch_ports[6], endpoint_b=switch_2.switch_ports[6]
|
||||
)
|
||||
|
||||
This produces:
|
||||
|
||||
**node_a NIC table**
|
||||
|
||||
+-------------------+--------------+---------------+-----------------+--------------+---------+
|
||||
| MAC Address | IP Address | Subnet Mask | Default Gateway | Speed (Mbps) | Status |
|
||||
+===================+==============+===============+=================+==============+=========+
|
||||
| 80:af:f2:f6:58:b7 | 102.169.0.10 | 255.255.255.0 | 192.168.0.1 | 100 | Enabled |
|
||||
+-------------------+--------------+---------------+-----------------+--------------+---------+
|
||||
|
||||
**node_a sys log**
|
||||
|
||||
.. code-block::
|
||||
|
||||
2023-08-08 15:50:08,355 INFO: Connected NIC 80:af:f2:f6:58:b7/192.168.0.10
|
||||
2023-08-08 15:50:08,355 INFO: Turned on
|
||||
2023-08-08 15:50:08,355 INFO: NIC 80:af:f2:f6:58:b7/192.168.0.10 enabled
|
||||
|
||||
**node_b NIC table**
|
||||
|
||||
+-------------------+--------------+---------------+-----------------+--------------+---------+
|
||||
| MAC Address | IP Address | Subnet Mask | Default Gateway | Speed (Mbps) | Status |
|
||||
+===================+==============+===============+=================+==============+=========+
|
||||
| 98:ad:eb:7c:dc:cb | 102.169.0.11 | 255.255.255.0 | 192.168.0.1 | 100 | Enabled |
|
||||
+-------------------+--------------+---------------+-----------------+--------------+---------+
|
||||
|
||||
**node_b sys log**
|
||||
|
||||
.. code-block::
|
||||
|
||||
2023-08-08 15:50:08,357 INFO: Connected NIC 98:ad:eb:7c:dc:cb/192.168.0.11
|
||||
2023-08-08 15:50:08,357 INFO: Turned on
|
||||
2023-08-08 15:50:08,357 INFO: NIC 98:ad:eb:7c:dc:cb/192.168.0.11 enabled
|
||||
|
||||
**node_c NIC table**
|
||||
|
||||
+-------------------+--------------+---------------+-----------------+--------------+---------+
|
||||
| MAC Address | IP Address | Subnet Mask | Default Gateway | Speed (Mbps) | Status |
|
||||
+===================+==============+===============+=================+==============+=========+
|
||||
| bc:72:82:5d:82:a4 | 102.169.0.12 | 255.255.255.0 | 192.168.0.1 | 100 | Enabled |
|
||||
+-------------------+--------------+---------------+-----------------+--------------+---------+
|
||||
|
||||
**node_c sys log**
|
||||
|
||||
.. code-block::
|
||||
|
||||
2023-08-08 15:50:08,358 INFO: Connected NIC bc:72:82:5d:82:a4/192.168.0.12
|
||||
2023-08-08 15:50:08,358 INFO: Turned on
|
||||
2023-08-08 15:50:08,358 INFO: NIC bc:72:82:5d:82:a4/192.168.0.12 enabled
|
||||
|
||||
**node_d NIC table**
|
||||
|
||||
+-------------------+--------------+---------------+-----------------+--------------+---------+
|
||||
| MAC Address | IP Address | Subnet Mask | Default Gateway | Speed (Mbps) | Status |
|
||||
+===================+==============+===============+=================+==============+=========+
|
||||
| 84:20:7c:ec:a5:c6 | 102.169.0.13 | 255.255.255.0 | 192.168.0.1 | 100 | Enabled |
|
||||
+-------------------+--------------+---------------+-----------------+--------------+---------+
|
||||
|
||||
**node_d sys log**
|
||||
|
||||
.. code-block::
|
||||
|
||||
2023-08-08 15:50:08,359 INFO: Connected NIC 84:20:7c:ec:a5:c6/192.168.0.13
|
||||
2023-08-08 15:50:08,360 INFO: Turned on
|
||||
2023-08-08 15:50:08,360 INFO: NIC 84:20:7c:ec:a5:c6/192.168.0.13 enabled
|
||||
|
||||
**switch_1 MAC table**
|
||||
|
||||
+------+-------------------+--------------+----------+
|
||||
| Port | MAC Address | Speed (Mbps) | Status |
|
||||
+======+===================+==============+==========+
|
||||
| 1 | 9d:ac:59:a0:05:13 | 100 | Enabled |
|
||||
+------+-------------------+--------------+----------+
|
||||
| 2 | 45:f5:8e:b6:f5:d3 | 100 | Enabled |
|
||||
+------+-------------------+--------------+----------+
|
||||
| 3 | ef:f5:b9:28:cb:ae | 100 | Disabled |
|
||||
+------+-------------------+--------------+----------+
|
||||
| 4 | 88:76:0a:72:fc:14 | 100 | Disabled |
|
||||
+------+-------------------+--------------+----------+
|
||||
| 5 | 79:de:da:bd:e2:ba | 100 | Disabled |
|
||||
+------+-------------------+--------------+----------+
|
||||
| 6 | 91:d5:83:a0:02:f2 | 100 | Enabled |
|
||||
+------+-------------------+--------------+----------+
|
||||
|
||||
|
||||
**switch_1 sys log**
|
||||
|
||||
.. code-block::
|
||||
|
||||
2023-08-08 15:50:08,373 INFO: Turned on
|
||||
2023-08-08 15:50:08,378 INFO: SwitchPort 9d:ac:59:a0:05:13 enabled
|
||||
2023-08-08 15:50:08,380 INFO: SwitchPort 45:f5:8e:b6:f5:d3 enabled
|
||||
2023-08-08 15:50:08,384 INFO: SwitchPort 91:d5:83:a0:02:f2 enabled
|
||||
|
||||
|
||||
**switch_2 MAC table**
|
||||
|
||||
+------+-------------------+--------------+----------+
|
||||
| Port | MAC Address | Speed (Mbps) | Status |
|
||||
+======+===================+==============+==========+
|
||||
| 1 | aa:58:fa:66:d7:be | 100 | Enabled |
|
||||
+------+-------------------+--------------+----------+
|
||||
| 2 | 72:d2:1e:88:e9:45 | 100 | Enabled |
|
||||
+------+-------------------+--------------+----------+
|
||||
| 3 | 8a:fc:2a:56:d5:c5 | 100 | Disabled |
|
||||
+------+-------------------+--------------+----------+
|
||||
| 4 | fb:b5:9a:04:4a:49 | 100 | Disabled |
|
||||
+------+-------------------+--------------+----------+
|
||||
| 5 | 88:aa:48:d0:21:9e | 100 | Disabled |
|
||||
+------+-------------------+--------------+----------+
|
||||
| 6 | 96:77:39:d1:de:44 | 100 | Enabled |
|
||||
+------+-------------------+--------------+----------+
|
||||
|
||||
|
||||
**switch_2 sys log**
|
||||
|
||||
.. code-block::
|
||||
|
||||
2023-08-08 15:50:08,374 INFO: Turned on
|
||||
2023-08-08 15:50:08,381 INFO: SwitchPort aa:58:fa:66:d7:be enabled
|
||||
2023-08-08 15:50:08,383 INFO: SwitchPort 72:d2:1e:88:e9:45 enabled
|
||||
2023-08-08 15:50:08,384 INFO: SwitchPort 96:77:39:d1:de:44 enabled
|
||||
|
||||
|
||||
Perform Ping
|
||||
************
|
||||
|
||||
Now with the network setup and operational, we can perform a ping to confirm that communication between nodes over a
|
||||
switched network is possible. In the below example, we ping 192.168.0.13 (node_d) from node_a:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
pc_a.ping("192.168.0.13")
|
||||
|
||||
|
||||
This produces:
|
||||
|
||||
**node_a sys log**
|
||||
|
||||
.. code-block::
|
||||
|
||||
2023-08-08 15:50:08,355 INFO: Connected NIC 80:af:f2:f6:58:b7/192.168.0.10
|
||||
2023-08-08 15:50:08,355 INFO: Turned on
|
||||
2023-08-08 15:50:08,355 INFO: NIC 80:af:f2:f6:58:b7/192.168.0.10 enabled
|
||||
2023-08-08 15:50:08,406 INFO: Attempting to ping 192.168.0.13
|
||||
2023-08-08 15:50:08,406 INFO: No entry in ARP cache for 192.168.0.13
|
||||
2023-08-08 15:50:08,406 INFO: Sending ARP request from NIC 80:af:f2:f6:58:b7/192.168.0.10 for ip 192.168.0.13
|
||||
2023-08-08 15:50:08,413 INFO: Received ARP response for 192.168.0.13 from 84:20:7c:ec:a5:c6 via NIC 80:af:f2:f6:58:b7/192.168.0.10
|
||||
2023-08-08 15:50:08,413 INFO: Adding ARP cache entry for 84:20:7c:ec:a5:c6/192.168.0.13 via NIC 80:af:f2:f6:58:b7/192.168.0.10
|
||||
2023-08-08 15:50:08,415 INFO: Sending echo request to 192.168.0.13
|
||||
2023-08-08 15:50:08,417 INFO: Received echo reply from 192.168.0.13
|
||||
2023-08-08 15:50:08,419 INFO: Sending echo request to 192.168.0.13
|
||||
2023-08-08 15:50:08,421 INFO: Received echo reply from 192.168.0.13
|
||||
2023-08-08 15:50:08,422 INFO: Sending echo request to 192.168.0.13
|
||||
2023-08-08 15:50:08,424 INFO: Received echo reply from 192.168.0.13
|
||||
2023-08-08 15:50:08,425 INFO: Sending echo request to 192.168.0.13
|
||||
2023-08-08 15:50:08,427 INFO: Received echo reply from 192.168.0.13
|
||||
|
||||
|
||||
**node_b sys log**
|
||||
|
||||
.. code-block::
|
||||
|
||||
2023-08-08 15:50:08,357 INFO: Connected NIC 98:ad:eb:7c:dc:cb/192.168.0.11
|
||||
2023-08-08 15:50:08,357 INFO: Turned on
|
||||
2023-08-08 15:50:08,357 INFO: NIC 98:ad:eb:7c:dc:cb/192.168.0.11 enabled
|
||||
2023-08-08 15:50:08,410 INFO: Received ARP request for 192.168.0.13 from 80:af:f2:f6:58:b7/192.168.0.10
|
||||
2023-08-08 15:50:08,410 INFO: Ignoring ARP request for 192.168.0.13
|
||||
|
||||
|
||||
**node_c sys log**
|
||||
|
||||
.. code-block::
|
||||
|
||||
2023-08-08 15:50:08,358 INFO: Connected NIC bc:72:82:5d:82:a4/192.168.0.12
|
||||
2023-08-08 15:50:08,358 INFO: Turned on
|
||||
2023-08-08 15:50:08,358 INFO: NIC bc:72:82:5d:82:a4/192.168.0.12 enabled
|
||||
2023-08-08 15:50:08,411 INFO: Received ARP request for 192.168.0.13 from 80:af:f2:f6:58:b7/192.168.0.10
|
||||
2023-08-08 15:50:08,411 INFO: Ignoring ARP request for 192.168.0.13
|
||||
|
||||
|
||||
**node_d sys log**
|
||||
|
||||
.. code-block::
|
||||
|
||||
2023-08-08 15:50:08,359 INFO: Connected NIC 84:20:7c:ec:a5:c6/192.168.0.13
|
||||
2023-08-08 15:50:08,360 INFO: Turned on
|
||||
2023-08-08 15:50:08,360 INFO: NIC 84:20:7c:ec:a5:c6/192.168.0.13 enabled
|
||||
2023-08-08 15:50:08,412 INFO: Received ARP request for 192.168.0.13 from 80:af:f2:f6:58:b7/192.168.0.10
|
||||
2023-08-08 15:50:08,412 INFO: Adding ARP cache entry for 80:af:f2:f6:58:b7/192.168.0.10 via NIC 84:20:7c:ec:a5:c6/192.168.0.13
|
||||
2023-08-08 15:50:08,412 INFO: Sending ARP reply from 84:20:7c:ec:a5:c6/192.168.0.13 to 192.168.0.10/80:af:f2:f6:58:b7
|
||||
2023-08-08 15:50:08,416 INFO: Received echo request from 192.168.0.10
|
||||
2023-08-08 15:50:08,417 INFO: Sending echo reply to 192.168.0.10
|
||||
2023-08-08 15:50:08,420 INFO: Received echo request from 192.168.0.10
|
||||
2023-08-08 15:50:08,420 INFO: Sending echo reply to 192.168.0.10
|
||||
2023-08-08 15:50:08,423 INFO: Received echo request from 192.168.0.10
|
||||
2023-08-08 15:50:08,423 INFO: Sending echo reply to 192.168.0.10
|
||||
2023-08-08 15:50:08,426 INFO: Received echo request from 192.168.0.10
|
||||
2023-08-08 15:50:08,426 INFO: Sending echo reply to 192.168.0.10
|
||||
|
||||
|
||||
|
||||
**switch_1 sys log**
|
||||
|
||||
.. code-block::
|
||||
|
||||
2023-08-08 15:50:08,373 INFO: Turned on
|
||||
2023-08-08 15:50:08,378 INFO: SwitchPort 9d:ac:59:a0:05:13 enabled
|
||||
2023-08-08 15:50:08,380 INFO: SwitchPort 45:f5:8e:b6:f5:d3 enabled
|
||||
2023-08-08 15:50:08,384 INFO: SwitchPort 91:d5:83:a0:02:f2 enabled
|
||||
2023-08-08 15:50:08,409 INFO: Added MAC table entry: Port 1 -> 80:af:f2:f6:58:b7
|
||||
2023-08-08 15:50:08,413 INFO: Added MAC table entry: Port 6 -> 84:20:7c:ec:a5:c6
|
||||
|
||||
|
||||
|
||||
**switch_2 sys log**
|
||||
|
||||
.. code-block::
|
||||
|
||||
2023-08-08 15:50:08,374 INFO: Turned on
|
||||
2023-08-08 15:50:08,381 INFO: SwitchPort aa:58:fa:66:d7:be enabled
|
||||
2023-08-08 15:50:08,383 INFO: SwitchPort 72:d2:1e:88:e9:45 enabled
|
||||
2023-08-08 15:50:08,384 INFO: SwitchPort 96:77:39:d1:de:44 enabled
|
||||
2023-08-08 15:50:08,411 INFO: Added MAC table entry: Port 6 -> 80:af:f2:f6:58:b7
|
||||
2023-08-08 15:50:08,412 INFO: Added MAC table entry: Port 2 -> 84:20:7c:ec:a5:c6
|
||||
@@ -1,75 +0,0 @@
|
||||
.. only:: comment
|
||||
|
||||
© Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
|
||||
|
||||
Physical Layer
|
||||
==============
|
||||
|
||||
The physical layer components are models of a ``NIC`` (Network Interface Card) and a ``Link``. These components allow
|
||||
modelling of layer 1 (physical layer) in the OSI model.
|
||||
|
||||
NIC
|
||||
###
|
||||
The ``NIC`` class is a realistic model of a Network Interface Card. The ``NIC`` acts as the interface between the
|
||||
``Node`` and the ``Link``.
|
||||
|
||||
NICs have the following attributes:
|
||||
|
||||
- **ip_address:** The IPv4 address assigned to the NIC.
|
||||
- **subnet_mask:** The subnet mask assigned to the NIC.
|
||||
- **gateway:** The default gateway IP address for forwarding network traffic to other networks.
|
||||
- **mac_address:** The MAC address of the NIC. Defaults to a randomly set MAC address.
|
||||
- **speed:** The speed of the NIC in Mbps (default is 100 Mbps).
|
||||
- **mtu:** The Maximum Transmission Unit (MTU) of the NIC in Bytes, representing the largest data packet size it can handle without fragmentation (default is 1500 B).
|
||||
- **wake_on_lan:** Indicates if the NIC supports Wake-on-LAN functionality.
|
||||
- **dns_servers:** List of IP addresses of DNS servers used for name resolution.
|
||||
- **connected_link:** The link to which the NIC is connected.
|
||||
- **enabled:** Indicates whether the NIC is enabled.
|
||||
|
||||
**Basic Example**
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
nic1 = NIC(
|
||||
ip_address="192.168.1.100",
|
||||
subnet_mask="255.255.255.0",
|
||||
gateway="192.168.1.1"
|
||||
)
|
||||
|
||||
Link
|
||||
####
|
||||
|
||||
The ``Link`` class represents a physical link between two network endpoints.
|
||||
|
||||
Links have the following attributes:
|
||||
|
||||
- **endpoint_a:** The first NIC connected to the Link.
|
||||
- **endpoint_b:** The second NIC connected to the Link.
|
||||
- **bandwidth:** The bandwidth of the Link in Mbps (default is 100 Mbps).
|
||||
- **current_load:** The current load on the link in Mbps.
|
||||
|
||||
**Basic Example**
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
nic1 = NIC(
|
||||
ip_address="192.168.1.100",
|
||||
subnet_mask="255.255.255.0",
|
||||
gateway="192.168.1.1"
|
||||
)
|
||||
nic1 = NIC(
|
||||
ip_address="192.168.1.101",
|
||||
subnet_mask="255.255.255.0",
|
||||
gateway="192.168.1.1"
|
||||
)
|
||||
|
||||
link = Link(
|
||||
endpoint_a=nic1,
|
||||
endpoint_b=nic2,
|
||||
bandwidth=1000
|
||||
)
|
||||
|
||||
Link, NIC, Node Interface
|
||||
#########################
|
||||
|
||||
.. image:: ../../../_static/node_nic_link_component_diagram.png
|
||||
@@ -34,7 +34,7 @@ specify the priority of IP packets for Quality of Service handling.
|
||||
**ICMPType:** Enumeration of common ICMP (Internet Control Message Protocol) types. It defines various types of ICMP
|
||||
messages used for network troubleshooting and error reporting.
|
||||
|
||||
**ICMPHeader:** Models an ICMP header and includes ICMP type, code, identifier, and sequence number. It is used to
|
||||
**ICMPPacket:** Models an ICMP header and includes ICMP type, code, identifier, and sequence number. It is used to
|
||||
create ICMP packets for network control and error reporting.
|
||||
|
||||
**IPPacket:** Represents the IP layer of a network frame. It includes source and destination IP addresses, protocol
|
||||
@@ -55,11 +55,24 @@ PrimAITE-specific metadata required for reinforcement learning (RL) purposes.
|
||||
Data Link Layer (Layer 2)
|
||||
#########################
|
||||
|
||||
**ARPEntry:** Represents an entry in the ARP cache. It consists of the following fields:
|
||||
|
||||
- **mac_address:** The MAC address associated with the IP address.
|
||||
- **nic_uuid:** The NIC (Network Interface Card) UUID through which the NIC with the IP address is reachable.
|
||||
|
||||
**ARPPacket:** Represents the ARP layer of a network frame, and it includes the following fields:
|
||||
|
||||
- **request:** ARP operation. Set to True for a request and False for a reply.
|
||||
- **sender_mac_addr:** Sender's MAC address.
|
||||
- **sender_ip:** Sender's IP address (IPv4 format).
|
||||
- **target_mac_addr:** Target's MAC address.
|
||||
- **target_ip:** Target's IP address (IPv4 format).
|
||||
|
||||
**EthernetHeader:** Represents the Ethernet layer of a network frame. It includes source and destination MAC addresses.
|
||||
This header is used to identify the physical hardware addresses of devices on a local network.
|
||||
|
||||
**Frame:** Represents a complete network frame with all layers. It includes an ``EthernetHeader``, an ``IPPacket``, an
|
||||
optional ``TCPHeader``, ``UDPHeader``, or ``ICMPHeader``, a ``PrimaiteHeader`` and an optional payload. This class
|
||||
optional ``TCPHeader``, ``UDPHeader``, or ``ICMPPacket``, a ``PrimaiteHeader`` and an optional payload. This class
|
||||
combines all the headers and data to create a complete network frame that can be sent over the network and used in the
|
||||
PrimAITE simulation.
|
||||
|
||||
|
||||
@@ -7,7 +7,67 @@ Simulation Structure
|
||||
====================
|
||||
|
||||
The simulation is made up of many smaller components which are related to each other in a tree-like structure. At the
|
||||
top level, there is an object called the ``SimulationController`` _(doesn't exist yet)_, which has a physical network
|
||||
and a software controller for managing software and users.
|
||||
top level, there is the :py:meth:`primaite.simulator.sim_container.Simulation`, which keeps track of the physical network
|
||||
and a domain controller for managing software and users.
|
||||
|
||||
Each node of the simulation 'tree' has responsibility for creating, deleting, and updating its direct descendants.
|
||||
Each node of the simulation 'tree' has responsibility for creating, deleting, and updating its direct descendants. Also,
|
||||
when a component's ``describe_state()`` method is called, it will include the state of its descendants. The
|
||||
``apply_action()`` method can be used to act on a component or one of its descendatnts. The diagram below shows the
|
||||
relationship between components.
|
||||
|
||||
.. image:: _static/component_relationship.png
|
||||
:width: 500
|
||||
:alt: The top level simulation object owns a NetworkContainer and a DomainController. The DomainController has a
|
||||
list of accounts. The network container has links and nodes. Nodes can own switchports, NICs, FileSystem,
|
||||
Application, Service, and Process.
|
||||
|
||||
|
||||
Actions
|
||||
=======
|
||||
Agents can interact with the simulation by using actions. Actions are standardised with the
|
||||
:py:class:`primaite.simulation.core.Action` class, which just holds a reference to two special functions.
|
||||
|
||||
1. The action function itself, it must accept a `request` parameters which is a list of strings that describe what the
|
||||
action should do. It must also accept a `context` dict which can house additional information surrounding the action.
|
||||
For example, the context will typically include information about which entity intiated the action.
|
||||
2. A validator function. This function should return a boolean value that decides if the request is permitted or not.
|
||||
It uses the same paramters as the action function.
|
||||
|
||||
Action Permissions
|
||||
------------------
|
||||
When an agent tries to perform an action on a simulation component, that action will only be executed if the request is
|
||||
validated. For example, some actions can require that an agent is logged into an admin account. Each action defines its
|
||||
own permissions using an instance of :py:class:`primaite.simulation.core.ActionPermissionValidator`. The below code
|
||||
snippet demonstrates usage of the ``ActionPermissionValidator``.
|
||||
|
||||
.. code:: python
|
||||
|
||||
from primaite.simulator.core import Action, ActionManager, SimComponent
|
||||
from primaite.simulator.domain.controller import AccountGroup, GroupMembershipValidator
|
||||
|
||||
class Smartphone(SimComponent):
|
||||
name: str
|
||||
apps = []
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
self.action_manager = ActionManager()
|
||||
|
||||
self.action_manager.add_action(
|
||||
"reset_factory_settings",
|
||||
Action(
|
||||
func = lambda request, context: self.reset_factory_settings(),
|
||||
validator = GroupMembershipValidator([AccountGroup.DOMAIN_ADMIN]),
|
||||
),
|
||||
)
|
||||
|
||||
def reset_factory_settings(self):
|
||||
self.apps = []
|
||||
|
||||
phone = Smartphone(name="phone1")
|
||||
|
||||
# try to wipe the phone as a domain user, this will have no effect
|
||||
phone.apply_action(["reset_factory_settings"], context={"request_source":{"groups":["DOMAIN_USER"]})
|
||||
|
||||
# try to wipe the phone as an admin user, this will wipe the phone
|
||||
phone.apply_action(["reset_factory_settings"], context={"request_source":{"groups":["DOMAIN_ADMIN"]})
|
||||
|
||||
@@ -33,11 +33,12 @@ dependencies = [
|
||||
"platformdirs==3.5.1",
|
||||
"plotly==5.15.0",
|
||||
"polars==0.18.4",
|
||||
"prettytable==3.8.0",
|
||||
"PyYAML==6.0",
|
||||
"stable-baselines3==1.6.2",
|
||||
"tensorflow==2.12.0",
|
||||
"typer[all]==0.9.0",
|
||||
"pydantic"
|
||||
"pydantic==2.1.1"
|
||||
]
|
||||
|
||||
[tool.setuptools.dynamic]
|
||||
|
||||
@@ -16,6 +16,9 @@ from platformdirs import PlatformDirs
|
||||
with open(Path(__file__).parent.resolve() / "VERSION", "r") as file:
|
||||
__version__ = file.readline().strip()
|
||||
|
||||
_PRIMAITE_ROOT: Path = Path(__file__).parent
|
||||
# TODO: Remove once we integrate the simulation into PrimAITE and it uses the primaite session path
|
||||
|
||||
|
||||
class _PrimaitePaths:
|
||||
"""
|
||||
|
||||
478
src/primaite/notebooks/create-simulation.ipynb
Normal file
478
src/primaite/notebooks/create-simulation.ipynb
Normal file
File diff suppressed because one or more lines are too long
@@ -0,0 +1,5 @@
|
||||
from primaite import _PRIMAITE_ROOT
|
||||
|
||||
TEMP_SIM_OUTPUT = _PRIMAITE_ROOT.parent.parent / "simulation_output"
|
||||
"A path at the repo root dir to use temporarily for sim output testing while in dev."
|
||||
# TODO: Remove once we integrate the simulation into PrimAITE and it uses the primaite session path
|
||||
|
||||
@@ -1,12 +1,143 @@
|
||||
"""Core of the PrimAITE Simulator."""
|
||||
from abc import abstractmethod
|
||||
from typing import Callable, Dict, List
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Callable, Dict, List, Optional, Union
|
||||
from uuid import uuid4
|
||||
|
||||
from pydantic import BaseModel
|
||||
from pydantic import BaseModel, ConfigDict, Extra
|
||||
|
||||
from primaite import getLogger
|
||||
|
||||
_LOGGER = getLogger(__name__)
|
||||
|
||||
|
||||
class ActionPermissionValidator(ABC):
|
||||
"""
|
||||
Base class for action validators.
|
||||
|
||||
The permissions manager is designed to be generic. So, although in the first instance the permissions
|
||||
are evaluated purely on membership to AccountGroup, this class can support validating permissions based on any
|
||||
arbitrary criteria.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def __call__(self, request: List[str], context: Dict) -> bool:
|
||||
"""Use the request and context paramters to decide whether the action should be permitted."""
|
||||
pass
|
||||
|
||||
|
||||
class AllowAllValidator(ActionPermissionValidator):
|
||||
"""Always allows the action."""
|
||||
|
||||
def __call__(self, request: List[str], context: Dict) -> bool:
|
||||
"""Always allow the action."""
|
||||
return True
|
||||
|
||||
|
||||
class Action:
|
||||
"""
|
||||
This object stores data related to a single action.
|
||||
|
||||
This includes the callable that can execute the action request, and the validator that will decide whether
|
||||
the action can be performed or not.
|
||||
"""
|
||||
|
||||
def __init__(self, func: Callable[[List[str], Dict], None], validator: ActionPermissionValidator) -> None:
|
||||
"""
|
||||
Save the functions that are for this action.
|
||||
|
||||
Here's a description for the intended use of both of these.
|
||||
|
||||
``func`` is a function that accepts a request and a context dict. Typically this would be a lambda function
|
||||
that invokes a class method of your SimComponent. For example if the component is a node and the action is for
|
||||
turning it off, then the SimComponent should have a turn_off(self) method that does not need to accept any args.
|
||||
Then, this Action will be given something like ``func = lambda request, context: self.turn_off()``.
|
||||
|
||||
``validator`` is an instance of a subclass of `ActionPermissionValidator`. This is essentially a callable that
|
||||
accepts `request` and `context` and returns a boolean to represent whether the permission is granted to perform
|
||||
the action.
|
||||
|
||||
:param func: Function that performs the request.
|
||||
:type func: Callable[[List[str], Dict], None]
|
||||
:param validator: Function that checks if the request is authenticated given the context.
|
||||
:type validator: ActionPermissionValidator
|
||||
"""
|
||||
self.func: Callable[[List[str], Dict], None] = func
|
||||
self.validator: ActionPermissionValidator = validator
|
||||
|
||||
|
||||
class ActionManager:
|
||||
"""
|
||||
ActionManager is used by `SimComponent` instances to keep track of actions.
|
||||
|
||||
Its main purpose is to be a lookup from action name to action function and corresponding validation function. This
|
||||
class is responsible for providing a consistent API for processing actions as well as helpful error messages.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
"""Initialise ActionManager with an empty action lookup."""
|
||||
self.actions: Dict[str, Action] = {}
|
||||
|
||||
def process_request(self, request: List[str], context: Dict) -> None:
|
||||
"""Process an action request.
|
||||
|
||||
:param request: A list of strings which specify what action to take. The first string must be one of the allowed
|
||||
actions, i.e. it must be a key of self.actions. The subsequent strings in the list are passed as parameters
|
||||
to the action function.
|
||||
:type request: List[str]
|
||||
:param context: Dictionary of additional information necessary to process or validate the request.
|
||||
:type context: Dict
|
||||
:raises RuntimeError: If the request parameter does not have a valid action identifier as the first item.
|
||||
"""
|
||||
action_key = request[0]
|
||||
|
||||
if action_key not in self.actions:
|
||||
msg = (
|
||||
f"Action request {request} could not be processed because {action_key} is not a valid action",
|
||||
"within this ActionManager",
|
||||
)
|
||||
_LOGGER.error(msg)
|
||||
raise RuntimeError(msg)
|
||||
|
||||
action = self.actions[action_key]
|
||||
action_options = request[1:]
|
||||
|
||||
if not action.validator(action_options, context):
|
||||
_LOGGER.debug(f"Action request {request} was denied due to insufficient permissions")
|
||||
return
|
||||
|
||||
action.func(action_options, context)
|
||||
|
||||
def add_action(self, name: str, action: Action) -> None:
|
||||
"""Add an action to this action manager.
|
||||
|
||||
:param name: The string associated to this action.
|
||||
:type name: str
|
||||
:param action: Action object.
|
||||
:type action: Action
|
||||
"""
|
||||
if name in self.actions:
|
||||
msg = f"Attempted to register an action but the action name {name} is already taken."
|
||||
_LOGGER.error(msg)
|
||||
raise RuntimeError(msg)
|
||||
|
||||
self.actions[name] = action
|
||||
|
||||
|
||||
class SimComponent(BaseModel):
|
||||
"""Extension of pydantic BaseModel with additional methods that must be defined by all classes in the simulator."""
|
||||
"""Extension of pydantic BaseModel with additional methods that must be defined by all classes in the simulator."""
|
||||
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True, extra=Extra.allow)
|
||||
"""Configure pydantic to allow arbitrary types and to let the instance have attributes not present in model."""
|
||||
|
||||
uuid: str
|
||||
"""The component UUID."""
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
if not kwargs.get("uuid"):
|
||||
kwargs["uuid"] = str(uuid4())
|
||||
super().__init__(**kwargs)
|
||||
self.action_manager: Optional[ActionManager] = None
|
||||
self._parent: Optional["SimComponent"] = None
|
||||
|
||||
@abstractmethod
|
||||
def describe_state(self) -> Dict:
|
||||
@@ -17,9 +148,12 @@ class SimComponent(BaseModel):
|
||||
object. If there are objects referenced by this object that are owned by something else, it is not included in
|
||||
this output.
|
||||
"""
|
||||
return {}
|
||||
state = {
|
||||
"uuid": self.uuid,
|
||||
}
|
||||
return state
|
||||
|
||||
def apply_action(self, action: List[str]) -> None:
|
||||
def apply_action(self, action: List[str], context: Dict = {}) -> None:
|
||||
"""
|
||||
Apply an action to a simulation component. Action data is passed in as a 'namespaced' list of strings.
|
||||
|
||||
@@ -34,16 +168,9 @@ class SimComponent(BaseModel):
|
||||
:param action: List describing the action to apply to this object.
|
||||
:type action: List[str]
|
||||
"""
|
||||
possible_actions = self._possible_actions()
|
||||
if action[0] in possible_actions:
|
||||
# take the first element off the action list and pass the remaining arguments to the corresponding action
|
||||
# function
|
||||
possible_actions[action.pop(0)](action)
|
||||
else:
|
||||
raise ValueError(f"{self.__class__.__name__} received invalid action {action}")
|
||||
|
||||
def _possible_actions(self) -> Dict[str, Callable[[List[str]], None]]:
|
||||
return {}
|
||||
if self.action_manager is None:
|
||||
return
|
||||
self.action_manager.process_request(action, context)
|
||||
|
||||
def apply_timestep(self, timestep: int) -> None:
|
||||
"""
|
||||
@@ -54,10 +181,31 @@ class SimComponent(BaseModel):
|
||||
"""
|
||||
pass
|
||||
|
||||
def reset_component_for_episode(self):
|
||||
def reset_component_for_episode(self, episode: int):
|
||||
"""
|
||||
Reset this component to its original state for a new episode.
|
||||
|
||||
Override this method with anything that needs to happen within the component for it to be reset.
|
||||
"""
|
||||
pass
|
||||
|
||||
@property
|
||||
def parent(self) -> "SimComponent":
|
||||
"""Reference to the parent object which manages this object.
|
||||
|
||||
:return: Parent object.
|
||||
:rtype: SimComponent
|
||||
"""
|
||||
return self._parent
|
||||
|
||||
@parent.setter
|
||||
def parent(self, new_parent: Union["SimComponent", None]) -> None:
|
||||
if self._parent and new_parent:
|
||||
msg = f"Overwriting parent of {self.uuid}. Old parent: {self._parent.uuid}, New parent: {new_parent.uuid}"
|
||||
_LOGGER.warn(msg)
|
||||
raise RuntimeWarning(msg)
|
||||
self._parent = new_parent
|
||||
|
||||
@parent.deleter
|
||||
def parent(self) -> None:
|
||||
self._parent = None
|
||||
|
||||
0
src/primaite/simulator/domain/__init__.py
Normal file
0
src/primaite/simulator/domain/__init__.py
Normal file
82
src/primaite/simulator/domain/account.py
Normal file
82
src/primaite/simulator/domain/account.py
Normal file
@@ -0,0 +1,82 @@
|
||||
"""User account simulation."""
|
||||
from enum import Enum
|
||||
from typing import Dict
|
||||
|
||||
from primaite import getLogger
|
||||
from primaite.simulator.core import SimComponent
|
||||
|
||||
_LOGGER = getLogger(__name__)
|
||||
|
||||
|
||||
class AccountType(Enum):
|
||||
"""Whether the account is intended for a user to log in or for a service to use."""
|
||||
|
||||
SERVICE = 1
|
||||
"Service accounts are used to grant permissions to software on nodes to perform actions"
|
||||
USER = 2
|
||||
"User accounts are used to allow agents to log in and perform actions"
|
||||
|
||||
|
||||
class PasswordPolicyLevel(Enum):
|
||||
"""Complexity requirements for account passwords."""
|
||||
|
||||
LOW = 1
|
||||
MEDIUM = 2
|
||||
HIGH = 3
|
||||
|
||||
|
||||
class Account(SimComponent):
|
||||
"""User accounts."""
|
||||
|
||||
num_logons: int = 0
|
||||
"The number of times this account was logged into since last reset."
|
||||
num_logoffs: int = 0
|
||||
"The number of times this account was logged out of since last reset."
|
||||
num_group_changes: int = 0
|
||||
"The number of times this account was moved in or out of an AccountGroup."
|
||||
username: str
|
||||
"Account username."
|
||||
password: str
|
||||
"Account password."
|
||||
account_type: AccountType
|
||||
"Account Type, currently this can be service account (used by apps) or user account."
|
||||
enabled: bool = True
|
||||
|
||||
def describe_state(self) -> Dict:
|
||||
"""
|
||||
Produce a dictionary describing the current state of this object.
|
||||
|
||||
Please see :py:meth:`primaite.simulator.core.SimComponent.describe_state` for a more detailed explanation.
|
||||
|
||||
:return: Current state of this object and child objects.
|
||||
:rtype: Dict
|
||||
"""
|
||||
state = super().describe_state()
|
||||
state.update(
|
||||
{
|
||||
"num_logons": self.num_logons,
|
||||
"num_logoffs": self.num_logoffs,
|
||||
"num_group_changes": self.num_group_changes,
|
||||
"username": self.username,
|
||||
"password": self.password,
|
||||
"account_type": self.account_type.name,
|
||||
"enabled": self.enabled,
|
||||
}
|
||||
)
|
||||
return state
|
||||
|
||||
def enable(self):
|
||||
"""Set the status to enabled."""
|
||||
self.enabled = True
|
||||
|
||||
def disable(self):
|
||||
"""Set the status to disabled."""
|
||||
self.enabled = False
|
||||
|
||||
def log_on(self) -> None:
|
||||
"""TODO. Once the accounts are integrated with nodes, populate this accordingly."""
|
||||
self.num_logons += 1
|
||||
|
||||
def log_off(self) -> None:
|
||||
"""TODO. Once the accounts are integrated with nodes, populate this accordingly."""
|
||||
self.num_logoffs += 1
|
||||
154
src/primaite/simulator/domain/controller.py
Normal file
154
src/primaite/simulator/domain/controller.py
Normal file
@@ -0,0 +1,154 @@
|
||||
from enum import Enum
|
||||
from typing import Dict, Final, List, Literal, Tuple
|
||||
|
||||
from primaite.simulator.core import Action, ActionManager, ActionPermissionValidator, SimComponent
|
||||
from primaite.simulator.domain.account import Account, AccountType
|
||||
|
||||
|
||||
# placeholder while these objects don't yet exist
|
||||
class temp_node:
|
||||
"""Placeholder for node class for type hinting purposes."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class temp_application:
|
||||
"""Placeholder for application class for type hinting purposes."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class temp_folder:
|
||||
"""Placeholder for folder class for type hinting purposes."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class temp_file:
|
||||
"""Placeholder for file class for type hinting purposes."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class AccountGroup(Enum):
|
||||
"""Permissions are set at group-level and accounts can belong to these groups."""
|
||||
|
||||
LOCAL_USER = 1
|
||||
"For performing basic actions on a node"
|
||||
DOMAIN_USER = 2
|
||||
"For performing basic actions to the domain"
|
||||
LOCAL_ADMIN = 3
|
||||
"For full access to actions on a node"
|
||||
DOMAIN_ADMIN = 4
|
||||
"For full access"
|
||||
|
||||
|
||||
class GroupMembershipValidator(ActionPermissionValidator):
|
||||
"""Permit actions based on group membership."""
|
||||
|
||||
def __init__(self, allowed_groups: List[AccountGroup]) -> None:
|
||||
"""Store a list of groups that should be granted permission.
|
||||
|
||||
:param allowed_groups: List of AccountGroups that are permitted to perform some action.
|
||||
:type allowed_groups: List[AccountGroup]
|
||||
"""
|
||||
self.allowed_groups = allowed_groups
|
||||
|
||||
def __call__(self, request: List[str], context: Dict) -> bool:
|
||||
"""Permit the action if the request comes from an account which belongs to the right group."""
|
||||
# if context request source is part of any groups mentioned in self.allow_groups, return true, otherwise false
|
||||
requestor_groups: List[str] = context["request_source"]["groups"]
|
||||
for allowed_group in self.allowed_groups:
|
||||
if allowed_group.name in requestor_groups:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
class DomainController(SimComponent):
|
||||
"""Main object for controlling the domain."""
|
||||
|
||||
# owned objects
|
||||
accounts: Dict[str, Account] = {}
|
||||
groups: Final[List[AccountGroup]] = list(AccountGroup)
|
||||
|
||||
domain_group_membership: Dict[Literal[AccountGroup.DOMAIN_ADMIN, AccountGroup.DOMAIN_USER], List[Account]] = {}
|
||||
local_group_membership: Dict[
|
||||
Tuple[temp_node, Literal[AccountGroup.LOCAL_ADMIN, AccountGroup.LOCAL_USER]], List[Account]
|
||||
] = {}
|
||||
|
||||
# references to non-owned objects. Not sure if all are needed here.
|
||||
nodes: Dict[str, temp_node] = {}
|
||||
applications: Dict[str, temp_application] = {}
|
||||
folders: List[temp_folder] = {}
|
||||
files: List[temp_file] = {}
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
|
||||
self.action_manager = ActionManager()
|
||||
# Action 'account' matches requests like:
|
||||
# ['account', '<account-uuid>', *account_action]
|
||||
self.action_manager.add_action(
|
||||
"account",
|
||||
Action(
|
||||
func=lambda request, context: self.accounts[request.pop(0)].apply_action(request, context),
|
||||
validator=GroupMembershipValidator([AccountGroup.DOMAIN_ADMIN]),
|
||||
),
|
||||
)
|
||||
|
||||
def describe_state(self) -> Dict:
|
||||
"""
|
||||
Produce a dictionary describing the current state of this object.
|
||||
|
||||
Please see :py:meth:`primaite.simulator.core.SimComponent.describe_state` for a more detailed explanation.
|
||||
|
||||
:return: Current state of this object and child objects.
|
||||
:rtype: Dict
|
||||
"""
|
||||
state = super().describe_state()
|
||||
state.update({"accounts": {uuid: acct.describe_state() for uuid, acct in self.accounts.items()}})
|
||||
return state
|
||||
|
||||
def _register_account(self, account: Account) -> None:
|
||||
"""TODO."""
|
||||
...
|
||||
|
||||
def _deregister_account(self, account: Account) -> None:
|
||||
"""TODO."""
|
||||
...
|
||||
|
||||
def create_account(self, username: str, password: str, account_type: AccountType) -> Account:
|
||||
"""TODO."""
|
||||
...
|
||||
|
||||
def delete_account(self, account: Account) -> None:
|
||||
"""TODO."""
|
||||
...
|
||||
|
||||
def rotate_all_credentials(self) -> None:
|
||||
"""TODO."""
|
||||
...
|
||||
|
||||
def rotate_account_credentials(self, account: Account) -> None:
|
||||
"""TODO."""
|
||||
...
|
||||
|
||||
def add_account_to_group(self, account: Account, group: AccountGroup) -> None:
|
||||
"""TODO."""
|
||||
...
|
||||
|
||||
def remove_account_from_group(self, account: Account, group: AccountGroup) -> None:
|
||||
"""TODO."""
|
||||
...
|
||||
|
||||
def check_account_permissions(self, account: Account, node: temp_node) -> List[AccountGroup]:
|
||||
"""Return a list of permission groups that this account has on this node."""
|
||||
...
|
||||
|
||||
def register_node(self, node: temp_node) -> None:
|
||||
"""TODO."""
|
||||
...
|
||||
|
||||
def deregister_node(self, node: temp_node) -> None:
|
||||
"""TODO."""
|
||||
...
|
||||
0
src/primaite/simulator/file_system/__init__.py
Normal file
0
src/primaite/simulator/file_system/__init__.py
Normal file
242
src/primaite/simulator/file_system/file_system.py
Normal file
242
src/primaite/simulator/file_system/file_system.py
Normal file
@@ -0,0 +1,242 @@
|
||||
from random import choice
|
||||
from typing import Dict, Optional
|
||||
|
||||
from primaite import getLogger
|
||||
from primaite.simulator.core import SimComponent
|
||||
from primaite.simulator.file_system.file_system_file import FileSystemFile
|
||||
from primaite.simulator.file_system.file_system_file_type import FileSystemFileType
|
||||
from primaite.simulator.file_system.file_system_folder import FileSystemFolder
|
||||
|
||||
_LOGGER = getLogger(__name__)
|
||||
|
||||
|
||||
class FileSystem(SimComponent):
|
||||
"""Class that contains all the simulation File System."""
|
||||
|
||||
folders: Dict[str, FileSystemFolder] = {}
|
||||
"""List containing all the folders in the file system."""
|
||||
|
||||
def describe_state(self) -> Dict:
|
||||
"""
|
||||
Produce a dictionary describing the current state of this object.
|
||||
|
||||
Please see :py:meth:`primaite.simulator.core.SimComponent.describe_state` for a more detailed explanation.
|
||||
|
||||
:return: Current state of this object and child objects.
|
||||
:rtype: Dict
|
||||
"""
|
||||
state = super().describe_state()
|
||||
state.update({"folders": {uuid: folder.describe_state() for uuid, folder in self.folders.items()}})
|
||||
return state
|
||||
|
||||
def get_folders(self) -> Dict:
|
||||
"""Returns the list of folders."""
|
||||
return self.folders
|
||||
|
||||
def create_file(
|
||||
self,
|
||||
file_name: str,
|
||||
size: Optional[float] = None,
|
||||
file_type: Optional[FileSystemFileType] = None,
|
||||
folder: Optional[FileSystemFolder] = None,
|
||||
folder_uuid: Optional[str] = None,
|
||||
) -> FileSystemFile:
|
||||
"""
|
||||
Creates a FileSystemFile and adds it to the list of files.
|
||||
|
||||
If no size or file_type are provided, one will be chosen randomly.
|
||||
If no folder_uuid or folder is provided, a new folder will be created.
|
||||
|
||||
:param: file_name: The file name
|
||||
:type: file_name: str
|
||||
|
||||
:param: size: The size the file takes on disk.
|
||||
:type: size: Optional[float]
|
||||
|
||||
:param: file_type: The type of the file
|
||||
:type: Optional[FileSystemFileType]
|
||||
|
||||
:param: folder: The folder to add the file to
|
||||
:type: folder: Optional[FileSystemFolder]
|
||||
|
||||
:param: folder_uuid: The uuid of the folder to add the file to
|
||||
:type: folder_uuid: Optional[str]
|
||||
"""
|
||||
file = None
|
||||
folder = None
|
||||
|
||||
if file_type is None:
|
||||
file_type = self.get_random_file_type()
|
||||
|
||||
# if no folder uuid provided, create a folder and add file to it
|
||||
if folder_uuid is not None:
|
||||
# otherwise check for existence and add file
|
||||
folder = self.get_folder_by_id(folder_uuid)
|
||||
|
||||
if folder is not None:
|
||||
# check if file with name already exists
|
||||
if folder.get_file_by_name(file_name):
|
||||
raise Exception(f'File with name "{file_name}" already exists.')
|
||||
|
||||
file = FileSystemFile(name=file_name, size=size, file_type=file_type)
|
||||
folder.add_file(file=file)
|
||||
else:
|
||||
# check if a "root" folder exists
|
||||
folder = self.get_folder_by_name("root")
|
||||
if folder is None:
|
||||
# create a root folder
|
||||
folder = FileSystemFolder(name="root")
|
||||
|
||||
# add file to root folder
|
||||
file = FileSystemFile(name=file_name, size=size, file_type=file_type)
|
||||
folder.add_file(file)
|
||||
self.folders[folder.uuid] = folder
|
||||
return file
|
||||
|
||||
def create_folder(
|
||||
self,
|
||||
folder_name: str,
|
||||
) -> FileSystemFolder:
|
||||
"""
|
||||
Creates a FileSystemFolder and adds it to the list of folders.
|
||||
|
||||
:param: folder_name: The name of the folder
|
||||
:type: folder_name: str
|
||||
"""
|
||||
# check if folder with name already exists
|
||||
if self.get_folder_by_name(folder_name):
|
||||
raise Exception(f'Folder with name "{folder_name}" already exists.')
|
||||
|
||||
folder = FileSystemFolder(name=folder_name)
|
||||
|
||||
self.folders[folder.uuid] = folder
|
||||
return folder
|
||||
|
||||
def delete_file(self, file: Optional[FileSystemFile] = None):
|
||||
"""
|
||||
Deletes a file and removes it from the files list.
|
||||
|
||||
:param file: The file to delete
|
||||
:type file: Optional[FileSystemFile]
|
||||
"""
|
||||
# iterate through folders to delete the item with the matching uuid
|
||||
for key in self.folders:
|
||||
self.get_folder_by_id(key).remove_file(file)
|
||||
|
||||
def delete_folder(self, folder: FileSystemFolder):
|
||||
"""
|
||||
Deletes a folder, removes it from the folders list and removes any child folders and files.
|
||||
|
||||
:param folder: The folder to remove
|
||||
:type folder: FileSystemFolder
|
||||
"""
|
||||
if folder is None or not isinstance(folder, FileSystemFolder):
|
||||
raise Exception(f"Invalid folder: {folder}")
|
||||
|
||||
if self.folders.get(folder.uuid):
|
||||
del self.folders[folder.uuid]
|
||||
else:
|
||||
_LOGGER.debug(f"File with UUID {folder.uuid} was not found.")
|
||||
|
||||
def move_file(self, file: FileSystemFile, src_folder: FileSystemFolder, target_folder: FileSystemFolder):
|
||||
"""
|
||||
Moves a file from one folder to another.
|
||||
|
||||
can provide
|
||||
|
||||
:param: file: The file to move
|
||||
:type: file: FileSystemFile
|
||||
|
||||
:param: src_folder: The folder where the file is located
|
||||
:type: FileSystemFolder
|
||||
|
||||
:param: target_folder: The folder where the file should be moved to
|
||||
:type: FileSystemFolder
|
||||
"""
|
||||
# check that the folders exist
|
||||
if src_folder is None:
|
||||
raise Exception("Source folder not provided")
|
||||
|
||||
if target_folder is None:
|
||||
raise Exception("Target folder not provided")
|
||||
|
||||
if file is None:
|
||||
raise Exception("File to be moved is None")
|
||||
|
||||
# check if file with name already exists
|
||||
if target_folder.get_file_by_name(file.name):
|
||||
raise Exception(f'Folder with name "{file.name}" already exists.')
|
||||
|
||||
# remove file from src
|
||||
src_folder.remove_file(file)
|
||||
|
||||
# add file to target
|
||||
target_folder.add_file(file)
|
||||
|
||||
def copy_file(self, file: FileSystemFile, src_folder: FileSystemFolder, target_folder: FileSystemFolder):
|
||||
"""
|
||||
Copies a file from one folder to another.
|
||||
|
||||
can provide
|
||||
|
||||
:param: file: The file to move
|
||||
:type: file: FileSystemFile
|
||||
|
||||
:param: src_folder: The folder where the file is located
|
||||
:type: FileSystemFolder
|
||||
|
||||
:param: target_folder: The folder where the file should be moved to
|
||||
:type: FileSystemFolder
|
||||
"""
|
||||
if src_folder is None:
|
||||
raise Exception("Source folder not provided")
|
||||
|
||||
if target_folder is None:
|
||||
raise Exception("Target folder not provided")
|
||||
|
||||
if file is None:
|
||||
raise Exception("File to be moved is None")
|
||||
|
||||
# check if file with name already exists
|
||||
if target_folder.get_file_by_name(file.name):
|
||||
raise Exception(f'Folder with name "{file.name}" already exists.')
|
||||
|
||||
# add file to target
|
||||
target_folder.add_file(file)
|
||||
|
||||
def get_file_by_id(self, file_id: str) -> FileSystemFile:
|
||||
"""Checks if the file exists in any file system folders."""
|
||||
for key in self.folders:
|
||||
file = self.folders[key].get_file_by_id(file_id=file_id)
|
||||
if file is not None:
|
||||
return file
|
||||
|
||||
def get_folder_by_name(self, folder_name: str) -> FileSystemFolder:
|
||||
"""
|
||||
Returns a the first folder with a matching name.
|
||||
|
||||
:return: Returns the first FileSydtemFolder with a matching name
|
||||
"""
|
||||
matching_folder = None
|
||||
for key in self.folders:
|
||||
if self.folders[key].name == folder_name:
|
||||
matching_folder = self.folders[key]
|
||||
break
|
||||
return matching_folder
|
||||
|
||||
def get_folder_by_id(self, folder_id: str) -> FileSystemFolder:
|
||||
"""
|
||||
Checks if the folder exists.
|
||||
|
||||
:param: folder_id: The id of the folder to find
|
||||
:type: folder_id: str
|
||||
"""
|
||||
return self.folders[folder_id]
|
||||
|
||||
def get_random_file_type(self) -> FileSystemFileType:
|
||||
"""
|
||||
Returns a random FileSystemFileTypeEnum.
|
||||
|
||||
:return: A random file type Enum
|
||||
"""
|
||||
return choice(list(FileSystemFileType))
|
||||
55
src/primaite/simulator/file_system/file_system_file.py
Normal file
55
src/primaite/simulator/file_system/file_system_file.py
Normal file
@@ -0,0 +1,55 @@
|
||||
from random import choice
|
||||
from typing import Dict
|
||||
|
||||
from primaite.simulator.file_system.file_system_file_type import file_type_sizes_KB, FileSystemFileType
|
||||
from primaite.simulator.file_system.file_system_item_abc import FileSystemItem
|
||||
|
||||
|
||||
class FileSystemFile(FileSystemItem):
|
||||
"""Class that represents a file in the simulation."""
|
||||
|
||||
file_type: FileSystemFileType = None
|
||||
"""The type of the FileSystemFile"""
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
"""
|
||||
Initialise FileSystemFile class.
|
||||
|
||||
:param name: The name of the file.
|
||||
:type name: str
|
||||
|
||||
:param file_type: The FileSystemFileType of the file
|
||||
:type file_type: Optional[FileSystemFileType]
|
||||
|
||||
:param size: The size of the FileSystemItem
|
||||
:type size: Optional[float]
|
||||
"""
|
||||
# set random file type if none provided
|
||||
|
||||
# set random file type if none provided
|
||||
if kwargs.get("file_type") is None:
|
||||
kwargs["file_type"] = choice(list(FileSystemFileType))
|
||||
|
||||
# set random file size if none provided
|
||||
if kwargs.get("size") is None:
|
||||
kwargs["size"] = file_type_sizes_KB[kwargs["file_type"]]
|
||||
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def describe_state(self) -> Dict:
|
||||
"""
|
||||
Produce a dictionary describing the current state of this object.
|
||||
|
||||
Please see :py:meth:`primaite.simulator.core.SimComponent.describe_state` for a more detailed explanation.
|
||||
|
||||
:return: Current state of this object and child objects.
|
||||
:rtype: Dict
|
||||
"""
|
||||
state = super().describe_state()
|
||||
state.update(
|
||||
{
|
||||
"uuid": self.uuid,
|
||||
"file_type": self.file_type.name,
|
||||
}
|
||||
)
|
||||
return state
|
||||
124
src/primaite/simulator/file_system/file_system_file_type.py
Normal file
124
src/primaite/simulator/file_system/file_system_file_type.py
Normal file
@@ -0,0 +1,124 @@
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class FileSystemFileType(str, Enum):
|
||||
"""An enumeration of common file types."""
|
||||
|
||||
UNKNOWN = 0
|
||||
"Unknown file type."
|
||||
|
||||
# Text formats
|
||||
TXT = 1
|
||||
"Plain text file."
|
||||
DOC = 2
|
||||
"Microsoft Word document (.doc)"
|
||||
DOCX = 3
|
||||
"Microsoft Word document (.docx)"
|
||||
PDF = 4
|
||||
"Portable Document Format."
|
||||
HTML = 5
|
||||
"HyperText Markup Language file."
|
||||
XML = 6
|
||||
"Extensible Markup Language file."
|
||||
CSV = 7
|
||||
"Comma-Separated Values file."
|
||||
|
||||
# Spreadsheet formats
|
||||
XLS = 8
|
||||
"Microsoft Excel file (.xls)"
|
||||
XLSX = 9
|
||||
"Microsoft Excel file (.xlsx)"
|
||||
|
||||
# Image formats
|
||||
JPEG = 10
|
||||
"JPEG image file."
|
||||
PNG = 11
|
||||
"PNG image file."
|
||||
GIF = 12
|
||||
"GIF image file."
|
||||
BMP = 13
|
||||
"Bitmap image file."
|
||||
|
||||
# Audio formats
|
||||
MP3 = 14
|
||||
"MP3 audio file."
|
||||
WAV = 15
|
||||
"WAV audio file."
|
||||
|
||||
# Video formats
|
||||
MP4 = 16
|
||||
"MP4 video file."
|
||||
AVI = 17
|
||||
"AVI video file."
|
||||
MKV = 18
|
||||
"MKV video file."
|
||||
FLV = 19
|
||||
"FLV video file."
|
||||
|
||||
# Presentation formats
|
||||
PPT = 20
|
||||
"Microsoft PowerPoint file (.ppt)"
|
||||
PPTX = 21
|
||||
"Microsoft PowerPoint file (.pptx)"
|
||||
|
||||
# Web formats
|
||||
JS = 22
|
||||
"JavaScript file."
|
||||
CSS = 23
|
||||
"Cascading Style Sheets file."
|
||||
|
||||
# Programming languages
|
||||
PY = 24
|
||||
"Python script file."
|
||||
C = 25
|
||||
"C source code file."
|
||||
CPP = 26
|
||||
"C++ source code file."
|
||||
JAVA = 27
|
||||
"Java source code file."
|
||||
|
||||
# Compressed file types
|
||||
RAR = 28
|
||||
"RAR archive file."
|
||||
ZIP = 29
|
||||
"ZIP archive file."
|
||||
TAR = 30
|
||||
"TAR archive file."
|
||||
GZ = 31
|
||||
"Gzip compressed file."
|
||||
|
||||
|
||||
file_type_sizes_KB = {
|
||||
FileSystemFileType.UNKNOWN: 0,
|
||||
FileSystemFileType.TXT: 4,
|
||||
FileSystemFileType.DOC: 50,
|
||||
FileSystemFileType.DOCX: 30,
|
||||
FileSystemFileType.PDF: 100,
|
||||
FileSystemFileType.HTML: 15,
|
||||
FileSystemFileType.XML: 10,
|
||||
FileSystemFileType.CSV: 15,
|
||||
FileSystemFileType.XLS: 100,
|
||||
FileSystemFileType.XLSX: 25,
|
||||
FileSystemFileType.JPEG: 100,
|
||||
FileSystemFileType.PNG: 40,
|
||||
FileSystemFileType.GIF: 30,
|
||||
FileSystemFileType.BMP: 300,
|
||||
FileSystemFileType.MP3: 5000,
|
||||
FileSystemFileType.WAV: 25000,
|
||||
FileSystemFileType.MP4: 25000,
|
||||
FileSystemFileType.AVI: 50000,
|
||||
FileSystemFileType.MKV: 50000,
|
||||
FileSystemFileType.FLV: 15000,
|
||||
FileSystemFileType.PPT: 200,
|
||||
FileSystemFileType.PPTX: 100,
|
||||
FileSystemFileType.JS: 10,
|
||||
FileSystemFileType.CSS: 5,
|
||||
FileSystemFileType.PY: 5,
|
||||
FileSystemFileType.C: 5,
|
||||
FileSystemFileType.CPP: 10,
|
||||
FileSystemFileType.JAVA: 10,
|
||||
FileSystemFileType.RAR: 1000,
|
||||
FileSystemFileType.ZIP: 1000,
|
||||
FileSystemFileType.TAR: 1000,
|
||||
FileSystemFileType.GZ: 800,
|
||||
}
|
||||
87
src/primaite/simulator/file_system/file_system_folder.py
Normal file
87
src/primaite/simulator/file_system/file_system_folder.py
Normal file
@@ -0,0 +1,87 @@
|
||||
from typing import Dict, Optional
|
||||
|
||||
from primaite import getLogger
|
||||
from primaite.simulator.file_system.file_system_file import FileSystemFile
|
||||
from primaite.simulator.file_system.file_system_item_abc import FileSystemItem
|
||||
|
||||
_LOGGER = getLogger(__name__)
|
||||
|
||||
|
||||
class FileSystemFolder(FileSystemItem):
|
||||
"""Simulation FileSystemFolder."""
|
||||
|
||||
files: Dict[str, FileSystemFile] = {}
|
||||
"""List of files stored in the folder."""
|
||||
|
||||
is_quarantined: bool = False
|
||||
"""Flag that marks the folder as quarantined if true."""
|
||||
|
||||
def describe_state(self) -> Dict:
|
||||
"""
|
||||
Produce a dictionary describing the current state of this object.
|
||||
|
||||
Please see :py:meth:`primaite.simulator.core.SimComponent.describe_state` for a more detailed explanation.
|
||||
|
||||
:return: Current state of this object and child objects.
|
||||
:rtype: Dict
|
||||
"""
|
||||
state = super().describe_state()
|
||||
state.update(
|
||||
{
|
||||
"files": {uuid: file.describe_state() for uuid, file in self.files.items()},
|
||||
"is_quarantined": self.is_quarantined,
|
||||
}
|
||||
)
|
||||
return state
|
||||
|
||||
def get_file_by_id(self, file_id: str) -> FileSystemFile:
|
||||
"""Return a FileSystemFile with the matching id."""
|
||||
return self.files.get(file_id)
|
||||
|
||||
def get_file_by_name(self, file_name: str) -> FileSystemFile:
|
||||
"""Return a FileSystemFile with the matching id."""
|
||||
return next((f for f in list(self.files) if f.name == file_name), None)
|
||||
|
||||
def add_file(self, file: FileSystemFile):
|
||||
"""Adds a file to the folder list."""
|
||||
if file is None or not isinstance(file, FileSystemFile):
|
||||
raise Exception(f"Invalid file: {file}")
|
||||
|
||||
# check if file with id already exists in folder
|
||||
if file.uuid in self.files:
|
||||
_LOGGER.debug(f"File with id {file.uuid} already exists in folder")
|
||||
else:
|
||||
# add to list
|
||||
self.files[file.uuid] = file
|
||||
self.size += file.size
|
||||
|
||||
def remove_file(self, file: Optional[FileSystemFile]):
|
||||
"""
|
||||
Removes a file from the folder list.
|
||||
|
||||
The method can take a FileSystemFile object or a file id.
|
||||
|
||||
:param: file: The file to remove
|
||||
:type: Optional[FileSystemFile]
|
||||
"""
|
||||
if file is None or not isinstance(file, FileSystemFile):
|
||||
raise Exception(f"Invalid file: {file}")
|
||||
|
||||
if self.files.get(file.uuid):
|
||||
del self.files[file.uuid]
|
||||
|
||||
self.size -= file.size
|
||||
else:
|
||||
_LOGGER.debug(f"File with UUID {file.uuid} was not found.")
|
||||
|
||||
def quarantine(self):
|
||||
"""Quarantines the File System Folder."""
|
||||
self.is_quarantined = True
|
||||
|
||||
def end_quarantine(self):
|
||||
"""Ends the quarantine of the File System Folder."""
|
||||
self.is_quarantined = False
|
||||
|
||||
def quarantine_status(self) -> bool:
|
||||
"""Returns true if the folder is being quarantined."""
|
||||
return self.is_quarantined
|
||||
31
src/primaite/simulator/file_system/file_system_item_abc.py
Normal file
31
src/primaite/simulator/file_system/file_system_item_abc.py
Normal file
@@ -0,0 +1,31 @@
|
||||
from typing import Dict
|
||||
|
||||
from primaite.simulator.core import SimComponent
|
||||
|
||||
|
||||
class FileSystemItem(SimComponent):
|
||||
"""Abstract base class for FileSystemItems used in the file system simulation."""
|
||||
|
||||
name: str
|
||||
"""The name of the FileSystemItem."""
|
||||
|
||||
size: float = 0
|
||||
"""The size the item takes up on disk."""
|
||||
|
||||
def describe_state(self) -> Dict:
|
||||
"""
|
||||
Produce a dictionary describing the current state of this object.
|
||||
|
||||
Please see :py:meth:`primaite.simulator.core.SimComponent.describe_state` for a more detailed explanation.
|
||||
|
||||
:return: Current state of this object and child objects.
|
||||
:rtype: Dict
|
||||
"""
|
||||
state = super().describe_state()
|
||||
state.update(
|
||||
{
|
||||
"name": self.name,
|
||||
"size": self.size,
|
||||
}
|
||||
)
|
||||
return state
|
||||
116
src/primaite/simulator/network/container.py
Normal file
116
src/primaite/simulator/network/container.py
Normal file
@@ -0,0 +1,116 @@
|
||||
from typing import Any, Dict, Union
|
||||
|
||||
from primaite import getLogger
|
||||
from primaite.simulator.core import Action, ActionManager, AllowAllValidator, SimComponent
|
||||
from primaite.simulator.network.hardware.base import Link, NIC, Node, SwitchPort
|
||||
|
||||
_LOGGER = getLogger(__name__)
|
||||
|
||||
|
||||
class Network(SimComponent):
|
||||
"""Top level container object representing the physical network."""
|
||||
|
||||
nodes: Dict[str, Node] = {}
|
||||
links: Dict[str, Link] = {}
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
"""Initialise the network."""
|
||||
super().__init__(**kwargs)
|
||||
|
||||
self.action_manager = ActionManager()
|
||||
self.action_manager.add_action(
|
||||
"node",
|
||||
Action(
|
||||
func=lambda request, context: self.nodes[request.pop(0)].apply_action(request, context),
|
||||
validator=AllowAllValidator(),
|
||||
),
|
||||
)
|
||||
|
||||
def describe_state(self) -> Dict:
|
||||
"""
|
||||
Produce a dictionary describing the current state of this object.
|
||||
|
||||
Please see :py:meth:`primaite.simulator.core.SimComponent.describe_state` for a more detailed explanation.
|
||||
|
||||
:return: Current state of this object and child objects.
|
||||
:rtype: Dict
|
||||
"""
|
||||
state = super().describe_state()
|
||||
state.update(
|
||||
{
|
||||
"nodes": {uuid: node.describe_state() for uuid, node in self.nodes.items()},
|
||||
"links": {uuid: link.describe_state() for uuid, link in self.links.items()},
|
||||
}
|
||||
)
|
||||
return state
|
||||
|
||||
def add_node(self, node: Node) -> None:
|
||||
"""
|
||||
Add an existing node to the network.
|
||||
|
||||
:param node: Node instance that the network should keep track of.
|
||||
:type node: Node
|
||||
"""
|
||||
if node in self:
|
||||
_LOGGER.warning(f"Can't add node {node.uuid}. It is already in the network.")
|
||||
return
|
||||
self.nodes[node.uuid] = node
|
||||
node.parent = self
|
||||
_LOGGER.info(f"Added node {node.uuid} to Network {self.uuid}")
|
||||
|
||||
def remove_node(self, node: Node) -> None:
|
||||
"""
|
||||
Remove a node from the network.
|
||||
|
||||
:param node: Node instance that is currently part of the network that should be removed.
|
||||
:type node: Node
|
||||
"""
|
||||
if node not in self:
|
||||
_LOGGER.warning(f"Can't remove node {node.uuid}. It's not in the network.")
|
||||
return
|
||||
self.nodes.pop(node.uuid)
|
||||
node.parent = None
|
||||
_LOGGER.info(f"Removed node {node.uuid} from network {self.uuid}")
|
||||
|
||||
def connect(self, endpoint_a: Union[NIC, SwitchPort], endpoint_b: Union[NIC, SwitchPort], **kwargs) -> None:
|
||||
"""Connect two nodes on the network by creating a link between an NIC/SwitchPort of each one.
|
||||
|
||||
:param endpoint_a: The endpoint to which to connect the link on the first node
|
||||
:type endpoint_a: Union[NIC, SwitchPort]
|
||||
:param endpoint_b: The endpoint to which to connct the link on the second node
|
||||
:type endpoint_b: Union[NIC, SwitchPort]
|
||||
:raises RuntimeError: _description_
|
||||
"""
|
||||
node_a = endpoint_a.parent
|
||||
node_b = endpoint_b.parent
|
||||
if node_a not in self:
|
||||
self.add_node(node_a)
|
||||
if node_b not in self:
|
||||
self.add_node(node_b)
|
||||
if node_a is node_b:
|
||||
_LOGGER.warn(f"Cannot link endpoint {endpoint_a} to {endpoint_b} because they belong to the same node.")
|
||||
return
|
||||
|
||||
link = Link(endpoint_a=endpoint_a, endpoint_b=endpoint_b, **kwargs)
|
||||
self.links[link.uuid] = link
|
||||
link.parent = self
|
||||
_LOGGER.info(f"Added link {link.uuid} to connect {endpoint_a} and {endpoint_b}")
|
||||
|
||||
def remove_link(self, link: Link) -> None:
|
||||
"""Disconnect a link from the network.
|
||||
|
||||
:param link: The link to be removed
|
||||
:type link: Link
|
||||
"""
|
||||
link.endpoint_a.disconnect_link()
|
||||
link.endpoint_b.disconnect_link()
|
||||
self.links.pop(link.uuid)
|
||||
link.parent = None
|
||||
_LOGGER.info(f"Removed link {link.uuid} from network {self.uuid}.")
|
||||
|
||||
def __contains__(self, item: Any) -> bool:
|
||||
if isinstance(item, Node):
|
||||
return item.uuid in self.nodes
|
||||
elif isinstance(item, Link):
|
||||
return item.uuid in self.links
|
||||
return False
|
||||
0
src/primaite/simulator/network/hardware/__init__.py
Normal file
0
src/primaite/simulator/network/hardware/__init__.py
Normal file
1091
src/primaite/simulator/network/hardware/base.py
Normal file
1091
src/primaite/simulator/network/hardware/base.py
Normal file
File diff suppressed because it is too large
Load Diff
69
src/primaite/simulator/network/protocols/arp.py
Normal file
69
src/primaite/simulator/network/protocols/arp.py
Normal file
@@ -0,0 +1,69 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from ipaddress import IPv4Address
|
||||
from typing import Optional
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
||||
class ARPEntry(BaseModel):
|
||||
"""
|
||||
Represents an entry in the ARP cache.
|
||||
|
||||
:param mac_address: The MAC address associated with the IP address.
|
||||
:param nic: The NIC through which the NIC with the IP address is reachable.
|
||||
"""
|
||||
|
||||
mac_address: str
|
||||
nic_uuid: str
|
||||
|
||||
|
||||
class ARPPacket(BaseModel):
|
||||
"""
|
||||
Represents the ARP layer of a network frame.
|
||||
|
||||
:param request: ARP operation. True if a request, False if a reply.
|
||||
:param sender_mac_addr: Sender MAC address.
|
||||
:param sender_ip: Sender IP address.
|
||||
:param target_mac_addr: Target MAC address.
|
||||
:param target_ip: Target IP address.
|
||||
|
||||
:Example:
|
||||
|
||||
>>> arp_request = ARPPacket(
|
||||
... sender_mac_addr="aa:bb:cc:dd:ee:ff",
|
||||
... sender_ip=IPv4Address("192.168.0.1"),
|
||||
... target_ip=IPv4Address("192.168.0.2")
|
||||
... )
|
||||
>>> arp_response = ARPPacket(
|
||||
... sender_mac_addr="aa:bb:cc:dd:ee:ff",
|
||||
... sender_ip=IPv4Address("192.168.0.1"),
|
||||
... target_ip=IPv4Address("192.168.0.2")
|
||||
... )
|
||||
"""
|
||||
|
||||
request: bool = True
|
||||
"ARP operation. True if a request, False if a reply."
|
||||
sender_mac_addr: str
|
||||
"Sender MAC address."
|
||||
sender_ip: IPv4Address
|
||||
"Sender IP address."
|
||||
target_mac_addr: Optional[str] = None
|
||||
"Target MAC address."
|
||||
target_ip: IPv4Address
|
||||
"Target IP address."
|
||||
|
||||
def generate_reply(self, mac_address: str) -> ARPPacket:
|
||||
"""
|
||||
Generate a new ARPPacket to be sent as a response with a given mac address.
|
||||
|
||||
:param mac_address: The mac_address that was being sought after from the original target IP address.
|
||||
:return: A new instance of ARPPacket.
|
||||
"""
|
||||
return ARPPacket(
|
||||
request=False,
|
||||
sender_ip=self.target_ip,
|
||||
sender_mac_addr=mac_address,
|
||||
target_ip=self.sender_ip,
|
||||
target_mac_addr=self.sender_mac_addr,
|
||||
)
|
||||
@@ -1,11 +1,14 @@
|
||||
from datetime import datetime
|
||||
from typing import Any, Optional
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from primaite import getLogger
|
||||
from primaite.simulator.network.transmission.network_layer import ICMPHeader, IPPacket, IPProtocol
|
||||
from primaite.simulator.network.protocols.arp import ARPPacket
|
||||
from primaite.simulator.network.transmission.network_layer import ICMPPacket, IPPacket, IPProtocol
|
||||
from primaite.simulator.network.transmission.primaite_layer import PrimaiteHeader
|
||||
from primaite.simulator.network.transmission.transport_layer import TCPHeader, UDPHeader
|
||||
from primaite.simulator.network.utils import convert_bytes_to_megabits
|
||||
|
||||
_LOGGER = getLogger(__name__)
|
||||
|
||||
@@ -74,9 +77,11 @@ class Frame(BaseModel):
|
||||
_LOGGER.error(msg)
|
||||
raise ValueError(msg)
|
||||
if kwargs["ip"].protocol == IPProtocol.ICMP and not kwargs.get("icmp"):
|
||||
msg = "Cannot build a Frame using the ICMP IP Protocol without a ICMPHeader"
|
||||
msg = "Cannot build a Frame using the ICMP IP Protocol without a ICMPPacket"
|
||||
_LOGGER.error(msg)
|
||||
raise ValueError(msg)
|
||||
kwargs["primaite"] = PrimaiteHeader()
|
||||
|
||||
super().__init__(**kwargs)
|
||||
|
||||
ethernet: EthernetHeader
|
||||
@@ -87,14 +92,44 @@ class Frame(BaseModel):
|
||||
"TCP header."
|
||||
udp: Optional[UDPHeader] = None
|
||||
"UDP header."
|
||||
icmp: Optional[ICMPHeader] = None
|
||||
icmp: Optional[ICMPPacket] = None
|
||||
"ICMP header."
|
||||
primaite: PrimaiteHeader = PrimaiteHeader()
|
||||
arp: Optional[ARPPacket] = None
|
||||
"ARP packet."
|
||||
primaite: PrimaiteHeader
|
||||
"PrimAITE header."
|
||||
payload: Optional[Any] = None
|
||||
"Raw data payload."
|
||||
sent_timestamp: Optional[datetime] = None
|
||||
"The time the Frame was sent from the original source NIC."
|
||||
received_timestamp: Optional[datetime] = None
|
||||
"The time the Frame was received at the final destination NIC."
|
||||
|
||||
def decrement_ttl(self):
|
||||
"""Decrement the IPPacket ttl by 1."""
|
||||
self.ip.ttl -= 1
|
||||
|
||||
@property
|
||||
def size(self) -> int:
|
||||
"""The size in Bytes."""
|
||||
return len(self.model_dump_json().encode("utf-8"))
|
||||
def can_transmit(self) -> bool:
|
||||
"""Informs whether the Frame can transmit based on the IPPacket tll being >= 1."""
|
||||
return self.ip.ttl >= 1
|
||||
|
||||
def set_sent_timestamp(self):
|
||||
"""Set the sent_timestamp."""
|
||||
if not self.sent_timestamp:
|
||||
self.sent_timestamp = datetime.now()
|
||||
|
||||
def set_received_timestamp(self):
|
||||
"""Set the received_timestamp."""
|
||||
if not self.received_timestamp:
|
||||
self.received_timestamp = datetime.now()
|
||||
|
||||
@property
|
||||
def size(self) -> float: # noqa - Keep it as MBits as this is how they're expressed
|
||||
"""The size of the Frame in Bytes."""
|
||||
return float(len(self.model_dump_json().encode("utf-8")))
|
||||
|
||||
@property
|
||||
def size_Mbits(self) -> float: # noqa - Keep it as MBits as this is how they're expressed
|
||||
"""The daa transfer size of the Frame in Mbits."""
|
||||
return convert_bytes_to_megabits(self.size)
|
||||
|
||||
@@ -120,18 +120,23 @@ def get_icmp_type_code_description(icmp_type: ICMPType, icmp_code: int) -> Union
|
||||
return icmp_code_descriptions[icmp_type].get(icmp_code)
|
||||
|
||||
|
||||
class ICMPHeader(BaseModel):
|
||||
"""Models an ICMP Header."""
|
||||
class ICMPPacket(BaseModel):
|
||||
"""Models an ICMP Packet."""
|
||||
|
||||
icmp_type: ICMPType = ICMPType.ECHO_REQUEST
|
||||
"ICMP Type."
|
||||
icmp_code: int = 0
|
||||
"ICMP Code."
|
||||
identifier: str = secrets.randbits(16)
|
||||
identifier: int
|
||||
"ICMP identifier (16 bits randomly generated)."
|
||||
sequence: int = 1
|
||||
sequence: int = 0
|
||||
"ICMP message sequence number."
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
if not kwargs.get("identifier"):
|
||||
kwargs["identifier"] = secrets.randbits(16)
|
||||
super().__init__(**kwargs)
|
||||
|
||||
@field_validator("icmp_code") # noqa
|
||||
@classmethod
|
||||
def _icmp_type_must_have_icmp_code(cls, v: int, info: FieldValidationInfo) -> int:
|
||||
|
||||
@@ -1,277 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
import secrets
|
||||
from ipaddress import IPv4Address, IPv4Network
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
from primaite import getLogger
|
||||
from primaite.exceptions import NetworkError
|
||||
from primaite.simulator.core import SimComponent
|
||||
from primaite.simulator.network.transmission.data_link_layer import Frame
|
||||
|
||||
_LOGGER = getLogger(__name__)
|
||||
|
||||
|
||||
def generate_mac_address(oui: Optional[str] = None) -> str:
|
||||
"""
|
||||
Generate a random MAC Address.
|
||||
|
||||
:Example:
|
||||
|
||||
>>> generate_mac_address()
|
||||
'ef:7e:97:c8:a8:ce'
|
||||
|
||||
>>> generate_mac_address(oui='aa:bb:cc')
|
||||
'aa:bb:cc:42:ba:41'
|
||||
|
||||
:param oui: The Organizationally Unique Identifier (OUI) portion of the MAC address. It should be a string with
|
||||
the first 3 bytes (24 bits) in the format "XX:XX:XX".
|
||||
:raises ValueError: If the 'oui' is not in the correct format (hexadecimal and 6 characters).
|
||||
"""
|
||||
random_bytes = [secrets.randbits(8) for _ in range(6)]
|
||||
|
||||
if oui:
|
||||
oui_pattern = re.compile(r"^([0-9A-Fa-f]{2}[:-]){2}[0-9A-Fa-f]{2}$")
|
||||
if not oui_pattern.match(oui):
|
||||
msg = f"Invalid oui. The oui should be in the format xx:xx:xx, where x is a hexadecimal digit, got '{oui}'"
|
||||
raise ValueError(msg)
|
||||
oui_bytes = [int(chunk, 16) for chunk in oui.split(":")]
|
||||
mac = oui_bytes + random_bytes[len(oui_bytes) :]
|
||||
else:
|
||||
mac = random_bytes
|
||||
|
||||
return ":".join(f"{b:02x}" for b in mac)
|
||||
|
||||
|
||||
class NIC(SimComponent):
|
||||
"""
|
||||
Models a Network Interface Card (NIC) in a computer or network device.
|
||||
|
||||
:param ip_address: The IPv4 address assigned to the NIC.
|
||||
:param subnet_mask: The subnet mask assigned to the NIC.
|
||||
:param gateway: The default gateway IP address for forwarding network traffic to other networks.
|
||||
:param mac_address: The MAC address of the NIC. Defaults to a randomly set MAC address.
|
||||
:param speed: The speed of the NIC in Mbps (default is 100 Mbps).
|
||||
:param mtu: The Maximum Transmission Unit (MTU) of the NIC in Bytes, representing the largest data packet size it
|
||||
can handle without fragmentation (default is 1500 B).
|
||||
:param wake_on_lan: Indicates if the NIC supports Wake-on-LAN functionality.
|
||||
:param dns_servers: List of IP addresses of DNS servers used for name resolution.
|
||||
"""
|
||||
|
||||
ip_address: IPv4Address
|
||||
"The IP address assigned to the NIC for communication on an IP-based network."
|
||||
subnet_mask: str
|
||||
"The subnet mask assigned to the NIC."
|
||||
gateway: IPv4Address
|
||||
"The default gateway IP address for forwarding network traffic to other networks. Randomly generated upon creation."
|
||||
mac_address: str = generate_mac_address()
|
||||
"The MAC address of the NIC. Defaults to a randomly set MAC address."
|
||||
speed: int = 100
|
||||
"The speed of the NIC in Mbps. Default is 100 Mbps."
|
||||
mtu: int = 1500
|
||||
"The Maximum Transmission Unit (MTU) of the NIC in Bytes. Default is 1500 B"
|
||||
wake_on_lan: bool = False
|
||||
"Indicates if the NIC supports Wake-on-LAN functionality."
|
||||
dns_servers: List[IPv4Address] = []
|
||||
"List of IP addresses of DNS servers used for name resolution."
|
||||
connected_link: Optional[Link] = None
|
||||
"The Link to which the NIC is connected."
|
||||
enabled: bool = False
|
||||
"Indicates whether the NIC is enabled."
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
"""
|
||||
NIC constructor.
|
||||
|
||||
Performs some type conversion the calls ``super().__init__()``. Then performs some checking on the ip_address
|
||||
and gateway just to check that it's all been configured correctly.
|
||||
|
||||
:raises ValueError: When the ip_address and gateway are the same. And when the ip_address/subnet mask are a
|
||||
network address.
|
||||
"""
|
||||
if not isinstance(kwargs["ip_address"], IPv4Address):
|
||||
kwargs["ip_address"] = IPv4Address(kwargs["ip_address"])
|
||||
if not isinstance(kwargs["gateway"], IPv4Address):
|
||||
kwargs["gateway"] = IPv4Address(kwargs["gateway"])
|
||||
super().__init__(**kwargs)
|
||||
|
||||
if self.ip_address == self.gateway:
|
||||
msg = f"NIC ip address {self.ip_address} cannot be the same as the gateway {self.gateway}"
|
||||
_LOGGER.error(msg)
|
||||
raise ValueError(msg)
|
||||
if self.ip_network.network_address == self.ip_address:
|
||||
msg = (
|
||||
f"Failed to set IP address {self.ip_address} and subnet mask {self.subnet_mask} as it is a "
|
||||
f"network address {self.ip_network.network_address}"
|
||||
)
|
||||
_LOGGER.error(msg)
|
||||
raise ValueError(msg)
|
||||
|
||||
@property
|
||||
def ip_network(self) -> IPv4Network:
|
||||
"""
|
||||
Return the IPv4Network of the NIC.
|
||||
|
||||
:return: The IPv4Network from the ip_address/subnet mask.
|
||||
"""
|
||||
return IPv4Network(f"{self.ip_address}/{self.subnet_mask}", strict=False)
|
||||
|
||||
def connect_link(self, link: Link):
|
||||
"""
|
||||
Connect the NIC to a link.
|
||||
|
||||
:param link: The link to which the NIC is connected.
|
||||
:type link: :class:`~primaite.simulator.network.transmission.physical_layer.Link`
|
||||
:raise NetworkError: When an attempt to connect a Link is made while the NIC has a connected Link.
|
||||
"""
|
||||
if not self.connected_link:
|
||||
if self.connected_link != link:
|
||||
# TODO: Inform the Node that a link has been connected
|
||||
self.connected_link = link
|
||||
else:
|
||||
_LOGGER.warning(f"Cannot connect link to NIC ({self.mac_address}) as it is already connected")
|
||||
else:
|
||||
msg = f"Cannot connect link to NIC ({self.mac_address}) as it already has a connection"
|
||||
_LOGGER.error(msg)
|
||||
raise NetworkError(msg)
|
||||
|
||||
def disconnect_link(self):
|
||||
"""Disconnect the NIC from the connected Link."""
|
||||
if self.connected_link.endpoint_a == self:
|
||||
self.connected_link.endpoint_a = None
|
||||
if self.connected_link.endpoint_b == self:
|
||||
self.connected_link.endpoint_b = None
|
||||
self.connected_link = None
|
||||
|
||||
def add_dns_server(self, ip_address: IPv4Address):
|
||||
"""
|
||||
Add a DNS server IP address.
|
||||
|
||||
:param ip_address: The IP address of the DNS server to be added.
|
||||
:type ip_address: ipaddress.IPv4Address
|
||||
"""
|
||||
pass
|
||||
|
||||
def remove_dns_server(self, ip_address: IPv4Address):
|
||||
"""
|
||||
Remove a DNS server IP Address.
|
||||
|
||||
:param ip_address: The IP address of the DNS server to be removed.
|
||||
:type ip_address: ipaddress.IPv4Address
|
||||
"""
|
||||
pass
|
||||
|
||||
def send_frame(self, frame: Frame):
|
||||
"""
|
||||
Send a network frame from the NIC to the connected link.
|
||||
|
||||
:param frame: The network frame to be sent.
|
||||
:type frame: :class:`~primaite.simulator.network.osi_layers.Frame`
|
||||
"""
|
||||
pass
|
||||
|
||||
def receive_frame(self, frame: Frame):
|
||||
"""
|
||||
Receive a network frame from the connected link.
|
||||
|
||||
The Frame is passed to the Node.
|
||||
|
||||
:param frame: The network frame being received.
|
||||
:type frame: :class:`~primaite.simulator.network.osi_layers.Frame`
|
||||
"""
|
||||
pass
|
||||
|
||||
def describe_state(self) -> Dict:
|
||||
"""
|
||||
Get the current state of the NIC as a dict.
|
||||
|
||||
:return: A dict containing the current state of the NIC.
|
||||
"""
|
||||
pass
|
||||
|
||||
def apply_action(self, action: str):
|
||||
"""
|
||||
Apply an action to the NIC.
|
||||
|
||||
:param action: The action to be applied.
|
||||
:type action: str
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
class Link(SimComponent):
|
||||
"""
|
||||
Represents a network link between two network interface cards (NICs).
|
||||
|
||||
:param endpoint_a: The first NIC connected to the Link.
|
||||
:type endpoint_a: NIC
|
||||
:param endpoint_b: The second NIC connected to the Link.
|
||||
:type endpoint_b: NIC
|
||||
:param bandwidth: The bandwidth of the Link in Mbps (default is 100 Mbps).
|
||||
:type bandwidth: int
|
||||
"""
|
||||
|
||||
endpoint_a: NIC
|
||||
"The first NIC connected to the Link."
|
||||
endpoint_b: NIC
|
||||
"The second NIC connected to the Link."
|
||||
bandwidth: int = 100
|
||||
"The bandwidth of the Link in Mbps (default is 100 Mbps)."
|
||||
current_load: int = 0
|
||||
"The current load on the link in Mbps."
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
"""
|
||||
Ensure that endpoint_a and endpoint_b are not the same NIC.
|
||||
|
||||
Connect the link to the NICs after creation.
|
||||
|
||||
:raises ValueError: If endpoint_a and endpoint_b are the same NIC.
|
||||
"""
|
||||
if kwargs["endpoint_a"] == kwargs["endpoint_b"]:
|
||||
msg = "endpoint_a and endpoint_b cannot be the same NIC"
|
||||
_LOGGER.error(msg)
|
||||
raise ValueError(msg)
|
||||
super().__init__(**kwargs)
|
||||
self.endpoint_a.connect_link(self)
|
||||
self.endpoint_b.connect_link(self)
|
||||
|
||||
def send_frame(self, sender_nic: NIC, frame: Frame):
|
||||
"""
|
||||
Send a network frame from one NIC to another connected NIC.
|
||||
|
||||
:param sender_nic: The NIC sending the frame.
|
||||
:type sender_nic: NIC
|
||||
:param frame: The network frame to be sent.
|
||||
:type frame: Frame
|
||||
"""
|
||||
pass
|
||||
|
||||
def receive_frame(self, sender_nic: NIC, frame: Frame):
|
||||
"""
|
||||
Receive a network frame from a connected NIC.
|
||||
|
||||
:param sender_nic: The NIC sending the frame.
|
||||
:type sender_nic: NIC
|
||||
:param frame: The network frame being received.
|
||||
:type frame: Frame
|
||||
"""
|
||||
pass
|
||||
|
||||
def describe_state(self) -> Dict:
|
||||
"""
|
||||
Get the current state of the Libk as a dict.
|
||||
|
||||
:return: A dict containing the current state of the Link.
|
||||
"""
|
||||
pass
|
||||
|
||||
def apply_action(self, action: str):
|
||||
"""
|
||||
Apply an action to the Link.
|
||||
|
||||
:param action: The action to be applied.
|
||||
:type action: str
|
||||
"""
|
||||
pass
|
||||
@@ -33,6 +33,8 @@ class Port(Enum):
|
||||
"Simple Network Management Protocol (SNMP) - Used for network device management."
|
||||
SNMP_TRAP = 162
|
||||
"SNMP Trap - Used for sending SNMP notifications (traps) to a network management system."
|
||||
ARP = 219
|
||||
"Address resolution Protocol - Used to connect a MAC address to an IP address."
|
||||
LDAP = 389
|
||||
"Lightweight Directory Access Protocol (LDAP) - Used for accessing and modifying directory information."
|
||||
HTTPS = 443
|
||||
@@ -114,6 +116,6 @@ class TCPHeader(BaseModel):
|
||||
... )
|
||||
"""
|
||||
|
||||
src_port: int
|
||||
dst_port: int
|
||||
src_port: Port
|
||||
dst_port: Port
|
||||
flags: List[TCPFlags] = [TCPFlags.SYN]
|
||||
|
||||
27
src/primaite/simulator/network/utils.py
Normal file
27
src/primaite/simulator/network/utils.py
Normal file
@@ -0,0 +1,27 @@
|
||||
from typing import Union
|
||||
|
||||
|
||||
def convert_bytes_to_megabits(B: Union[int, float]) -> float: # noqa - Keep it as B as this is how Bytes are expressed
|
||||
"""
|
||||
Convert Bytes (file size) to Megabits (data transfer).
|
||||
|
||||
:param B: The file size in Bytes.
|
||||
:return: File bits to transfer in Megabits.
|
||||
"""
|
||||
if isinstance(B, int):
|
||||
B = float(B)
|
||||
bits = B * 8.0
|
||||
return bits / 1024.0**2.0
|
||||
|
||||
|
||||
def convert_megabits_to_bytes(Mbits: Union[int, float]) -> float: # noqa - The same for Mbits
|
||||
"""
|
||||
Convert Megabits (data transfer) to Bytes (file size).
|
||||
|
||||
:param Mbits bits to transfer in Megabits.
|
||||
:return: The file size in Bytes.
|
||||
"""
|
||||
if isinstance(Mbits, int):
|
||||
Mbits = float(Mbits)
|
||||
bits = Mbits * 1024.0**2.0
|
||||
return bits / 8
|
||||
56
src/primaite/simulator/sim_container.py
Normal file
56
src/primaite/simulator/sim_container.py
Normal file
@@ -0,0 +1,56 @@
|
||||
from typing import Dict
|
||||
|
||||
from primaite.simulator.core import Action, ActionManager, AllowAllValidator, SimComponent
|
||||
from primaite.simulator.domain.controller import DomainController
|
||||
from primaite.simulator.network.container import Network
|
||||
|
||||
|
||||
class Simulation(SimComponent):
|
||||
"""Top-level simulation object which holds a reference to all other parts of the simulation."""
|
||||
|
||||
network: Network
|
||||
domain: DomainController
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
"""Initialise the Simulation."""
|
||||
if not kwargs.get("network"):
|
||||
kwargs["network"] = Network()
|
||||
|
||||
if not kwargs.get("domain"):
|
||||
kwargs["domain"] = DomainController()
|
||||
|
||||
super().__init__(**kwargs)
|
||||
|
||||
self.action_manager = ActionManager()
|
||||
# pass through network actions to the network objects
|
||||
self.action_manager.add_action(
|
||||
"network",
|
||||
Action(
|
||||
func=lambda request, context: self.network.apply_action(request, context), validator=AllowAllValidator()
|
||||
),
|
||||
)
|
||||
# pass through domain actions to the domain object
|
||||
self.action_manager.add_action(
|
||||
"domain",
|
||||
Action(
|
||||
func=lambda request, context: self.domain.apply_action(request, context), validator=AllowAllValidator()
|
||||
),
|
||||
)
|
||||
|
||||
def describe_state(self) -> Dict:
|
||||
"""
|
||||
Produce a dictionary describing the current state of this object.
|
||||
|
||||
Please see :py:meth:`primaite.simulator.core.SimComponent.describe_state` for a more detailed explanation.
|
||||
|
||||
:return: Current state of this object and child objects.
|
||||
:rtype: Dict
|
||||
"""
|
||||
state = super().describe_state()
|
||||
state.update(
|
||||
{
|
||||
"network": self.network.describe_state(),
|
||||
"domain": self.domain.describe_state(),
|
||||
}
|
||||
)
|
||||
return state
|
||||
0
src/primaite/simulator/system/__init__.py
Normal file
0
src/primaite/simulator/system/__init__.py
Normal file
95
src/primaite/simulator/system/applications/application.py
Normal file
95
src/primaite/simulator/system/applications/application.py
Normal file
@@ -0,0 +1,95 @@
|
||||
from abc import abstractmethod
|
||||
from enum import Enum
|
||||
from typing import Any, Dict, List, Set
|
||||
|
||||
from primaite.simulator.system.software import IOSoftware
|
||||
|
||||
|
||||
class ApplicationOperatingState(Enum):
|
||||
"""Enumeration of Application Operating States."""
|
||||
|
||||
RUNNING = 1
|
||||
"The application is running."
|
||||
CLOSED = 2
|
||||
"The application is closed or not running."
|
||||
INSTALLING = 3
|
||||
"The application is being installed or updated."
|
||||
|
||||
|
||||
class Application(IOSoftware):
|
||||
"""
|
||||
Represents an Application in the simulation environment.
|
||||
|
||||
Applications are user-facing programs that may perform input/output operations.
|
||||
"""
|
||||
|
||||
operating_state: ApplicationOperatingState
|
||||
"The current operating state of the Application."
|
||||
execution_control_status: str
|
||||
"Control status of the application's execution. It could be 'manual' or 'automatic'."
|
||||
num_executions: int = 0
|
||||
"The number of times the application has been executed. Default is 0."
|
||||
groups: Set[str] = set()
|
||||
"The set of groups to which the application belongs."
|
||||
|
||||
@abstractmethod
|
||||
def describe_state(self) -> Dict:
|
||||
"""
|
||||
Produce a dictionary describing the current state of this object.
|
||||
|
||||
Please see :py:meth:`primaite.simulator.core.SimComponent.describe_state` for a more detailed explanation.
|
||||
|
||||
:return: Current state of this object and child objects.
|
||||
:rtype: Dict
|
||||
"""
|
||||
state = super().describe_state()
|
||||
state.update(
|
||||
{
|
||||
"opearting_state": self.operating_state.name,
|
||||
"execution_control_status": self.execution_control_status,
|
||||
"num_executions": self.num_executions,
|
||||
"groups": list(self.groups),
|
||||
}
|
||||
)
|
||||
return state
|
||||
|
||||
def apply_action(self, action: List[str]) -> None:
|
||||
"""
|
||||
Applies a list of actions to the Application.
|
||||
|
||||
:param action: A list of actions to apply.
|
||||
"""
|
||||
pass
|
||||
|
||||
def reset_component_for_episode(self, episode: int):
|
||||
"""
|
||||
Resets the Application component for a new episode.
|
||||
|
||||
This method ensures the Application is ready for a new episode, including resetting any
|
||||
stateful properties or statistics, and clearing any message queues.
|
||||
"""
|
||||
pass
|
||||
|
||||
def send(self, payload: Any, session_id: str, **kwargs) -> bool:
|
||||
"""
|
||||
Sends a payload to 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 send.
|
||||
:return: True if successful, False otherwise.
|
||||
"""
|
||||
pass
|
||||
|
||||
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.
|
||||
:return: True if successful, False otherwise.
|
||||
"""
|
||||
pass
|
||||
0
src/primaite/simulator/system/core/__init__.py
Normal file
0
src/primaite/simulator/system/core/__init__.py
Normal file
76
src/primaite/simulator/system/core/packet_capture.py
Normal file
76
src/primaite/simulator/system/core/packet_capture.py
Normal file
@@ -0,0 +1,76 @@
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
from primaite.simulator import TEMP_SIM_OUTPUT
|
||||
|
||||
|
||||
class _JSONFilter(logging.Filter):
|
||||
def filter(self, record: logging.LogRecord) -> bool:
|
||||
"""Filter logs that start and end with '{' and '}' (JSON-like messages)."""
|
||||
return record.getMessage().startswith("{") and record.getMessage().endswith("}")
|
||||
|
||||
|
||||
class PacketCapture:
|
||||
"""
|
||||
Represents a PacketCapture component on a Node in the simulation environment.
|
||||
|
||||
PacketCapture is a service that logs Frames as json strings; It's Wireshark for PrimAITE.
|
||||
|
||||
The PCAPs are logged to: <simulation output directory>/<hostname>/<hostname>_<ip address>_pcap.log
|
||||
"""
|
||||
|
||||
def __init__(self, hostname: str, ip_address: Optional[str] = None, switch_port_number: Optional[int] = None):
|
||||
"""
|
||||
Initialize the PacketCapture process.
|
||||
|
||||
:param hostname: The hostname for which PCAP logs are being recorded.
|
||||
:param ip_address: The IP address associated with the PCAP logs.
|
||||
"""
|
||||
self.hostname: str = hostname
|
||||
"The hostname for which PCAP logs are being recorded."
|
||||
self.ip_address: str = ip_address
|
||||
"The IP address associated with the PCAP logs."
|
||||
self.switch_port_number = switch_port_number
|
||||
"The SwitchPort number."
|
||||
self._setup_logger()
|
||||
|
||||
def _setup_logger(self):
|
||||
"""Set up the logger configuration."""
|
||||
log_path = self._get_log_path()
|
||||
|
||||
file_handler = logging.FileHandler(filename=log_path)
|
||||
file_handler.setLevel(60) # Custom log level > CRITICAL to prevent any unwanted standard DEBUG-CRITICAL logs
|
||||
|
||||
log_format = "%(message)s"
|
||||
file_handler.setFormatter(logging.Formatter(log_format))
|
||||
|
||||
self.logger = logging.getLogger(self._logger_name)
|
||||
self.logger.setLevel(60) # Custom log level > CRITICAL to prevent any unwanted standard DEBUG-CRITICAL logs
|
||||
self.logger.addHandler(file_handler)
|
||||
|
||||
self.logger.addFilter(_JSONFilter())
|
||||
|
||||
@property
|
||||
def _logger_name(self) -> str:
|
||||
"""Get PCAP the logger name."""
|
||||
if self.ip_address:
|
||||
return f"{self.hostname}_{self.ip_address}_pcap"
|
||||
if self.switch_port_number:
|
||||
return f"{self.hostname}_port-{self.switch_port_number}_pcap"
|
||||
return f"{self.hostname}_pcap"
|
||||
|
||||
def _get_log_path(self) -> Path:
|
||||
"""Get the path for the log file."""
|
||||
root = TEMP_SIM_OUTPUT / self.hostname
|
||||
root.mkdir(exist_ok=True, parents=True)
|
||||
return root / f"{self._logger_name}.log"
|
||||
|
||||
def capture(self, frame): # noqa - I'll have a circular import and cant use if TYPE_CHECKING ;(
|
||||
"""
|
||||
Capture a Frame and log it.
|
||||
|
||||
:param frame: The PCAP frame to capture.
|
||||
"""
|
||||
msg = frame.model_dump_json()
|
||||
self.logger.log(level=60, msg=msg) # Log at custom log level > CRITICAL
|
||||
183
src/primaite/simulator/system/core/session_manager.py
Normal file
183
src/primaite/simulator/system/core/session_manager.py
Normal file
@@ -0,0 +1,183 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from ipaddress import IPv4Address
|
||||
from typing import Any, Dict, Optional, Tuple, TYPE_CHECKING
|
||||
|
||||
from primaite.simulator.core import SimComponent
|
||||
from primaite.simulator.network.transmission.data_link_layer import Frame
|
||||
from primaite.simulator.network.transmission.network_layer import IPProtocol
|
||||
from primaite.simulator.network.transmission.transport_layer import Port
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from primaite.simulator.network.hardware.base import ARPCache
|
||||
from primaite.simulator.system.core.software_manager import SoftwareManager
|
||||
from primaite.simulator.system.core.sys_log import SysLog
|
||||
|
||||
|
||||
class Session(SimComponent):
|
||||
"""
|
||||
Models a network session.
|
||||
|
||||
Encapsulates information related to communication between two network endpoints, including the protocol,
|
||||
source and destination IPs and ports.
|
||||
|
||||
:param protocol: The IP protocol used in the session.
|
||||
:param src_ip: The source IP address.
|
||||
:param dst_ip: The destination IP address.
|
||||
:param src_port: The source port number (optional).
|
||||
:param dst_port: The destination port number (optional).
|
||||
:param connected: A flag indicating whether the session is connected.
|
||||
"""
|
||||
|
||||
protocol: IPProtocol
|
||||
src_ip: IPv4Address
|
||||
dst_ip: IPv4Address
|
||||
src_port: Optional[Port]
|
||||
dst_port: Optional[Port]
|
||||
connected: bool = False
|
||||
|
||||
@classmethod
|
||||
def from_session_key(
|
||||
cls, session_key: Tuple[IPProtocol, IPv4Address, IPv4Address, Optional[Port], Optional[Port]]
|
||||
) -> Session:
|
||||
"""
|
||||
Create a Session instance from a session key tuple.
|
||||
|
||||
:param session_key: Tuple containing the session details.
|
||||
:return: A Session instance.
|
||||
"""
|
||||
protocol, src_ip, dst_ip, src_port, dst_port = session_key
|
||||
return Session(protocol=protocol, src_ip=src_ip, dst_ip=dst_ip, src_port=src_port, dst_port=dst_port)
|
||||
|
||||
def describe_state(self) -> Dict:
|
||||
"""
|
||||
Produce a dictionary describing the current state of this object.
|
||||
|
||||
Please see :py:meth:`primaite.simulator.core.SimComponent.describe_state` for a more detailed explanation.
|
||||
|
||||
:return: Current state of this object and child objects.
|
||||
:rtype: Dict
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
class SessionManager:
|
||||
"""
|
||||
Manages network sessions, including session creation, lookup, and communication with other components.
|
||||
|
||||
:param sys_log: A reference to the system log component.
|
||||
:param arp_cache: A reference to the ARP cache component.
|
||||
"""
|
||||
|
||||
def __init__(self, sys_log: SysLog, arp_cache: "ARPCache"):
|
||||
self.sessions_by_key: Dict[
|
||||
Tuple[IPProtocol, IPv4Address, IPv4Address, Optional[Port], Optional[Port]], Session
|
||||
] = {}
|
||||
self.sessions_by_uuid: Dict[str, Session] = {}
|
||||
self.sys_log: SysLog = sys_log
|
||||
self.software_manager: SoftwareManager = None # Noqa
|
||||
self.arp_cache: "ARPCache" = arp_cache
|
||||
|
||||
def describe_state(self) -> Dict:
|
||||
"""
|
||||
Produce a dictionary describing the current state of this object.
|
||||
|
||||
Please see :py:meth:`primaite.simulator.core.SimComponent.describe_state` for a more detailed explanation.
|
||||
|
||||
:return: Current state of this object and child objects.
|
||||
:rtype: Dict
|
||||
"""
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def _get_session_key(
|
||||
frame: Frame, from_source: bool = True
|
||||
) -> Tuple[IPProtocol, IPv4Address, IPv4Address, Optional[Port], Optional[Port]]:
|
||||
"""
|
||||
Extracts the session key from the given frame.
|
||||
|
||||
The session key is a tuple containing the following elements:
|
||||
- IPProtocol: The transport protocol (e.g. TCP, UDP, ICMP).
|
||||
- IPv4Address: The source IP address.
|
||||
- IPv4Address: The destination IP address.
|
||||
- Optional[Port]: The source port number (if applicable).
|
||||
- Optional[Port]: The destination port number (if applicable).
|
||||
|
||||
:param frame: The network frame from which to extract the session key.
|
||||
:param from_source: A flag to indicate if the key should be extracted from the source or destination.
|
||||
:return: A tuple containing the session key.
|
||||
"""
|
||||
protocol = frame.ip.protocol
|
||||
src_ip = frame.ip.src_ip
|
||||
dst_ip = frame.ip.dst_ip
|
||||
if protocol == IPProtocol.TCP:
|
||||
if from_source:
|
||||
src_port = frame.tcp.src_port
|
||||
dst_port = frame.tcp.dst_port
|
||||
else:
|
||||
dst_port = frame.tcp.src_port
|
||||
src_port = frame.tcp.dst_port
|
||||
elif protocol == IPProtocol.UDP:
|
||||
if from_source:
|
||||
src_port = frame.udp.src_port
|
||||
dst_port = frame.udp.dst_port
|
||||
else:
|
||||
dst_port = frame.udp.src_port
|
||||
src_port = frame.udp.dst_port
|
||||
else:
|
||||
src_port = None
|
||||
dst_port = None
|
||||
return protocol, src_ip, dst_ip, src_port, dst_port
|
||||
|
||||
def receive_payload_from_software_manager(self, payload: Any, session_id: Optional[int] = None):
|
||||
"""
|
||||
Receive a payload from the SoftwareManager.
|
||||
|
||||
If no session_id, a Session is established. Once established, the payload is sent to ``send_payload_to_nic``.
|
||||
|
||||
:param payload: The payload to be sent.
|
||||
:param session_id: The Session ID the payload is to originate from. Optional. If None, one will be created.
|
||||
"""
|
||||
# TODO: Implement session creation and
|
||||
|
||||
self.send_payload_to_nic(payload, session_id)
|
||||
|
||||
def send_payload_to_software_manager(self, payload: Any, session_id: int):
|
||||
"""
|
||||
Send a payload to the software manager.
|
||||
|
||||
:param payload: The payload to be sent.
|
||||
:param session_id: The Session ID the payload originates from.
|
||||
"""
|
||||
self.software_manager.receive_payload_from_session_manger()
|
||||
|
||||
def send_payload_to_nic(self, payload: Any, session_id: int):
|
||||
"""
|
||||
Send a payload across the Network.
|
||||
|
||||
Takes a payload and a session_id. Builds a Frame and sends it across the network via a NIC.
|
||||
|
||||
:param payload: The payload to be sent.
|
||||
:param session_id: The Session ID the payload originates from
|
||||
"""
|
||||
# TODO: Implement frame construction and sent to NIC.
|
||||
pass
|
||||
|
||||
def receive_payload_from_nic(self, frame: Frame):
|
||||
"""
|
||||
Receive a Frame from the NIC.
|
||||
|
||||
Extract the session key using the _get_session_key method, and forward the payload to the appropriate
|
||||
session. If the session does not exist, a new one is created.
|
||||
|
||||
:param frame: The frame being received.
|
||||
"""
|
||||
session_key = self._get_session_key(frame)
|
||||
session = self.sessions_by_key.get(session_key)
|
||||
if not session:
|
||||
# Create new session
|
||||
session = Session.from_session_key(session_key)
|
||||
self.sessions_by_key[session_key] = session
|
||||
self.sessions_by_uuid[session.uuid] = session
|
||||
self.software_manager.receive_payload_from_session_manger(payload=frame, session=session)
|
||||
# TODO: Implement the frame deconstruction and send to SoftwareManager.
|
||||
99
src/primaite/simulator/system/core/software_manager.py
Normal file
99
src/primaite/simulator/system/core/software_manager.py
Normal file
@@ -0,0 +1,99 @@
|
||||
from typing import Any, Dict, Optional, Tuple, TYPE_CHECKING, Union
|
||||
|
||||
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.session_manager import Session
|
||||
from primaite.simulator.system.core.sys_log import SysLog
|
||||
from primaite.simulator.system.services.service import Service
|
||||
from primaite.simulator.system.software import SoftwareType
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from primaite.simulator.system.core.session_manager import SessionManager
|
||||
from primaite.simulator.system.core.sys_log import SysLog
|
||||
|
||||
|
||||
class SoftwareManager:
|
||||
"""A class that manages all running Services and Applications on a Node and facilitates their communication."""
|
||||
|
||||
def __init__(self, session_manager: "SessionManager", sys_log: "SysLog"):
|
||||
"""
|
||||
Initialize a new instance of SoftwareManager.
|
||||
|
||||
:param session_manager: The session manager handling network communications.
|
||||
"""
|
||||
self.session_manager = session_manager
|
||||
self.services: Dict[str, Service] = {}
|
||||
self.applications: Dict[str, Application] = {}
|
||||
self.port_protocol_mapping: Dict[Tuple[Port, IPProtocol], Union[Service, Application]] = {}
|
||||
self.sys_log: SysLog = sys_log
|
||||
|
||||
def add_service(self, name: str, service: Service, port: Port, protocol: IPProtocol):
|
||||
"""
|
||||
Add a Service to the manager.
|
||||
|
||||
:param name: The name of the service.
|
||||
:param service: The service instance.
|
||||
:param port: The port used by the service.
|
||||
:param protocol: The network protocol used by the service.
|
||||
"""
|
||||
service.software_manager = self
|
||||
self.services[name] = service
|
||||
self.port_protocol_mapping[(port, protocol)] = service
|
||||
|
||||
def add_application(self, name: str, application: Application, port: Port, protocol: IPProtocol):
|
||||
"""
|
||||
Add an Application to the manager.
|
||||
|
||||
:param name: The name of the application.
|
||||
:param application: The application instance.
|
||||
:param port: The port used by the application.
|
||||
:param protocol: The network protocol used by the application.
|
||||
"""
|
||||
application.software_manager = self
|
||||
self.applications[name] = application
|
||||
self.port_protocol_mapping[(port, protocol)] = application
|
||||
|
||||
def send_internal_payload(self, target_software: str, target_software_type: SoftwareType, payload: Any):
|
||||
"""
|
||||
Send a payload to a specific service or application.
|
||||
|
||||
:param target_software: The name of the target service or application.
|
||||
:param target_software_type: The type of software (Service, Application, Process).
|
||||
:param payload: The data to be sent.
|
||||
:param receiver_type: The type of the target, either 'service' or 'application'.
|
||||
"""
|
||||
if target_software_type is SoftwareType.SERVICE:
|
||||
receiver = self.services.get(target_software)
|
||||
elif target_software_type is SoftwareType.APPLICATION:
|
||||
receiver = self.applications.get(target_software)
|
||||
else:
|
||||
raise ValueError(f"Invalid receiver type {target_software_type}")
|
||||
|
||||
if receiver:
|
||||
receiver.receive_payload(payload)
|
||||
else:
|
||||
raise ValueError(f"No {target_software_type.name.lower()} found with the name {target_software}")
|
||||
|
||||
def send_payload_to_session_manger(self, payload: Any, session_id: Optional[int] = None):
|
||||
"""
|
||||
Send a payload to the SessionManager.
|
||||
|
||||
:param payload: The payload to be sent.
|
||||
:param session_id: The Session ID the payload is to originate from. Optional.
|
||||
"""
|
||||
self.session_manager.receive_payload_from_software_manager(payload, session_id)
|
||||
|
||||
def receive_payload_from_session_manger(self, payload: Any, session: Session):
|
||||
"""
|
||||
Receive a payload from the SessionManager and forward it to the corresponding service or application.
|
||||
|
||||
:param payload: The payload being received.
|
||||
:param session: The transport session the payload originates from.
|
||||
"""
|
||||
# receiver: Optional[Union[Service, Application]] = self.port_protocol_mapping.get((port, protocol), None)
|
||||
# if receiver:
|
||||
# receiver.receive_payload(None, payload)
|
||||
# else:
|
||||
# raise ValueError(f"No service or application found for port {port} and protocol {protocol}")
|
||||
pass
|
||||
103
src/primaite/simulator/system/core/sys_log.py
Normal file
103
src/primaite/simulator/system/core/sys_log.py
Normal file
@@ -0,0 +1,103 @@
|
||||
import logging
|
||||
from pathlib import Path
|
||||
|
||||
from primaite.simulator import TEMP_SIM_OUTPUT
|
||||
|
||||
|
||||
class _NotJSONFilter(logging.Filter):
|
||||
def filter(self, record: logging.LogRecord) -> bool:
|
||||
"""
|
||||
Determines if a log message does not start and end with '{' and '}' (i.e., it is not a JSON-like message).
|
||||
|
||||
:param record: LogRecord object containing all the information pertinent to the event being logged.
|
||||
:return: True if log message is not JSON-like, False otherwise.
|
||||
"""
|
||||
return not record.getMessage().startswith("{") and not record.getMessage().endswith("}")
|
||||
|
||||
|
||||
class SysLog:
|
||||
"""
|
||||
A SysLog class is a simple logger dedicated to managing and writing system logs for a Node.
|
||||
|
||||
Each log message is written to a file located at: <simulation output directory>/<hostname>/<hostname>_sys.log
|
||||
"""
|
||||
|
||||
def __init__(self, hostname: str):
|
||||
"""
|
||||
Constructs a SysLog instance for a given hostname.
|
||||
|
||||
:param hostname: The hostname associated with the system logs being recorded.
|
||||
"""
|
||||
self.hostname = hostname
|
||||
self._setup_logger()
|
||||
|
||||
def _setup_logger(self):
|
||||
"""
|
||||
Configures the logger for this SysLog instance.
|
||||
|
||||
The logger is set to the DEBUG level, and is equipped with a handler that writes to a file and filters out
|
||||
JSON-like messages.
|
||||
"""
|
||||
log_path = self._get_log_path()
|
||||
|
||||
file_handler = logging.FileHandler(filename=log_path)
|
||||
file_handler.setLevel(logging.DEBUG)
|
||||
|
||||
log_format = "%(asctime)s %(levelname)s: %(message)s"
|
||||
file_handler.setFormatter(logging.Formatter(log_format))
|
||||
|
||||
self.logger = logging.getLogger(f"{self.hostname}_sys_log")
|
||||
self.logger.setLevel(logging.DEBUG)
|
||||
self.logger.addHandler(file_handler)
|
||||
|
||||
self.logger.addFilter(_NotJSONFilter())
|
||||
|
||||
def _get_log_path(self) -> Path:
|
||||
"""
|
||||
Constructs the path for the log file based on the hostname.
|
||||
|
||||
:return: Path object representing the location of the log file.
|
||||
"""
|
||||
root = TEMP_SIM_OUTPUT / self.hostname
|
||||
root.mkdir(exist_ok=True, parents=True)
|
||||
return root / f"{self.hostname}_sys.log"
|
||||
|
||||
def debug(self, msg: str):
|
||||
"""
|
||||
Logs a message with the DEBUG level.
|
||||
|
||||
:param msg: The message to be logged.
|
||||
"""
|
||||
self.logger.debug(msg)
|
||||
|
||||
def info(self, msg: str):
|
||||
"""
|
||||
Logs a message with the INFO level.
|
||||
|
||||
:param msg: The message to be logged.
|
||||
"""
|
||||
self.logger.info(msg)
|
||||
|
||||
def warning(self, msg: str):
|
||||
"""
|
||||
Logs a message with the WARNING level.
|
||||
|
||||
:param msg: The message to be logged.
|
||||
"""
|
||||
self.logger.warning(msg)
|
||||
|
||||
def error(self, msg: str):
|
||||
"""
|
||||
Logs a message with the ERROR level.
|
||||
|
||||
:param msg: The message to be logged.
|
||||
"""
|
||||
self.logger.error(msg)
|
||||
|
||||
def critical(self, msg: str):
|
||||
"""
|
||||
Logs a message with the CRITICAL level.
|
||||
|
||||
:param msg: The message to be logged.
|
||||
"""
|
||||
self.logger.critical(msg)
|
||||
0
src/primaite/simulator/system/processes/__init__.py
Normal file
0
src/primaite/simulator/system/processes/__init__.py
Normal file
39
src/primaite/simulator/system/processes/process.py
Normal file
39
src/primaite/simulator/system/processes/process.py
Normal file
@@ -0,0 +1,39 @@
|
||||
from abc import abstractmethod
|
||||
from enum import Enum
|
||||
from typing import Dict
|
||||
|
||||
from primaite.simulator.system.software import Software
|
||||
|
||||
|
||||
class ProcessOperatingState(Enum):
|
||||
"""Enumeration of Process Operating States."""
|
||||
|
||||
RUNNING = 1
|
||||
"The process is running."
|
||||
PAUSED = 2
|
||||
"The process is temporarily paused."
|
||||
|
||||
|
||||
class Process(Software):
|
||||
"""
|
||||
Represents a Process, a program in execution, in the simulation environment.
|
||||
|
||||
Processes are executed by a Node and do not have the ability to performing input/output operations.
|
||||
"""
|
||||
|
||||
operating_state: ProcessOperatingState
|
||||
"The current operating state of the Process."
|
||||
|
||||
@abstractmethod
|
||||
def describe_state(self) -> Dict:
|
||||
"""
|
||||
Produce a dictionary describing the current state of this object.
|
||||
|
||||
Please see :py:meth:`primaite.simulator.core.SimComponent.describe_state` for a more detailed explanation.
|
||||
|
||||
:return: Current state of this object and child objects.
|
||||
:rtype: Dict
|
||||
"""
|
||||
state = super().describe_state()
|
||||
state.update({"operating_state": self.operating_state.name})
|
||||
return state
|
||||
0
src/primaite/simulator/system/services/__init__.py
Normal file
0
src/primaite/simulator/system/services/__init__.py
Normal file
88
src/primaite/simulator/system/services/service.py
Normal file
88
src/primaite/simulator/system/services/service.py
Normal file
@@ -0,0 +1,88 @@
|
||||
from abc import abstractmethod
|
||||
from enum import Enum
|
||||
from typing import Any, Dict, List
|
||||
|
||||
from primaite.simulator.system.software import IOSoftware
|
||||
|
||||
|
||||
class ServiceOperatingState(Enum):
|
||||
"""Enumeration of Service Operating States."""
|
||||
|
||||
RUNNING = 1
|
||||
"The service is currently running."
|
||||
STOPPED = 2
|
||||
"The service is not running."
|
||||
INSTALLING = 3
|
||||
"The service is being installed or updated."
|
||||
RESTARTING = 4
|
||||
"The service is in the process of restarting."
|
||||
PAUSED = 5
|
||||
"The service is temporarily paused."
|
||||
DISABLED = 6
|
||||
"The service is disabled and cannot be started."
|
||||
|
||||
|
||||
class Service(IOSoftware):
|
||||
"""
|
||||
Represents a Service in the simulation environment.
|
||||
|
||||
Services are programs that run in the background and may perform input/output operations.
|
||||
"""
|
||||
|
||||
operating_state: ServiceOperatingState
|
||||
"The current operating state of the Service."
|
||||
|
||||
@abstractmethod
|
||||
def describe_state(self) -> Dict:
|
||||
"""
|
||||
Produce a dictionary describing the current state of this object.
|
||||
|
||||
Please see :py:meth:`primaite.simulator.core.SimComponent.describe_state` for a more detailed explanation.
|
||||
|
||||
:return: Current state of this object and child objects.
|
||||
:rtype: Dict
|
||||
"""
|
||||
state = super().describe_state()
|
||||
state.update({"operating_state": self.operating_state.name})
|
||||
return state
|
||||
|
||||
def apply_action(self, action: List[str]) -> None:
|
||||
"""
|
||||
Applies a list of actions to the Service.
|
||||
|
||||
:param action: A list of actions to apply.
|
||||
"""
|
||||
pass
|
||||
|
||||
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 send(self, payload: Any, session_id: str, **kwargs) -> bool:
|
||||
"""
|
||||
Sends a payload to 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 send.
|
||||
:return: True if successful, False otherwise.
|
||||
"""
|
||||
pass
|
||||
|
||||
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.
|
||||
:return: True if successful, False otherwise.
|
||||
"""
|
||||
pass
|
||||
193
src/primaite/simulator/system/software.py
Normal file
193
src/primaite/simulator/system/software.py
Normal file
@@ -0,0 +1,193 @@
|
||||
from abc import abstractmethod
|
||||
from enum import Enum
|
||||
from typing import Any, Dict, List, Set
|
||||
|
||||
from primaite.simulator.core import SimComponent
|
||||
from primaite.simulator.network.transmission.transport_layer import Port
|
||||
|
||||
|
||||
class SoftwareType(Enum):
|
||||
"""
|
||||
An enumeration representing the different types of software within a simulated environment.
|
||||
|
||||
Members:
|
||||
- APPLICATION: User-facing programs that may perform input/output operations.
|
||||
- SERVICE: Represents programs that run in the background and may perform input/output operations.
|
||||
- PROCESS: Software executed by a Node that does not have the ability to performing input/output operations.
|
||||
"""
|
||||
|
||||
APPLICATION = 1
|
||||
"User-facing software that may perform input/output operations."
|
||||
SERVICE = 2
|
||||
"Software that runs in the background and may perform input/output operations."
|
||||
PROCESS = 3
|
||||
"Software executed by a Node that does not have the ability to performing input/output operations."
|
||||
|
||||
|
||||
class SoftwareHealthState(Enum):
|
||||
"""Enumeration of the Software Health States."""
|
||||
|
||||
GOOD = 1
|
||||
"The software is in a good and healthy condition."
|
||||
COMPROMISED = 2
|
||||
"The software's security has been compromised."
|
||||
OVERWHELMED = 3
|
||||
"he software is overwhelmed and not functioning properly."
|
||||
PATCHING = 4
|
||||
"The software is undergoing patching or updates."
|
||||
|
||||
|
||||
class SoftwareCriticality(Enum):
|
||||
"""Enumeration of Software Criticality Levels."""
|
||||
|
||||
LOWEST = 1
|
||||
"The lowest level of criticality."
|
||||
LOW = 2
|
||||
"A low level of criticality."
|
||||
MEDIUM = 3
|
||||
"A medium level of criticality."
|
||||
HIGH = 4
|
||||
"A high level of criticality."
|
||||
HIGHEST = 5
|
||||
"The highest level of criticality."
|
||||
|
||||
|
||||
class Software(SimComponent):
|
||||
"""
|
||||
A base class representing software in a simulator environment.
|
||||
|
||||
This class is intended to be subclassed by specific types of software entities.
|
||||
It outlines the fundamental attributes and behaviors expected of any software in the simulation.
|
||||
"""
|
||||
|
||||
name: str
|
||||
"The name of the software."
|
||||
health_state_actual: SoftwareHealthState
|
||||
"The actual health state of the software."
|
||||
health_state_visible: SoftwareHealthState
|
||||
"The health state of the software visible to the red agent."
|
||||
criticality: SoftwareCriticality
|
||||
"The criticality level of the software."
|
||||
patching_count: int = 0
|
||||
"The count of patches applied to the software, defaults to 0."
|
||||
scanning_count: int = 0
|
||||
"The count of times the software has been scanned, defaults to 0."
|
||||
revealed_to_red: bool = False
|
||||
"Indicates if the software has been revealed to red agent, defaults is False."
|
||||
|
||||
@abstractmethod
|
||||
def describe_state(self) -> Dict:
|
||||
"""
|
||||
Produce a dictionary describing the current state of this object.
|
||||
|
||||
Please see :py:meth:`primaite.simulator.core.SimComponent.describe_state` for a more detailed explanation.
|
||||
|
||||
:return: Current state of this object and child objects.
|
||||
:rtype: Dict
|
||||
"""
|
||||
state = super().describe_state()
|
||||
state.update(
|
||||
{
|
||||
"health_state": self.health_state_actual.name,
|
||||
"health_state_red_view": self.health_state_visible.name,
|
||||
"criticality": self.criticality.name,
|
||||
"patching_count": self.patching_count,
|
||||
"scanning_count": self.scanning_count,
|
||||
"revealed_to_red": self.revealed_to_red,
|
||||
}
|
||||
)
|
||||
return state
|
||||
|
||||
def apply_action(self, action: List[str]) -> None:
|
||||
"""
|
||||
Applies a list of actions to the software.
|
||||
|
||||
The specifics of how these actions are applied should be implemented in subclasses.
|
||||
|
||||
:param action: A list of actions to apply.
|
||||
:type action: List[str]
|
||||
"""
|
||||
pass
|
||||
|
||||
def reset_component_for_episode(self, episode: int):
|
||||
"""
|
||||
Resets the software component for a new episode.
|
||||
|
||||
This method should ensure the software is ready for a new episode, including resetting any
|
||||
stateful properties or statistics, and clearing any message queues. The specifics of what constitutes a
|
||||
"reset" should be implemented in subclasses.
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
class IOSoftware(Software):
|
||||
"""
|
||||
Represents software in a simulator environment that is capable of input/output operations.
|
||||
|
||||
This base class is meant to be sub-classed by Application and Service classes. It provides the blueprint for
|
||||
Applications and Services that can receive payloads from a Node's SessionManager (corresponding to layer 5 in the
|
||||
OSI Model), process them according to their internals, and send a response payload back to the SessionManager if
|
||||
required.
|
||||
"""
|
||||
|
||||
installing_count: int = 0
|
||||
"The number of times the software has been installed. Default is 0."
|
||||
max_sessions: int = 1
|
||||
"The maximum number of sessions that the software can handle simultaneously. Default is 0."
|
||||
tcp: bool = True
|
||||
"Indicates if the software uses TCP protocol for communication. Default is True."
|
||||
udp: bool = True
|
||||
"Indicates if the software uses UDP protocol for communication. Default is True."
|
||||
ports: Set[Port]
|
||||
"The set of ports to which the software is connected."
|
||||
|
||||
@abstractmethod
|
||||
def describe_state(self) -> Dict:
|
||||
"""
|
||||
Produce a dictionary describing the current state of this object.
|
||||
|
||||
Please see :py:meth:`primaite.simulator.core.SimComponent.describe_state` for a more detailed explanation.
|
||||
|
||||
:return: Current state of this object and child objects.
|
||||
:rtype: Dict
|
||||
"""
|
||||
state = super().describe_state()
|
||||
state.update(
|
||||
{
|
||||
"installing_count": self.installing_count,
|
||||
"max_sessions": self.max_sessions,
|
||||
"tcp": self.tcp,
|
||||
"udp": self.udp,
|
||||
"ports": [port.name for port in self.ports], # TODO: not sure if this should be port.name or port.value
|
||||
}
|
||||
)
|
||||
return state
|
||||
|
||||
def send(self, payload: Any, session_id: str, **kwargs) -> bool:
|
||||
"""
|
||||
Sends a payload to 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 send.
|
||||
: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 sent, False otherwise.
|
||||
"""
|
||||
pass
|
||||
|
||||
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.
|
||||
"""
|
||||
pass
|
||||
@@ -112,7 +112,7 @@ def temp_session_path() -> Path:
|
||||
session_timestamp = datetime.now()
|
||||
date_dir = session_timestamp.strftime("%Y-%m-%d")
|
||||
session_path = session_timestamp.strftime("%Y-%m-%d_%H-%M-%S")
|
||||
session_path = Path(tempfile.gettempdir()) / "primaite" / date_dir / session_path
|
||||
session_path = Path(tempfile.gettempdir()) / "_primaite" / date_dir / session_path
|
||||
session_path.mkdir(exist_ok=True, parents=True)
|
||||
|
||||
return session_path
|
||||
|
||||
@@ -0,0 +1,190 @@
|
||||
from enum import Enum
|
||||
from typing import Dict, List, Literal
|
||||
|
||||
import pytest
|
||||
|
||||
from primaite.simulator.core import Action, ActionManager, AllowAllValidator, SimComponent
|
||||
from primaite.simulator.domain.controller import AccountGroup, GroupMembershipValidator
|
||||
|
||||
|
||||
def test_group_action_validation() -> None:
|
||||
"""
|
||||
Check that actions are denied when an unauthorised request is made.
|
||||
|
||||
This test checks the integration between SimComponent and the permissions validation system. First, we create a
|
||||
basic node and folder class. We configure the node so that only admins can create a folder. Then, we try to create
|
||||
a folder as both an admin user and a non-admin user.
|
||||
"""
|
||||
|
||||
class Folder(SimComponent):
|
||||
name: str
|
||||
|
||||
def describe_state(self) -> Dict:
|
||||
return super().describe_state()
|
||||
|
||||
class Node(SimComponent):
|
||||
name: str
|
||||
folders: List[Folder] = []
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
self.action_manager = ActionManager()
|
||||
|
||||
self.action_manager.add_action(
|
||||
"create_folder",
|
||||
Action(
|
||||
func=lambda request, context: self.create_folder(request[0]),
|
||||
validator=GroupMembershipValidator([AccountGroup.LOCAL_ADMIN, AccountGroup.DOMAIN_ADMIN]),
|
||||
),
|
||||
)
|
||||
|
||||
def describe_state(self) -> Dict:
|
||||
return super().describe_state()
|
||||
|
||||
def create_folder(self, folder_name: str) -> None:
|
||||
new_folder = Folder(uuid="0000-0000-0001", name=folder_name)
|
||||
self.folders.append(new_folder)
|
||||
|
||||
def remove_folder(self, folder: Folder) -> None:
|
||||
self.folders = [x for x in self.folders if x is not folder]
|
||||
|
||||
# check that the folder is created when a local admin tried to do it
|
||||
permitted_context = {"request_source": {"agent": "BLUE", "account": "User1", "groups": ["LOCAL_ADMIN"]}}
|
||||
my_node = Node(uuid="0000-0000-1234", name="pc")
|
||||
my_node.apply_action(["create_folder", "memes"], context=permitted_context)
|
||||
assert len(my_node.folders) == 1
|
||||
assert my_node.folders[0].name == "memes"
|
||||
|
||||
# check that the number of folders is still 1 even after attempting to create a second one without permissions
|
||||
invalid_context = {"request_source": {"agent": "BLUE", "account": "User1", "groups": ["LOCAL_USER", "DOMAIN_USER"]}}
|
||||
my_node.apply_action(["create_folder", "memes2"], context=invalid_context)
|
||||
assert len(my_node.folders) == 1
|
||||
assert my_node.folders[0].name == "memes"
|
||||
|
||||
|
||||
def test_hierarchical_action_with_validation() -> None:
|
||||
"""
|
||||
Check that validation works with sub-objects.
|
||||
|
||||
This test creates a parent object (Node) and a child object (Application) which both accept actions. The node allows
|
||||
action passthrough to applications. The purpose of this test is to check that after an action is passed through to
|
||||
a child object, that the permission system still works as intended.
|
||||
"""
|
||||
|
||||
class Application(SimComponent):
|
||||
name: str
|
||||
state: Literal["on", "off", "disabled"] = "off"
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
self.action_manager = ActionManager()
|
||||
|
||||
self.action_manager.add_action(
|
||||
"turn_on",
|
||||
Action(
|
||||
func=lambda request, context: self.turn_on(),
|
||||
validator=AllowAllValidator(),
|
||||
),
|
||||
)
|
||||
self.action_manager.add_action(
|
||||
"turn_off",
|
||||
Action(
|
||||
func=lambda request, context: self.turn_off(),
|
||||
validator=AllowAllValidator(),
|
||||
),
|
||||
)
|
||||
self.action_manager.add_action(
|
||||
"disable",
|
||||
Action(
|
||||
func=lambda request, context: self.disable(),
|
||||
validator=GroupMembershipValidator([AccountGroup.LOCAL_ADMIN, AccountGroup.DOMAIN_ADMIN]),
|
||||
),
|
||||
)
|
||||
self.action_manager.add_action(
|
||||
"enable",
|
||||
Action(
|
||||
func=lambda request, context: self.enable(),
|
||||
validator=GroupMembershipValidator([AccountGroup.LOCAL_ADMIN, AccountGroup.DOMAIN_ADMIN]),
|
||||
),
|
||||
)
|
||||
|
||||
def describe_state(self) -> Dict:
|
||||
return super().describe_state()
|
||||
|
||||
def disable(self) -> None:
|
||||
self.state = "disabled"
|
||||
|
||||
def enable(self) -> None:
|
||||
if self.state == "disabled":
|
||||
self.state = "off"
|
||||
|
||||
def turn_on(self) -> None:
|
||||
if self.state == "off":
|
||||
self.state = "on"
|
||||
|
||||
def turn_off(self) -> None:
|
||||
if self.state == "on":
|
||||
self.state = "off"
|
||||
|
||||
class Node(SimComponent):
|
||||
name: str
|
||||
state: Literal["on", "off"] = "on"
|
||||
apps: List[Application] = []
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
self.action_manager = ActionManager()
|
||||
|
||||
self.action_manager.add_action(
|
||||
"apps",
|
||||
Action(
|
||||
func=lambda request, context: self.send_action_to_app(request.pop(0), request, context),
|
||||
validator=AllowAllValidator(),
|
||||
),
|
||||
)
|
||||
|
||||
def describe_state(self) -> Dict:
|
||||
return super().describe_state()
|
||||
|
||||
def install_app(self, app_name: str) -> None:
|
||||
new_app = Application(name=app_name)
|
||||
self.apps.append(new_app)
|
||||
|
||||
def send_action_to_app(self, app_name: str, options: List[str], context: Dict):
|
||||
for app in self.apps:
|
||||
if app_name == app.name:
|
||||
app.apply_action(options, context)
|
||||
break
|
||||
else:
|
||||
msg = f"Node has no app with name {app_name}"
|
||||
raise LookupError(msg)
|
||||
|
||||
my_node = Node(name="pc")
|
||||
my_node.install_app("Chrome")
|
||||
my_node.install_app("Firefox")
|
||||
|
||||
non_admin_context = {
|
||||
"request_source": {"agent": "BLUE", "account": "User1", "groups": ["LOCAL_USER", "DOMAIN_USER"]}
|
||||
}
|
||||
|
||||
admin_context = {
|
||||
"request_source": {
|
||||
"agent": "BLUE",
|
||||
"account": "User1",
|
||||
"groups": ["LOCAL_ADMIN", "DOMAIN_ADMIN", "LOCAL_USER", "DOMAIN_USER"],
|
||||
}
|
||||
}
|
||||
|
||||
# check that a non-admin can't disable this app
|
||||
my_node.apply_action(["apps", "Chrome", "disable"], non_admin_context)
|
||||
assert my_node.apps[0].name == "Chrome" # if failure occurs on this line, the test itself is broken
|
||||
assert my_node.apps[0].state == "off"
|
||||
|
||||
# check that a non-admin can turn this app on
|
||||
my_node.apply_action(["apps", "Firefox", "turn_on"], non_admin_context)
|
||||
assert my_node.apps[1].name == "Firefox" # if failure occurs on this line, the test itself is broken
|
||||
assert my_node.apps[1].state == "on"
|
||||
|
||||
# check that an admin can disable this app
|
||||
my_node.apply_action(["apps", "Chrome", "disable"], admin_context)
|
||||
assert my_node.apps[0].state == "disabled"
|
||||
86
tests/integration_tests/network/test_frame_transmission.py
Normal file
86
tests/integration_tests/network/test_frame_transmission.py
Normal file
@@ -0,0 +1,86 @@
|
||||
from primaite.simulator.network.hardware.base import Link, NIC, Node, Switch
|
||||
|
||||
|
||||
def test_node_to_node_ping():
|
||||
"""Tests two Nodes are able to ping each other."""
|
||||
# TODO Add actual checks. Manual check performed for now.
|
||||
node_a = Node(hostname="node_a")
|
||||
nic_a = NIC(ip_address="192.168.0.10", subnet_mask="255.255.255.0", gateway="192.168.0.1")
|
||||
node_a.connect_nic(nic_a)
|
||||
node_a.power_on()
|
||||
|
||||
node_b = Node(hostname="node_b")
|
||||
nic_b = NIC(ip_address="192.168.0.11", subnet_mask="255.255.255.0", gateway="192.168.0.1")
|
||||
node_b.connect_nic(nic_b)
|
||||
node_b.power_on()
|
||||
|
||||
Link(endpoint_a=nic_a, endpoint_b=nic_b)
|
||||
|
||||
assert node_a.ping("192.168.0.11")
|
||||
|
||||
|
||||
def test_multi_nic():
|
||||
"""Tests that Nodes with multiple NICs can ping each other and the data go across the correct links."""
|
||||
# TODO Add actual checks. Manual check performed for now.
|
||||
node_a = Node(hostname="node_a")
|
||||
nic_a = NIC(ip_address="192.168.0.10", subnet_mask="255.255.255.0", gateway="192.168.0.1")
|
||||
node_a.connect_nic(nic_a)
|
||||
node_a.power_on()
|
||||
|
||||
node_b = Node(hostname="node_b")
|
||||
nic_b1 = NIC(ip_address="192.168.0.11", subnet_mask="255.255.255.0", gateway="192.168.0.1")
|
||||
nic_b2 = NIC(ip_address="10.0.0.12", subnet_mask="255.0.0.0", gateway="10.0.0.1")
|
||||
node_b.connect_nic(nic_b1)
|
||||
node_b.connect_nic(nic_b2)
|
||||
node_b.power_on()
|
||||
|
||||
node_c = Node(hostname="node_c")
|
||||
nic_c = NIC(ip_address="10.0.0.13", subnet_mask="255.0.0.0", gateway="10.0.0.1")
|
||||
node_c.connect_nic(nic_c)
|
||||
node_c.power_on()
|
||||
|
||||
Link(endpoint_a=nic_a, endpoint_b=nic_b1)
|
||||
|
||||
Link(endpoint_a=nic_b2, endpoint_b=nic_c)
|
||||
|
||||
node_a.ping("192.168.0.11")
|
||||
|
||||
node_c.ping("10.0.0.12")
|
||||
|
||||
|
||||
def test_switched_network():
|
||||
"""Tests a larges network of Nodes and Switches with one node pinging another."""
|
||||
# TODO Add actual checks. Manual check performed for now.
|
||||
pc_a = Node(hostname="pc_a")
|
||||
nic_a = NIC(ip_address="192.168.0.10", subnet_mask="255.255.255.0", gateway="192.168.0.1")
|
||||
pc_a.connect_nic(nic_a)
|
||||
pc_a.power_on()
|
||||
|
||||
pc_b = Node(hostname="pc_b")
|
||||
nic_b = NIC(ip_address="192.168.0.11", subnet_mask="255.255.255.0", gateway="192.168.0.1")
|
||||
pc_b.connect_nic(nic_b)
|
||||
pc_b.power_on()
|
||||
|
||||
pc_c = Node(hostname="pc_c")
|
||||
nic_c = NIC(ip_address="192.168.0.12", subnet_mask="255.255.255.0", gateway="192.168.0.1")
|
||||
pc_c.connect_nic(nic_c)
|
||||
pc_c.power_on()
|
||||
|
||||
pc_d = Node(hostname="pc_d")
|
||||
nic_d = NIC(ip_address="192.168.0.13", subnet_mask="255.255.255.0", gateway="192.168.0.1")
|
||||
pc_d.connect_nic(nic_d)
|
||||
pc_d.power_on()
|
||||
|
||||
switch_1 = Switch(hostname="switch_1", num_ports=6)
|
||||
switch_1.power_on()
|
||||
|
||||
switch_2 = Switch(hostname="switch_2", num_ports=6)
|
||||
switch_2.power_on()
|
||||
|
||||
link_nic_a_switch_1 = Link(endpoint_a=nic_a, endpoint_b=switch_1.switch_ports[1])
|
||||
link_nic_b_switch_1 = Link(endpoint_a=nic_b, endpoint_b=switch_1.switch_ports[2])
|
||||
link_nic_c_switch_2 = Link(endpoint_a=nic_c, endpoint_b=switch_2.switch_ports[1])
|
||||
link_nic_d_switch_2 = Link(endpoint_a=nic_d, endpoint_b=switch_2.switch_ports[2])
|
||||
link_switch_1_switch_2 = Link(endpoint_a=switch_1.switch_ports[6], endpoint_b=switch_2.switch_ports[6])
|
||||
|
||||
pc_a.ping("192.168.0.13")
|
||||
21
tests/integration_tests/network/test_link_connection.py
Normal file
21
tests/integration_tests/network/test_link_connection.py
Normal file
@@ -0,0 +1,21 @@
|
||||
from primaite.simulator.network.hardware.base import Link, NIC, Node
|
||||
|
||||
|
||||
def test_link_up():
|
||||
"""Tests Nodes, NICs, and Links can all be connected and be in an enabled/up state."""
|
||||
node_a = Node(hostname="node_a")
|
||||
nic_a = NIC(ip_address="192.168.0.10", subnet_mask="255.255.255.0", gateway="192.168.0.1")
|
||||
node_a.connect_nic(nic_a)
|
||||
node_a.power_on()
|
||||
assert nic_a.enabled
|
||||
|
||||
node_b = Node(hostname="node_b")
|
||||
nic_b = NIC(ip_address="192.168.0.11", subnet_mask="255.255.255.0", gateway="192.168.0.1")
|
||||
node_b.connect_nic(nic_b)
|
||||
node_b.power_on()
|
||||
|
||||
assert nic_b.enabled
|
||||
|
||||
link = Link(endpoint_a=nic_a, endpoint_b=nic_b)
|
||||
|
||||
assert link.is_up
|
||||
97
tests/integration_tests/network/test_network_creation.py
Normal file
97
tests/integration_tests/network/test_network_creation.py
Normal file
@@ -0,0 +1,97 @@
|
||||
import pytest
|
||||
|
||||
from primaite.simulator.network.container import Network
|
||||
from primaite.simulator.network.hardware.base import NIC, Node
|
||||
|
||||
|
||||
def test_adding_removing_nodes():
|
||||
"""Check that we can create and add a node to a network."""
|
||||
net = Network()
|
||||
n1 = Node(hostname="computer")
|
||||
net.add_node(n1)
|
||||
assert n1.parent is net
|
||||
assert n1 in net
|
||||
|
||||
net.remove_node(n1)
|
||||
assert n1.parent is None
|
||||
assert n1 not in net
|
||||
|
||||
|
||||
def test_readding_node():
|
||||
"""Check that warning is raised when readding a node."""
|
||||
net = Network()
|
||||
n1 = Node(hostname="computer")
|
||||
net.add_node(n1)
|
||||
net.add_node(n1)
|
||||
assert n1.parent is net
|
||||
assert n1 in net
|
||||
|
||||
|
||||
def test_removing_nonexistent_node():
|
||||
"""Check that warning is raised when trying to remove a node that is not in the network."""
|
||||
net = Network()
|
||||
n1 = Node(hostname="computer")
|
||||
net.remove_node(n1)
|
||||
assert n1.parent is None
|
||||
assert n1 not in net
|
||||
|
||||
|
||||
def test_connecting_nodes():
|
||||
"""Check that two nodes on the network can be connected."""
|
||||
net = Network()
|
||||
n1 = Node(hostname="computer")
|
||||
n1_nic = NIC(ip_address="120.30.0.1", gateway="192.168.0.1", subnet_mask="255.255.255.0")
|
||||
n1.connect_nic(n1_nic)
|
||||
n2 = Node(hostname="server")
|
||||
n2_nic = NIC(ip_address="120.30.0.2", gateway="192.168.0.1", subnet_mask="255.255.255.0")
|
||||
n2.connect_nic(n2_nic)
|
||||
|
||||
net.add_node(n1)
|
||||
net.add_node(n2)
|
||||
|
||||
net.connect(n1.nics[n1_nic.uuid], n2.nics[n2_nic.uuid], bandwidth=30)
|
||||
|
||||
assert len(net.links) == 1
|
||||
link = list(net.links.values())[0]
|
||||
assert link in net
|
||||
assert link.parent is net
|
||||
|
||||
|
||||
def test_connecting_node_to_itself():
|
||||
net = Network()
|
||||
node = Node(hostname="computer")
|
||||
nic1 = NIC(ip_address="120.30.0.1", gateway="192.168.0.1", subnet_mask="255.255.255.0")
|
||||
node.connect_nic(nic1)
|
||||
nic2 = NIC(ip_address="120.30.0.2", gateway="192.168.0.1", subnet_mask="255.255.255.0")
|
||||
node.connect_nic(nic2)
|
||||
|
||||
net.add_node(node)
|
||||
|
||||
net.connect(node.nics[nic1.uuid], node.nics[nic2.uuid], bandwidth=30)
|
||||
|
||||
assert node in net
|
||||
assert nic1.connected_link is None
|
||||
assert nic2.connected_link is None
|
||||
assert len(net.links) == 0
|
||||
|
||||
|
||||
def test_disconnecting_nodes():
|
||||
net = Network()
|
||||
|
||||
n1 = Node(hostname="computer")
|
||||
n1_nic = NIC(ip_address="120.30.0.1", gateway="192.168.0.1", subnet_mask="255.255.255.0")
|
||||
n1.connect_nic(n1_nic)
|
||||
net.add_node(n1)
|
||||
|
||||
n2 = Node(hostname="server")
|
||||
n2_nic = NIC(ip_address="120.30.0.2", gateway="192.168.0.1", subnet_mask="255.255.255.0")
|
||||
n2.connect_nic(n2_nic)
|
||||
net.add_node(n2)
|
||||
|
||||
net.connect(n1.nics[n1_nic.uuid], n2.nics[n2_nic.uuid], bandwidth=30)
|
||||
assert len(net.links) == 1
|
||||
|
||||
link = list(net.links.values())[0]
|
||||
net.remove_link(link)
|
||||
assert link not in net
|
||||
assert len(net.links) == 0
|
||||
@@ -1,6 +1,6 @@
|
||||
import pytest
|
||||
|
||||
from primaite.simulator.network.transmission.physical_layer import Link, NIC
|
||||
from primaite.simulator.network.hardware.base import Link, NIC
|
||||
|
||||
|
||||
def test_link_fails_with_same_nic():
|
||||
|
||||
@@ -0,0 +1,18 @@
|
||||
"""Test the account module of the simulator."""
|
||||
from primaite.simulator.domain.account import Account, AccountType
|
||||
|
||||
|
||||
def test_account_serialise():
|
||||
"""Test that an account can be serialised. If pydantic throws error then this test fails."""
|
||||
acct = Account(username="Jake", password="JakePass1!", account_type=AccountType.USER)
|
||||
serialised = acct.model_dump_json()
|
||||
print(serialised)
|
||||
|
||||
|
||||
def test_account_deserialise():
|
||||
"""Test that an account can be deserialised. The test fails if pydantic throws an error."""
|
||||
acct_json = (
|
||||
'{"uuid":"dfb2bcaa-d3a1-48fd-af3f-c943354622b4","num_logons":0,"num_logoffs":0,"num_group_changes":0,'
|
||||
'"username":"Jake","password":"JakePass1!","account_type":2,"status":2,"action_manager":null}'
|
||||
)
|
||||
acct = Account.model_validate_json(acct_json)
|
||||
@@ -0,0 +1,134 @@
|
||||
from primaite.simulator.file_system.file_system import FileSystem
|
||||
from primaite.simulator.file_system.file_system_file import FileSystemFile
|
||||
from primaite.simulator.file_system.file_system_folder import FileSystemFolder
|
||||
|
||||
|
||||
def test_create_folder_and_file():
|
||||
"""Test creating a folder and a file."""
|
||||
file_system = FileSystem()
|
||||
folder = file_system.create_folder(folder_name="test_folder")
|
||||
assert len(file_system.folders) is 1
|
||||
|
||||
file = file_system.create_file(file_name="test_file", size=10, folder_uuid=folder.uuid)
|
||||
assert len(file_system.get_folder_by_id(folder.uuid).files) is 1
|
||||
|
||||
assert file_system.get_file_by_id(file.uuid).name is "test_file"
|
||||
assert file_system.get_file_by_id(file.uuid).size == 10
|
||||
|
||||
|
||||
def test_create_file():
|
||||
"""Tests that creating a file without a folder creates a folder and sets that as the file's parent."""
|
||||
file_system = FileSystem()
|
||||
|
||||
file = file_system.create_file(file_name="test_file", size=10)
|
||||
assert len(file_system.folders) is 1
|
||||
assert file_system.get_folder_by_name("root").get_file_by_id(file.uuid) is file
|
||||
|
||||
|
||||
def test_delete_file():
|
||||
"""Tests that a file can be deleted."""
|
||||
file_system = FileSystem()
|
||||
|
||||
file = file_system.create_file(file_name="test_file", size=10)
|
||||
assert len(file_system.folders) is 1
|
||||
|
||||
folder_id = list(file_system.folders.keys())[0]
|
||||
folder = file_system.get_folder_by_id(folder_id)
|
||||
assert folder.get_file_by_id(file.uuid) is file
|
||||
|
||||
file_system.delete_file(file=file)
|
||||
assert len(file_system.folders) is 1
|
||||
assert len(folder.files) is 0
|
||||
|
||||
|
||||
def test_delete_non_existent_file():
|
||||
"""Tests deleting a non existent file."""
|
||||
file_system = FileSystem()
|
||||
|
||||
file = file_system.create_file(file_name="test_file", size=10)
|
||||
not_added_file = FileSystemFile(name="not_added")
|
||||
# folder should be created
|
||||
assert len(file_system.folders) is 1
|
||||
# should only have 1 file in the file system
|
||||
folder_id = list(file_system.folders.keys())[0]
|
||||
folder = file_system.get_folder_by_id(folder_id)
|
||||
assert len(list(folder.files)) is 1
|
||||
|
||||
assert folder.get_file_by_id(file.uuid) is file
|
||||
|
||||
# deleting should not change how many files are in folder
|
||||
file_system.delete_file(file=not_added_file)
|
||||
assert len(file_system.folders) is 1
|
||||
assert len(list(folder.files)) is 1
|
||||
|
||||
|
||||
def test_delete_folder():
|
||||
file_system = FileSystem()
|
||||
folder = file_system.create_folder(folder_name="test_folder")
|
||||
assert len(file_system.folders) is 1
|
||||
|
||||
file_system.delete_folder(folder)
|
||||
assert len(file_system.folders) is 0
|
||||
|
||||
|
||||
def test_deleting_a_non_existent_folder():
|
||||
file_system = FileSystem()
|
||||
folder = file_system.create_folder(folder_name="test_folder")
|
||||
not_added_folder = FileSystemFolder(name="fake_folder")
|
||||
assert len(file_system.folders) is 1
|
||||
|
||||
file_system.delete_folder(not_added_folder)
|
||||
assert len(file_system.folders) is 1
|
||||
|
||||
|
||||
def test_move_file():
|
||||
"""Tests the file move function."""
|
||||
file_system = FileSystem()
|
||||
src_folder = file_system.create_folder(folder_name="test_folder_1")
|
||||
assert len(file_system.folders) is 1
|
||||
|
||||
target_folder = file_system.create_folder(folder_name="test_folder_2")
|
||||
assert len(file_system.folders) is 2
|
||||
|
||||
file = file_system.create_file(file_name="test_file", size=10, folder_uuid=src_folder.uuid)
|
||||
assert len(file_system.get_folder_by_id(src_folder.uuid).files) is 1
|
||||
assert len(file_system.get_folder_by_id(target_folder.uuid).files) is 0
|
||||
|
||||
file_system.move_file(file=file, src_folder=src_folder, target_folder=target_folder)
|
||||
|
||||
assert len(file_system.get_folder_by_id(src_folder.uuid).files) is 0
|
||||
assert len(file_system.get_folder_by_id(target_folder.uuid).files) is 1
|
||||
|
||||
|
||||
def test_copy_file():
|
||||
"""Tests the file copy function."""
|
||||
file_system = FileSystem()
|
||||
src_folder = file_system.create_folder(folder_name="test_folder_1")
|
||||
assert len(file_system.folders) is 1
|
||||
|
||||
target_folder = file_system.create_folder(folder_name="test_folder_2")
|
||||
assert len(file_system.folders) is 2
|
||||
|
||||
file = file_system.create_file(file_name="test_file", size=10, folder_uuid=src_folder.uuid)
|
||||
assert len(file_system.get_folder_by_id(src_folder.uuid).files) is 1
|
||||
assert len(file_system.get_folder_by_id(target_folder.uuid).files) is 0
|
||||
|
||||
file_system.copy_file(file=file, src_folder=src_folder, target_folder=target_folder)
|
||||
|
||||
assert len(file_system.get_folder_by_id(src_folder.uuid).files) is 1
|
||||
assert len(file_system.get_folder_by_id(target_folder.uuid).files) is 1
|
||||
|
||||
|
||||
def test_serialisation():
|
||||
"""Test to check that the object serialisation works correctly."""
|
||||
file_system = FileSystem()
|
||||
folder = file_system.create_folder(folder_name="test_folder")
|
||||
assert len(file_system.folders) is 1
|
||||
|
||||
file_system.create_file(file_name="test_file", size=10, folder_uuid=folder.uuid)
|
||||
assert file_system.get_folder_by_id(folder.uuid) is folder
|
||||
|
||||
serialised_file_sys = file_system.model_dump_json()
|
||||
deserialised_file_sys = FileSystem.model_validate_json(serialised_file_sys)
|
||||
|
||||
assert file_system.model_dump_json() == deserialised_file_sys.model_dump_json()
|
||||
@@ -0,0 +1,23 @@
|
||||
from primaite.simulator.file_system.file_system_file import FileSystemFile
|
||||
from primaite.simulator.file_system.file_system_file_type import FileSystemFileType
|
||||
|
||||
|
||||
def test_file_type():
|
||||
"""Tests tha the FileSystemFile type is set correctly."""
|
||||
file = FileSystemFile(name="test", file_type=FileSystemFileType.DOC)
|
||||
assert file.file_type is FileSystemFileType.DOC
|
||||
|
||||
|
||||
def test_get_size():
|
||||
"""Tests that the file size is being returned properly."""
|
||||
file = FileSystemFile(name="test", size=1.5)
|
||||
assert file.size == 1.5
|
||||
|
||||
|
||||
def test_serialisation():
|
||||
"""Test to check that the object serialisation works correctly."""
|
||||
file = FileSystemFile(name="test", size=1.5, file_type=FileSystemFileType.DOC)
|
||||
serialised_file = file.model_dump_json()
|
||||
deserialised_file = FileSystemFile.model_validate_json(serialised_file)
|
||||
|
||||
assert file.model_dump_json() == deserialised_file.model_dump_json()
|
||||
@@ -0,0 +1,75 @@
|
||||
from primaite.simulator.file_system.file_system_file import FileSystemFile
|
||||
from primaite.simulator.file_system.file_system_file_type import FileSystemFileType
|
||||
from primaite.simulator.file_system.file_system_folder import FileSystemFolder
|
||||
|
||||
|
||||
def test_adding_removing_file():
|
||||
"""Test the adding and removing of a file from a folder."""
|
||||
folder = FileSystemFolder(name="test")
|
||||
|
||||
file = FileSystemFile(name="test_file", size=10, file_type=FileSystemFileType.DOC)
|
||||
|
||||
folder.add_file(file)
|
||||
assert folder.size == 10
|
||||
assert len(folder.files) is 1
|
||||
|
||||
folder.remove_file(file)
|
||||
assert folder.size == 0
|
||||
assert len(folder.files) is 0
|
||||
|
||||
|
||||
def test_remove_non_existent_file():
|
||||
"""Test the removing of a file that does not exist."""
|
||||
folder = FileSystemFolder(name="test")
|
||||
|
||||
file = FileSystemFile(name="test_file", size=10, file_type=FileSystemFileType.DOC)
|
||||
not_added_file = FileSystemFile(name="fake_file", size=10, file_type=FileSystemFileType.DOC)
|
||||
|
||||
folder.add_file(file)
|
||||
assert folder.size == 10
|
||||
assert len(folder.files) is 1
|
||||
|
||||
folder.remove_file(not_added_file)
|
||||
assert folder.size == 10
|
||||
assert len(folder.files) is 1
|
||||
|
||||
|
||||
def test_get_file_by_id():
|
||||
"""Test to make sure that the correct file is returned."""
|
||||
folder = FileSystemFolder(name="test")
|
||||
|
||||
file = FileSystemFile(name="test_file", size=10, file_type=FileSystemFileType.DOC)
|
||||
file2 = FileSystemFile(name="test_file_2", size=10, file_type=FileSystemFileType.DOC)
|
||||
|
||||
folder.add_file(file)
|
||||
folder.add_file(file2)
|
||||
assert folder.size == 20
|
||||
assert len(folder.files) is 2
|
||||
|
||||
assert folder.get_file_by_id(file_id=file.uuid) is file
|
||||
|
||||
|
||||
def test_folder_quarantine_state():
|
||||
"""Tests the changing of folder quarantine status."""
|
||||
folder = FileSystemFolder(name="test")
|
||||
|
||||
assert folder.quarantine_status() is False
|
||||
|
||||
folder.quarantine()
|
||||
assert folder.quarantine_status() is True
|
||||
|
||||
folder.end_quarantine()
|
||||
assert folder.quarantine_status() is False
|
||||
|
||||
|
||||
def test_serialisation():
|
||||
"""Test to check that the object serialisation works correctly."""
|
||||
folder = FileSystemFolder(name="test")
|
||||
file = FileSystemFile(name="test_file", size=10, file_type=FileSystemFileType.DOC)
|
||||
folder.add_file(file)
|
||||
|
||||
serialised_folder = folder.model_dump_json()
|
||||
|
||||
deserialised_folder = FileSystemFolder.model_validate_json(serialised_folder)
|
||||
|
||||
assert folder.model_dump_json() == deserialised_folder.model_dump_json()
|
||||
@@ -3,7 +3,7 @@ from ipaddress import IPv4Address
|
||||
|
||||
import pytest
|
||||
|
||||
from primaite.simulator.network.transmission.physical_layer import generate_mac_address, NIC
|
||||
from primaite.simulator.network.hardware.base import generate_mac_address, NIC
|
||||
|
||||
|
||||
def test_mac_address_generation():
|
||||
@@ -0,0 +1,10 @@
|
||||
import re
|
||||
from ipaddress import IPv4Address
|
||||
|
||||
import pytest
|
||||
|
||||
from primaite.simulator.network.hardware.base import Node
|
||||
|
||||
|
||||
def test_node_creation():
|
||||
node = Node(hostname="host_1")
|
||||
@@ -1,7 +1,7 @@
|
||||
import pytest
|
||||
|
||||
from primaite.simulator.network.transmission.data_link_layer import EthernetHeader, Frame
|
||||
from primaite.simulator.network.transmission.network_layer import ICMPHeader, IPPacket, IPProtocol, Precedence
|
||||
from primaite.simulator.network.transmission.network_layer import ICMPPacket, IPPacket, IPProtocol, Precedence
|
||||
from primaite.simulator.network.transmission.primaite_layer import AgentSource, DataStatus
|
||||
from primaite.simulator.network.transmission.transport_layer import Port, TCPFlags, TCPHeader, UDPHeader
|
||||
|
||||
@@ -76,7 +76,7 @@ def test_icmp_frame_creation():
|
||||
frame = Frame(
|
||||
ethernet=EthernetHeader(src_mac_addr="aa:bb:cc:dd:ee:ff", dst_mac_addr="11:22:33:44:55:66"),
|
||||
ip=IPPacket(src_ip="192.168.0.10", dst_ip="192.168.0.20", protocol=IPProtocol.ICMP),
|
||||
icmp=ICMPHeader(),
|
||||
icmp=ICMPPacket(),
|
||||
)
|
||||
assert frame
|
||||
|
||||
|
||||
@@ -1,24 +1,24 @@
|
||||
import pytest
|
||||
|
||||
from primaite.simulator.network.transmission.network_layer import ICMPHeader, ICMPType
|
||||
from primaite.simulator.network.transmission.network_layer import ICMPPacket, ICMPType
|
||||
|
||||
|
||||
def test_icmp_minimal_header_creation():
|
||||
"""Checks the minimal ICMPHeader (ping 1 request) creation using default values."""
|
||||
ping = ICMPHeader()
|
||||
"""Checks the minimal ICMPPacket (ping 1 request) creation using default values."""
|
||||
ping = ICMPPacket()
|
||||
|
||||
assert ping.icmp_type == ICMPType.ECHO_REQUEST
|
||||
assert ping.icmp_code == 0
|
||||
assert ping.identifier
|
||||
assert ping.sequence == 1
|
||||
assert ping.sequence == 0
|
||||
|
||||
|
||||
def test_valid_icmp_type_code_pairing():
|
||||
"""Tests ICMPHeader creation with valid type and code pairing."""
|
||||
assert ICMPHeader(icmp_type=ICMPType.DESTINATION_UNREACHABLE, icmp_code=6)
|
||||
"""Tests ICMPPacket creation with valid type and code pairing."""
|
||||
assert ICMPPacket(icmp_type=ICMPType.DESTINATION_UNREACHABLE, icmp_code=6)
|
||||
|
||||
|
||||
def test_invalid_icmp_type_code_pairing():
|
||||
"""Tests ICMPHeader creation fails with invalid type and code pairing."""
|
||||
"""Tests ICMPPacket creation fails with invalid type and code pairing."""
|
||||
with pytest.raises(ValueError):
|
||||
assert ICMPHeader(icmp_type=ICMPType.DESTINATION_UNREACHABLE, icmp_code=16)
|
||||
assert ICMPPacket(icmp_type=ICMPType.DESTINATION_UNREACHABLE, icmp_code=16)
|
||||
|
||||
@@ -0,0 +1,17 @@
|
||||
import json
|
||||
|
||||
from primaite.simulator.network.container import Network
|
||||
|
||||
|
||||
def test_creating_container():
|
||||
"""Check that we can create a network container"""
|
||||
net = Network()
|
||||
assert net.nodes == {}
|
||||
assert net.links == {}
|
||||
|
||||
|
||||
def test_describe_state():
|
||||
"""Check that we can describe network state without raising errors, and that the result is JSON serialisable."""
|
||||
net = Network()
|
||||
state = net.describe_state()
|
||||
json.dumps(state) # if this function call raises an error, the test fails, state was not JSON-serialisable
|
||||
@@ -42,38 +42,5 @@ class TestIsolatedSimComponent:
|
||||
return {}
|
||||
|
||||
comp = TestComponent(name="computer", size=(5, 10))
|
||||
dump = comp.model_dump()
|
||||
assert dump == {"name": "computer", "size": (5, 10)}
|
||||
|
||||
def test_apply_action(self):
|
||||
"""Validate that we can override apply_action behaviour and it updates the state of the component."""
|
||||
|
||||
class TestComponent(SimComponent):
|
||||
name: str
|
||||
status: Literal["on", "off"] = "off"
|
||||
|
||||
def describe_state(self) -> Dict:
|
||||
return {}
|
||||
|
||||
def _possible_actions(self) -> Dict[str, Callable[[List[str]], None]]:
|
||||
return {
|
||||
"turn_off": self._turn_off,
|
||||
"turn_on": self._turn_on,
|
||||
}
|
||||
|
||||
def _turn_off(self, options: List[str]) -> None:
|
||||
assert len(options) == 0, "This action does not support options."
|
||||
self.status = "off"
|
||||
|
||||
def _turn_on(self, options: List[str]) -> None:
|
||||
assert len(options) == 0, "This action does not support options."
|
||||
self.status = "on"
|
||||
|
||||
comp = TestComponent(name="computer", status="off")
|
||||
|
||||
assert comp.status == "off"
|
||||
comp.apply_action(["turn_on"])
|
||||
assert comp.status == "on"
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
comp.apply_action(["do_nothing"])
|
||||
dump = comp.model_dump_json()
|
||||
assert comp == TestComponent.model_validate_json(dump)
|
||||
|
||||
16
tests/unit_tests/_primaite/_simulator/test_sim_conatiner.py
Normal file
16
tests/unit_tests/_primaite/_simulator/test_sim_conatiner.py
Normal file
@@ -0,0 +1,16 @@
|
||||
from primaite.simulator.sim_container import Simulation
|
||||
|
||||
|
||||
def test_creating_empty_simulation():
|
||||
"""Check that no errors occur when trying to setup a simulation without providing parameters"""
|
||||
empty_sim = Simulation()
|
||||
|
||||
|
||||
def test_empty_sim_state():
|
||||
"""Check that describe_state has the right subcomponents."""
|
||||
empty_sim = Simulation()
|
||||
sim_state = empty_sim.describe_state()
|
||||
network_state = empty_sim.network.describe_state()
|
||||
domain_state = empty_sim.domain.describe_state()
|
||||
assert sim_state["network"] == network_state
|
||||
assert sim_state["domain"] == domain_state
|
||||
Reference in New Issue
Block a user