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

#2736 - Fix up broken reward tests

Browse Source
This commit is contained in:
Marek Wolan
2024-08-19 13:59:35 +01:00
parent 05f9751fa8
commit f344d292db
5 changed files with 39 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
src/primaite/game/agent/rewards.py
View File

@@ -255,29 +255,26 @@ class WebpageUnavailablePenalty(AbstractReward):
"execute", "execute",
] ]
if ( # skip calculating if sticky and no new codes, reusing last step value
not request_attempted and self.sticky if not request_attempted and self.sticky:
): # skip calculating if sticky and no new codes, reusing last step value
return self.reward return self.reward
if last_action_response.response.status != "success": if last_action_response.response.status != "success":
self.reward = -1.0 self.reward = -1.0
# elif web_browser_state is NOT_PRESENT_IN_STATE or not web_browser_state["history"]:
elif web_browser_state is NOT_PRESENT_IN_STATE or "history" not in web_browser_state:
_LOGGER.debug( _LOGGER.debug(
"Web browser reward could not be calculated because the web browser history on node", "Web browser reward could not be calculated because the web browser history on node",
f"{self._node} was not reported in the simulation state. Returning 0.0", f"{self._node} was not reported in the simulation state. Returning 0.0",
) )
self.reward = 0.0 self.reward = 0.0
elif not web_browser_state["history"]: else:
self.reward = 0.0 # 0 if no requests have been attempted yet outcome = web_browser_state["history"][-1]["outcome"]
outcome = web_browser_state["history"][-1]["outcome"] if outcome == "PENDING":
if outcome == "PENDING": self.reward = 0.0 # 0 if a request was attempted but not yet resolved
self.reward = 0.0 # 0 if a request was attempted but not yet resolved elif outcome == 200:
elif outcome == 200: self.reward = 1.0 # 1 for successful request
self.reward = 1.0 # 1 for successful request else: # includes failure codes and SERVER_UNREACHABLE
else: # includes failure codes and SERVER_UNREACHABLE self.reward = -1.0 # -1 for failure
self.reward = -1.0 # -1 for failure
return self.reward return self.reward
@@ -325,7 +322,7 @@ class GreenAdminDatabaseUnreachablePenalty(AbstractReward):
db_state = access_from_nested_dict(state, self.location_in_state) db_state = access_from_nested_dict(state, self.location_in_state)
# If the last request was actually sent, then check if the connection was established. # If the last request was actually sent, then check if the connection was established.
if db_state is NOT_PRESENT_IN_STATE or "last_connection_successful" not in db_state: if db_state is NOT_PRESENT_IN_STATE:
_LOGGER.debug(f"Can't calculate reward for {self.__class__.__name__}") _LOGGER.debug(f"Can't calculate reward for {self.__class__.__name__}")
self.reward = 0.0 self.reward = 0.0
@@ -338,9 +335,8 @@ class GreenAdminDatabaseUnreachablePenalty(AbstractReward):
"execute", "execute",
] ]
if ( # skip calculating if sticky and no new codes, reusing last step value
not request_attempted and self.sticky if not request_attempted and self.sticky:
): # skip calculating if sticky and no new codes, reusing last step value
return self.reward return self.reward
self.reward = 1.0 if last_action_response.response.status == "success" else -1.0 self.reward = 1.0 if last_action_response.response.status == "success" else -1.0

2
src/primaite/simulator/system/applications/database_client.py
View File

@@ -75,8 +75,6 @@ class DatabaseClient(Application, identifier="DatabaseClient"):
server_password: Optional[str] = None server_password: Optional[str] = None
_query_success_tracker: Dict[str, bool] = {} _query_success_tracker: Dict[str, bool] = {}
"""Keep track of connections that were established or verified during this step. Used for rewards.""" """Keep track of connections that were established or verified during this step. Used for rewards."""
last_query_response: Optional[Dict] = None
"""Keep track of the latest query response. Used to determine rewards."""
_server_connection_id: Optional[str] = None _server_connection_id: Optional[str] = None
"""Connection ID to the Database Server.""" """Connection ID to the Database Server."""
client_connections: Dict[str, DatabaseClientConnection] = {} client_connections: Dict[str, DatabaseClientConnection] = {}

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

@@ -12,6 +12,7 @@ from primaite.simulator.network.hardware.nodes.network.router import ACLAction,
from primaite.simulator.network.transmission.network_layer import IPProtocol 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
from primaite.simulator.system.applications.database_client import DatabaseClient 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.simulator.system.services.database.database_service import DatabaseService
from tests import TEST_ASSETS_ROOT from tests import TEST_ASSETS_ROOT
from tests.conftest import ControlledAgent from tests.conftest import ControlledAgent
@@ -19,32 +20,30 @@ from tests.conftest import ControlledAgent
def test_WebpageUnavailablePenalty(game_and_agent): def test_WebpageUnavailablePenalty(game_and_agent):
"""Test that we get the right reward for failing to fetch a website.""" """Test that we get the right reward for failing to fetch a website."""
# set up the scenario, configure the web browser to the correct url
game, agent = game_and_agent game, agent = game_and_agent
agent: ControlledAgent agent: ControlledAgent
comp = WebpageUnavailablePenalty(node_hostname="client_1") comp = WebpageUnavailablePenalty(node_hostname="client_1")
agent.reward_function.register_component(comp, 0.7)
action = ("DONOTHING", {})
agent.store_action(action)
game.step()
# client 1 has not attempted to fetch webpage yet!
assert agent.reward_function.current_reward == 0.0
client_1 = game.simulation.network.get_node_by_hostname("client_1") client_1 = game.simulation.network.get_node_by_hostname("client_1")
browser = client_1.software_manager.software.get("WebBrowser") browser: WebBrowser = client_1.software_manager.software.get("WebBrowser")
browser.run() browser.run()
browser.target_url = "http://www.example.com" browser.target_url = "http://www.example.com"
assert browser.get_webpage() agent.reward_function.register_component(comp, 0.7)
action = ("DONOTHING", {})
agent.store_action(action) # Check that before trying to fetch the webpage, the reward is 0.0
agent.store_action(("DONOTHING", {}))
game.step()
assert agent.reward_function.current_reward == 0.0
# Check that successfully fetching the webpage yields a reward of 0.7
agent.store_action(("NODE_APPLICATION_EXECUTE", {"node_id": 0, "application_id": 0}))
game.step() game.step()
assert agent.reward_function.current_reward == 0.7 assert agent.reward_function.current_reward == 0.7
# 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: 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)
assert not browser.get_webpage() agent.store_action(("NODE_APPLICATION_EXECUTE", {"node_id": 0, "application_id": 0}))
agent.store_action(action)
game.step() game.step()
assert agent.reward_function.current_reward == -0.7 assert agent.reward_function.current_reward == -0.7
@@ -70,32 +69,25 @@ def test_uc2_rewards(game_and_agent):
comp = GreenAdminDatabaseUnreachablePenalty("client_1") comp = GreenAdminDatabaseUnreachablePenalty("client_1")
response = db_client.apply_request( request = ["network", "node", "client_1", "application", "DatabaseClient", "execute"]
[ response = game.simulation.apply_request(request)
"execute",
]
)
state = game.get_sim_state() state = game.get_sim_state()
reward_value = comp.calculate( reward_value = comp.calculate(
state, state,
last_action_response=AgentHistoryItem( last_action_response=AgentHistoryItem(
timestep=0, action="NODE_APPLICATION_EXECUTE", parameters={}, request=["execute"], response=response timestep=0, action="NODE_APPLICATION_EXECUTE", parameters={}, request=request, response=response
), ),
) )
assert reward_value == 1.0 assert reward_value == 1.0
router.acl.remove_rule(position=2) router.acl.remove_rule(position=2)
db_client.apply_request( response = game.simulation.apply_request(request)
[
"execute",
]
)
state = game.get_sim_state() state = game.get_sim_state()
reward_value = comp.calculate( reward_value = comp.calculate(
state, state,
last_action_response=AgentHistoryItem( last_action_response=AgentHistoryItem(
timestep=0, action="NODE_APPLICATION_EXECUTE", parameters={}, request=["execute"], response=response timestep=0, action="NODE_APPLICATION_EXECUTE", parameters={}, request=request, response=response
), ),
) )
assert reward_value == -1.0 assert reward_value == -1.0

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

@@ -146,7 +146,6 @@ def test_data_manipulation_disrupts_green_agent_connection(data_manipulation_db_
assert db_server_service.db_file.health_status is FileSystemItemHealthStatus.GOOD assert db_server_service.db_file.health_status is FileSystemItemHealthStatus.GOOD
assert green_db_connection.query("SELECT") assert green_db_connection.query("SELECT")
assert green_db_client.last_query_response.get("status_code") == 200
data_manipulation_bot.port_scan_p_of_success = 1 data_manipulation_bot.port_scan_p_of_success = 1
data_manipulation_bot.data_manipulation_p_of_success = 1 data_manipulation_bot.data_manipulation_p_of_success = 1
@@ -155,4 +154,3 @@ def test_data_manipulation_disrupts_green_agent_connection(data_manipulation_db_
assert db_server_service.db_file.health_status is FileSystemItemHealthStatus.COMPROMISED assert db_server_service.db_file.health_status is FileSystemItemHealthStatus.COMPROMISED
assert green_db_connection.query("SELECT") is False assert green_db_connection.query("SELECT") is False
assert green_db_client.last_query_response.get("status_code") != 200

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

@@ -103,7 +103,7 @@ def test_ransomware_script_attack(ransomware_script_and_db_server):
def test_ransomware_disrupts_green_agent_connection(ransomware_script_db_server_green_client): def test_ransomware_disrupts_green_agent_connection(ransomware_script_db_server_green_client):
"""Test to see show that the database service still operate""" """Test to show that the database service still operates after corruption"""
network: Network = ransomware_script_db_server_green_client network: Network = ransomware_script_db_server_green_client
client_1: Computer = network.get_node_by_hostname("client_1") client_1: Computer = network.get_node_by_hostname("client_1")
@@ -111,17 +111,18 @@ def test_ransomware_disrupts_green_agent_connection(ransomware_script_db_server_
client_2: Computer = network.get_node_by_hostname("client_2") client_2: Computer = network.get_node_by_hostname("client_2")
green_db_client: DatabaseClient = client_2.software_manager.software.get("DatabaseClient") green_db_client: DatabaseClient = client_2.software_manager.software.get("DatabaseClient")
green_db_client.connect()
green_db_client_connection: DatabaseClientConnection = green_db_client.get_new_connection() green_db_client_connection: DatabaseClientConnection = green_db_client.get_new_connection()
server: Server = network.get_node_by_hostname("server_1") server: Server = network.get_node_by_hostname("server_1")
db_server_service: DatabaseService = server.software_manager.software.get("DatabaseService") db_server_service: DatabaseService = server.software_manager.software.get("DatabaseService")
assert db_server_service.db_file.health_status is FileSystemItemHealthStatus.GOOD assert db_server_service.db_file.health_status is FileSystemItemHealthStatus.GOOD
assert green_db_client_connection.query("SELECT") assert green_db_client.query("SELECT") is True
assert green_db_client.last_query_response.get("status_code") == 200
ransomware_script_application.attack() ransomware_script_application.attack()
network.apply_timestep(0)
assert db_server_service.db_file.health_status is FileSystemItemHealthStatus.CORRUPT assert db_server_service.db_file.health_status is FileSystemItemHealthStatus.CORRUPT
assert green_db_client_connection.query("SELECT") is True assert green_db_client.query("SELECT") is True # Still operates but now the data field of response is empty
assert green_db_client.last_query_response.get("status_code") == 200