diff --git a/.azure/azure-ci-build-pipeline.yaml b/.azure/azure-ci-build-pipeline.yaml index 0bb03594..9070270a 100644 --- a/.azure/azure-ci-build-pipeline.yaml +++ b/.azure/azure-ci-build-pipeline.yaml @@ -86,5 +86,5 @@ stages: displayName: 'Perform PrimAITE Setup' - script: | - pytest -n 4 + pytest -n auto displayName: 'Run tests' diff --git a/.azuredevops/pull_request_template.md b/.azuredevops/pull_request_template.md index fd28ed57..f7533b37 100644 --- a/.azuredevops/pull_request_template.md +++ b/.azuredevops/pull_request_template.md @@ -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 diff --git a/.gitignore b/.gitignore index 60f5f54c..ff86b65f 100644 --- a/.gitignore +++ b/.gitignore @@ -150,3 +150,4 @@ src/primaite/outputs/ # benchmark session outputs benchmark/output +src/primaite/notebooks/scratch.ipynb diff --git a/CHANGELOG.md b/CHANGELOG.md index d66257b5..2b495c09 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/docs/_static/component_relationship.png b/docs/_static/component_relationship.png new file mode 100644 index 00000000..c2dd1102 Binary files /dev/null and b/docs/_static/component_relationship.png differ diff --git a/docs/_static/four_node_two_switch_network.png b/docs/_static/four_node_two_switch_network.png new file mode 100644 index 00000000..42839107 Binary files /dev/null and b/docs/_static/four_node_two_switch_network.png differ diff --git a/docs/source/simulation.rst b/docs/source/simulation.rst index 81476998..a2784628 100644 --- a/docs/source/simulation.rst +++ b/docs/source/simulation.rst @@ -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 diff --git a/docs/source/simulation_components/network/base_hardware.rst b/docs/source/simulation_components/network/base_hardware.rst new file mode 100644 index 00000000..452667d2 --- /dev/null +++ b/docs/source/simulation_components/network/base_hardware.rst @@ -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 diff --git a/docs/source/simulation_components/network/physical_layer.rst b/docs/source/simulation_components/network/physical_layer.rst deleted file mode 100644 index 1e87b72e..00000000 --- a/docs/source/simulation_components/network/physical_layer.rst +++ /dev/null @@ -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 diff --git a/docs/source/simulation_components/network/transport_to_data_link_layer.rst b/docs/source/simulation_components/network/transport_to_data_link_layer.rst index 8273339c..4961d337 100644 --- a/docs/source/simulation_components/network/transport_to_data_link_layer.rst +++ b/docs/source/simulation_components/network/transport_to_data_link_layer.rst @@ -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. diff --git a/docs/source/simulation_structure.rst b/docs/source/simulation_structure.rst index 65373a72..7630ae0f 100644 --- a/docs/source/simulation_structure.rst +++ b/docs/source/simulation_structure.rst @@ -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"]}) diff --git a/pyproject.toml b/pyproject.toml index 80094b7c..d7e71a28 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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] diff --git a/src/primaite/__init__.py b/src/primaite/__init__.py index ad157c9c..30fc9ab9 100644 --- a/src/primaite/__init__.py +++ b/src/primaite/__init__.py @@ -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: """ diff --git a/src/primaite/notebooks/create-simulation.ipynb b/src/primaite/notebooks/create-simulation.ipynb new file mode 100644 index 00000000..baf7bd2c --- /dev/null +++ b/src/primaite/notebooks/create-simulation.ipynb @@ -0,0 +1,478 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Build a simulation using the Python API\n", + "\n", + "Currently, this notbook manipulates the simulation by directly placing objects inside of the attributes of the network and domain. It should be refactored when proper methods exist for adding these objects.\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Import the Simulation class" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [], + "source": [ + "from primaite.simulator.sim_container import Simulation\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Create an empty simulation. By default this has a network with no nodes or links, and a domain controller with no accounts.\n", + "\n", + "Let's use the simulation's `describe_state()` method to verify that it is empty." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "{'uuid': '2ef348c6-32e5-4c5c-83b7-3b82d0b6123b',\n", + " 'network': {'uuid': 'dd2d1a02-d461-4505-8bbd-fd0681750175',\n", + " 'nodes': {},\n", + " 'links': {}},\n", + " 'domain': {'uuid': 'ae0423ee-51fa-41e7-be80-c642b39707f6', 'accounts': {}}}" + ] + }, + "execution_count": 2, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "my_sim = Simulation()\n", + "net = my_sim.network\n", + "my_sim.describe_state()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Add nodes" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [], + "source": [ + "from primaite.simulator.network.hardware.base import Node\n" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [], + "source": [ + "my_pc = Node(hostname=\"primaite_pc\",)\n", + "net.add_node(my_pc)\n", + "my_server = Node(hostname=\"google_server\")\n", + "net.add_node(my_server)\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Connect the nodes" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [], + "source": [ + "from primaite.simulator.network.hardware.base import NIC, Link, Switch\n" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "2023-08-24 13:06:28,617: NIC cc:be:ec:43:a6:4c/130.1.1.1 connected to Link cc:be:ec:43:a6:4c/130.1.1.1<-->79:2b:4a:70:c3:50\n", + "2023-08-24 13:06:28,618: SwitchPort 79:2b:4a:70:c3:50 connected to Link cc:be:ec:43:a6:4c/130.1.1.1<-->79:2b:4a:70:c3:50\n", + "2023-08-24 13:06:28,619: NIC c2:1e:48:e1:a4:ad/130.1.1.2 connected to Link c2:1e:48:e1:a4:ad/130.1.1.2<-->1a:2d:12:38:80:2f\n", + "2023-08-24 13:06:28,620: SwitchPort 1a:2d:12:38:80:2f connected to Link c2:1e:48:e1:a4:ad/130.1.1.2<-->1a:2d:12:38:80:2f\n" + ] + } + ], + "source": [ + "my_swtich = Switch(hostname=\"switch1\", num_ports=12)\n", + "net.add_node(my_swtich)\n", + "\n", + "pc_nic = NIC(ip_address=\"130.1.1.1\", gateway=\"130.1.1.255\", subnet_mask=\"255.255.255.0\")\n", + "my_pc.connect_nic(pc_nic)\n", + "\n", + "\n", + "server_nic = NIC(ip_address=\"130.1.1.2\", gateway=\"130.1.1.255\", subnet_mask=\"255.255.255.0\")\n", + "my_server.connect_nic(server_nic)\n", + "\n", + "\n", + "net.connect(pc_nic, my_swtich.switch_ports[1])\n", + "net.connect(server_nic, my_swtich.switch_ports[2])\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Add files and folders to nodes\n" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [], + "source": [ + "from primaite.simulator.file_system.file_system_file_type import FileSystemFileType\n", + "from primaite.simulator.file_system.file_system_file import FileSystemFile" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": {}, + "outputs": [], + "source": [ + "my_pc_downloads_folder = my_pc.file_system.create_folder(\"downloads\")\n", + "my_pc_downloads_folder.add_file(FileSystemFile(name=\"firefox_installer.zip\",file_type=FileSystemFileType.ZIP))" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "FileSystemFile(uuid='7d56a563-ecc0-4011-8c97-240dd6c885c0', name='favicon.ico', size=40.0, file_type=, action_manager=None)" + ] + }, + "execution_count": 9, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "my_server_folder = my_server.file_system.create_folder(\"static\")\n", + "my_server.file_system.create_file(\"favicon.ico\", file_type=FileSystemFileType.PNG)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Add applications to nodes" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "metadata": {}, + "outputs": [], + "source": [ + "from primaite.simulator.system.applications.application import Application, ApplicationOperatingState\n", + "from primaite.simulator.system.software import SoftwareHealthState, SoftwareCriticality\n", + "from primaite.simulator.network.transmission.transport_layer import Port\n", + "\n", + "# no applications exist yet so we will create our own.\n", + "class MSPaint(Application):\n", + " def describe_state(self):\n", + " return super().describe_state()" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "metadata": {}, + "outputs": [], + "source": [ + "mspaint = MSPaint(name = \"mspaint\", health_state_actual=SoftwareHealthState.GOOD, health_state_visible=SoftwareHealthState.GOOD, criticality=SoftwareCriticality.MEDIUM, ports={Port.HTTP}, operating_state=ApplicationOperatingState.RUNNING,execution_control_status='manual')" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "metadata": {}, + "outputs": [], + "source": [ + "my_pc.applications[mspaint.uuid] = mspaint" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Create a domain account" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "metadata": {}, + "outputs": [], + "source": [ + "from primaite.simulator.domain.account import Account, AccountType\n" + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "metadata": {}, + "outputs": [], + "source": [ + "acct = Account(username=\"admin\", password=\"admin12\", account_type=AccountType.USER)\n", + "my_sim.domain.accounts[acct.uuid] = acct" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Verify that the state dictionary contains no non-serialisable objects." + ] + }, + { + "cell_type": "code", + "execution_count": 15, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "{'uuid': '2ef348c6-32e5-4c5c-83b7-3b82d0b6123b',\n", + " 'network': {'uuid': 'dd2d1a02-d461-4505-8bbd-fd0681750175',\n", + " 'nodes': {'2f03b32b-7290-4921-8670-faebe4a19d63': {'uuid': '2f03b32b-7290-4921-8670-faebe4a19d63',\n", + " 'hostname': 'primaite_pc',\n", + " 'operating_state': 0,\n", + " 'NICs': {'e07e2a7f-b09f-4bd8-8e92-cffbf1f2270b': {'uuid': 'e07e2a7f-b09f-4bd8-8e92-cffbf1f2270b',\n", + " 'ip_adress': '130.1.1.1',\n", + " 'subnet_mask': '255.255.255.0',\n", + " 'gateway': '130.1.1.255',\n", + " 'mac_address': 'cc:be:ec:43:a6:4c',\n", + " 'speed': 100,\n", + " 'mtu': 1500,\n", + " 'wake_on_lan': False,\n", + " 'dns_servers': [],\n", + " 'enabled': False}},\n", + " 'file_system': {'uuid': '0b7206af-3e0a-41b0-8115-ae9e0dbbcd81',\n", + " 'folders': {'c161bc7c-9abd-4666-9b49-2745fdb65ebe': {'uuid': 'c161bc7c-9abd-4666-9b49-2745fdb65ebe',\n", + " 'name': 'downloads',\n", + " 'size': 1000.0,\n", + " 'files': {'f807d777-d167-4f37-9f9b-ced634af6ed5': {'uuid': 'f807d777-d167-4f37-9f9b-ced634af6ed5',\n", + " 'name': 'firefox_installer.zip',\n", + " 'size': 1000.0,\n", + " 'file_type': 'ZIP'}},\n", + " 'is_quarantined': False}}},\n", + " 'applications': {'ea466b2f-1ed5-49fd-9579-44852bff684d': {'uuid': 'ea466b2f-1ed5-49fd-9579-44852bff684d',\n", + " 'health_state': 'GOOD',\n", + " 'health_state_red_view': 'GOOD',\n", + " 'criticality': 'MEDIUM',\n", + " 'patching_count': 0,\n", + " 'scanning_count': 0,\n", + " 'revealed_to_red': False,\n", + " 'installing_count': 0,\n", + " 'max_sessions': 1,\n", + " 'tcp': True,\n", + " 'udp': True,\n", + " 'ports': ['HTTP'],\n", + " 'opearting_state': 'RUNNING',\n", + " 'execution_control_status': 'manual',\n", + " 'num_executions': 0,\n", + " 'groups': []}},\n", + " 'services': {},\n", + " 'process': {}},\n", + " 'e9afc0bc-fb21-48a3-9868-2ede6a3181dc': {'uuid': 'e9afc0bc-fb21-48a3-9868-2ede6a3181dc',\n", + " 'hostname': 'google_server',\n", + " 'operating_state': 0,\n", + " 'NICs': {'956ce240-8fb3-4fde-8635-ac4ea601a582': {'uuid': '956ce240-8fb3-4fde-8635-ac4ea601a582',\n", + " 'ip_adress': '130.1.1.2',\n", + " 'subnet_mask': '255.255.255.0',\n", + " 'gateway': '130.1.1.255',\n", + " 'mac_address': 'c2:1e:48:e1:a4:ad',\n", + " 'speed': 100,\n", + " 'mtu': 1500,\n", + " 'wake_on_lan': False,\n", + " 'dns_servers': [],\n", + " 'enabled': False}},\n", + " 'file_system': {'uuid': 'c3f99c30-b493-4fb6-b13e-d2005d851b59',\n", + " 'folders': {'869eda49-21f2-4fc1-8681-78725cdd5c70': {'uuid': '869eda49-21f2-4fc1-8681-78725cdd5c70',\n", + " 'name': 'static',\n", + " 'size': 0,\n", + " 'files': {},\n", + " 'is_quarantined': False},\n", + " '9fbe0e41-0d6a-4142-9c73-9c0de2dbde6e': {'uuid': '9fbe0e41-0d6a-4142-9c73-9c0de2dbde6e',\n", + " 'name': 'root',\n", + " 'size': 40.0,\n", + " 'files': {'7d56a563-ecc0-4011-8c97-240dd6c885c0': {'uuid': '7d56a563-ecc0-4011-8c97-240dd6c885c0',\n", + " 'name': 'favicon.ico',\n", + " 'size': 40.0,\n", + " 'file_type': 'PNG'}},\n", + " 'is_quarantined': False}}},\n", + " 'applications': {},\n", + " 'services': {},\n", + " 'process': {}},\n", + " '47814452-ef47-4e6b-9087-796c438d4698': {'uuid': '47814452-ef47-4e6b-9087-796c438d4698',\n", + " 'num_ports': 12,\n", + " 'ports': {1: {'uuid': 'b76fe86f-bb92-4346-8e83-217a2fb0bc67',\n", + " 'mac_address': '79:2b:4a:70:c3:50',\n", + " 'speed': 100,\n", + " 'mtu': 1500,\n", + " 'enabled': False},\n", + " 2: {'uuid': '6f8fc6e7-76a4-441a-b7af-441edbdcc6ac',\n", + " 'mac_address': '1a:2d:12:38:80:2f',\n", + " 'speed': 100,\n", + " 'mtu': 1500,\n", + " 'enabled': False},\n", + " 3: {'uuid': '1aa75a3c-01f1-4293-9894-5396fa412690',\n", + " 'mac_address': 'd1:7b:36:c1:82:c1',\n", + " 'speed': 100,\n", + " 'mtu': 1500,\n", + " 'enabled': False},\n", + " 4: {'uuid': 'fe6c9f44-59d5-403e-973a-6f19fce7b9b9',\n", + " 'mac_address': 'e3:6b:cc:0c:98:9b',\n", + " 'speed': 100,\n", + " 'mtu': 1500,\n", + " 'enabled': False},\n", + " 5: {'uuid': 'e9e83e37-8537-4884-98a6-87017540078f',\n", + " 'mac_address': '32:09:c0:4a:f1:20',\n", + " 'speed': 100,\n", + " 'mtu': 1500,\n", + " 'enabled': False},\n", + " 6: {'uuid': '747f2cd3-8902-4da8-8829-b0b53fe79735',\n", + " 'mac_address': 'e8:20:0b:04:b8:76',\n", + " 'speed': 100,\n", + " 'mtu': 1500,\n", + " 'enabled': False},\n", + " 7: {'uuid': '88ed129e-0ddb-4d29-ba3c-58d81efe240e',\n", + " 'mac_address': '7f:b4:f4:2e:b6:71',\n", + " 'speed': 100,\n", + " 'mtu': 1500,\n", + " 'enabled': False},\n", + " 8: {'uuid': '6c1a4c3c-25d8-46f6-98a8-54073d0ca0d3',\n", + " 'mac_address': 'f6:22:2d:24:b9:71',\n", + " 'speed': 100,\n", + " 'mtu': 1500,\n", + " 'enabled': False},\n", + " 9: {'uuid': 'b2bfc006-6a6b-4701-a75a-27954592d429',\n", + " 'mac_address': 'b6:a5:92:a5:aa:1b',\n", + " 'speed': 100,\n", + " 'mtu': 1500,\n", + " 'enabled': False},\n", + " 10: {'uuid': '3c607386-87a2-4d0b-ac04-449416ca5b1f',\n", + " 'mac_address': 'b3:75:7d:ce:88:0a',\n", + " 'speed': 100,\n", + " 'mtu': 1500,\n", + " 'enabled': False},\n", + " 11: {'uuid': '590002c8-27fa-4c31-b17b-7b89dbf8cdf8',\n", + " 'mac_address': 'c0:25:a6:64:52:8e',\n", + " 'speed': 100,\n", + " 'mtu': 1500,\n", + " 'enabled': False},\n", + " 12: {'uuid': 'b7e25eed-547a-4c17-8cb9-8b976ce4bbd9',\n", + " 'mac_address': '98:50:96:47:ca:bc',\n", + " 'speed': 100,\n", + " 'mtu': 1500,\n", + " 'enabled': False}},\n", + " 'mac_address_table': {}}},\n", + " 'links': {'a51a4435-20ae-43cf-a151-26e824968b3d': {'uuid': 'a51a4435-20ae-43cf-a151-26e824968b3d',\n", + " 'endpoint_a': 'e07e2a7f-b09f-4bd8-8e92-cffbf1f2270b',\n", + " 'endpoint_b': 'b76fe86f-bb92-4346-8e83-217a2fb0bc67',\n", + " 'bandwidth': 100.0,\n", + " 'current_load': 0.0},\n", + " 'ae3486e5-f78e-4092-96d1-d7e8176f2b7d': {'uuid': 'ae3486e5-f78e-4092-96d1-d7e8176f2b7d',\n", + " 'endpoint_a': '956ce240-8fb3-4fde-8635-ac4ea601a582',\n", + " 'endpoint_b': '6f8fc6e7-76a4-441a-b7af-441edbdcc6ac',\n", + " 'bandwidth': 100.0,\n", + " 'current_load': 0.0}}},\n", + " 'domain': {'uuid': 'ae0423ee-51fa-41e7-be80-c642b39707f6',\n", + " 'accounts': {'917eda28-9a67-4449-bddd-87e2141a3162': {'uuid': '917eda28-9a67-4449-bddd-87e2141a3162',\n", + " 'num_logons': 0,\n", + " 'num_logoffs': 0,\n", + " 'num_group_changes': 0,\n", + " 'username': 'admin',\n", + " 'password': 'admin12',\n", + " 'account_type': 'USER',\n", + " 'enabled': True}}}}" + ] + }, + "execution_count": 15, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "my_sim.describe_state()" + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "'{\"uuid\": \"2ef348c6-32e5-4c5c-83b7-3b82d0b6123b\", \"network\": {\"uuid\": \"dd2d1a02-d461-4505-8bbd-fd0681750175\", \"nodes\": {\"2f03b32b-7290-4921-8670-faebe4a19d63\": {\"uuid\": \"2f03b32b-7290-4921-8670-faebe4a19d63\", \"hostname\": \"primaite_pc\", \"operating_state\": 0, \"NICs\": {\"e07e2a7f-b09f-4bd8-8e92-cffbf1f2270b\": {\"uuid\": \"e07e2a7f-b09f-4bd8-8e92-cffbf1f2270b\", \"ip_adress\": \"130.1.1.1\", \"subnet_mask\": \"255.255.255.0\", \"gateway\": \"130.1.1.255\", \"mac_address\": \"cc:be:ec:43:a6:4c\", \"speed\": 100, \"mtu\": 1500, \"wake_on_lan\": false, \"dns_servers\": [], \"enabled\": false}}, \"file_system\": {\"uuid\": \"0b7206af-3e0a-41b0-8115-ae9e0dbbcd81\", \"folders\": {\"c161bc7c-9abd-4666-9b49-2745fdb65ebe\": {\"uuid\": \"c161bc7c-9abd-4666-9b49-2745fdb65ebe\", \"name\": \"downloads\", \"size\": 1000.0, \"files\": {\"f807d777-d167-4f37-9f9b-ced634af6ed5\": {\"uuid\": \"f807d777-d167-4f37-9f9b-ced634af6ed5\", \"name\": \"firefox_installer.zip\", \"size\": 1000.0, \"file_type\": \"ZIP\"}}, \"is_quarantined\": false}}}, \"applications\": {\"ea466b2f-1ed5-49fd-9579-44852bff684d\": {\"uuid\": \"ea466b2f-1ed5-49fd-9579-44852bff684d\", \"health_state\": \"GOOD\", \"health_state_red_view\": \"GOOD\", \"criticality\": \"MEDIUM\", \"patching_count\": 0, \"scanning_count\": 0, \"revealed_to_red\": false, \"installing_count\": 0, \"max_sessions\": 1, \"tcp\": true, \"udp\": true, \"ports\": [\"HTTP\"], \"opearting_state\": \"RUNNING\", \"execution_control_status\": \"manual\", \"num_executions\": 0, \"groups\": []}}, \"services\": {}, \"process\": {}}, \"e9afc0bc-fb21-48a3-9868-2ede6a3181dc\": {\"uuid\": \"e9afc0bc-fb21-48a3-9868-2ede6a3181dc\", \"hostname\": \"google_server\", \"operating_state\": 0, \"NICs\": {\"956ce240-8fb3-4fde-8635-ac4ea601a582\": {\"uuid\": \"956ce240-8fb3-4fde-8635-ac4ea601a582\", \"ip_adress\": \"130.1.1.2\", \"subnet_mask\": \"255.255.255.0\", \"gateway\": \"130.1.1.255\", \"mac_address\": \"c2:1e:48:e1:a4:ad\", \"speed\": 100, \"mtu\": 1500, \"wake_on_lan\": false, \"dns_servers\": [], \"enabled\": false}}, \"file_system\": {\"uuid\": \"c3f99c30-b493-4fb6-b13e-d2005d851b59\", \"folders\": {\"869eda49-21f2-4fc1-8681-78725cdd5c70\": {\"uuid\": \"869eda49-21f2-4fc1-8681-78725cdd5c70\", \"name\": \"static\", \"size\": 0, \"files\": {}, \"is_quarantined\": false}, \"9fbe0e41-0d6a-4142-9c73-9c0de2dbde6e\": {\"uuid\": \"9fbe0e41-0d6a-4142-9c73-9c0de2dbde6e\", \"name\": \"root\", \"size\": 40.0, \"files\": {\"7d56a563-ecc0-4011-8c97-240dd6c885c0\": {\"uuid\": \"7d56a563-ecc0-4011-8c97-240dd6c885c0\", \"name\": \"favicon.ico\", \"size\": 40.0, \"file_type\": \"PNG\"}}, \"is_quarantined\": false}}}, \"applications\": {}, \"services\": {}, \"process\": {}}, \"47814452-ef47-4e6b-9087-796c438d4698\": {\"uuid\": \"47814452-ef47-4e6b-9087-796c438d4698\", \"num_ports\": 12, \"ports\": {\"1\": {\"uuid\": \"b76fe86f-bb92-4346-8e83-217a2fb0bc67\", \"mac_address\": \"79:2b:4a:70:c3:50\", \"speed\": 100, \"mtu\": 1500, \"enabled\": false}, \"2\": {\"uuid\": \"6f8fc6e7-76a4-441a-b7af-441edbdcc6ac\", \"mac_address\": \"1a:2d:12:38:80:2f\", \"speed\": 100, \"mtu\": 1500, \"enabled\": false}, \"3\": {\"uuid\": \"1aa75a3c-01f1-4293-9894-5396fa412690\", \"mac_address\": \"d1:7b:36:c1:82:c1\", \"speed\": 100, \"mtu\": 1500, \"enabled\": false}, \"4\": {\"uuid\": \"fe6c9f44-59d5-403e-973a-6f19fce7b9b9\", \"mac_address\": \"e3:6b:cc:0c:98:9b\", \"speed\": 100, \"mtu\": 1500, \"enabled\": false}, \"5\": {\"uuid\": \"e9e83e37-8537-4884-98a6-87017540078f\", \"mac_address\": \"32:09:c0:4a:f1:20\", \"speed\": 100, \"mtu\": 1500, \"enabled\": false}, \"6\": {\"uuid\": \"747f2cd3-8902-4da8-8829-b0b53fe79735\", \"mac_address\": \"e8:20:0b:04:b8:76\", \"speed\": 100, \"mtu\": 1500, \"enabled\": false}, \"7\": {\"uuid\": \"88ed129e-0ddb-4d29-ba3c-58d81efe240e\", \"mac_address\": \"7f:b4:f4:2e:b6:71\", \"speed\": 100, \"mtu\": 1500, \"enabled\": false}, \"8\": {\"uuid\": \"6c1a4c3c-25d8-46f6-98a8-54073d0ca0d3\", \"mac_address\": \"f6:22:2d:24:b9:71\", \"speed\": 100, \"mtu\": 1500, \"enabled\": false}, \"9\": {\"uuid\": \"b2bfc006-6a6b-4701-a75a-27954592d429\", \"mac_address\": \"b6:a5:92:a5:aa:1b\", \"speed\": 100, \"mtu\": 1500, \"enabled\": false}, \"10\": {\"uuid\": \"3c607386-87a2-4d0b-ac04-449416ca5b1f\", \"mac_address\": \"b3:75:7d:ce:88:0a\", \"speed\": 100, \"mtu\": 1500, \"enabled\": false}, \"11\": {\"uuid\": \"590002c8-27fa-4c31-b17b-7b89dbf8cdf8\", \"mac_address\": \"c0:25:a6:64:52:8e\", \"speed\": 100, \"mtu\": 1500, \"enabled\": false}, \"12\": {\"uuid\": \"b7e25eed-547a-4c17-8cb9-8b976ce4bbd9\", \"mac_address\": \"98:50:96:47:ca:bc\", \"speed\": 100, \"mtu\": 1500, \"enabled\": false}}, \"mac_address_table\": {}}}, \"links\": {\"a51a4435-20ae-43cf-a151-26e824968b3d\": {\"uuid\": \"a51a4435-20ae-43cf-a151-26e824968b3d\", \"endpoint_a\": \"e07e2a7f-b09f-4bd8-8e92-cffbf1f2270b\", \"endpoint_b\": \"b76fe86f-bb92-4346-8e83-217a2fb0bc67\", \"bandwidth\": 100.0, \"current_load\": 0.0}, \"ae3486e5-f78e-4092-96d1-d7e8176f2b7d\": {\"uuid\": \"ae3486e5-f78e-4092-96d1-d7e8176f2b7d\", \"endpoint_a\": \"956ce240-8fb3-4fde-8635-ac4ea601a582\", \"endpoint_b\": \"6f8fc6e7-76a4-441a-b7af-441edbdcc6ac\", \"bandwidth\": 100.0, \"current_load\": 0.0}}}, \"domain\": {\"uuid\": \"ae0423ee-51fa-41e7-be80-c642b39707f6\", \"accounts\": {\"917eda28-9a67-4449-bddd-87e2141a3162\": {\"uuid\": \"917eda28-9a67-4449-bddd-87e2141a3162\", \"num_logons\": 0, \"num_logoffs\": 0, \"num_group_changes\": 0, \"username\": \"admin\", \"password\": \"admin12\", \"account_type\": \"USER\", \"enabled\": true}}}}'" + ] + }, + "execution_count": 16, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "import json\n", + "json.dumps(my_sim.describe_state())" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.12" + }, + "orig_nbformat": 4 + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/src/primaite/simulator/__init__.py b/src/primaite/simulator/__init__.py index e69de29b..1cfe7f49 100644 --- a/src/primaite/simulator/__init__.py +++ b/src/primaite/simulator/__init__.py @@ -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 diff --git a/src/primaite/simulator/core.py b/src/primaite/simulator/core.py index c3130116..b7dfcf72 100644 --- a/src/primaite/simulator/core.py +++ b/src/primaite/simulator/core.py @@ -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 diff --git a/src/primaite/simulator/domain/__init__.py b/src/primaite/simulator/domain/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/primaite/simulator/domain/account.py b/src/primaite/simulator/domain/account.py new file mode 100644 index 00000000..d235c00e --- /dev/null +++ b/src/primaite/simulator/domain/account.py @@ -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 diff --git a/src/primaite/simulator/domain/controller.py b/src/primaite/simulator/domain/controller.py new file mode 100644 index 00000000..f772ab22 --- /dev/null +++ b/src/primaite/simulator/domain/controller.py @@ -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_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.""" + ... diff --git a/src/primaite/simulator/file_system/__init__.py b/src/primaite/simulator/file_system/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/primaite/simulator/file_system/file_system.py b/src/primaite/simulator/file_system/file_system.py new file mode 100644 index 00000000..440b7dc5 --- /dev/null +++ b/src/primaite/simulator/file_system/file_system.py @@ -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)) diff --git a/src/primaite/simulator/file_system/file_system_file.py b/src/primaite/simulator/file_system/file_system_file.py new file mode 100644 index 00000000..c25f5973 --- /dev/null +++ b/src/primaite/simulator/file_system/file_system_file.py @@ -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 diff --git a/src/primaite/simulator/file_system/file_system_file_type.py b/src/primaite/simulator/file_system/file_system_file_type.py new file mode 100644 index 00000000..7e2d8706 --- /dev/null +++ b/src/primaite/simulator/file_system/file_system_file_type.py @@ -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, +} diff --git a/src/primaite/simulator/file_system/file_system_folder.py b/src/primaite/simulator/file_system/file_system_folder.py new file mode 100644 index 00000000..4e461a3a --- /dev/null +++ b/src/primaite/simulator/file_system/file_system_folder.py @@ -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 diff --git a/src/primaite/simulator/file_system/file_system_item_abc.py b/src/primaite/simulator/file_system/file_system_item_abc.py new file mode 100644 index 00000000..3b368819 --- /dev/null +++ b/src/primaite/simulator/file_system/file_system_item_abc.py @@ -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 diff --git a/src/primaite/simulator/network/container.py b/src/primaite/simulator/network/container.py new file mode 100644 index 00000000..85676034 --- /dev/null +++ b/src/primaite/simulator/network/container.py @@ -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 diff --git a/src/primaite/simulator/network/hardware/__init__.py b/src/primaite/simulator/network/hardware/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py new file mode 100644 index 00000000..9acdf0b4 --- /dev/null +++ b/src/primaite/simulator/network/hardware/base.py @@ -0,0 +1,1091 @@ +from __future__ import annotations + +import re +import secrets +from enum import Enum +from ipaddress import IPv4Address, IPv4Network +from typing import Dict, List, Optional, Tuple, Union + +from prettytable import PrettyTable + +from primaite import getLogger +from primaite.exceptions import NetworkError +from primaite.simulator.core import SimComponent +from primaite.simulator.domain.account import Account +from primaite.simulator.file_system.file_system import FileSystem +from primaite.simulator.network.protocols.arp import ARPEntry, ARPPacket +from primaite.simulator.network.transmission.data_link_layer import EthernetHeader, Frame +from primaite.simulator.network.transmission.network_layer import ICMPPacket, ICMPType, IPPacket, IPProtocol +from primaite.simulator.network.transmission.transport_layer import Port, TCPHeader +from primaite.simulator.system.applications.application import Application +from primaite.simulator.system.core.packet_capture import PacketCapture +from primaite.simulator.system.core.session_manager import SessionManager +from primaite.simulator.system.core.software_manager import SoftwareManager +from primaite.simulator.system.core.sys_log import SysLog +from primaite.simulator.system.processes.process import Process +from primaite.simulator.system.services.service import Service + +_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}'" + _LOGGER.error(msg) + 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 + "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_node: Optional[Node] = None + "The Node to which the NIC is connected." + connected_link: Optional[Link] = None + "The Link to which the NIC is connected." + enabled: bool = False + "Indicates whether the NIC is enabled." + pcap: Optional[PacketCapture] = None + + 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"]) + if "mac_address" not in kwargs: + kwargs["mac_address"] = generate_mac_address() + 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) + + 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( + { + "ip_adress": str(self.ip_address), + "subnet_mask": str(self.subnet_mask), + "gateway": str(self.gateway), + "mac_address": self.mac_address, + "speed": self.speed, + "mtu": self.mtu, + "wake_on_lan": self.wake_on_lan, + "dns_servers": self.dns_servers, + "enabled": self.enabled, + } + ) + return state + + @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 enable(self): + """Attempt to enable the NIC.""" + if self.enabled: + return + if not self.connected_node: + _LOGGER.error(f"NIC {self} cannot be enabled as it is not connected to a Node") + return + if self.connected_node.operating_state != NodeOperatingState.ON: + self.connected_node.sys_log.error(f"NIC {self} cannot be enabled as the endpoint is not turned on") + return + + self.enabled = True + self.connected_node.sys_log.info(f"NIC {self} enabled") + self.pcap = PacketCapture(hostname=self.connected_node.hostname, ip_address=self.ip_address) + if self.connected_link: + self.connected_link.endpoint_up() + + def disable(self): + """Disable the NIC.""" + if not self.enabled: + return + + self.enabled = False + if self.connected_node: + self.connected_node.sys_log.info(f"NIC {self} disabled") + else: + _LOGGER.info(f"NIC {self} disabled") + if self.connected_link: + self.connected_link.endpoint_down() + + 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` + """ + if self.connected_link: + _LOGGER.error(f"Cannot connect Link to NIC ({self.mac_address}) as it already has a connection") + return + + if self.connected_link == link: + _LOGGER.error(f"Cannot connect Link to NIC ({self.mac_address}) as it is already connected") + return + + # TODO: Inform the Node that a link has been connected + self.connected_link = link + _LOGGER.info(f"NIC {self} connected to Link {link}") + + 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) -> bool: + """ + 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` + """ + if self.enabled: + frame.set_sent_timestamp() + self.pcap.capture(frame) + self.connected_link.transmit_frame(sender_nic=self, frame=frame) + return True + # Cannot send Frame as the NIC is not enabled + return False + + def receive_frame(self, frame: Frame) -> bool: + """ + Receive a network frame from the connected link if the NIC is enabled. + + The Frame is passed to the Node. + + :param frame: The network frame being received. + :type frame: :class:`~primaite.simulator.network.osi_layers.Frame` + """ + if self.enabled: + frame.decrement_ttl() + frame.set_received_timestamp() + self.pcap.capture(frame) + self.connected_node.receive_frame(frame=frame, from_nic=self) + return True + return False + + def __str__(self) -> str: + return f"{self.mac_address}/{self.ip_address}" + + +class SwitchPort(SimComponent): + """ + Models a switch port in a network switch device. + + :param mac_address: The MAC address of the SwitchPort. Defaults to a randomly set MAC address. + :param speed: The speed of the SwitchPort in Mbps (default is 100 Mbps). + :param mtu: The Maximum Transmission Unit (MTU) of the SwitchPort in Bytes, representing the largest data packet + size it can handle without fragmentation (default is 1500 B). + """ + + port_num: int = 1 + mac_address: str + "The MAC address of the SwitchPort. Defaults to a randomly set MAC address." + speed: int = 100 + "The speed of the SwitchPort in Mbps. Default is 100 Mbps." + mtu: int = 1500 + "The Maximum Transmission Unit (MTU) of the SwitchPort in Bytes. Default is 1500 B" + connected_node: Optional[Switch] = None + "The Node to which the SwitchPort is connected." + connected_link: Optional[Link] = None + "The Link to which the SwitchPort is connected." + enabled: bool = False + "Indicates whether the SwitchPort is enabled." + pcap: Optional[PacketCapture] = None + + def __init__(self, **kwargs): + """The SwitchPort constructor.""" + if "mac_address" not in kwargs: + kwargs["mac_address"] = generate_mac_address() + 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( + { + "mac_address": self.mac_address, + "speed": self.speed, + "mtu": self.mtu, + "enabled": self.enabled, + } + ) + return state + + def enable(self): + """Attempt to enable the SwitchPort.""" + if self.enabled: + return + + if not self.connected_node: + _LOGGER.error(f"SwitchPort {self} cannot be enabled as it is not connected to a Node") + return + + if self.connected_node.operating_state != NodeOperatingState.ON: + self.connected_node.sys_log.info(f"SwitchPort {self} cannot be enabled as the endpoint is not turned on") + return + + self.enabled = True + self.connected_node.sys_log.info(f"SwitchPort {self} enabled") + self.pcap = PacketCapture(hostname=self.connected_node.hostname, switch_port_number=self.port_num) + if self.connected_link: + self.connected_link.endpoint_up() + + def disable(self): + """Disable the SwitchPort.""" + if not self.enabled: + return + self.enabled = False + if self.connected_node: + self.connected_node.sys_log.info(f"SwitchPort {self} disabled") + else: + _LOGGER.info(f"SwitchPort {self} disabled") + if self.connected_link: + self.connected_link.endpoint_down() + + def connect_link(self, link: Link): + """ + Connect the SwitchPort to a link. + + :param link: The link to which the SwitchPort is connected. + """ + if self.connected_link: + _LOGGER.error(f"Cannot connect link to SwitchPort {self.mac_address} as it already has a connection") + return + + if self.connected_link == link: + _LOGGER.error(f"Cannot connect Link to SwitchPort {self.mac_address} as it is already connected") + return + + # TODO: Inform the Switch that a link has been connected + self.connected_link = link + _LOGGER.info(f"SwitchPort {self} connected to Link {link}") + self.enable() + + def disconnect_link(self): + """Disconnect the SwitchPort 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 send_frame(self, frame: Frame) -> bool: + """ + Send a network frame from the SwitchPort to the connected link. + + :param frame: The network frame to be sent. + """ + if self.enabled: + self.pcap.capture(frame) + self.connected_link.transmit_frame(sender_nic=self, frame=frame) + return True + # Cannot send Frame as the SwitchPort is not enabled + return False + + def receive_frame(self, frame: Frame) -> bool: + """ + Receive a network frame from the connected link if the SwitchPort is enabled. + + The Frame is passed to the Node. + + :param frame: The network frame being received. + """ + if self.enabled: + frame.decrement_ttl() + self.pcap.capture(frame) + self.connected_node.forward_frame(frame=frame, incoming_port=self) + return True + return False + + def __str__(self) -> str: + return f"{self.mac_address}" + + +class Link(SimComponent): + """ + Represents a network link between NIC<-->NIC, NIC<-->SwitchPort, or SwitchPort<-->SwitchPort. + + :param endpoint_a: The first NIC or SwitchPort connected to the Link. + :param endpoint_b: The second NIC or SwitchPort connected to the Link. + :param bandwidth: The bandwidth of the Link in Mbps (default is 100 Mbps). + """ + + endpoint_a: Union[NIC, SwitchPort] + "The first NIC or SwitchPort connected to the Link." + endpoint_b: Union[NIC, SwitchPort] + "The second NIC or SwitchPort connected to the Link." + bandwidth: float = 100.0 + "The bandwidth of the Link in Mbps (default is 100 Mbps)." + current_load: float = 0.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 or SwitchPort" + _LOGGER.error(msg) + raise ValueError(msg) + super().__init__(**kwargs) + self.endpoint_a.connect_link(self) + self.endpoint_b.connect_link(self) + self.endpoint_up() + + 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( + { + "endpoint_a": self.endpoint_a.uuid, + "endpoint_b": self.endpoint_b.uuid, + "bandwidth": self.bandwidth, + "current_load": self.current_load, + } + ) + return state + + @property + def current_load_percent(self) -> str: + """Get the current load formatted as a percentage string.""" + return f"{self.current_load / self.bandwidth:.5f}%" + + def endpoint_up(self): + """Let the Link know and endpoint has been brought up.""" + if self.is_up: + _LOGGER.info(f"Link {self} up") + + def endpoint_down(self): + """Let the Link know and endpoint has been brought down.""" + if not self.is_up: + self.current_load = 0.0 + _LOGGER.info(f"Link {self} down") + + @property + def is_up(self) -> bool: + """ + Informs whether the link is up. + + This is based upon both NIC endpoints being enabled. + """ + return self.endpoint_a.enabled and self.endpoint_b.enabled + + def _can_transmit(self, frame: Frame) -> bool: + if self.is_up: + frame_size_Mbits = frame.size_Mbits # noqa - Leaving it as Mbits as this is how they're expressed + return self.current_load + frame_size_Mbits <= self.bandwidth + return False + + def transmit_frame(self, sender_nic: Union[NIC, SwitchPort], frame: Frame) -> bool: + """ + Send a network frame from one NIC or SwitchPort to another connected NIC or SwitchPort. + + :param sender_nic: The NIC or SwitchPort sending the frame. + :param frame: The network frame to be sent. + :return: True if the Frame can be sent, otherwise False. + """ + can_transmit = self._can_transmit(frame) + if not can_transmit: + _LOGGER.info(f"Cannot transmit frame as {self} is at capacity") + return False + + receiver = self.endpoint_a + if receiver == sender_nic: + receiver = self.endpoint_b + frame_size = frame.size_Mbits + + if receiver.receive_frame(frame): + # Frame transmitted successfully + # Load the frame size on the link + self.current_load += frame_size + _LOGGER.info( + f"Added {frame_size:.3f} Mbits to {self}, current load {self.current_load:.3f} Mbits " + f"({self.current_load_percent})" + ) + return True + return False + + def reset_component_for_episode(self, episode: int): + """ + Link reset function. + + Reset: + - returns the link current_load to 0. + """ + self.current_load = 0 + + def __str__(self) -> str: + return f"{self.endpoint_a}<-->{self.endpoint_b}" + + +class ARPCache: + """ + The ARPCache (Address Resolution Protocol) class. + + Responsible for maintaining a mapping between IP addresses and MAC addresses (ARP cache) for the network. It + provides methods for looking up, adding, and removing entries, and for processing ARPPackets. + """ + + def __init__(self, sys_log: "SysLog"): + """ + Initialize an ARP (Address Resolution Protocol) cache. + + :param sys_log: The nodes sys log. + """ + self.sys_log: "SysLog" = sys_log + self.arp: Dict[IPv4Address, ARPEntry] = {} + self.nics: Dict[str, "NIC"] = {} + + def _add_arp_cache_entry(self, ip_address: IPv4Address, mac_address: str, nic: NIC): + """ + Add an ARP entry to the cache. + + :param ip_address: The IP address to be added to the 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. + """ + self.sys_log.info(f"Adding ARP cache entry for {mac_address}/{ip_address} via NIC {nic}") + arp_entry = ARPEntry(mac_address=mac_address, nic_uuid=nic.uuid) + self.arp[ip_address] = arp_entry + + def _remove_arp_cache_entry(self, ip_address: IPv4Address): + """ + Remove an ARP entry from the cache. + + :param ip_address: The IP address to be removed from the cache. + """ + if ip_address in self.arp: + del self.arp[ip_address] + + def get_arp_cache_mac_address(self, ip_address: IPv4Address) -> Optional[str]: + """ + Get the MAC address associated with an IP address. + + :param ip_address: The IP address to look up in the cache. + :return: The MAC address associated with the IP address, or None if not found. + """ + arp_entry = self.arp.get(ip_address) + if arp_entry: + return arp_entry.mac_address + + def get_arp_cache_nic(self, ip_address: IPv4Address) -> Optional[NIC]: + """ + Get the NIC associated with an IP address. + + :param ip_address: The IP address to look up in the cache. + :return: The NIC associated with the IP address, or None if not found. + """ + arp_entry = self.arp.get(ip_address) + if arp_entry: + return self.nics[arp_entry.nic_uuid] + + def clear_arp_cache(self): + """Clear the entire ARP cache, removing all stored entries.""" + self.arp.clear() + + def send_arp_request(self, target_ip_address: Union[IPv4Address, str]): + """ + Perform a standard ARP request for a given target IP address. + + Broadcasts the request through all enabled NICs to determine the MAC address corresponding to the target IP + address. + + :param target_ip_address: The target IP address to send an ARP request for. + """ + for nic in self.nics.values(): + if nic.enabled: + self.sys_log.info(f"Sending ARP request from NIC {nic} for ip {target_ip_address}") + tcp_header = TCPHeader(src_port=Port.ARP, dst_port=Port.ARP) + + # Network Layer + ip_packet = IPPacket( + src_ip=nic.ip_address, + dst_ip=target_ip_address, + ) + # Data Link Layer + ethernet_header = EthernetHeader(src_mac_addr=nic.mac_address, dst_mac_addr="ff:ff:ff:ff:ff:ff") + arp_packet = ARPPacket( + sender_ip=nic.ip_address, sender_mac_addr=nic.mac_address, target_ip=target_ip_address + ) + frame = Frame(ethernet=ethernet_header, ip=ip_packet, tcp=tcp_header, arp=arp_packet) + nic.send_frame(frame) + + def process_arp_packet(self, from_nic: NIC, arp_packet: ARPPacket): + """ + Process a received ARP packet, handling both ARP requests and responses. + + If an ARP request is received for the local IP, a response is sent back. + If an ARP response is received, the ARP cache is updated with the new entry. + + :param from_nic: The NIC that received the ARP packet. + :param arp_packet: The ARP packet to be processed. + """ + # ARP Reply + if not arp_packet.request: + self.sys_log.info( + f"Received ARP response for {arp_packet.sender_ip} from {arp_packet.sender_mac_addr} via NIC {from_nic}" + ) + self._add_arp_cache_entry( + ip_address=arp_packet.sender_ip, mac_address=arp_packet.sender_mac_addr, nic=from_nic + ) + return + + # ARP Request + self.sys_log.info( + f"Received ARP request for {arp_packet.target_ip} from " + f"{arp_packet.sender_mac_addr}/{arp_packet.sender_ip} " + ) + + # Unmatched ARP Request + if arp_packet.target_ip != from_nic.ip_address: + self.sys_log.info(f"Ignoring ARP request for {arp_packet.target_ip}") + return + + # Matched ARP request + self._add_arp_cache_entry(ip_address=arp_packet.sender_ip, mac_address=arp_packet.sender_mac_addr, nic=from_nic) + arp_packet = arp_packet.generate_reply(from_nic.mac_address) + self.sys_log.info( + f"Sending ARP reply from {arp_packet.sender_mac_addr}/{arp_packet.sender_ip} " + f"to {arp_packet.target_ip}/{arp_packet.target_mac_addr} " + ) + + tcp_header = TCPHeader(src_port=Port.ARP, dst_port=Port.ARP) + + # Network Layer + ip_packet = IPPacket( + src_ip=arp_packet.sender_ip, + dst_ip=arp_packet.target_ip, + ) + # Data Link Layer + ethernet_header = EthernetHeader( + src_mac_addr=arp_packet.sender_mac_addr, dst_mac_addr=arp_packet.target_mac_addr + ) + frame = Frame(ethernet=ethernet_header, ip=ip_packet, tcp=tcp_header, arp=arp_packet) + from_nic.send_frame(frame) + + +class ICMP: + """ + The ICMP (Internet Control Message Protocol) class. + + Provides functionalities for managing and handling ICMP packets, including echo requests and replies. + """ + + def __init__(self, sys_log: SysLog, arp_cache: ARPCache): + """ + Initialize the ICMP (Internet Control Message Protocol) service. + + :param sys_log: The system log to store system messages and information. + :param arp_cache: The ARP cache for resolving IP to MAC address mappings. + """ + self.sys_log: SysLog = sys_log + self.arp: ARPCache = arp_cache + + def process_icmp(self, frame: Frame): + """ + Process an ICMP packet, including handling echo requests and replies. + + :param frame: The Frame containing the ICMP packet to process. + """ + if frame.icmp.icmp_type == ICMPType.ECHO_REQUEST: + self.sys_log.info(f"Received echo request from {frame.ip.src_ip}") + target_mac_address = self.arp.get_arp_cache_mac_address(frame.ip.src_ip) + src_nic = self.arp.get_arp_cache_nic(frame.ip.src_ip) + tcp_header = TCPHeader(src_port=Port.ARP, dst_port=Port.ARP) + + # Network Layer + ip_packet = IPPacket(src_ip=src_nic.ip_address, dst_ip=frame.ip.src_ip, protocol=IPProtocol.ICMP) + # Data Link Layer + ethernet_header = EthernetHeader(src_mac_addr=src_nic.mac_address, dst_mac_addr=target_mac_address) + icmp_reply_packet = ICMPPacket( + icmp_type=ICMPType.ECHO_REPLY, + icmp_code=0, + identifier=frame.icmp.identifier, + sequence=frame.icmp.sequence + 1, + ) + frame = Frame(ethernet=ethernet_header, ip=ip_packet, tcp=tcp_header, icmp=icmp_reply_packet) + self.sys_log.info(f"Sending echo reply to {frame.ip.dst_ip}") + src_nic.send_frame(frame) + elif frame.icmp.icmp_type == ICMPType.ECHO_REPLY: + self.sys_log.info(f"Received echo reply from {frame.ip.src_ip}") + + def ping( + self, target_ip_address: IPv4Address, sequence: int = 0, identifier: Optional[int] = None + ) -> Tuple[int, Union[int, None]]: + """ + Send an ICMP echo request (ping) to a target IP address and manage the sequence and identifier. + + :param target_ip_address: The target IP address to send the ping. + :param sequence: The sequence number of the echo request. Defaults to 0. + :param identifier: An optional identifier for the ICMP packet. If None, a default will be used. + :return: A tuple containing the next sequence number and the identifier, or (0, None) if the target IP address + was not found in the ARP cache. + """ + nic = self.arp.get_arp_cache_nic(target_ip_address) + # TODO: Eventually this ARP request needs to be done elsewhere. It's not the resonsibility of the + # ping function to handle ARP lookups + # No existing ARP entry + if not nic: + self.sys_log.info(f"No entry in ARP cache for {target_ip_address}") + self.arp.send_arp_request(target_ip_address) + return 0, None + + # ARP entry exists + sequence += 1 + target_mac_address = self.arp.get_arp_cache_mac_address(target_ip_address) + src_nic = self.arp.get_arp_cache_nic(target_ip_address) + tcp_header = TCPHeader(src_port=Port.ARP, dst_port=Port.ARP) + + # Network Layer + ip_packet = IPPacket( + src_ip=nic.ip_address, + dst_ip=target_ip_address, + protocol=IPProtocol.ICMP, + ) + # Data Link Layer + ethernet_header = EthernetHeader(src_mac_addr=src_nic.mac_address, dst_mac_addr=target_mac_address) + icmp_packet = ICMPPacket(identifier=identifier, sequence=sequence) + frame = Frame(ethernet=ethernet_header, ip=ip_packet, tcp=tcp_header, icmp=icmp_packet) + self.sys_log.info(f"Sending echo request to {target_ip_address}") + nic.send_frame(frame) + return sequence, icmp_packet.identifier + + +class NodeOperatingState(Enum): + """Enumeration of Node Operating States.""" + + OFF = 0 + "The node is powered off." + ON = 1 + "The node is powered on." + SHUTTING_DOWN = 2 + "The node is in the process of shutting down." + BOOTING = 3 + "The node is in the process of booting up." + + +class Node(SimComponent): + """ + A basic Node class that represents a node on the network. + + This class manages the state of the node, including the NICs (Network Interface Cards), accounts, applications, + services, processes, file system, and various managers like ARP, ICMP, SessionManager, and SoftwareManager. + + :param hostname: The node hostname on the network. + :param operating_state: The node operating state, either ON or OFF. + """ + + hostname: str + "The node hostname on the network." + operating_state: NodeOperatingState = NodeOperatingState.OFF + "The hardware state of the node." + nics: Dict[str, NIC] = {} + "The NICs on the node." + + accounts: Dict[str, Account] = {} + "All accounts on the node." + applications: Dict[str, Application] = {} + "All applications on the node." + services: Dict[str, Service] = {} + "All services on the node." + processes: Dict[str, Process] = {} + "All processes on the node." + file_system: FileSystem + "The nodes file system." + sys_log: SysLog + arp: ARPCache + icmp: ICMP + session_manager: SessionManager + software_manager: SoftwareManager + + revealed_to_red: bool = False + "Informs whether the node has been revealed to a red agent." + + def __init__(self, **kwargs): + """ + Initialize the Node with various components and managers. + + This method initializes the ARP cache, ICMP handler, session manager, and software manager if they are not + provided. + """ + if not kwargs.get("sys_log"): + kwargs["sys_log"] = SysLog(kwargs["hostname"]) + if not kwargs.get("arp_cache"): + kwargs["arp"] = ARPCache(sys_log=kwargs.get("sys_log")) + if not kwargs.get("icmp"): + kwargs["icmp"] = ICMP(sys_log=kwargs.get("sys_log"), arp_cache=kwargs.get("arp")) + if not kwargs.get("session_manager"): + kwargs["session_manager"] = SessionManager(sys_log=kwargs.get("sys_log"), arp_cache=kwargs.get("arp")) + if not kwargs.get("software_manager"): + kwargs["software_manager"] = SoftwareManager( + sys_log=kwargs.get("sys_log"), session_manager=kwargs.get("session_manager") + ) + if not kwargs.get("file_system"): + kwargs["file_system"] = FileSystem() + super().__init__(**kwargs) + self.arp.nics = self.nics + + 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( + { + "hostname": self.hostname, + "operating_state": self.operating_state.value, + "NICs": {uuid: nic.describe_state() for uuid, nic in self.nics.items()}, + # "switch_ports": {uuid, sp for uuid, sp in self.switch_ports.items()}, + "file_system": self.file_system.describe_state(), + "applications": {uuid: app.describe_state() for uuid, app in self.applications.items()}, + "services": {uuid: svc.describe_state() for uuid, svc in self.services.items()}, + "process": {uuid: proc.describe_state() for uuid, proc in self.processes.items()}, + } + ) + return state + + def show(self): + """Prints a table of the NICs on the Node..""" + from prettytable import PrettyTable + + table = PrettyTable(["MAC Address", "Address", "Default Gateway", "Speed", "Status"]) + + for nic in self.nics.values(): + table.add_row( + [ + nic.mac_address, + f"{nic.ip_address}/{nic.ip_network.prefixlen}", + nic.gateway, + nic.speed, + "Enabled" if nic.enabled else "Disabled", + ] + ) + print(table) + + def power_on(self): + """Power on the Node, enabling its NICs if it is in the OFF state.""" + if self.operating_state == NodeOperatingState.OFF: + self.operating_state = NodeOperatingState.ON + self.sys_log.info("Turned on") + for nic in self.nics.values(): + nic.enable() + + def power_off(self): + """Power off the Node, disabling its NICs if it is in the ON state.""" + if self.operating_state == NodeOperatingState.ON: + for nic in self.nics.values(): + nic.disable() + self.operating_state = NodeOperatingState.OFF + self.sys_log.info("Turned off") + + def connect_nic(self, nic: NIC): + """ + Connect a NIC (Network Interface Card) to the node. + + :param nic: The NIC to connect. + :raise NetworkError: If the NIC is already connected. + """ + if nic.uuid not in self.nics: + self.nics[nic.uuid] = nic + nic.connected_node = self + nic.parent = self + self.sys_log.info(f"Connected NIC {nic}") + if self.operating_state == NodeOperatingState.ON: + nic.enable() + else: + msg = f"Cannot connect NIC {nic} as it is already connected" + self.sys_log.logger.error(msg) + _LOGGER.error(msg) + raise NetworkError(msg) + + def disconnect_nic(self, nic: Union[NIC, str]): + """ + Disconnect a NIC (Network Interface Card) from the node. + + :param nic: The NIC to Disconnect, or its UUID. + :raise NetworkError: If the NIC is not connected. + """ + if isinstance(nic, str): + nic = self.nics.get(nic) + if nic or nic.uuid in self.nics: + self.nics.pop(nic.uuid) + nic.parent = None + nic.disable() + self.sys_log.info(f"Disconnected NIC {nic}") + else: + msg = f"Cannot disconnect NIC {nic} as it is not connected" + self.sys_log.logger.error(msg) + _LOGGER.error(msg) + raise NetworkError(msg) + + def ping(self, target_ip_address: Union[IPv4Address, str], pings: int = 4) -> bool: + """ + Ping an IP address, performing a standard ICMP echo request/response. + + :param target_ip_address: The target IP address to ping. + :param pings: The number of pings to attempt, default is 4. + :return: True if the ping is successful, otherwise False. + """ + if not isinstance(target_ip_address, IPv4Address): + target_ip_address = IPv4Address(target_ip_address) + if self.operating_state == NodeOperatingState.ON: + self.sys_log.info(f"Attempting to ping {target_ip_address}") + sequence, identifier = 0, None + while sequence < pings: + sequence, identifier = self.icmp.ping(target_ip_address, sequence, identifier) + return True + self.sys_log.info("Ping failed as the node is turned off") + return False + + def send_frame(self, frame: Frame): + """ + Send a Frame from the Node to the connected NIC. + + :param frame: The Frame to be sent. + """ + nic: NIC = self._get_arp_cache_nic(frame.ip.dst_ip) + nic.send_frame(frame) + + def receive_frame(self, frame: Frame, from_nic: NIC): + """ + Receive a Frame from the connected NIC and process it. + + Depending on the protocol, the frame is passed to the appropriate handler such as ARP or ICMP, or up to the + SessionManager if no code manager exists. + + :param frame: The Frame being received. + :param from_nic: The NIC that received the frame. + """ + if frame.ip.protocol == IPProtocol.TCP: + if frame.tcp.src_port == Port.ARP: + self.arp.process_arp_packet(from_nic=from_nic, arp_packet=frame.arp) + elif frame.ip.protocol == IPProtocol.UDP: + pass + elif frame.ip.protocol == IPProtocol.ICMP: + self.icmp.process_icmp(frame=frame) + + +class Switch(Node): + """A class representing a Layer 2 network switch.""" + + num_ports: int = 24 + "The number of ports on the switch." + switch_ports: Dict[int, SwitchPort] = {} + "The SwitchPorts on the switch." + mac_address_table: Dict[str, SwitchPort] = {} + "A MAC address table mapping destination MAC addresses to corresponding SwitchPorts." + + def __init__(self, **kwargs): + super().__init__(**kwargs) + if not self.switch_ports: + self.switch_ports = {i: SwitchPort() for i in range(1, self.num_ports + 1)} + for port_num, port in self.switch_ports.items(): + port.connected_node = self + port.parent = self + port.port_num = port_num + + def show(self): + """Prints a table of the SwitchPorts on the Switch.""" + table = PrettyTable(["Port", "MAC Address", "Speed", "Status"]) + + for port_num, port in self.switch_ports.items(): + table.add_row([port_num, port.mac_address, port.speed, "Enabled" if port.enabled else "Disabled"]) + print(table) + + 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 + """ + return { + "uuid": self.uuid, + "num_ports": self.num_ports, # redundant? + "ports": {port_num: port.describe_state() for port_num, port in self.switch_ports.items()}, + "mac_address_table": {mac: port for mac, port in self.mac_address_table.items()}, + } + + def _add_mac_table_entry(self, mac_address: str, switch_port: SwitchPort): + mac_table_port = self.mac_address_table.get(mac_address) + if not mac_table_port: + self.mac_address_table[mac_address] = switch_port + self.sys_log.info(f"Added MAC table entry: Port {switch_port.port_num} -> {mac_address}") + else: + if mac_table_port != switch_port: + self.mac_address_table.pop(mac_address) + self.sys_log.info(f"Removed MAC table entry: Port {mac_table_port.port_num} -> {mac_address}") + self._add_mac_table_entry(mac_address, switch_port) + + def forward_frame(self, frame: Frame, incoming_port: SwitchPort): + """ + Forward a frame to the appropriate port based on the destination MAC address. + + :param frame: The Frame to be forwarded. + :param incoming_port: The port number from which the frame was received. + """ + src_mac = frame.ethernet.src_mac_addr + dst_mac = frame.ethernet.dst_mac_addr + self._add_mac_table_entry(src_mac, incoming_port) + + outgoing_port = self.mac_address_table.get(dst_mac) + if outgoing_port or dst_mac != "ff:ff:ff:ff:ff:ff": + outgoing_port.send_frame(frame) + else: + # If the destination MAC is not in the table, flood to all ports except incoming + for port in self.switch_ports.values(): + if port != incoming_port: + port.send_frame(frame) + + def disconnect_link_from_port(self, link: Link, port_number: int): + """ + Disconnect a given link from the specified port number on the switch. + + :param link: The Link object to be disconnected. + :param port_number: The port number on the switch from where the link should be disconnected. + :raise NetworkError: When an invalid port number is provided or the link does not match the connection. + """ + port = self.switch_ports.get(port_number) + if port is None: + msg = f"Invalid port number {port_number} on the switch" + _LOGGER.error(msg) + raise NetworkError(msg) + + if port.connected_link != link: + msg = f"The link does not match the connection at port number {port_number}" + _LOGGER.error(msg) + raise NetworkError(msg) + + port.disconnect_link() diff --git a/src/primaite/simulator/network/hardware/nodes/__init__.py b/src/primaite/simulator/network/hardware/nodes/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/primaite/simulator/network/protocols/__init__.py b/src/primaite/simulator/network/protocols/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/primaite/simulator/network/protocols/arp.py b/src/primaite/simulator/network/protocols/arp.py new file mode 100644 index 00000000..bae14d28 --- /dev/null +++ b/src/primaite/simulator/network/protocols/arp.py @@ -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, + ) diff --git a/src/primaite/simulator/network/transmission/data_link_layer.py b/src/primaite/simulator/network/transmission/data_link_layer.py index b9d969bd..1b7ccf7d 100644 --- a/src/primaite/simulator/network/transmission/data_link_layer.py +++ b/src/primaite/simulator/network/transmission/data_link_layer.py @@ -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) diff --git a/src/primaite/simulator/network/transmission/network_layer.py b/src/primaite/simulator/network/transmission/network_layer.py index 69b682cc..afd1ecef 100644 --- a/src/primaite/simulator/network/transmission/network_layer.py +++ b/src/primaite/simulator/network/transmission/network_layer.py @@ -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: diff --git a/src/primaite/simulator/network/transmission/physical_layer.py b/src/primaite/simulator/network/transmission/physical_layer.py deleted file mode 100644 index ee2297b6..00000000 --- a/src/primaite/simulator/network/transmission/physical_layer.py +++ /dev/null @@ -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 diff --git a/src/primaite/simulator/network/transmission/transport_layer.py b/src/primaite/simulator/network/transmission/transport_layer.py index c8e6b89d..b95b4a74 100644 --- a/src/primaite/simulator/network/transmission/transport_layer.py +++ b/src/primaite/simulator/network/transmission/transport_layer.py @@ -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] diff --git a/src/primaite/simulator/network/utils.py b/src/primaite/simulator/network/utils.py new file mode 100644 index 00000000..496f5e13 --- /dev/null +++ b/src/primaite/simulator/network/utils.py @@ -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 diff --git a/src/primaite/simulator/sim_container.py b/src/primaite/simulator/sim_container.py new file mode 100644 index 00000000..319defe4 --- /dev/null +++ b/src/primaite/simulator/sim_container.py @@ -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 diff --git a/src/primaite/simulator/system/__init__.py b/src/primaite/simulator/system/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/primaite/simulator/system/applications/__init__.py b/src/primaite/simulator/system/applications/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/primaite/simulator/system/applications/application.py b/src/primaite/simulator/system/applications/application.py new file mode 100644 index 00000000..37748560 --- /dev/null +++ b/src/primaite/simulator/system/applications/application.py @@ -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 diff --git a/src/primaite/simulator/system/core/__init__.py b/src/primaite/simulator/system/core/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/primaite/simulator/system/core/packet_capture.py b/src/primaite/simulator/system/core/packet_capture.py new file mode 100644 index 00000000..c985af1f --- /dev/null +++ b/src/primaite/simulator/system/core/packet_capture.py @@ -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: //__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 diff --git a/src/primaite/simulator/system/core/session_manager.py b/src/primaite/simulator/system/core/session_manager.py new file mode 100644 index 00000000..fe7b06b2 --- /dev/null +++ b/src/primaite/simulator/system/core/session_manager.py @@ -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. diff --git a/src/primaite/simulator/system/core/software_manager.py b/src/primaite/simulator/system/core/software_manager.py new file mode 100644 index 00000000..411fb6e9 --- /dev/null +++ b/src/primaite/simulator/system/core/software_manager.py @@ -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 diff --git a/src/primaite/simulator/system/core/sys_log.py b/src/primaite/simulator/system/core/sys_log.py new file mode 100644 index 00000000..4b858c2e --- /dev/null +++ b/src/primaite/simulator/system/core/sys_log.py @@ -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: //_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) diff --git a/src/primaite/simulator/system/processes/__init__.py b/src/primaite/simulator/system/processes/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/primaite/simulator/system/processes/process.py b/src/primaite/simulator/system/processes/process.py new file mode 100644 index 00000000..c4e94845 --- /dev/null +++ b/src/primaite/simulator/system/processes/process.py @@ -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 diff --git a/src/primaite/simulator/system/services/__init__.py b/src/primaite/simulator/system/services/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/primaite/simulator/system/services/service.py b/src/primaite/simulator/system/services/service.py new file mode 100644 index 00000000..eafff3f0 --- /dev/null +++ b/src/primaite/simulator/system/services/service.py @@ -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 diff --git a/src/primaite/simulator/system/software.py b/src/primaite/simulator/system/software.py new file mode 100644 index 00000000..a2acd9fb --- /dev/null +++ b/src/primaite/simulator/system/software.py @@ -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 diff --git a/tests/conftest.py b/tests/conftest.py index f40b0b94..f1c05187 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -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 diff --git a/tests/integration_tests/component_creation/__init__.py b/tests/integration_tests/component_creation/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/integration_tests/component_creation/test_permission_system.py b/tests/integration_tests/component_creation/test_permission_system.py new file mode 100644 index 00000000..6816ba84 --- /dev/null +++ b/tests/integration_tests/component_creation/test_permission_system.py @@ -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" diff --git a/tests/integration_tests/network/test_frame_transmission.py b/tests/integration_tests/network/test_frame_transmission.py new file mode 100644 index 00000000..3840c302 --- /dev/null +++ b/tests/integration_tests/network/test_frame_transmission.py @@ -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") diff --git a/tests/integration_tests/network/test_link_connection.py b/tests/integration_tests/network/test_link_connection.py new file mode 100644 index 00000000..e08e40b9 --- /dev/null +++ b/tests/integration_tests/network/test_link_connection.py @@ -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 diff --git a/tests/integration_tests/network/test_network_creation.py b/tests/integration_tests/network/test_network_creation.py new file mode 100644 index 00000000..356eb1db --- /dev/null +++ b/tests/integration_tests/network/test_network_creation.py @@ -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 diff --git a/tests/integration_tests/network/test_nic_link_connection.py b/tests/integration_tests/network/test_nic_link_connection.py index 6bca3c0a..52a0c735 100644 --- a/tests/integration_tests/network/test_nic_link_connection.py +++ b/tests/integration_tests/network/test_nic_link_connection.py @@ -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(): diff --git a/tests/unit_tests/_primaite/_simulator/_domain/__init__.py b/tests/unit_tests/_primaite/_simulator/_domain/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit_tests/_primaite/_simulator/_domain/test_account.py b/tests/unit_tests/_primaite/_simulator/_domain/test_account.py new file mode 100644 index 00000000..b5632ea7 --- /dev/null +++ b/tests/unit_tests/_primaite/_simulator/_domain/test_account.py @@ -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) diff --git a/tests/unit_tests/_primaite/_simulator/_domain/test_controller.py b/tests/unit_tests/_primaite/_simulator/_domain/test_controller.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit_tests/_primaite/_simulator/_file_system/__init__.py b/tests/unit_tests/_primaite/_simulator/_file_system/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit_tests/_primaite/_simulator/_file_system/test_file_system.py b/tests/unit_tests/_primaite/_simulator/_file_system/test_file_system.py new file mode 100644 index 00000000..348eb440 --- /dev/null +++ b/tests/unit_tests/_primaite/_simulator/_file_system/test_file_system.py @@ -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() diff --git a/tests/unit_tests/_primaite/_simulator/_file_system/test_file_system_file.py b/tests/unit_tests/_primaite/_simulator/_file_system/test_file_system_file.py new file mode 100644 index 00000000..629b9bb9 --- /dev/null +++ b/tests/unit_tests/_primaite/_simulator/_file_system/test_file_system_file.py @@ -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() diff --git a/tests/unit_tests/_primaite/_simulator/_file_system/test_file_system_folder.py b/tests/unit_tests/_primaite/_simulator/_file_system/test_file_system_folder.py new file mode 100644 index 00000000..1940e886 --- /dev/null +++ b/tests/unit_tests/_primaite/_simulator/_file_system/test_file_system_folder.py @@ -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() diff --git a/tests/unit_tests/_primaite/_simulator/_network/_hardware/__init__.py b/tests/unit_tests/_primaite/_simulator/_network/_hardware/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit_tests/_primaite/_simulator/_network/_transmission/test_physical_layer.py b/tests/unit_tests/_primaite/_simulator/_network/_hardware/test_nic.py similarity index 95% rename from tests/unit_tests/_primaite/_simulator/_network/_transmission/test_physical_layer.py rename to tests/unit_tests/_primaite/_simulator/_network/_hardware/test_nic.py index 5a33e723..dc508508 100644 --- a/tests/unit_tests/_primaite/_simulator/_network/_transmission/test_physical_layer.py +++ b/tests/unit_tests/_primaite/_simulator/_network/_hardware/test_nic.py @@ -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(): diff --git a/tests/unit_tests/_primaite/_simulator/_network/_hardware/test_node.py b/tests/unit_tests/_primaite/_simulator/_network/_hardware/test_node.py new file mode 100644 index 00000000..0e5fb4c7 --- /dev/null +++ b/tests/unit_tests/_primaite/_simulator/_network/_hardware/test_node.py @@ -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") diff --git a/tests/unit_tests/_primaite/_simulator/_network/_transmission/test_data_link_layer.py b/tests/unit_tests/_primaite/_simulator/_network/_transmission/test_data_link_layer.py index 83b215ca..8a78d1bc 100644 --- a/tests/unit_tests/_primaite/_simulator/_network/_transmission/test_data_link_layer.py +++ b/tests/unit_tests/_primaite/_simulator/_network/_transmission/test_data_link_layer.py @@ -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 diff --git a/tests/unit_tests/_primaite/_simulator/_network/_transmission/test_network_layer.py b/tests/unit_tests/_primaite/_simulator/_network/_transmission/test_network_layer.py index 584ff25d..a7189452 100644 --- a/tests/unit_tests/_primaite/_simulator/_network/_transmission/test_network_layer.py +++ b/tests/unit_tests/_primaite/_simulator/_network/_transmission/test_network_layer.py @@ -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) diff --git a/tests/unit_tests/_primaite/_simulator/_network/test_container.py b/tests/unit_tests/_primaite/_simulator/_network/test_container.py new file mode 100644 index 00000000..290e7cc3 --- /dev/null +++ b/tests/unit_tests/_primaite/_simulator/_network/test_container.py @@ -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 diff --git a/tests/unit_tests/_primaite/_simulator/test_core.py b/tests/unit_tests/_primaite/_simulator/test_core.py index de0732f9..0d227633 100644 --- a/tests/unit_tests/_primaite/_simulator/test_core.py +++ b/tests/unit_tests/_primaite/_simulator/test_core.py @@ -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) diff --git a/tests/unit_tests/_primaite/_simulator/test_sim_conatiner.py b/tests/unit_tests/_primaite/_simulator/test_sim_conatiner.py new file mode 100644 index 00000000..4543259d --- /dev/null +++ b/tests/unit_tests/_primaite/_simulator/test_sim_conatiner.py @@ -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