From 5f9c2ef0ad96b748f6e43bffc31996207ae75f58 Mon Sep 17 00:00:00 2001 From: Ondrej Lukas Date: Wed, 13 Nov 2024 12:55:21 +0100 Subject: [PATCH 1/9] Add firewall check --- env/worlds/network_security_game.py | 45 +++++++++++++++-------------- 1 file changed, 24 insertions(+), 21 deletions(-) diff --git a/env/worlds/network_security_game.py b/env/worlds/network_security_game.py index b4056b94..b784b5a0 100755 --- a/env/worlds/network_security_game.py +++ b/env/worlds/network_security_game.py @@ -690,31 +690,34 @@ def _execute_block_ip_action(self, current_state:components.GameState, action:co if "target_host" in action.parameters.keys() and action.parameters["target_host"] in current_state.controlled_hosts: # For now there is only one FW in the main router, but this should change in the future. # This means we ignore the 'target_host' that would be the router where this is applied. - - # Stop the blocked host to connect _to_ any other IP - try: - self._firewall[blocked_host] = set() - self.logger.debug(f"Removing all allowed connections from {blocked_host}") - except KeyError: - # The blocked_host host was not in the list - pass - # Stop the other hosts to connect _to the blocked_host_ - for host in self._firewall.keys(): + if self._firewall_check(action.parameters["source_host"], action.parameters["target_host"]): + + # Stop the blocked host to connect _to_ any other IP try: - self._firewall[host].remove(blocked_host) - self.logger.debug(f"Removing {blocked_host} from allowed connections from {host}") + self._firewall[blocked_host] = set() + self.logger.debug(f"Removing all allowed connections from {blocked_host}") except KeyError: # The blocked_host host was not in the list pass - # Update the state of blocked ips. It is a dict with key target_host and a set with blocked hosts inside - new_blocked = set() - # Store the blocked host IP in the set of blocked hosts - new_blocked.add(action.parameters["blocked_host"]) - if len(new_blocked) > 0: - if action.parameters["target_host"] not in next_blocked.keys(): - next_blocked[action.parameters["target_host"]] = new_blocked - else: - next_blocked[action.parameters["target_host"]] = next_blocked[action.parameters["target_host"]].union(new_blocked) + # Stop the other hosts to connect _to the blocked_host_ + for host in self._firewall.keys(): + try: + self._firewall[host].remove(blocked_host) + self.logger.debug(f"Removing {blocked_host} from allowed connections from {host}") + except KeyError: + # The blocked_host host was not in the list + pass + # Update the state of blocked ips. It is a dict with key target_host and a set with blocked hosts inside + new_blocked = set() + # Store the blocked host IP in the set of blocked hosts + new_blocked.add(action.parameters["blocked_host"]) + if len(new_blocked) > 0: + if action.parameters["target_host"] not in next_blocked.keys(): + next_blocked[action.parameters["target_host"]] = new_blocked + else: + next_blocked[action.parameters["target_host"]] = next_blocked[action.parameters["target_host"]].union(new_blocked) + else: + self.logger.debug(f"\t\t\t Connection from '{action.parameters['source_host']}->'{action.parameters['target_host']} is blocked blocked by FW") else: self.logger.info(f"\t\t\t Invalid target_host:'{action.parameters['target_host']}'") else: From 19bf38ca632bfd65b9556947016861501d512e3c Mon Sep 17 00:00:00 2001 From: Ondrej Lukas Date: Wed, 13 Nov 2024 16:22:53 +0100 Subject: [PATCH 2/9] Re-work blocking action --- env/worlds/network_security_game.py | 49 +++++++++++++++-------------- 1 file changed, 25 insertions(+), 24 deletions(-) diff --git a/env/worlds/network_security_game.py b/env/worlds/network_security_game.py index b784b5a0..2d8ae8ab 100755 --- a/env/worlds/network_security_game.py +++ b/env/worlds/network_security_game.py @@ -27,6 +27,7 @@ def __init__(self, task_config_file, world_name="NetSecEnv") -> None: self._services = {} # Dict of all services in the environment. Keys: hostname (`str`), values: `set` of `Service` objetcs. self._data = {} # Dict of all services in the environment. Keys: hostname (`str`), values `set` of `Service` objetcs. self._firewall = {} # dict of all the allowed connections in the environment. Keys `IP` ,values: `set` of `IP` objects. + self._fw_blocks = {} self._data_content = {} #content of each datapoint from self._data # All exploits in the environment self._exploits = {} @@ -680,10 +681,8 @@ def _execute_block_ip_action(self, current_state:components.GameState, action:co - Add the rule to the FW list - Update the state """ - blocked_host = action.parameters['blocked_host'] - next_nets, next_known_h, next_controlled_h, next_services, next_data, next_blocked = self._state_parts_deep_copy(current_state) - self.logger.info(f"\t\tBlockIP {action.parameters['target_host']}") + self.logger.info(f"\t\tBlockConnection {action.parameters['target_host']} <-> {action.parameters['blocked_host']}") # Is the src in the controlled hosts? if "source_host" in action.parameters.keys() and action.parameters["source_host"] in current_state.controlled_hosts: # Is the target in the controlled hosts? @@ -691,31 +690,33 @@ def _execute_block_ip_action(self, current_state:components.GameState, action:co # For now there is only one FW in the main router, but this should change in the future. # This means we ignore the 'target_host' that would be the router where this is applied. if self._firewall_check(action.parameters["source_host"], action.parameters["target_host"]): - - # Stop the blocked host to connect _to_ any other IP - try: - self._firewall[blocked_host] = set() - self.logger.debug(f"Removing all allowed connections from {blocked_host}") - except KeyError: - # The blocked_host host was not in the list - pass - # Stop the other hosts to connect _to the blocked_host_ - for host in self._firewall.keys(): + if action.parameters["target_host"] != action.parameters['blocked_host']: + try: + #remove connection target_host -> blocked_host + self._firewall[action.parameters["target_host"]].pop(action.parameters["blocked_host"]) + except KeyError: + pass try: - self._firewall[host].remove(blocked_host) - self.logger.debug(f"Removing {blocked_host} from allowed connections from {host}") + #remove blocked_host -> target_host + self._firewall[action.parameters["blocked_host"]].pop(action.parameters["blocked_host"]) except KeyError: - # The blocked_host host was not in the list pass - # Update the state of blocked ips. It is a dict with key target_host and a set with blocked hosts inside - new_blocked = set() - # Store the blocked host IP in the set of blocked hosts - new_blocked.add(action.parameters["blocked_host"]) - if len(new_blocked) > 0: + + #Update the FW_Rules visible to agents + if action.parameters["target_host"] not in self._fw_blocks.keys(): + self._fw_blocks[action.parameters["target_host"]] = set() + self._fw_blocks[action.parameters["target_host"]].add(action.parameters["blocked_host"]) + if action.parameters["blocked_host"] not in self._fw_blocks.keys(): + self._fw_blocks[action.parameters["blocked_host"]] = set() + self._fw_blocks[action.parameters["blocked_host"]].add(action.parameters["target_host"]) + + # update the state if action.parameters["target_host"] not in next_blocked.keys(): - next_blocked[action.parameters["target_host"]] = new_blocked - else: - next_blocked[action.parameters["target_host"]] = next_blocked[action.parameters["target_host"]].union(new_blocked) + next_blocked[action.parameters["target_host"]] = set() + if action.parameters["blocked_host"] not in next_blocked.keys(): + next_blocked[action.parameters["blocked_host"]] = set() + next_blocked[action.parameters["target_host"]].add(action.parameters["blocked_host"]) + next_blocked[action.parameters["blocked_host"]].add(action.parameters["target_host"]) else: self.logger.debug(f"\t\t\t Connection from '{action.parameters['source_host']}->'{action.parameters['target_host']} is blocked blocked by FW") else: From 593b80a2f8dcf02a69f8e1627ffc0e0ad2bcb46e Mon Sep 17 00:00:00 2001 From: Ondrej Lukas Date: Wed, 13 Nov 2024 16:28:48 +0100 Subject: [PATCH 3/9] Discover FW blocks with FindData action --- env/worlds/network_security_game.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/env/worlds/network_security_game.py b/env/worlds/network_security_game.py index 2d8ae8ab..a0c49f43 100755 --- a/env/worlds/network_security_game.py +++ b/env/worlds/network_security_game.py @@ -458,6 +458,16 @@ def _get_data_in_host(self, host_ip:str, controlled_hosts:set)->set: else: self.logger.debug("\t\t\tCan't get data in host. The host is not controlled.") return data + + def _get_known_blocks_in_host(self, host_ip:str, controlled_hosts:set)->set: + known_blocks = set() + if host_ip in controlled_hosts: #only return data if the agent controls the host + if host_ip in self._ip_to_hostname: + if host_ip in self._fw_blocks: + known_blocks = self._fw_blocks[host_ip] + else: + self.logger.debug("\t\t\tCan't get data in host. The host is not controlled.") + return known_blocks def _get_data_content(self, host_ip:str, data_id:str)->str: """ @@ -580,6 +590,13 @@ def _execute_find_data_action(self, current:components.GameState, action:compone next_data[action.parameters["target_host"]] = new_data else: next_data[action.parameters["target_host"]] = next_data[action.parameters["target_host"]].union(new_data) + # ADD KNOWN FW BLOCKS + new_blocks = self._get_known_blocks_in_host(action.parameters["target_host"], current.controlled_hosts) + if len(new_blocks) > 0: + if action.parameters["target_host"] not in next_blocked.keys(): + next_blocked[action.parameters["target_host"]] = new_blocks + else: + next_blocked[action.parameters["target_host"]] = next_blocked[action.parameters["target_host"]].union(new_blocks) else: self.logger.debug(f"\t\t\tConnection {action.parameters['source_host']} -> {action.parameters['target_host']} blocked by FW. Skipping") else: From b826da3ac96683e685a53331a19e7ac75b1cd7d1 Mon Sep 17 00:00:00 2001 From: Ondrej Lukas Date: Wed, 13 Nov 2024 16:37:41 +0100 Subject: [PATCH 4/9] repalce pop() with discard --- env/worlds/network_security_game.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/env/worlds/network_security_game.py b/env/worlds/network_security_game.py index a0c49f43..7374354b 100755 --- a/env/worlds/network_security_game.py +++ b/env/worlds/network_security_game.py @@ -710,12 +710,14 @@ def _execute_block_ip_action(self, current_state:components.GameState, action:co if action.parameters["target_host"] != action.parameters['blocked_host']: try: #remove connection target_host -> blocked_host - self._firewall[action.parameters["target_host"]].pop(action.parameters["blocked_host"]) + self._firewall[action.parameters["target_host"]].discard(action.parameters["blocked_host"]) + self.logger.debug(f"\t\t\t Removed rule:'{action.parameters['target_host']}' -> {action.parameters['blocked_host']}") except KeyError: pass try: #remove blocked_host -> target_host - self._firewall[action.parameters["blocked_host"]].pop(action.parameters["blocked_host"]) + self._firewall[action.parameters["blocked_host"]].discard(action.parameters["target_host"]) + self.logger.debug(f"\t\t\t Removed rule:'{action.parameters['blocked_host']}' -> {action.parameters['target_host']}") except KeyError: pass From 7e7e5d274349e95a20481df3cc558095704de31e Mon Sep 17 00:00:00 2001 From: Ondrej Lukas Date: Wed, 13 Nov 2024 16:38:57 +0100 Subject: [PATCH 5/9] remove all fw_blocks when reset --- env/worlds/network_security_game.py | 1 + 1 file changed, 1 insertion(+) diff --git a/env/worlds/network_security_game.py b/env/worlds/network_security_game.py index 7374354b..7dad58e8 100755 --- a/env/worlds/network_security_game.py +++ b/env/worlds/network_security_game.py @@ -893,6 +893,7 @@ def reset(self)->None: self._data = copy.deepcopy(self._data_original) # reset self._data_content to orignal state self._data_content_original = copy.deepcopy(self._data_content_original) + self._fw_blocks = {} self._actions_played = [] From 99d6e90205e3367019e6d19a43e2202b3ec67d01 Mon Sep 17 00:00:00 2001 From: Ondrej Lukas Date: Wed, 13 Nov 2024 16:41:23 +0100 Subject: [PATCH 6/9] reset firewall properly --- env/worlds/network_security_game.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/env/worlds/network_security_game.py b/env/worlds/network_security_game.py index 7dad58e8..58275216 100755 --- a/env/worlds/network_security_game.py +++ b/env/worlds/network_security_game.py @@ -77,6 +77,7 @@ def __init__(self, task_config_file, world_name="NetSecEnv") -> None: # Make a copy of data placements so it is possible to reset to it when episode ends self._data_original = copy.deepcopy(self._data) self._data_content_original = copy.deepcopy(self._data_content) + self._firewall_original = copy.deepcopy(self._firewall) self._actions_played = [] self.logger.info("Environment initialization finished") @@ -893,6 +894,7 @@ def reset(self)->None: self._data = copy.deepcopy(self._data_original) # reset self._data_content to orignal state self._data_content_original = copy.deepcopy(self._data_content_original) + self._firewall = copy.deepcopy(self._firewall_original) self._fw_blocks = {} From 44fb8bc7bd42ff4051f2be92431cc0451de454f5 Mon Sep 17 00:00:00 2001 From: Ondrej Lukas Date: Wed, 13 Nov 2024 16:47:25 +0100 Subject: [PATCH 7/9] better logging --- env/worlds/network_security_game.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/env/worlds/network_security_game.py b/env/worlds/network_security_game.py index 58275216..c2bd3532 100755 --- a/env/worlds/network_security_game.py +++ b/env/worlds/network_security_game.py @@ -700,7 +700,6 @@ def _execute_block_ip_action(self, current_state:components.GameState, action:co - Update the state """ next_nets, next_known_h, next_controlled_h, next_services, next_data, next_blocked = self._state_parts_deep_copy(current_state) - self.logger.info(f"\t\tBlockConnection {action.parameters['target_host']} <-> {action.parameters['blocked_host']}") # Is the src in the controlled hosts? if "source_host" in action.parameters.keys() and action.parameters["source_host"] in current_state.controlled_hosts: # Is the target in the controlled hosts? @@ -709,6 +708,7 @@ def _execute_block_ip_action(self, current_state:components.GameState, action:co # This means we ignore the 'target_host' that would be the router where this is applied. if self._firewall_check(action.parameters["source_host"], action.parameters["target_host"]): if action.parameters["target_host"] != action.parameters['blocked_host']: + self.logger.info(f"\t\tBlockConnection {action.parameters['target_host']} <-> {action.parameters['blocked_host']}") try: #remove connection target_host -> blocked_host self._firewall[action.parameters["target_host"]].discard(action.parameters["blocked_host"]) @@ -737,6 +737,8 @@ def _execute_block_ip_action(self, current_state:components.GameState, action:co next_blocked[action.parameters["blocked_host"]] = set() next_blocked[action.parameters["target_host"]].add(action.parameters["blocked_host"]) next_blocked[action.parameters["blocked_host"]].add(action.parameters["target_host"]) + else: + self.logger.info(f"\t\t\t Cant block connection form :'{action.parameters['target_host']}' to '{action.parameters['blocked_host']}'") else: self.logger.debug(f"\t\t\t Connection from '{action.parameters['source_host']}->'{action.parameters['target_host']} is blocked blocked by FW") else: From 9a04f79c09c813d2e4752d18ee6c33da13ad0ecb Mon Sep 17 00:00:00 2001 From: Ondrej Lukas Date: Fri, 15 Nov 2024 16:24:05 +0100 Subject: [PATCH 8/9] add tests for blockAction --- tests/test_actions.py | 53 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 53 insertions(+) diff --git a/tests/test_actions.py b/tests/test_actions.py index c0a3e66d..d08d9f2e 100644 --- a/tests/test_actions.py +++ b/tests/test_actions.py @@ -92,6 +92,21 @@ def env_obs_found_data2(env_obs_exploited_service2): new_state = env.step(state=state, action=action, agent_id=None) return (env, new_state) +@pytest.fixture +def env_obs_blocked_connection(env_obs_exploited_service): + "After blocking" + env, state = env_obs_exploited_service + source_host = components.IP('192.168.2.2') + target_host = components.IP('192.168.1.3') + blocked_host = components.IP('192.168.2.2') + parameters = { + "target_host":target_host, + "source_host":source_host, + "blocked_host": blocked_host} + action = components.Action(components.ActionType.BlockIP, parameters) + new_state = env.step(state=state, action=action, agent_id=None) + return (env, new_state) + class TestActionsNoDefender: def test_scan_network_not_exist(self, env_obs): """ @@ -303,3 +318,41 @@ def test_exploit_service_witout_find_service_in_host(self, env_obs_scan): new_state = env.step(state=state, action=action, agent_id=None) assert state == new_state assert components.IP('192.168.1.3') not in new_state.known_services + + def test_block_ip_same_host(self, env_obs_exploited_service): + env, state = env_obs_exploited_service + target_host = components.IP('192.168.2.2') + blocked_host = components.IP("1.1.1.1") + parameters = { + "target_host":target_host, + "source_host":target_host, + "blocked_host": blocked_host} + action = components.Action(components.ActionType.BlockIP, parameters) + new_state = env.step(state=state, action=action, agent_id=None) + assert target_host in new_state.known_blocks.keys() + assert blocked_host in new_state.known_blocks[target_host] + + def test_block_ip_same_different_source(self, env_obs_exploited_service): + env, state = env_obs_exploited_service + source_host = components.IP('192.168.2.2') + target_host = components.IP("192.168.1.3") + blocked_host = components.IP("1.1.1.1") + parameters = { + "target_host":target_host, + "source_host":source_host, + "blocked_host": blocked_host} + action = components.Action(components.ActionType.BlockIP, parameters) + new_state = env.step(state=state, action=action, agent_id=None) + assert target_host in new_state.known_blocks.keys() + assert blocked_host in new_state.known_blocks[target_host] + + def test_block_ip_self_block(self, env_obs_exploited_service): + env, state = env_obs_exploited_service + target_host = components.IP('192.168.2.2') + parameters = { + "target_host":target_host, + "source_host":components.IP('192.168.2.2'), + "blocked_host": target_host} + action = components.Action(components.ActionType.BlockIP, parameters) + new_state = env.step(state=state, action=action, agent_id=None) + assert target_host not in new_state.known_blocks.keys() \ No newline at end of file From 517947111ab6b13845bac99a267b4417eafe083e Mon Sep 17 00:00:00 2001 From: Ondrej Lukas Date: Fri, 15 Nov 2024 16:30:25 +0100 Subject: [PATCH 9/9] more test cases --- tests/test_actions.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/tests/test_actions.py b/tests/test_actions.py index d08d9f2e..e3864207 100644 --- a/tests/test_actions.py +++ b/tests/test_actions.py @@ -331,7 +331,9 @@ def test_block_ip_same_host(self, env_obs_exploited_service): new_state = env.step(state=state, action=action, agent_id=None) assert target_host in new_state.known_blocks.keys() assert blocked_host in new_state.known_blocks[target_host] - + assert target_host in env._fw_blocks.keys() + assert blocked_host in env._fw_blocks[target_host] + def test_block_ip_same_different_source(self, env_obs_exploited_service): env, state = env_obs_exploited_service source_host = components.IP('192.168.2.2') @@ -345,6 +347,9 @@ def test_block_ip_same_different_source(self, env_obs_exploited_service): new_state = env.step(state=state, action=action, agent_id=None) assert target_host in new_state.known_blocks.keys() assert blocked_host in new_state.known_blocks[target_host] + assert target_host in env._fw_blocks.keys() + assert blocked_host in env._fw_blocks[target_host] + def test_block_ip_self_block(self, env_obs_exploited_service): env, state = env_obs_exploited_service @@ -355,4 +360,5 @@ def test_block_ip_self_block(self, env_obs_exploited_service): "blocked_host": target_host} action = components.Action(components.ActionType.BlockIP, parameters) new_state = env.step(state=state, action=action, agent_id=None) - assert target_host not in new_state.known_blocks.keys() \ No newline at end of file + assert target_host not in new_state.known_blocks.keys() + assert target_host not in env._fw_blocks.keys()