Add port and protocol custom validators

2024-09-20 11:21:28 +01:00
parent 08f1cf1fbd
commit 695891f55c
77 changed files with 616 additions and 613 deletions
--- a/tests/integration_tests/network/test_broadcast.py
+++ b/tests/integration_tests/network/test_broadcast.py
@@ -8,10 +8,10 @@ from primaite.simulator.network.container import Network
 from primaite.simulator.network.hardware.nodes.host.computer import Computer
 from primaite.simulator.network.hardware.nodes.host.server import Server
 from primaite.simulator.network.hardware.nodes.network.switch import Switch
-from primaite.simulator.network.transmission.network_layer import IPProtocol
-from primaite.simulator.network.transmission.transport_layer import Port
+from primaite.simulator.network.transmission.transport_layer import PORT_LOOKUP
 from primaite.simulator.system.applications.application import Application
 from primaite.simulator.system.services.service import Service
+from primaite.utils.validators import PROTOCOL_LOOKUP


 class BroadcastTestService(Service):
@@ -20,8 +20,8 @@ class BroadcastTestService(Service):
    def __init__(self, **kwargs):
        # Set default service properties for broadcasting
        kwargs["name"] = "BroadcastService"
-        kwargs["port"] = Port["HTTP"]
-        kwargs["protocol"] = IPProtocol["TCP"]
+        kwargs["port"] = PORT_LOOKUP["HTTP"]
+        kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"]
        super().__init__(**kwargs)

    def describe_state(self) -> Dict:
@@ -33,12 +33,14 @@ class BroadcastTestService(Service):
        super().send(
            payload="unicast",
            dest_ip_address=ip_address,
-            dest_port=Port["HTTP"],
+            dest_port=PORT_LOOKUP["HTTP"],
        )

    def broadcast(self, ip_network: IPv4Network):
        # Send a broadcast payload to an entire IP network
-        super().send(payload="broadcast", dest_ip_address=ip_network, dest_port=Port["HTTP"], ip_protocol=self.protocol)
+        super().send(
+            payload="broadcast", dest_ip_address=ip_network, dest_port=PORT_LOOKUP["HTTP"], ip_protocol=self.protocol
+        )


 class BroadcastTestClient(Application, identifier="BroadcastTestClient"):
@@ -49,8 +51,8 @@ class BroadcastTestClient(Application, identifier="BroadcastTestClient"):
    def __init__(self, **kwargs):
        # Set default client properties
        kwargs["name"] = "BroadcastTestClient"
-        kwargs["port"] = Port["HTTP"]
-        kwargs["protocol"] = IPProtocol["TCP"]
+        kwargs["port"] = PORT_LOOKUP["HTTP"]
+        kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"]
        super().__init__(**kwargs)

    def describe_state(self) -> Dict:
--- a/tests/integration_tests/network/test_firewall.py
+++ b/tests/integration_tests/network/test_firewall.py
@@ -7,10 +7,10 @@ from primaite.simulator.network.container import Network
 from primaite.simulator.network.hardware.nodes.host.computer import Computer
 from primaite.simulator.network.hardware.nodes.network.firewall import Firewall
 from primaite.simulator.network.hardware.nodes.network.router import ACLAction
-from primaite.simulator.network.transmission.network_layer import IPProtocol
-from primaite.simulator.network.transmission.transport_layer import Port
+from primaite.simulator.network.transmission.transport_layer import PORT_LOOKUP
 from primaite.simulator.system.services.ntp.ntp_client import NTPClient
 from primaite.simulator.system.services.ntp.ntp_server import NTPServer
+from primaite.utils.validators import PROTOCOL_LOOKUP


@pytest.fixture(scope="function")
@@ -53,31 +53,31 @@ def dmz_external_internal_network() -> Network:
    )

    # Allow ICMP
-    firewall_node.internal_inbound_acl.add_rule(action=ACLAction.PERMIT, protocol=IPProtocol["ICMP"], position=23)
-    firewall_node.internal_outbound_acl.add_rule(action=ACLAction.PERMIT, protocol=IPProtocol["ICMP"], position=23)
-    firewall_node.external_inbound_acl.add_rule(action=ACLAction.PERMIT, protocol=IPProtocol["ICMP"], position=23)
-    firewall_node.external_outbound_acl.add_rule(action=ACLAction.PERMIT, protocol=IPProtocol["ICMP"], position=23)
-    firewall_node.dmz_inbound_acl.add_rule(action=ACLAction.PERMIT, protocol=IPProtocol["ICMP"], position=23)
-    firewall_node.dmz_outbound_acl.add_rule(action=ACLAction.PERMIT, protocol=IPProtocol["ICMP"], position=23)
+    firewall_node.internal_inbound_acl.add_rule(action=ACLAction.PERMIT, protocol=PROTOCOL_LOOKUP["ICMP"], position=23)
+    firewall_node.internal_outbound_acl.add_rule(action=ACLAction.PERMIT, protocol=PROTOCOL_LOOKUP["ICMP"], position=23)
+    firewall_node.external_inbound_acl.add_rule(action=ACLAction.PERMIT, protocol=PROTOCOL_LOOKUP["ICMP"], position=23)
+    firewall_node.external_outbound_acl.add_rule(action=ACLAction.PERMIT, protocol=PROTOCOL_LOOKUP["ICMP"], position=23)
+    firewall_node.dmz_inbound_acl.add_rule(action=ACLAction.PERMIT, protocol=PROTOCOL_LOOKUP["ICMP"], position=23)
+    firewall_node.dmz_outbound_acl.add_rule(action=ACLAction.PERMIT, protocol=PROTOCOL_LOOKUP["ICMP"], position=23)

    # Allow ARP
    firewall_node.internal_inbound_acl.add_rule(
-        action=ACLAction.PERMIT, src_port=Port["ARP"], dst_port=Port["ARP"], position=22
+        action=ACLAction.PERMIT, src_port=PORT_LOOKUP["ARP"], dst_port=PORT_LOOKUP["ARP"], position=22
    )
    firewall_node.internal_outbound_acl.add_rule(
-        action=ACLAction.PERMIT, src_port=Port["ARP"], dst_port=Port["ARP"], position=22
+        action=ACLAction.PERMIT, src_port=PORT_LOOKUP["ARP"], dst_port=PORT_LOOKUP["ARP"], position=22
    )
    firewall_node.external_inbound_acl.add_rule(
-        action=ACLAction.PERMIT, src_port=Port["ARP"], dst_port=Port["ARP"], position=22
+        action=ACLAction.PERMIT, src_port=PORT_LOOKUP["ARP"], dst_port=PORT_LOOKUP["ARP"], position=22
    )
    firewall_node.external_outbound_acl.add_rule(
-        action=ACLAction.PERMIT, src_port=Port["ARP"], dst_port=Port["ARP"], position=22
+        action=ACLAction.PERMIT, src_port=PORT_LOOKUP["ARP"], dst_port=PORT_LOOKUP["ARP"], position=22
    )
    firewall_node.dmz_inbound_acl.add_rule(
-        action=ACLAction.PERMIT, src_port=Port["ARP"], dst_port=Port["ARP"], position=22
+        action=ACLAction.PERMIT, src_port=PORT_LOOKUP["ARP"], dst_port=PORT_LOOKUP["ARP"], position=22
    )
    firewall_node.dmz_outbound_acl.add_rule(
-        action=ACLAction.PERMIT, src_port=Port["ARP"], dst_port=Port["ARP"], position=22
+        action=ACLAction.PERMIT, src_port=PORT_LOOKUP["ARP"], dst_port=PORT_LOOKUP["ARP"], position=22
    )

    # external node
@@ -267,10 +267,10 @@ def test_service_allowed_with_rule(dmz_external_internal_network):
    assert not internal_ntp_client.time

    firewall.internal_outbound_acl.add_rule(
-        action=ACLAction.PERMIT, src_port=Port["NTP"], dst_port=Port["NTP"], position=1
+        action=ACLAction.PERMIT, src_port=PORT_LOOKUP["NTP"], dst_port=PORT_LOOKUP["NTP"], position=1
    )
    firewall.internal_inbound_acl.add_rule(
-        action=ACLAction.PERMIT, src_port=Port["NTP"], dst_port=Port["NTP"], position=1
+        action=ACLAction.PERMIT, src_port=PORT_LOOKUP["NTP"], dst_port=PORT_LOOKUP["NTP"], position=1
    )

    internal_ntp_client.request_time()
@@ -279,8 +279,12 @@ def test_service_allowed_with_rule(dmz_external_internal_network):

    assert not dmz_ntp_client.time

-    firewall.dmz_outbound_acl.add_rule(action=ACLAction.PERMIT, src_port=Port["NTP"], dst_port=Port["NTP"], position=1)
-    firewall.dmz_inbound_acl.add_rule(action=ACLAction.PERMIT, src_port=Port["NTP"], dst_port=Port["NTP"], position=1)
+    firewall.dmz_outbound_acl.add_rule(
+        action=ACLAction.PERMIT, src_port=PORT_LOOKUP["NTP"], dst_port=PORT_LOOKUP["NTP"], position=1
+    )
+    firewall.dmz_inbound_acl.add_rule(
+        action=ACLAction.PERMIT, src_port=PORT_LOOKUP["NTP"], dst_port=PORT_LOOKUP["NTP"], position=1
+    )

    dmz_ntp_client.request_time()

--- a/tests/integration_tests/network/test_routing.py
+++ b/tests/integration_tests/network/test_routing.py
@@ -6,10 +6,10 @@ import pytest
 from primaite.simulator.network.container import Network
 from primaite.simulator.network.hardware.nodes.host.computer import Computer
 from primaite.simulator.network.hardware.nodes.network.router import ACLAction, Router
-from primaite.simulator.network.transmission.network_layer import IPProtocol
-from primaite.simulator.network.transmission.transport_layer import Port
+from primaite.simulator.network.transmission.transport_layer import PORT_LOOKUP
 from primaite.simulator.system.services.ntp.ntp_client import NTPClient
 from primaite.simulator.system.services.ntp.ntp_server import NTPServer
+from primaite.utils.validators import PROTOCOL_LOOKUP


@pytest.fixture(scope="function")
@@ -73,8 +73,10 @@ def multi_hop_network() -> Network:
    router_1.enable_port(2)

    # Configure Router 1 ACLs
-    router_1.acl.add_rule(action=ACLAction.PERMIT, src_port=Port["ARP"], dst_port=Port["ARP"], position=22)
-    router_1.acl.add_rule(action=ACLAction.PERMIT, protocol=IPProtocol["ICMP"], position=23)
+    router_1.acl.add_rule(
+        action=ACLAction.PERMIT, src_port=PORT_LOOKUP["ARP"], dst_port=PORT_LOOKUP["ARP"], position=22
+    )
+    router_1.acl.add_rule(action=ACLAction.PERMIT, protocol=PROTOCOL_LOOKUP["ICMP"], position=23)

    # Configure PC B
    pc_b = Computer(
@@ -197,8 +199,12 @@ def test_routing_services(multi_hop_network):
    router_1: Router = multi_hop_network.get_node_by_hostname("router_1")  # noqa
    router_2: Router = multi_hop_network.get_node_by_hostname("router_2")  # noqa

-    router_1.acl.add_rule(action=ACLAction.PERMIT, src_port=Port["NTP"], dst_port=Port["NTP"], position=21)
-    router_2.acl.add_rule(action=ACLAction.PERMIT, src_port=Port["NTP"], dst_port=Port["NTP"], position=21)
+    router_1.acl.add_rule(
+        action=ACLAction.PERMIT, src_port=PORT_LOOKUP["NTP"], dst_port=PORT_LOOKUP["NTP"], position=21
+    )
+    router_2.acl.add_rule(
+        action=ACLAction.PERMIT, src_port=PORT_LOOKUP["NTP"], dst_port=PORT_LOOKUP["NTP"], position=21
+    )

    assert ntp_client.time is None
    ntp_client.request_time()
--- a/tests/integration_tests/network/test_wireless_router.py
+++ b/tests/integration_tests/network/test_wireless_router.py
@@ -7,8 +7,8 @@ from primaite.simulator.network.container import Network
 from primaite.simulator.network.hardware.nodes.host.computer import Computer
 from primaite.simulator.network.hardware.nodes.network.router import ACLAction
 from primaite.simulator.network.hardware.nodes.network.wireless_router import WirelessRouter
-from primaite.simulator.network.transmission.network_layer import IPProtocol
-from primaite.simulator.network.transmission.transport_layer import Port
+from primaite.simulator.network.transmission.transport_layer import PORT_LOOKUP
+from primaite.utils.validators import PROTOCOL_LOOKUP
 from tests import TEST_ASSETS_ROOT


@@ -37,8 +37,10 @@ def wireless_wan_network():
    network.connect(pc_a.network_interface[1], router_1.network_interface[2])

    # Configure Router 1 ACLs
-    router_1.acl.add_rule(action=ACLAction.PERMIT, src_port=Port["ARP"], dst_port=Port["ARP"], position=22)
-    router_1.acl.add_rule(action=ACLAction.PERMIT, protocol=IPProtocol["ICMP"], position=23)
+    router_1.acl.add_rule(
+        action=ACLAction.PERMIT, src_port=PORT_LOOKUP["ARP"], dst_port=PORT_LOOKUP["ARP"], position=22
+    )
+    router_1.acl.add_rule(action=ACLAction.PERMIT, protocol=PROTOCOL_LOOKUP["ICMP"], position=23)

    # Configure PC B
    pc_b = Computer(