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

Add port and protocol custom validators

Browse Source
This commit is contained in:
Marek Wolan
2024-09-20 11:21:28 +01:00
parent 08f1cf1fbd
commit 695891f55c
77 changed files with 616 additions and 613 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -11,7 +11,7 @@ from primaite.simulator.network.hardware.base import UserManager
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.router import ACLAction
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.network.transmission.transport_layer import PORT_LOOKUP
from primaite.simulator.system.applications.red_applications.c2.c2_beacon import C2Beacon
from primaite.simulator.system.applications.red_applications.c2.c2_server import C2Command, C2Server
from primaite.simulator.system.services.database.database_service import DatabaseService
@@ -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_LOOKUP["HTTP"], dst_port=PORT_LOOKUP["HTTP"], position=4)
router.acl.add_rule(action=ACLAction.PERMIT, src_port=PORT_LOOKUP["DNS"], dst_port=PORT_LOOKUP["DNS"], position=5)
router.acl.add_rule(action=ACLAction.PERMIT, src_port=PORT_LOOKUP["FTP"], dst_port=PORT_LOOKUP["FTP"], position=6)
c2_server_host = game.simulation.network.get_node_by_hostname("client_1")
c2_server_host.software_manager.install(software_class=C2Server)

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

@@ -11,7 +11,7 @@ from primaite.game.agent.actions import (
)
from primaite.session.environment import PrimaiteGymEnv
from primaite.simulator.file_system.file_system_item_abc import FileSystemItemHealthStatus
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 ApplicationOperatingState
from primaite.simulator.system.applications.database_client import DatabaseClient
from primaite.simulator.system.applications.red_applications.dos_bot import DoSBot
@@ -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_LOOKUP["POSTGRES_SERVER"]
assert dos_bot.payload == "HACC"
assert not dos_bot.repeat
assert dos_bot.port_scan_p_of_success == 0.875

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

@@ -9,7 +9,7 @@ from primaite.simulator.network.hardware.base import UserManager
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.router import ACLAction
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.network.transmission.transport_layer import PORT_LOOKUP
from primaite.simulator.system.services.service import ServiceOperatingState
from primaite.simulator.system.services.terminal.terminal import RemoteTerminalConnection
@@ -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_LOOKUP["SSH"], dst_port=PORT_LOOKUP["SSH"], position=4)
return (game, agent)

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

@@ -4,7 +4,7 @@ import pytest
from primaite.game.agent.observations.acl_observation import ACLObservation
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.transport_layer import Port
from primaite.simulator.network.transmission.transport_layer import PORT_LOOKUP
from primaite.simulator.sim_container import Simulation
from primaite.simulator.system.services.ntp.ntp_client import NTPClient
from primaite.simulator.system.services.ntp.ntp_server import NTPServer
@@ -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_LOOKUP["NTP"], src_port=PORT_LOOKUP["NTP"], position=1)
acl_obs = ACLObservation(
where=["network", "nodes", router.hostname, "acl", "acl"],

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

@@ -5,8 +5,8 @@ from primaite.simulator.network.hardware.node_operating_state import NodeOperati
from primaite.simulator.network.hardware.nodes.network.firewall import Firewall
from primaite.simulator.network.hardware.nodes.network.router import ACLAction
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.utils.validators import PROTOCOL_LOOKUP
def check_default_rules(acl_obs):
@@ -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=PROTOCOL_LOOKUP["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_LOOKUP["HTTP"],
dst_port=PORT_LOOKUP["HTTP"],
position=5,
)

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

@@ -8,9 +8,9 @@ from primaite.simulator.network.container import Network
from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
from primaite.simulator.network.hardware.nodes.network.router import ACLAction, Router
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.sim_container import Simulation
from primaite.utils.validators import PROTOCOL_LOOKUP
def test_router_observation():
@@ -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=PROTOCOL_LOOKUP["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_LOOKUP["HTTP"],
dst_port=PORT_LOOKUP["HTTP"],
position=5,
)
# Observe the state using the RouterObservation instance

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

@@ -3,7 +3,7 @@ import pytest
from primaite.session.environment import PrimaiteGymEnv
from primaite.simulator.network.hardware.nodes.network.router import ACLAction, Router
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.network.transmission.transport_layer import PORT_LOOKUP
from tests import TEST_ASSETS_ROOT
DATA_MANIPULATION_CONFIG = TEST_ASSETS_ROOT / "configs" / "data_manipulation.yaml"
@@ -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_LOOKUP["SSH"], dst_port=PORT_LOOKUP["SSH"], position=3)
return env

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

@@ -21,11 +21,11 @@ from primaite.game.agent.interface import ProxyAgent
from primaite.game.game import PrimaiteGame
from primaite.session.environment import PrimaiteGymEnv
from primaite.simulator.file_system.file_system_item_abc import FileSystemItemHealthStatus
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 ApplicationOperatingState
from primaite.simulator.system.applications.web_browser import WebBrowser
from primaite.simulator.system.software import SoftwareHealthState
from primaite.utils.validators import PROTOCOL_LOOKUP
from tests import TEST_ASSETS_ROOT
FIREWALL_ACTIONS_NETWORK = TEST_ASSETS_ROOT / "configs/firewall_actions_network.yaml"
@@ -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_LOOKUP["DNS"]
assert firewall.internal_outbound_acl.acl[1].src_port == PORT_LOOKUP["ARP"]
assert firewall.internal_outbound_acl.acl[1].protocol == PROTOCOL_LOOKUP["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_LOOKUP["HTTP"]
assert firewall.dmz_inbound_acl.acl[1].src_port == PORT_LOOKUP["HTTP"]
assert firewall.dmz_inbound_acl.acl[1].protocol == PROTOCOL_LOOKUP["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_LOOKUP["HTTP"]
assert firewall.dmz_outbound_acl.acl[2].src_port == PORT_LOOKUP["HTTP"]
assert firewall.dmz_outbound_acl.acl[2].protocol == PROTOCOL_LOOKUP["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_LOOKUP["POSTGRES_SERVER"]
assert firewall.external_inbound_acl.acl[10].src_port == PORT_LOOKUP["POSTGRES_SERVER"]
assert firewall.external_inbound_acl.acl[10].protocol == PROTOCOL_LOOKUP["ICMP"]
env.step(10) # Remove ACL rule from External Inbound
assert firewall.external_inbound_acl.num_rules == 1

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

@@ -9,11 +9,11 @@ from primaite.interface.request import RequestResponse
from primaite.session.environment import PrimaiteGymEnv
from primaite.simulator.network.hardware.nodes.host.server import Server
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.applications.database_client import DatabaseClient
from primaite.simulator.system.applications.web_browser import WebBrowser
from primaite.simulator.system.services.database.database_service import DatabaseService
from primaite.utils.validators import PROTOCOL_LOOKUP
from tests import TEST_ASSETS_ROOT
from tests.conftest import ControlledAgent
@@ -42,7 +42,12 @@ 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=PROTOCOL_LOOKUP["TCP"],
src_port=PORT_LOOKUP["HTTP"],
dst_port=PORT_LOOKUP["HTTP"],
)
agent.store_action(("NODE_APPLICATION_EXECUTE", {"node_id": 0, "application_id": 0}))
game.step()
assert agent.reward_function.current_reward == -0.7
@@ -66,7 +71,7 @@ def test_uc2_rewards(game_and_agent):
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
ACLAction.PERMIT, src_port=PORT_LOOKUP["POSTGRES_SERVER"], dst_port=PORT_LOOKUP["POSTGRES_SERVER"], position=2
)
comp = GreenAdminDatabaseUnreachablePenalty("client_1")