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

port refactor - eod commit, airspace still broken

Browse Source
This commit is contained in:
Marek Wolan
2024-09-18 16:02:25 +01:00
parent 310876cd3b
commit dd931d900b
92 changed files with 957 additions and 682 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
tests/conftest.py
View File

@@ -45,8 +45,8 @@ class DummyService(Service):
def __init__(self, **kwargs):
kwargs["name"] = "DummyService"
kwargs["port"] = Port.HTTP
kwargs["protocol"] = IPProtocol.TCP
kwargs["port"] = Port["HTTP"]
kwargs["protocol"] = IPProtocol["TCP"]
super().__init__(**kwargs)
def receive(self, payload: Any, session_id: str, **kwargs) -> bool:
@@ -58,8 +58,8 @@ class DummyApplication(Application, identifier="DummyApplication"):
def __init__(self, **kwargs):
kwargs["name"] = "DummyApplication"
kwargs["port"] = Port.HTTP
kwargs["protocol"] = IPProtocol.TCP
kwargs["port"] = Port["HTTP"]
kwargs["protocol"] = IPProtocol["TCP"]
super().__init__(**kwargs)
def describe_state(self) -> Dict:
@@ -77,7 +77,7 @@ def uc2_network() -> Network:
@pytest.fixture(scope="function")
def service(file_system) -> DummyService:
return DummyService(
name="DummyService", port=Port.ARP, file_system=file_system, sys_log=SysLog(hostname="dummy_service")
name="DummyService", port=Port["ARP"], file_system=file_system, sys_log=SysLog(hostname="dummy_service")
)
@@ -90,7 +90,7 @@ def service_class():
def application(file_system) -> DummyApplication:
return DummyApplication(
name="DummyApplication",
port=Port.ARP,
port=Port["ARP"],
file_system=file_system,
sys_log=SysLog(hostname="dummy_application"),
)
@@ -350,10 +350,10 @@ def install_stuff_to_sim(sim: Simulation):
network.connect(endpoint_a=server_2.network_interface[1], endpoint_b=switch_2.network_interface[2])
# 2: Configure base ACL
router.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.ARP, dst_port=Port.ARP, position=22)
router.acl.add_rule(action=ACLAction.PERMIT, protocol=IPProtocol.ICMP, position=23)
router.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.DNS, dst_port=Port.DNS, position=1)
router.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.HTTP, dst_port=Port.HTTP, position=3)
router.acl.add_rule(action=ACLAction.PERMIT, src_port=Port["ARP"], dst_port=Port["ARP"], position=22)
router.acl.add_rule(action=ACLAction.PERMIT, protocol=IPProtocol["ICMP"], position=23)
router.acl.add_rule(action=ACLAction.PERMIT, src_port=Port["DNS"], dst_port=Port["DNS"], position=1)
router.acl.add_rule(action=ACLAction.PERMIT, src_port=Port["HTTP"], dst_port=Port["HTTP"], position=3)
# 3: Install server software
server_1.software_manager.install(DNSServer)
@@ -379,13 +379,13 @@ def install_stuff_to_sim(sim: Simulation):
r = sim.network.router_nodes[0]
for i, acl_rule in enumerate(r.acl.acl):
if i == 1:
assert acl_rule.src_port == acl_rule.dst_port == Port.DNS
assert acl_rule.src_port == acl_rule.dst_port == Port["DNS"]
elif i == 3:
assert acl_rule.src_port == acl_rule.dst_port == Port.HTTP
assert acl_rule.src_port == acl_rule.dst_port == Port["HTTP"]
elif i == 22:
assert acl_rule.src_port == acl_rule.dst_port == Port.ARP
assert acl_rule.src_port == acl_rule.dst_port == Port["ARP"]
elif i == 23:
assert acl_rule.protocol == IPProtocol.ICMP
assert acl_rule.protocol == IPProtocol["ICMP"]
elif i == 24:
...
else:

32
tests/integration_tests/configuration_file_parsing/nodes/network/test_firewall_config.py
View File

@@ -68,44 +68,44 @@ def test_firewall_acl_rules_correctly_added(dmz_config):
# ICMP and ARP should be allowed internal_inbound
assert firewall.internal_inbound_acl.num_rules == 2
assert firewall.internal_inbound_acl.acl[22].action == ACLAction.PERMIT
assert firewall.internal_inbound_acl.acl[22].src_port == Port.ARP
assert firewall.internal_inbound_acl.acl[22].dst_port == Port.ARP
assert firewall.internal_inbound_acl.acl[22].src_port == Port["ARP"]
assert firewall.internal_inbound_acl.acl[22].dst_port == Port["ARP"]
assert firewall.internal_inbound_acl.acl[23].action == ACLAction.PERMIT
assert firewall.internal_inbound_acl.acl[23].protocol == IPProtocol.ICMP
assert firewall.internal_inbound_acl.acl[23].protocol == IPProtocol["ICMP"]
assert firewall.internal_inbound_acl.implicit_action == ACLAction.DENY
# ICMP and ARP should be allowed internal_outbound
assert firewall.internal_outbound_acl.num_rules == 2
assert firewall.internal_outbound_acl.acl[22].action == ACLAction.PERMIT
assert firewall.internal_outbound_acl.acl[22].src_port == Port.ARP
assert firewall.internal_outbound_acl.acl[22].dst_port == Port.ARP
assert firewall.internal_outbound_acl.acl[22].src_port == Port["ARP"]
assert firewall.internal_outbound_acl.acl[22].dst_port == Port["ARP"]
assert firewall.internal_outbound_acl.acl[23].action == ACLAction.PERMIT
assert firewall.internal_outbound_acl.acl[23].protocol == IPProtocol.ICMP
assert firewall.internal_outbound_acl.acl[23].protocol == IPProtocol["ICMP"]
assert firewall.internal_outbound_acl.implicit_action == ACLAction.DENY
# ICMP and ARP should be allowed dmz_inbound
assert firewall.dmz_inbound_acl.num_rules == 2
assert firewall.dmz_inbound_acl.acl[22].action == ACLAction.PERMIT
assert firewall.dmz_inbound_acl.acl[22].src_port == Port.ARP
assert firewall.dmz_inbound_acl.acl[22].dst_port == Port.ARP
assert firewall.dmz_inbound_acl.acl[22].src_port == Port["ARP"]
assert firewall.dmz_inbound_acl.acl[22].dst_port == Port["ARP"]
assert firewall.dmz_inbound_acl.acl[23].action == ACLAction.PERMIT
assert firewall.dmz_inbound_acl.acl[23].protocol == IPProtocol.ICMP
assert firewall.dmz_inbound_acl.acl[23].protocol == IPProtocol["ICMP"]
assert firewall.dmz_inbound_acl.implicit_action == ACLAction.DENY
# ICMP and ARP should be allowed dmz_outbound
assert firewall.dmz_outbound_acl.num_rules == 2
assert firewall.dmz_outbound_acl.acl[22].action == ACLAction.PERMIT
assert firewall.dmz_outbound_acl.acl[22].src_port == Port.ARP
assert firewall.dmz_outbound_acl.acl[22].dst_port == Port.ARP
assert firewall.dmz_outbound_acl.acl[22].src_port == Port["ARP"]
assert firewall.dmz_outbound_acl.acl[22].dst_port == Port["ARP"]
assert firewall.dmz_outbound_acl.acl[23].action == ACLAction.PERMIT
assert firewall.dmz_outbound_acl.acl[23].protocol == IPProtocol.ICMP
assert firewall.dmz_outbound_acl.acl[23].protocol == IPProtocol["ICMP"]
assert firewall.dmz_outbound_acl.implicit_action == ACLAction.DENY
# ICMP and ARP should be allowed external_inbound
assert firewall.external_inbound_acl.num_rules == 1
assert firewall.external_inbound_acl.acl[22].action == ACLAction.PERMIT
assert firewall.external_inbound_acl.acl[22].src_port == Port.ARP
assert firewall.external_inbound_acl.acl[22].dst_port == Port.ARP
assert firewall.external_inbound_acl.acl[22].src_port == Port["ARP"]
assert firewall.external_inbound_acl.acl[22].dst_port == Port["ARP"]
# external_inbound should have implicit action PERMIT
# ICMP does not have a provided ACL Rule but implicit action should allow anything
assert firewall.external_inbound_acl.implicit_action == ACLAction.PERMIT
@@ -113,8 +113,8 @@ def test_firewall_acl_rules_correctly_added(dmz_config):
# ICMP and ARP should be allowed external_outbound
assert firewall.external_outbound_acl.num_rules == 1
assert firewall.external_outbound_acl.acl[22].action == ACLAction.PERMIT
assert firewall.external_outbound_acl.acl[22].src_port == Port.ARP
assert firewall.external_outbound_acl.acl[22].dst_port == Port.ARP
assert firewall.external_outbound_acl.acl[22].src_port == Port["ARP"]
assert firewall.external_outbound_acl.acl[22].dst_port == Port["ARP"]
# external_outbound should have implicit action PERMIT
# ICMP does not have a provided ACL Rule but implicit action should allow anything
assert firewall.external_outbound_acl.implicit_action == ACLAction.PERMIT

6
tests/integration_tests/configuration_file_parsing/nodes/network/test_router_config.py
View File

@@ -63,8 +63,8 @@ def test_router_acl_rules_correctly_added(dmz_config):
# ICMP and ARP should be allowed
assert router_1.acl.num_rules == 2
assert router_1.acl.acl[22].action == ACLAction.PERMIT
assert router_1.acl.acl[22].src_port == Port.ARP
assert router_1.acl.acl[22].dst_port == Port.ARP
assert router_1.acl.acl[22].src_port == Port["ARP"]
assert router_1.acl.acl[22].dst_port == Port["ARP"]
assert router_1.acl.acl[23].action == ACLAction.PERMIT
assert router_1.acl.acl[23].protocol == IPProtocol.ICMP
assert router_1.acl.acl[23].protocol == IPProtocol["ICMP"]
assert router_1.acl.implicit_action == ACLAction.DENY

8
tests/integration_tests/extensions/applications/extended_application.py
View File

@@ -44,10 +44,10 @@ class ExtendedApplication(Application, identifier="ExtendedApplication"):
def __init__(self, **kwargs):
kwargs["name"] = "ExtendedApplication"
kwargs["protocol"] = IPProtocol.TCP
kwargs["protocol"] = IPProtocol["TCP"]
# default for web is port 80
if kwargs.get("port") is None:
kwargs["port"] = Port.HTTP
kwargs["port"] = Port["HTTP"]
super().__init__(**kwargs)
self.run()
@@ -127,7 +127,7 @@ class ExtendedApplication(Application, identifier="ExtendedApplication"):
if self.send(
payload=payload,
dest_ip_address=self.domain_name_ip_address,
dest_port=parsed_url.port if parsed_url.port else Port.HTTP,
dest_port=parsed_url.port if parsed_url.port else Port["HTTP"],
):
self.sys_log.info(
f"{self.name}: Received HTTP {payload.request_method.name} "
@@ -155,7 +155,7 @@ class ExtendedApplication(Application, identifier="ExtendedApplication"):
self,
payload: HttpRequestPacket,
dest_ip_address: Optional[IPv4Address] = None,
dest_port: Optional[Port] = Port.HTTP,
dest_port: Optional[int] = Port["HTTP"],
session_id: Optional[str] = None,
**kwargs,
) -> bool:

4
tests/integration_tests/extensions/services/extended_service.py
View File

@@ -38,8 +38,8 @@ class ExtendedService(Service, identifier='extendedservice'):
def __init__(self, **kwargs):
kwargs["name"] = "ExtendedService"
kwargs["port"] = Port.POSTGRES_SERVER
kwargs["protocol"] = IPProtocol.TCP
kwargs["port"] = Port["POSTGRES_SERVER"]
kwargs["protocol"] = IPProtocol["TCP"]
super().__init__(**kwargs)
self._create_db_file()
if kwargs.get('options'):

6
tests/integration_tests/game_layer/actions/test_c2_suite_actions.py
View File

@@ -26,9 +26,9 @@ def game_and_agent_fixture(game_and_agent):
game, agent = game_and_agent
router = game.simulation.network.get_node_by_hostname("router")
router.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.HTTP, dst_port=Port.HTTP, position=4)
router.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.DNS, dst_port=Port.DNS, position=5)
router.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.FTP, dst_port=Port.FTP, position=6)
router.acl.add_rule(action=ACLAction.PERMIT, src_port=Port["HTTP"], dst_port=Port["HTTP"], position=4)
router.acl.add_rule(action=ACLAction.PERMIT, src_port=Port["DNS"], dst_port=Port["DNS"], position=5)
router.acl.add_rule(action=ACLAction.PERMIT, src_port=Port["FTP"], dst_port=Port["FTP"], position=6)
c2_server_host = game.simulation.network.get_node_by_hostname("client_1")
c2_server_host.software_manager.install(software_class=C2Server)

2
tests/integration_tests/game_layer/actions/test_configure_actions.py
View File

@@ -200,7 +200,7 @@ class TestConfigureDoSBot:
game.step()
assert dos_bot.target_ip_address == IPv4Address("192.168.1.99")
assert dos_bot.target_port == Port.POSTGRES_SERVER
assert dos_bot.target_port == Port["POSTGRES_SERVER"]
assert dos_bot.payload == "HACC"
assert not dos_bot.repeat
assert dos_bot.port_scan_p_of_success == 0.875

2
tests/integration_tests/game_layer/actions/test_terminal_actions.py
View File

@@ -20,7 +20,7 @@ def game_and_agent_fixture(game_and_agent):
game, agent = game_and_agent
router = game.simulation.network.get_node_by_hostname("router")
router.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.SSH, dst_port=Port.SSH, position=4)
router.acl.add_rule(action=ACLAction.PERMIT, src_port=Port["SSH"], dst_port=Port["SSH"], position=4)
return (game, agent)

2
tests/integration_tests/game_layer/observations/test_acl_observations.py
View File

@@ -33,7 +33,7 @@ def test_acl_observations(simulation):
server.software_manager.install(NTPServer)
# add router acl rule
router.acl.add_rule(action=ACLAction.PERMIT, dst_port=Port.NTP, src_port=Port.NTP, position=1)
router.acl.add_rule(action=ACLAction.PERMIT, dst_port=Port["NTP"], src_port=Port["NTP"], position=1)
acl_obs = ACLObservation(
where=["network", "nodes", router.hostname, "acl", "acl"],

6
tests/integration_tests/game_layer/observations/test_firewall_observation.py
View File

@@ -62,13 +62,13 @@ def test_firewall_observation():
# add a rule to the internal inbound and check that the observation is correct
firewall.internal_inbound_acl.add_rule(
action=ACLAction.DENY,
protocol=IPProtocol.TCP,
protocol=IPProtocol["TCP"],
src_ip_address="10.0.0.1",
src_wildcard_mask="0.0.0.1",
dst_ip_address="10.0.0.2",
dst_wildcard_mask="0.0.0.1",
src_port=Port.HTTP,
dst_port=Port.HTTP,
src_port=Port["HTTP"],
dst_port=Port["HTTP"],
position=5,
)

2
tests/integration_tests/game_layer/observations/test_nic_observations.py
View File

@@ -152,7 +152,7 @@ def test_config_nic_categories(simulation):
def test_nic_monitored_traffic(simulation):
monitored_traffic = {"icmp": ["NONE"], "tcp": ["DNS"]}
monitored_traffic = {"icmp": ["NONE"], "tcp": [53,]}
pc: Computer = simulation.network.get_node_by_hostname("client_1")
pc2: Computer = simulation.network.get_node_by_hostname("client_2")

6
tests/integration_tests/game_layer/observations/test_router_observation.py
View File

@@ -39,13 +39,13 @@ def test_router_observation():
# Add an ACL rule to the router
router.acl.add_rule(
action=ACLAction.DENY,
protocol=IPProtocol.TCP,
protocol=IPProtocol["TCP"],
src_ip_address="10.0.0.1",
src_wildcard_mask="0.0.0.1",
dst_ip_address="10.0.0.2",
dst_wildcard_mask="0.0.0.1",
src_port=Port.HTTP,
dst_port=Port.HTTP,
src_port=Port["HTTP"],
dst_port=Port["HTTP"],
position=5,
)
# Observe the state using the RouterObservation instance

2
tests/integration_tests/game_layer/observations/test_user_observations.py
View File

@@ -15,7 +15,7 @@ def env_with_ssh() -> PrimaiteGymEnv:
env = PrimaiteGymEnv(DATA_MANIPULATION_CONFIG)
env.agent.flatten_obs = False
router: Router = env.game.simulation.network.get_node_by_hostname("router_1")
router.acl.add_rule(ACLAction.PERMIT, src_port=Port.SSH, dst_port=Port.SSH, position=3)
router.acl.add_rule(ACLAction.PERMIT, src_port=Port["SSH"], dst_port=Port["SSH"], position=3)
return env

24
tests/integration_tests/game_layer/test_actions.py
View File

@@ -608,9 +608,9 @@ def test_firewall_acl_add_remove_rule_integration():
assert firewall.internal_outbound_acl.acl[1].action.name == "DENY"
assert firewall.internal_outbound_acl.acl[1].src_ip_address == IPv4Address("192.168.0.10")
assert firewall.internal_outbound_acl.acl[1].dst_ip_address is None
assert firewall.internal_outbound_acl.acl[1].dst_port == Port.DNS
assert firewall.internal_outbound_acl.acl[1].src_port == Port.ARP
assert firewall.internal_outbound_acl.acl[1].protocol == IPProtocol.ICMP
assert firewall.internal_outbound_acl.acl[1].dst_port == Port["DNS"]
assert firewall.internal_outbound_acl.acl[1].src_port == Port["ARP"]
assert firewall.internal_outbound_acl.acl[1].protocol == IPProtocol["ICMP"]
env.step(4) # Remove ACL rule from Internal Outbound
assert firewall.internal_outbound_acl.num_rules == 2
@@ -620,9 +620,9 @@ def test_firewall_acl_add_remove_rule_integration():
assert firewall.dmz_inbound_acl.acl[1].action.name == "DENY"
assert firewall.dmz_inbound_acl.acl[1].src_ip_address == IPv4Address("192.168.10.10")
assert firewall.dmz_inbound_acl.acl[1].dst_ip_address == IPv4Address("192.168.0.10")
assert firewall.dmz_inbound_acl.acl[1].dst_port == Port.HTTP
assert firewall.dmz_inbound_acl.acl[1].src_port == Port.HTTP
assert firewall.dmz_inbound_acl.acl[1].protocol == IPProtocol.UDP
assert firewall.dmz_inbound_acl.acl[1].dst_port == Port["HTTP"]
assert firewall.dmz_inbound_acl.acl[1].src_port == Port["HTTP"]
assert firewall.dmz_inbound_acl.acl[1].protocol == IPProtocol["UDP"]
env.step(6) # Remove ACL rule from DMZ Inbound
assert firewall.dmz_inbound_acl.num_rules == 2
@@ -632,9 +632,9 @@ def test_firewall_acl_add_remove_rule_integration():
assert firewall.dmz_outbound_acl.acl[2].action.name == "DENY"
assert firewall.dmz_outbound_acl.acl[2].src_ip_address == IPv4Address("192.168.10.10")
assert firewall.dmz_outbound_acl.acl[2].dst_ip_address == IPv4Address("192.168.0.10")
assert firewall.dmz_outbound_acl.acl[2].dst_port == Port.HTTP
assert firewall.dmz_outbound_acl.acl[2].src_port == Port.HTTP
assert firewall.dmz_outbound_acl.acl[2].protocol == IPProtocol.TCP
assert firewall.dmz_outbound_acl.acl[2].dst_port == Port["HTTP"]
assert firewall.dmz_outbound_acl.acl[2].src_port == Port["HTTP"]
assert firewall.dmz_outbound_acl.acl[2].protocol == IPProtocol["TCP"]
env.step(8) # Remove ACL rule from DMZ Outbound
assert firewall.dmz_outbound_acl.num_rules == 2
@@ -644,9 +644,9 @@ def test_firewall_acl_add_remove_rule_integration():
assert firewall.external_inbound_acl.acl[10].action.name == "DENY"
assert firewall.external_inbound_acl.acl[10].src_ip_address == IPv4Address("192.168.20.10")
assert firewall.external_inbound_acl.acl[10].dst_ip_address == IPv4Address("192.168.10.10")
assert firewall.external_inbound_acl.acl[10].dst_port == Port.POSTGRES_SERVER
assert firewall.external_inbound_acl.acl[10].src_port == Port.POSTGRES_SERVER
assert firewall.external_inbound_acl.acl[10].protocol == IPProtocol.ICMP
assert firewall.external_inbound_acl.acl[10].dst_port == Port["POSTGRES_SERVER"]
assert firewall.external_inbound_acl.acl[10].src_port == Port["POSTGRES_SERVER"]
assert firewall.external_inbound_acl.acl[10].protocol == IPProtocol["ICMP"]
env.step(10) # Remove ACL rule from External Inbound
assert firewall.external_inbound_acl.num_rules == 1

4
tests/integration_tests/game_layer/test_rewards.py
View File

@@ -42,7 +42,7 @@ def test_WebpageUnavailablePenalty(game_and_agent):
# Block the web traffic, check that failing to fetch the webpage yields a reward of -0.7
router: Router = game.simulation.network.get_node_by_hostname("router")
router.acl.add_rule(action=ACLAction.DENY, protocol=IPProtocol.TCP, src_port=Port.HTTP, dst_port=Port.HTTP)
router.acl.add_rule(action=ACLAction.DENY, protocol=IPProtocol["TCP"], src_port=Port["HTTP"], dst_port=Port["HTTP"])
agent.store_action(("NODE_APPLICATION_EXECUTE", {"node_id": 0, "application_id": 0}))
game.step()
assert agent.reward_function.current_reward == -0.7
@@ -65,7 +65,7 @@ def test_uc2_rewards(game_and_agent):
db_client.run()
router: Router = game.simulation.network.get_node_by_hostname("router")
router.acl.add_rule(ACLAction.PERMIT, src_port=Port.POSTGRES_SERVER, dst_port=Port.POSTGRES_SERVER, position=2)
router.acl.add_rule(ACLAction.PERMIT, src_port=Port["POSTGRES_SERVER"], dst_port=Port["POSTGRES_SERVER"], position=2)
comp = GreenAdminDatabaseUnreachablePenalty("client_1")

8
tests/integration_tests/network/test_airspace_config.py
View File

@@ -13,8 +13,8 @@ def test_override_freq_max_capacity_mbps():
config_dict = yaml.safe_load(f)
network = PrimaiteGame.from_config(cfg=config_dict).simulation.network
assert network.airspace.get_frequency_max_capacity_mbps(AirSpaceFrequency.WIFI_2_4) == 123.45
assert network.airspace.get_frequency_max_capacity_mbps(AirSpaceFrequency.WIFI_5) == 0.0
assert network.airspace.get_frequency_max_capacity_mbps(AirSpaceFrequency["WIFI_2_4"]) == 123.45
assert network.airspace.get_frequency_max_capacity_mbps(AirSpaceFrequency["WIFI_5"]) == 0.0
pc_a = network.get_node_by_hostname("pc_a")
pc_b = network.get_node_by_hostname("pc_b")
@@ -32,8 +32,8 @@ def test_override_freq_max_capacity_mbps_blocked():
config_dict = yaml.safe_load(f)
network = PrimaiteGame.from_config(cfg=config_dict).simulation.network
assert network.airspace.get_frequency_max_capacity_mbps(AirSpaceFrequency.WIFI_2_4) == 0.0
assert network.airspace.get_frequency_max_capacity_mbps(AirSpaceFrequency.WIFI_5) == 0.0
assert network.airspace.get_frequency_max_capacity_mbps(AirSpaceFrequency["WIFI_2_4"]) == 0.0
assert network.airspace.get_frequency_max_capacity_mbps(AirSpaceFrequency["WIFI_5"]) == 0.0
pc_a = network.get_node_by_hostname("pc_a")
pc_b = network.get_node_by_hostname("pc_b")

12
tests/integration_tests/network/test_broadcast.py
View File

@@ -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["HTTP"]
kwargs["protocol"] = IPProtocol["TCP"]
super().__init__(**kwargs)
def describe_state(self) -> Dict:
@@ -33,12 +33,12 @@ class BroadcastTestService(Service):
super().send(
payload="unicast",
dest_ip_address=ip_address,
dest_port=Port.HTTP,
dest_port=Port["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["HTTP"], ip_protocol=self.protocol)
class BroadcastTestClient(Application, identifier="BroadcastTestClient"):
@@ -49,8 +49,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["HTTP"]
kwargs["protocol"] = IPProtocol["TCP"]
super().__init__(**kwargs)
def describe_state(self) -> Dict:

32
tests/integration_tests/network/test_firewall.py
View File

@@ -53,28 +53,28 @@ 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=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)
# 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["ARP"], dst_port=Port["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["ARP"], dst_port=Port["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["ARP"], dst_port=Port["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["ARP"], dst_port=Port["ARP"], position=22
)
firewall_node.dmz_inbound_acl.add_rule(action=ACLAction.PERMIT, src_port=Port.ARP, dst_port=Port.ARP, position=22)
firewall_node.dmz_outbound_acl.add_rule(action=ACLAction.PERMIT, src_port=Port.ARP, dst_port=Port.ARP, position=22)
firewall_node.dmz_inbound_acl.add_rule(action=ACLAction.PERMIT, src_port=Port["ARP"], dst_port=Port["ARP"], position=22)
firewall_node.dmz_outbound_acl.add_rule(action=ACLAction.PERMIT, src_port=Port["ARP"], dst_port=Port["ARP"], position=22)
# external node
external_node = Computer(
@@ -262,8 +262,8 @@ 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)
firewall.internal_inbound_acl.add_rule(action=ACLAction.PERMIT, src_port=Port.NTP, dst_port=Port.NTP, position=1)
firewall.internal_outbound_acl.add_rule(action=ACLAction.PERMIT, src_port=Port["NTP"], dst_port=Port["NTP"], position=1)
firewall.internal_inbound_acl.add_rule(action=ACLAction.PERMIT, src_port=Port["NTP"], dst_port=Port["NTP"], position=1)
internal_ntp_client.request_time()
@@ -271,8 +271,8 @@ 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["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)
dmz_ntp_client.request_time()

8
tests/integration_tests/network/test_routing.py
View File

@@ -73,8 +73,8 @@ 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["ARP"], dst_port=Port["ARP"], position=22)
router_1.acl.add_rule(action=ACLAction.PERMIT, protocol=IPProtocol["ICMP"], position=23)
# Configure PC B
pc_b = Computer(
@@ -197,8 +197,8 @@ 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["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)
assert ntp_client.time is None
ntp_client.request_time()

4
tests/integration_tests/network/test_wireless_router.py
View File

@@ -37,8 +37,8 @@ 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["ARP"], dst_port=Port["ARP"], position=22)
router_1.acl.add_rule(action=ACLAction.PERMIT, protocol=IPProtocol["ICMP"], position=23)
# Configure PC B
pc_b = Computer(

20
tests/integration_tests/system/red_applications/test_c2_suite_integration.py
View File

@@ -227,7 +227,7 @@ def test_c2_suite_acl_block(basic_network):
assert c2_beacon.c2_connection_active == True
# Now we add a HTTP blocking acl (Thus preventing a keep alive)
router.acl.add_rule(action=ACLAction.DENY, src_port=Port.HTTP, dst_port=Port.HTTP, position=0)
router.acl.add_rule(action=ACLAction.DENY, src_port=Port["HTTP"], dst_port=Port["HTTP"], position=0)
c2_beacon.apply_timestep(2)
c2_beacon.apply_timestep(3)
@@ -322,8 +322,8 @@ def test_c2_suite_acl_bypass(basic_network):
################ Confirm Default Setup #########################
# Permitting all HTTP & FTP traffic
router.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.HTTP, dst_port=Port.HTTP, position=0)
router.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.FTP, dst_port=Port.FTP, position=1)
router.acl.add_rule(action=ACLAction.PERMIT, src_port=Port["HTTP"], dst_port=Port["HTTP"], position=0)
router.acl.add_rule(action=ACLAction.PERMIT, src_port=Port["FTP"], dst_port=Port["FTP"], position=1)
c2_beacon.apply_timestep(0)
assert c2_beacon.keep_alive_inactivity == 1
@@ -337,7 +337,7 @@ def test_c2_suite_acl_bypass(basic_network):
################ Denying HTTP Traffic #########################
# Now we add a HTTP blocking acl (Thus preventing a keep alive)
router.acl.add_rule(action=ACLAction.DENY, src_port=Port.HTTP, dst_port=Port.HTTP, position=0)
router.acl.add_rule(action=ACLAction.DENY, src_port=Port["HTTP"], dst_port=Port["HTTP"], position=0)
blocking_acl: AccessControlList = router.acl.acl[0]
# Asserts to show the C2 Suite is unable to maintain connection:
@@ -359,8 +359,8 @@ def test_c2_suite_acl_bypass(basic_network):
c2_beacon.configure(
c2_server_ip_address="192.168.0.2",
keep_alive_frequency=2,
masquerade_port=Port.FTP,
masquerade_protocol=IPProtocol.TCP,
masquerade_port=Port["FTP"],
masquerade_protocol=IPProtocol["TCP"],
)
c2_beacon.establish()
@@ -407,8 +407,8 @@ def test_c2_suite_acl_bypass(basic_network):
################ Denying FTP Traffic & Enable HTTP #########################
# Blocking FTP and re-permitting HTTP:
router.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.HTTP, dst_port=Port.HTTP, position=0)
router.acl.add_rule(action=ACLAction.DENY, src_port=Port.FTP, dst_port=Port.FTP, position=1)
router.acl.add_rule(action=ACLAction.PERMIT, src_port=Port["HTTP"], dst_port=Port["HTTP"], position=0)
router.acl.add_rule(action=ACLAction.DENY, src_port=Port["FTP"], dst_port=Port["FTP"], position=1)
blocking_acl: AccessControlList = router.acl.acl[1]
# Asserts to show the C2 Suite is unable to maintain connection:
@@ -430,8 +430,8 @@ def test_c2_suite_acl_bypass(basic_network):
c2_beacon.configure(
c2_server_ip_address="192.168.0.2",
keep_alive_frequency=2,
masquerade_port=Port.HTTP,
masquerade_protocol=IPProtocol.TCP,
masquerade_port=Port["HTTP"],
masquerade_protocol=IPProtocol["TCP"],
)
c2_beacon.establish()

2
tests/integration_tests/system/red_applications/test_data_manipulation_bot_and_server.py
View File

@@ -52,7 +52,7 @@ def data_manipulation_db_server_green_client(example_network) -> Network:
router_1: Router = example_network.get_node_by_hostname("router_1")
router_1.acl.add_rule(
action=ACLAction.PERMIT, src_port=Port.POSTGRES_SERVER, dst_port=Port.POSTGRES_SERVER, position=0
action=ACLAction.PERMIT, src_port=Port["POSTGRES_SERVER"], dst_port=Port["POSTGRES_SERVER"], position=0
)
client_1: Computer = network.get_node_by_hostname("client_1")

6
tests/integration_tests/system/red_applications/test_dos_bot_and_server.py
View File

@@ -26,7 +26,7 @@ def dos_bot_and_db_server(client_server) -> Tuple[DoSBot, Computer, DatabaseServ
dos_bot: DoSBot = computer.software_manager.software.get("DoSBot")
dos_bot.configure(
target_ip_address=IPv4Address(server.network_interface[1].ip_address),
target_port=Port.POSTGRES_SERVER,
target_port=Port["POSTGRES_SERVER"],
)
# Install DB Server service on server
@@ -43,7 +43,7 @@ def dos_bot_db_server_green_client(example_network) -> Network:
router_1: Router = example_network.get_node_by_hostname("router_1")
router_1.acl.add_rule(
action=ACLAction.PERMIT, src_port=Port.POSTGRES_SERVER, dst_port=Port.POSTGRES_SERVER, position=0
action=ACLAction.PERMIT, src_port=Port["POSTGRES_SERVER"], dst_port=Port["POSTGRES_SERVER"], position=0
)
client_1: Computer = network.get_node_by_hostname("client_1")
@@ -56,7 +56,7 @@ def dos_bot_db_server_green_client(example_network) -> Network:
dos_bot: DoSBot = client_1.software_manager.software.get("DoSBot")
dos_bot.configure(
target_ip_address=IPv4Address(server.network_interface[1].ip_address),
target_port=Port.POSTGRES_SERVER,
target_port=Port["POSTGRES_SERVER"],
)
# install db server service on server

2
tests/integration_tests/system/red_applications/test_ransomware_script.py
View File

@@ -47,7 +47,7 @@ def ransomware_script_db_server_green_client(example_network) -> Network:
router_1: Router = example_network.get_node_by_hostname("router_1")
router_1.acl.add_rule(
action=ACLAction.PERMIT, src_port=Port.POSTGRES_SERVER, dst_port=Port.POSTGRES_SERVER, position=0
action=ACLAction.PERMIT, src_port=Port["POSTGRES_SERVER"], dst_port=Port["POSTGRES_SERVER"], position=0
)
client_1: Computer = network.get_node_by_hostname("client_1")

16
tests/integration_tests/system/test_nmap.py
View File

@@ -73,10 +73,10 @@ def test_port_scan_one_node_one_port(example_network):
client_2 = network.get_node_by_hostname("client_2")
actual_result = client_1_nmap.port_scan(
target_ip_address=client_2.network_interface[1].ip_address, target_port=Port.DNS, target_protocol=IPProtocol.TCP
target_ip_address=client_2.network_interface[1].ip_address, target_port=Port["DNS"], target_protocol=IPProtocol["TCP"]
)
expected_result = {IPv4Address("192.168.10.22"): {IPProtocol.TCP: [Port.DNS]}}
expected_result = {IPv4Address("192.168.10.22"): {IPProtocol["TCP"]: [Port["DNS"]]}}
assert actual_result == expected_result
@@ -101,14 +101,14 @@ def test_port_scan_full_subnet_all_ports_and_protocols(example_network):
actual_result = client_1_nmap.port_scan(
target_ip_address=IPv4Network("192.168.10.0/24"),
target_port=[Port.ARP, Port.HTTP, Port.FTP, Port.DNS, Port.NTP],
target_port=[Port["ARP"], Port["HTTP"], Port["FTP"], Port["DNS"], Port["NTP"]],
)
expected_result = {
IPv4Address("192.168.10.1"): {IPProtocol.UDP: [Port.ARP]},
IPv4Address("192.168.10.1"): {IPProtocol["UDP"]: [Port["ARP"]]},
IPv4Address("192.168.10.22"): {
IPProtocol.TCP: [Port.HTTP, Port.FTP, Port.DNS],
IPProtocol.UDP: [Port.ARP, Port.NTP],
IPProtocol["TCP"]: [Port["HTTP"], Port["FTP"], Port["DNS"]],
IPProtocol["UDP"]: [Port["ARP"], Port["NTP"]],
},
}
@@ -122,10 +122,10 @@ def test_network_service_recon_all_ports_and_protocols(example_network):
client_1_nmap: NMAP = client_1.software_manager.software["NMAP"] # noqa
actual_result = client_1_nmap.network_service_recon(
target_ip_address=IPv4Network("192.168.10.0/24"), target_port=Port.HTTP, target_protocol=IPProtocol.TCP
target_ip_address=IPv4Network("192.168.10.0/24"), target_port=Port["HTTP"], target_protocol=IPProtocol["TCP"]
)
expected_result = {IPv4Address("192.168.10.22"): {IPProtocol.TCP: [Port.HTTP]}}
expected_result = {IPv4Address("192.168.10.22"): {IPProtocol["TCP"]: [Port["HTTP"]]}}
assert sort_dict(actual_result) == sort_dict(expected_result)

16
tests/integration_tests/system/test_service_listening_on_ports.py
View File

@@ -16,9 +16,9 @@ from tests import TEST_ASSETS_ROOT
class _DatabaseListener(Service):
name: str = "DatabaseListener"
protocol: IPProtocol = IPProtocol.TCP
port: Port = Port.NONE
listen_on_ports: Set[Port] = {Port.POSTGRES_SERVER}
protocol: str = IPProtocol["TCP"]
port: int = Port["NONE"]
listen_on_ports: Set[int] = {Port["POSTGRES_SERVER"]}
payloads_received: List[Any] = Field(default_factory=list)
def receive(self, payload: Any, session_id: str, **kwargs) -> bool:
@@ -51,8 +51,8 @@ def test_http_listener(client_server):
computer.session_manager.receive_payload_from_software_manager(
payload="masquerade as Database traffic",
dst_ip_address=server.network_interface[1].ip_address,
dst_port=Port.POSTGRES_SERVER,
ip_protocol=IPProtocol.TCP,
dst_port=Port["POSTGRES_SERVER"],
ip_protocol=IPProtocol["TCP"],
)
assert len(server_db_listener.payloads_received) == 1
@@ -76,9 +76,9 @@ def test_set_listen_on_ports_from_config():
network = PrimaiteGame.from_config(cfg=config_dict).simulation.network
client: Computer = network.get_node_by_hostname("client")
assert Port.SMB in client.software_manager.get_open_ports()
assert Port.IPP in client.software_manager.get_open_ports()
assert Port["SMB"] in client.software_manager.get_open_ports()
assert Port["IPP"] in client.software_manager.get_open_ports()
web_browser = client.software_manager.software["WebBrowser"]
assert not web_browser.listen_on_ports.difference({Port.SMB, Port.IPP})
assert not web_browser.listen_on_ports.difference({Port["SMB"], Port["IPP"]})

12
tests/integration_tests/system/test_web_client_server_and_database.py
View File

@@ -24,17 +24,17 @@ def web_client_web_server_database(example_network) -> Tuple[Network, Computer,
# add rules to network router
router_1: Router = example_network.get_node_by_hostname("router_1")
router_1.acl.add_rule(
action=ACLAction.PERMIT, src_port=Port.POSTGRES_SERVER, dst_port=Port.POSTGRES_SERVER, position=0
action=ACLAction.PERMIT, src_port=Port["POSTGRES_SERVER"], dst_port=Port["POSTGRES_SERVER"], position=0
)
# Allow DNS requests
router_1.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.DNS, dst_port=Port.DNS, position=1)
router_1.acl.add_rule(action=ACLAction.PERMIT, src_port=Port["DNS"], dst_port=Port["DNS"], position=1)
# Allow FTP requests
router_1.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.FTP, dst_port=Port.FTP, position=2)
router_1.acl.add_rule(action=ACLAction.PERMIT, src_port=Port["FTP"], dst_port=Port["FTP"], position=2)
# Open port 80 for web server
router_1.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.HTTP, dst_port=Port.HTTP, position=3)
router_1.acl.add_rule(action=ACLAction.PERMIT, src_port=Port["HTTP"], dst_port=Port["HTTP"], position=3)
# Create Computer
computer: Computer = example_network.get_node_by_hostname("client_1")
@@ -148,7 +148,7 @@ class TestWebBrowserHistory:
assert web_browser.history[-1].response_code == 200
router = network.get_node_by_hostname("router_1")
router.acl.add_rule(action=ACLAction.DENY, src_port=Port.HTTP, dst_port=Port.HTTP, position=0)
router.acl.add_rule(action=ACLAction.DENY, src_port=Port["HTTP"], dst_port=Port["HTTP"], position=0)
assert not web_browser.get_webpage()
assert len(web_browser.history) == 3
# with current NIC behaviour, even if you block communication, you won't get SERVER_UNREACHABLE because
@@ -166,7 +166,7 @@ class TestWebBrowserHistory:
web_browser.get_webpage()
router = network.get_node_by_hostname("router_1")
router.acl.add_rule(action=ACLAction.DENY, src_port=Port.HTTP, dst_port=Port.HTTP, position=0)
router.acl.add_rule(action=ACLAction.DENY, src_port=Port["HTTP"], dst_port=Port["HTTP"], position=0)
web_browser.get_webpage()
state = computer.describe_state()

4
tests/integration_tests/test_simulation/test_request_response.py
View File

@@ -171,7 +171,7 @@ class TestDataManipulationGreenRequests:
assert client_1_browser_execute.status == "success"
assert client_2_browser_execute.status == "success"
router.acl.add_rule(ACLAction.DENY, src_port=Port.HTTP, dst_port=Port.HTTP, position=3)
router.acl.add_rule(ACLAction.DENY, src_port=Port["HTTP"], dst_port=Port["HTTP"], position=3)
client_1_browser_execute = net.apply_request(["node", "client_1", "application", "WebBrowser", "execute"])
client_2_browser_execute = net.apply_request(["node", "client_2", "application", "WebBrowser", "execute"])
assert client_1_browser_execute.status == "failure"
@@ -182,7 +182,7 @@ class TestDataManipulationGreenRequests:
assert client_1_db_client_execute.status == "success"
assert client_2_db_client_execute.status == "success"
router.acl.add_rule(ACLAction.DENY, src_port=Port.POSTGRES_SERVER, dst_port=Port.POSTGRES_SERVER)
router.acl.add_rule(ACLAction.DENY, src_port=Port["POSTGRES_SERVER"], dst_port=Port["POSTGRES_SERVER"])
client_1_db_client_execute = net.apply_request(["node", "client_1", "application", "DatabaseClient", "execute"])
client_2_db_client_execute = net.apply_request(["node", "client_2", "application", "DatabaseClient", "execute"])
assert client_1_db_client_execute.status == "failure"

70
tests/unit_tests/_primaite/_simulator/_network/_hardware/nodes/test_acl.py
View File

@@ -28,20 +28,20 @@ def router_with_acl_rules():
# Add rules here as needed
acl.add_rule(
action=ACLAction.PERMIT,
protocol=IPProtocol.TCP,
protocol=IPProtocol["TCP"],
src_ip_address="192.168.1.1",
src_port=Port.HTTPS,
src_port=Port["HTTPS"],
dst_ip_address="192.168.1.2",
dst_port=Port.HTTP,
dst_port=Port["HTTP"],
position=1,
)
acl.add_rule(
action=ACLAction.DENY,
protocol=IPProtocol.TCP,
protocol=IPProtocol["TCP"],
src_ip_address="192.168.1.3",
src_port=Port(8080),
src_port=8080,
dst_ip_address="192.168.1.4",
dst_port=Port(80),
dst_port=80,
position=2,
)
return router
@@ -65,21 +65,21 @@ def router_with_wildcard_acl():
# Rule to permit traffic from a specific source IP and port to a specific destination IP and port
acl.add_rule(
action=ACLAction.PERMIT,
protocol=IPProtocol.TCP,
protocol=IPProtocol["TCP"],
src_ip_address="192.168.1.1",
src_port=Port(8080),
src_port=8080,
dst_ip_address="10.1.1.2",
dst_port=Port(80),
dst_port=80,
position=1,
)
# Rule to deny traffic from an IP range to a specific destination IP and port
acl.add_rule(
action=ACLAction.DENY,
protocol=IPProtocol.TCP,
protocol=IPProtocol["TCP"],
src_ip_address="192.168.1.0",
src_wildcard_mask="0.0.0.255",
dst_ip_address="10.1.1.3",
dst_port=Port(443),
dst_port=443,
position=2,
)
# Rule to permit any traffic to a range of destination IPs
@@ -109,11 +109,11 @@ def test_add_rule(router_with_acl_rules):
acl = router_with_acl_rules.acl
assert acl.acl[1].action == ACLAction.PERMIT
assert acl.acl[1].protocol == IPProtocol.TCP
assert acl.acl[1].protocol == IPProtocol["TCP"]
assert acl.acl[1].src_ip_address == IPv4Address("192.168.1.1")
assert acl.acl[1].src_port == Port.HTTPS
assert acl.acl[1].src_port == Port["HTTPS"]
assert acl.acl[1].dst_ip_address == IPv4Address("192.168.1.2")
assert acl.acl[1].dst_port == Port.HTTP
assert acl.acl[1].dst_port == Port["HTTP"]
def test_remove_rule(router_with_acl_rules):
@@ -136,8 +136,8 @@ def test_traffic_permitted_by_specific_rule(router_with_acl_rules):
acl = router_with_acl_rules.acl
permitted_frame = Frame(
ethernet=EthernetHeader(src_mac_addr=generate_mac_address(), dst_mac_addr=generate_mac_address()),
ip=IPPacket(src_ip_address="192.168.1.1", dst_ip_address="192.168.1.2", protocol=IPProtocol.TCP),
tcp=TCPHeader(src_port=Port.HTTPS, dst_port=Port.HTTP),
ip=IPPacket(src_ip_address="192.168.1.1", dst_ip_address="192.168.1.2", protocol=IPProtocol["TCP"]),
tcp=TCPHeader(src_port=Port["HTTPS"], dst_port=Port["HTTP"]),
)
is_permitted, _ = acl.is_permitted(permitted_frame)
assert is_permitted
@@ -153,8 +153,8 @@ def test_traffic_denied_by_specific_rule(router_with_acl_rules):
acl = router_with_acl_rules.acl
not_permitted_frame = Frame(
ethernet=EthernetHeader(src_mac_addr=generate_mac_address(), dst_mac_addr=generate_mac_address()),
ip=IPPacket(src_ip_address="192.168.1.3", dst_ip_address="192.168.1.4", protocol=IPProtocol.TCP),
tcp=TCPHeader(src_port=Port(8080), dst_port=Port(80)),
ip=IPPacket(src_ip_address="192.168.1.3", dst_ip_address="192.168.1.4", protocol=IPProtocol["TCP"]),
tcp=TCPHeader(src_port=8080, dst_port=80),
)
is_permitted, _ = acl.is_permitted(not_permitted_frame)
assert not is_permitted
@@ -173,8 +173,8 @@ def test_default_rule(router_with_acl_rules):
acl = router_with_acl_rules.acl
not_permitted_frame = Frame(
ethernet=EthernetHeader(src_mac_addr=generate_mac_address(), dst_mac_addr=generate_mac_address()),
ip=IPPacket(src_ip_address="192.168.1.5", dst_ip_address="192.168.1.12", protocol=IPProtocol.UDP),
udp=UDPHeader(src_port=Port.HTTPS, dst_port=Port.HTTP),
ip=IPPacket(src_ip_address="192.168.1.5", dst_ip_address="192.168.1.12", protocol=IPProtocol["UDP"]),
udp=UDPHeader(src_port=Port["HTTPS"], dst_port=Port["HTTP"]),
)
is_permitted, rule = acl.is_permitted(not_permitted_frame)
assert not is_permitted
@@ -189,8 +189,8 @@ def test_direct_ip_match_with_acl(router_with_wildcard_acl):
acl = router_with_wildcard_acl.acl
frame = Frame(
ethernet=EthernetHeader(src_mac_addr=generate_mac_address(), dst_mac_addr=generate_mac_address()),
ip=IPPacket(src_ip_address="192.168.1.1", dst_ip_address="10.1.1.2", protocol=IPProtocol.TCP),
tcp=TCPHeader(src_port=Port(8080), dst_port=Port(80)),
ip=IPPacket(src_ip_address="192.168.1.1", dst_ip_address="10.1.1.2", protocol=IPProtocol["TCP"]),
tcp=TCPHeader(src_port=8080, dst_port=80),
)
assert acl.is_permitted(frame)[0], "Direct IP match should be permitted."
@@ -204,8 +204,8 @@ def test_ip_range_match_denied_with_acl(router_with_wildcard_acl):
acl = router_with_wildcard_acl.acl
frame = Frame(
ethernet=EthernetHeader(src_mac_addr=generate_mac_address(), dst_mac_addr=generate_mac_address()),
ip=IPPacket(src_ip_address="192.168.1.100", dst_ip_address="10.1.1.3", protocol=IPProtocol.TCP),
tcp=TCPHeader(src_port=Port(8080), dst_port=Port(443)),
ip=IPPacket(src_ip_address="192.168.1.100", dst_ip_address="10.1.1.3", protocol=IPProtocol["TCP"]),
tcp=TCPHeader(src_port=8080, dst_port=443),
)
assert not acl.is_permitted(frame)[0], "IP range match with wildcard mask should be denied."
@@ -219,8 +219,8 @@ def test_traffic_permitted_to_destination_range_with_acl(router_with_wildcard_ac
acl = router_with_wildcard_acl.acl
frame = Frame(
ethernet=EthernetHeader(src_mac_addr=generate_mac_address(), dst_mac_addr=generate_mac_address()),
ip=IPPacket(src_ip_address="192.168.1.50", dst_ip_address="10.2.200.200", protocol=IPProtocol.UDP),
udp=UDPHeader(src_port=Port(1433), dst_port=Port(1433)),
ip=IPPacket(src_ip_address="192.168.1.50", dst_ip_address="10.2.200.200", protocol=IPProtocol["UDP"]),
udp=UDPHeader(src_port=1433, dst_port=1433),
)
assert acl.is_permitted(frame)[0], "Traffic to destination IP range should be permitted."
@@ -253,23 +253,23 @@ def test_ip_traffic_from_specific_subnet():
permitted_frame_1 = Frame(
ethernet=EthernetHeader(src_mac_addr=generate_mac_address(), dst_mac_addr=generate_mac_address()),
ip=IPPacket(src_ip_address="192.168.1.50", dst_ip_address="10.2.200.200", protocol=IPProtocol.TCP),
tcp=TCPHeader(src_port=Port.POSTGRES_SERVER, dst_port=Port.POSTGRES_SERVER),
ip=IPPacket(src_ip_address="192.168.1.50", dst_ip_address="10.2.200.200", protocol=IPProtocol["TCP"]),
tcp=TCPHeader(src_port=Port["POSTGRES_SERVER"], dst_port=Port["POSTGRES_SERVER"]),
)
assert acl.is_permitted(permitted_frame_1)[0]
permitted_frame_2 = Frame(
ethernet=EthernetHeader(src_mac_addr=generate_mac_address(), dst_mac_addr=generate_mac_address()),
ip=IPPacket(src_ip_address="192.168.1.10", dst_ip_address="85.199.214.101", protocol=IPProtocol.UDP),
udp=UDPHeader(src_port=Port.NTP, dst_port=Port.NTP),
ip=IPPacket(src_ip_address="192.168.1.10", dst_ip_address="85.199.214.101", protocol=IPProtocol["UDP"]),
udp=UDPHeader(src_port=Port["NTP"], dst_port=Port["NTP"]),
)
assert acl.is_permitted(permitted_frame_2)[0]
permitted_frame_3 = Frame(
ethernet=EthernetHeader(src_mac_addr=generate_mac_address(), dst_mac_addr=generate_mac_address()),
ip=IPPacket(src_ip_address="192.168.1.200", dst_ip_address="192.168.1.1", protocol=IPProtocol.ICMP),
ip=IPPacket(src_ip_address="192.168.1.200", dst_ip_address="192.168.1.1", protocol=IPProtocol["ICMP"]),
icmp=ICMPPacket(identifier=1),
)
@@ -277,16 +277,16 @@ def test_ip_traffic_from_specific_subnet():
not_permitted_frame_1 = Frame(
ethernet=EthernetHeader(src_mac_addr=generate_mac_address(), dst_mac_addr=generate_mac_address()),
ip=IPPacket(src_ip_address="192.168.0.50", dst_ip_address="10.2.200.200", protocol=IPProtocol.TCP),
tcp=TCPHeader(src_port=Port.POSTGRES_SERVER, dst_port=Port.POSTGRES_SERVER),
ip=IPPacket(src_ip_address="192.168.0.50", dst_ip_address="10.2.200.200", protocol=IPProtocol["TCP"]),
tcp=TCPHeader(src_port=Port["POSTGRES_SERVER"], dst_port=Port["POSTGRES_SERVER"]),
)
assert not acl.is_permitted(not_permitted_frame_1)[0]
not_permitted_frame_2 = Frame(
ethernet=EthernetHeader(src_mac_addr=generate_mac_address(), dst_mac_addr=generate_mac_address()),
ip=IPPacket(src_ip_address="192.168.2.10", dst_ip_address="85.199.214.101", protocol=IPProtocol.UDP),
udp=UDPHeader(src_port=Port.NTP, dst_port=Port.NTP),
ip=IPPacket(src_ip_address="192.168.2.10", dst_ip_address="85.199.214.101", protocol=IPProtocol["UDP"]),
udp=UDPHeader(src_port=Port["NTP"], dst_port=Port["NTP"]),
)
assert not acl.is_permitted(not_permitted_frame_2)[0]

4
tests/unit_tests/_primaite/_simulator/_network/_hardware/nodes/test_router.py
View File

@@ -67,12 +67,12 @@ def test_wireless_router_from_config():
r0 = rt.acl.acl[0]
assert r0.action == ACLAction.PERMIT
assert r0.src_port == r0.dst_port == Port.POSTGRES_SERVER
assert r0.src_port == r0.dst_port == Port["POSTGRES_SERVER"]
assert r0.src_ip_address == r0.dst_ip_address == r0.dst_wildcard_mask == r0.src_wildcard_mask == r0.protocol == None
r1 = rt.acl.acl[1]
assert r1.action == ACLAction.PERMIT
assert r1.protocol == IPProtocol.ICMP
assert r1.protocol == IPProtocol["ICMP"]
assert (
r1.src_ip_address
== r1.dst_ip_address

14
tests/unit_tests/_primaite/_simulator/_network/_transmission/test_data_link_layer.py
View File

@@ -20,7 +20,7 @@ def test_frame_minimal_instantiation():
)
# Check network layer default values
assert frame.ip.protocol == IPProtocol.TCP
assert frame.ip.protocol == IPProtocol["TCP"]
assert frame.ip.ttl == 64
assert frame.ip.precedence == Precedence.ROUTINE
@@ -40,7 +40,7 @@ def test_frame_creation_fails_tcp_without_header():
with pytest.raises(ValueError):
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_address="192.168.0.10", dst_ip_address="192.168.0.20", protocol=IPProtocol.TCP),
ip=IPPacket(src_ip_address="192.168.0.10", dst_ip_address="192.168.0.20", protocol=IPProtocol["TCP"]),
)
@@ -49,7 +49,7 @@ def test_frame_creation_fails_udp_without_header():
with pytest.raises(ValueError):
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_address="192.168.0.10", dst_ip_address="192.168.0.20", protocol=IPProtocol.UDP),
ip=IPPacket(src_ip_address="192.168.0.10", dst_ip_address="192.168.0.20", protocol=IPProtocol["UDP"]),
)
@@ -58,7 +58,7 @@ def test_frame_creation_fails_tcp_with_udp_header():
with pytest.raises(ValueError):
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_address="192.168.0.10", dst_ip_address="192.168.0.20", protocol=IPProtocol.TCP),
ip=IPPacket(src_ip_address="192.168.0.10", dst_ip_address="192.168.0.20", protocol=IPProtocol["TCP"]),
udp=UDPHeader(src_port=8080, dst_port=80),
)
@@ -68,7 +68,7 @@ def test_frame_creation_fails_udp_with_tcp_header():
with pytest.raises(ValueError):
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_address="192.168.0.10", dst_ip_address="192.168.0.20", protocol=IPProtocol.UDP),
ip=IPPacket(src_ip_address="192.168.0.10", dst_ip_address="192.168.0.20", protocol=IPProtocol["UDP"]),
udp=TCPHeader(src_port=8080, dst_port=80),
)
@@ -77,7 +77,7 @@ def test_icmp_frame_creation():
"""Tests Frame creation for ICMP."""
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_address="192.168.0.10", dst_ip_address="192.168.0.20", protocol=IPProtocol.ICMP),
ip=IPPacket(src_ip_address="192.168.0.10", dst_ip_address="192.168.0.20", protocol=IPProtocol["ICMP"]),
icmp=ICMPPacket(),
)
assert frame
@@ -88,5 +88,5 @@ def test_icmp_frame_creation_fails_without_icmp_header():
with pytest.raises(ValueError):
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_address="192.168.0.10", dst_ip_address="192.168.0.20", protocol=IPProtocol.ICMP),
ip=IPPacket(src_ip_address="192.168.0.10", dst_ip_address="192.168.0.20", protocol=IPProtocol["ICMP"]),
)

20
tests/unit_tests/_primaite/_simulator/_system/_applications/_red_applications/test_c2_suite.py
View File

@@ -129,19 +129,19 @@ def test_c2_handle_switching_port(basic_c2_network):
# Assert to confirm that both the C2 server and the C2 beacon are configured correctly.
assert c2_beacon.c2_config.keep_alive_frequency is 2
assert c2_beacon.c2_config.masquerade_port is Port.HTTP
assert c2_beacon.c2_config.masquerade_protocol is IPProtocol.TCP
assert c2_beacon.c2_config.masquerade_port is Port["HTTP"]
assert c2_beacon.c2_config.masquerade_protocol is IPProtocol["TCP"]
assert c2_server.c2_config.keep_alive_frequency is 2
assert c2_server.c2_config.masquerade_port is Port.HTTP
assert c2_server.c2_config.masquerade_protocol is IPProtocol.TCP
assert c2_server.c2_config.masquerade_port is Port["HTTP"]
assert c2_server.c2_config.masquerade_protocol is IPProtocol["TCP"]
# Configuring the C2 Beacon.
c2_beacon.configure(
c2_server_ip_address="192.168.0.1",
keep_alive_frequency=2,
masquerade_port=Port.FTP,
masquerade_protocol=IPProtocol.TCP,
masquerade_port=Port["FTP"],
masquerade_protocol=IPProtocol["TCP"],
)
# Asserting that the c2 applications have established a c2 connection
@@ -150,11 +150,11 @@ def test_c2_handle_switching_port(basic_c2_network):
# Assert to confirm that both the C2 server and the C2 beacon
# Have reconfigured their C2 settings.
assert c2_beacon.c2_config.masquerade_port is Port.FTP
assert c2_beacon.c2_config.masquerade_protocol is IPProtocol.TCP
assert c2_beacon.c2_config.masquerade_port is Port["FTP"]
assert c2_beacon.c2_config.masquerade_protocol is IPProtocol["TCP"]
assert c2_server.c2_config.masquerade_port is Port.FTP
assert c2_server.c2_config.masquerade_protocol is IPProtocol.TCP
assert c2_server.c2_config.masquerade_port is Port["FTP"]
assert c2_server.c2_config.masquerade_protocol is IPProtocol["TCP"]
def test_c2_handle_switching_frequency(basic_c2_network):

4
tests/unit_tests/_primaite/_simulator/_system/_applications/_red_applications/test_data_manipulation_bot.py
View File

@@ -27,8 +27,8 @@ def test_create_dm_bot(dm_client):
data_manipulation_bot: DataManipulationBot = dm_client.software_manager.software.get("DataManipulationBot")
assert data_manipulation_bot.name == "DataManipulationBot"
assert data_manipulation_bot.port == Port.NONE
assert data_manipulation_bot.protocol == IPProtocol.NONE
assert data_manipulation_bot.port == Port["NONE"]
assert data_manipulation_bot.protocol == IPProtocol["NONE"]
assert data_manipulation_bot.payload == "DELETE"

4
tests/unit_tests/_primaite/_simulator/_system/_applications/test_web_browser.py
View File

@@ -39,8 +39,8 @@ def test_create_web_client():
# Web Browser should be pre-installed in computer
web_browser: WebBrowser = computer.software_manager.software.get("WebBrowser")
assert web_browser.name is "WebBrowser"
assert web_browser.port is Port.HTTP
assert web_browser.protocol is IPProtocol.TCP
assert web_browser.port is Port["HTTP"]
assert web_browser.protocol is IPProtocol["TCP"]
def test_receive_invalid_payload(web_browser):

4
tests/unit_tests/_primaite/_simulator/_system/_services/test_dns_client.py
View File

@@ -28,8 +28,8 @@ def test_create_dns_client(dns_client):
assert dns_client is not None
dns_client_service: DNSClient = dns_client.software_manager.software.get("DNSClient")
assert dns_client_service.name is "DNSClient"
assert dns_client_service.port is Port.DNS
assert dns_client_service.protocol is IPProtocol.TCP
assert dns_client_service.port is Port["DNS"]
assert dns_client_service.protocol is IPProtocol["TCP"]
def test_dns_client_add_domain_to_cache_when_not_running(dns_client):

4
tests/unit_tests/_primaite/_simulator/_system/_services/test_dns_server.py
View File

@@ -32,8 +32,8 @@ def test_create_dns_server(dns_server):
assert dns_server is not None
dns_server_service: DNSServer = dns_server.software_manager.software.get("DNSServer")
assert dns_server_service.name is "DNSServer"
assert dns_server_service.port is Port.DNS
assert dns_server_service.protocol is IPProtocol.TCP
assert dns_server_service.port is Port["DNS"]
assert dns_server_service.protocol is IPProtocol["TCP"]
def test_dns_server_domain_name_registration(dns_server):

10
tests/unit_tests/_primaite/_simulator/_system/_services/test_ftp_client.py
View File

@@ -31,8 +31,8 @@ def test_create_ftp_client(ftp_client):
assert ftp_client is not None
ftp_client_service: FTPClient = ftp_client.software_manager.software.get("FTPClient")
assert ftp_client_service.name is "FTPClient"
assert ftp_client_service.port is Port.FTP
assert ftp_client_service.protocol is IPProtocol.TCP
assert ftp_client_service.port is Port["FTP"]
assert ftp_client_service.protocol is IPProtocol["TCP"]
def test_ftp_client_store_file(ftp_client):
@@ -61,7 +61,7 @@ def test_ftp_should_not_process_commands_if_service_not_running(ftp_client):
"""Method _process_ftp_command should return false if service is not running."""
payload: FTPPacket = FTPPacket(
ftp_command=FTPCommand.PORT,
ftp_command_args=Port.FTP,
ftp_command_args=Port["FTP"],
status_code=FTPStatusCode.OK,
)
@@ -102,7 +102,7 @@ def test_offline_ftp_client_receives_request(ftp_client):
payload: FTPPacket = FTPPacket(
ftp_command=FTPCommand.PORT,
ftp_command_args=Port.FTP,
ftp_command_args=Port["FTP"],
status_code=FTPStatusCode.OK,
)
@@ -119,7 +119,7 @@ def test_receive_should_ignore_payload_with_none_status_code(ftp_client):
"""Receive should ignore payload with no set status code to prevent infinite send/receive loops."""
payload: FTPPacket = FTPPacket(
ftp_command=FTPCommand.PORT,
ftp_command_args=Port.FTP,
ftp_command_args=Port["FTP"],
status_code=None,
)
ftp_client_service: FTPClient = ftp_client.software_manager.software.get("FTPClient")

4
tests/unit_tests/_primaite/_simulator/_system/_services/test_ftp_server.py
View File

@@ -30,8 +30,8 @@ def test_create_ftp_server(ftp_server):
assert ftp_server is not None
ftp_server_service: FTPServer = ftp_server.software_manager.software.get("FTPServer")
assert ftp_server_service.name is "FTPServer"
assert ftp_server_service.port is Port.FTP
assert ftp_server_service.protocol is IPProtocol.TCP
assert ftp_server_service.port is Port["FTP"]
assert ftp_server_service.protocol is IPProtocol["TCP"]
def test_ftp_server_store_file(ftp_server):

8
tests/unit_tests/_primaite/_simulator/_system/_services/test_terminal.py
View File

@@ -77,11 +77,11 @@ 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["ARP"], dst_port=Port["ARP"], position=22)
router_1.acl.add_rule(action=ACLAction.PERMIT, protocol=IPProtocol["ICMP"], position=23)
# add ACL rule to allow SSH traffic
router_1.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.SSH, dst_port=Port.SSH, position=21)
router_1.acl.add_rule(action=ACLAction.PERMIT, src_port=Port["SSH"], dst_port=Port["SSH"], position=21)
# Configure PC B
pc_b = Computer(
@@ -329,7 +329,7 @@ def test_SSH_across_network(wireless_wan_network):
terminal_a: Terminal = pc_a.software_manager.software.get("Terminal")
terminal_b: Terminal = pc_b.software_manager.software.get("Terminal")
router_2.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.SSH, dst_port=Port.SSH, position=21)
router_2.acl.add_rule(action=ACLAction.PERMIT, src_port=Port["SSH"], dst_port=Port["SSH"], position=21)
assert len(terminal_a._connections) == 0

4
tests/unit_tests/_primaite/_simulator/_system/_services/test_web_server.py
View File

@@ -33,8 +33,8 @@ def test_create_web_server(web_server):
assert web_server is not None
web_server_service: WebServer = web_server.software_manager.software.get("WebServer")
assert web_server_service.name is "WebServer"
assert web_server_service.port is Port.HTTP
assert web_server_service.protocol is IPProtocol.TCP
assert web_server_service.port is Port["HTTP"]
assert web_server_service.protocol is IPProtocol["TCP"]
def test_handling_get_request_not_found_path(web_server):

4
tests/unit_tests/_primaite/_simulator/_system/test_software.py
View File

@@ -19,10 +19,10 @@ class TestSoftware(Service):
def software(file_system):
return TestSoftware(
name="TestSoftware",
port=Port.ARP,
port=Port["ARP"],
file_system=file_system,
sys_log=SysLog(hostname="test_service"),
protocol=IPProtocol.TCP,
protocol=IPProtocol["TCP"],
)

12
tests/unit_tests/_primaite/_utils/test_dict_enum_keys_conversion.py
View File

@@ -11,7 +11,7 @@ def test_simple_conversion():
The original dictionary contains one level of nested dictionary with enums as keys.
The expected output should have string values of enums as keys.
"""
original_dict = {IPProtocol.UDP: {Port.ARP: {"inbound": 0, "outbound": 1016.0}}}
original_dict = {IPProtocol["UDP"]: {Port["ARP"]: {"inbound": 0, "outbound": 1016.0}}}
expected_dict = {"udp": {219: {"inbound": 0, "outbound": 1016.0}}}
assert convert_dict_enum_keys_to_enum_values(original_dict) == expected_dict
@@ -36,8 +36,8 @@ def test_mixed_keys():
The expected output should have string values of enums and original string keys.
"""
original_dict = {
IPProtocol.TCP: {"port": {"inbound": 0, "outbound": 1016.0}},
"protocol": {Port.HTTP: {"inbound": 10, "outbound": 2020.0}},
IPProtocol["TCP"]: {"port": {"inbound": 0, "outbound": 1016.0}},
"protocol": {Port["HTTP"]: {"inbound": 10, "outbound": 2020.0}},
}
expected_dict = {
"tcp": {"port": {"inbound": 0, "outbound": 1016.0}},
@@ -66,7 +66,7 @@ def test_nested_dicts():
The expected output should have string values of enums as keys at all levels.
"""
original_dict = {
IPProtocol.UDP: {Port.ARP: {"inbound": 0, "outbound": 1016.0, "details": {IPProtocol.TCP: {"latency": "low"}}}}
IPProtocol["UDP"]: {Port["ARP"]: {"inbound": 0, "outbound": 1016.0, "details": {IPProtocol["TCP"]: {"latency": "low"}}}}
}
expected_dict = {"udp": {219: {"inbound": 0, "outbound": 1016.0, "details": {"tcp": {"latency": "low"}}}}}
assert convert_dict_enum_keys_to_enum_values(original_dict) == expected_dict
@@ -79,6 +79,6 @@ def test_non_dict_values():
The original dictionary contains lists and tuples as values.
The expected output should preserve these non-dictionary values while converting enum keys to string values.
"""
original_dict = {IPProtocol.UDP: [Port.ARP, Port.HTTP], "protocols": (IPProtocol.TCP, IPProtocol.UDP)}
expected_dict = {"udp": [Port.ARP, Port.HTTP], "protocols": (IPProtocol.TCP, IPProtocol.UDP)}
original_dict = {IPProtocol["UDP"]: [Port["ARP"], Port["HTTP"]], "protocols": (IPProtocol["TCP"], IPProtocol["UDP"])}
expected_dict = {"udp": [Port["ARP"], Port["HTTP"]], "protocols": (IPProtocol["TCP"], IPProtocol["UDP"])}
assert convert_dict_enum_keys_to_enum_values(original_dict) == expected_dict