From 309100383439abba5e961885116e31742c5f791c Mon Sep 17 00:00:00 2001 From: William Blum Date: Fri, 9 Sep 2022 14:20:43 -0700 Subject: [PATCH] Fix regresssion #68 where nodes were incorrectly interpreted as owned when calculating rewards Regression was introduced in commit 2f6a045d4133bd225ffaf8a6a80d4904bd572e18 --- cyberbattle/simulation/actions.py | 71 ++++++++++--------- cyberbattle/simulation/actions_test.py | 2 +- cyberbattle/simulation/commandcontrol_test.py | 2 +- 3 files changed, 38 insertions(+), 37 deletions(-) diff --git a/cyberbattle/simulation/actions.py b/cyberbattle/simulation/actions.py index 7987d2d1..0f6c8f7a 100644 --- a/cyberbattle/simulation/actions.py +++ b/cyberbattle/simulation/actions.py @@ -96,8 +96,8 @@ class NodeTrackingInformation: # Map (vulnid, local_or_remote) to time of last attack. # local_or_remote is true for local, false for remote last_attack: Dict[Tuple[model.VulnerabilityID, bool], time] = dataclasses.field(default_factory=dict) - # Last time another node connected to this node - last_connection: Optional[time] = None + # Last time the node got owned by the attacker agent + last_owned_at: Optional[time] = None # All node properties discovered so far discovered_properties: Set[int] = dataclasses.field(default_factory=set) @@ -218,15 +218,26 @@ def __mark_allnodeproperties_as_discovered(self, node_id: model.NodeID): def __mark_node_as_owned(self, node_id: model.NodeID, - privilege: PrivilegeLevel = model.PrivilegeLevel.LocalUser) -> None: - if node_id not in self._discovered_nodes: - self._discovered_nodes[node_id] = NodeTrackingInformation() + privilege: PrivilegeLevel = model.PrivilegeLevel.LocalUser) -> Tuple[Optional[time], bool]: + """Mark a node as owned. + Return the time it was previously own (or None) and whether it was already owned.""" node_info = self._environment.get_node(node_id) - node_info.agent_installed = True - node_info.privilege_level = model.escalate(node_info.privilege_level, privilege) - self._environment.network.nodes[node_id].update({'data': node_info}) - self.__mark_allnodeproperties_as_discovered(node_id) + last_owned_at, is_currently_owned = self.__is_node_owned_history(node_id, node_info) + + if not is_currently_owned: + if node_id not in self._discovered_nodes: + self._discovered_nodes[node_id] = NodeTrackingInformation() + node_info.agent_installed = True + node_info.privilege_level = model.escalate(node_info.privilege_level, privilege) + self._environment.network.nodes[node_id].update({'data': node_info}) + + self.__mark_allnodeproperties_as_discovered(node_id) + + # Record that the node just got owned at the current time + self._discovered_nodes[node_id].last_owned_at = time() + + return last_owned_at, is_currently_owned def __mark_discovered_entities(self, reference_node: model.NodeID, outcome: model.VulnerabilityOutcome) -> Tuple[int, float, int]: """Mark discovered entities as such and return @@ -313,8 +324,6 @@ def __process_outcome(self, reward = 0 - was_previously_owned_at, is_currently_owned = self.__is_node_owned_history(node_id, node_info) - # if the vulnerability type is a privilege escalation # and if the escalation level is not already reached on that node, # then add the escalation tag to the node properties @@ -322,19 +331,18 @@ def __process_outcome(self, if outcome.tag in node_info.properties: return False, ActionResult(reward=Penalty.REPEAT, outcome=outcome) - if not was_previously_owned_at: - reward += float(node_info.value) + last_owned_at, is_currently_owned = self.__mark_node_as_owned(node_id, outcome.level) - self.__mark_node_as_owned(node_id, outcome.level) + if not last_owned_at: + reward += float(node_info.value) node_info.properties.append(outcome.tag) elif isinstance(outcome, model.LateralMove): - if not was_previously_owned_at: - reward += float(node_info.value) + last_owned_at, is_currently_owned = self.__mark_node_as_owned(node_id) - if not is_currently_owned: - self.__mark_node_as_owned(node_id) + if not last_owned_at: + reward += float(node_info.value) elif isinstance(outcome, model.ProbeSucceeded): for p in outcome.discovered_properties: @@ -454,14 +462,12 @@ def __is_passing_firewall_rules(self, rules: List[model.FirewallRule], port_name return False def __is_node_owned_history(self, target_node_id, target_node_data): - """ Returns whether the node was previously owned and if it's still currently owned.""" - was_previously_owned_at = self._discovered_nodes[target_node_id].last_connection - self._discovered_nodes[target_node_id].last_connection = time() + """ Returns the last time the node got owned and whether it is still currently owned.""" + last_previously_owned_at = self._discovered_nodes[target_node_id].last_owned_at if target_node_id in self._discovered_nodes else None - is_currently_owned = was_previously_owned_at is not None and \ - target_node_data.last_reimaging is not None and \ - was_previously_owned_at >= target_node_data.last_reimaging - return was_previously_owned_at, is_currently_owned + is_currently_owned = last_previously_owned_at is not None and \ + (target_node_data.last_reimaging is None or last_previously_owned_at >= target_node_data.last_reimaging) + return last_previously_owned_at, is_currently_owned def connect_to_remote_machine( self, @@ -525,27 +531,22 @@ def connect_to_remote_machine( return ActionResult(reward=Penalty.WRONG_PASSWORD, outcome=None) - is_already_owned = target_node_data.agent_installed + last_owned_at, is_already_owned = self.__mark_node_as_owned(target_node_id) + if is_already_owned: - return ActionResult(reward=Penalty.REPEAT, - outcome=model.LateralMove()) + return ActionResult(reward=Penalty.REPEAT, outcome=model.LateralMove()) if target_node_id not in self._discovered_nodes: self._discovered_nodes[target_node_id] = NodeTrackingInformation() - was_previously_owned_at, is_currently_owned = self.__is_node_owned_history(target_node_id, target_node_data) - - if is_currently_owned: - return ActionResult(reward=Penalty.REPEAT, outcome=model.LateralMove()) - self.__annotate_edge(source_node_id, target_node_id, EdgeAnnotation.LATERAL_MOVE) - self.__mark_node_as_owned(target_node_id) + logger.info(f"Infected node '{target_node_id}' from '{source_node_id}'" + f" via {port_name} with credential '{credential}'") if target_node.owned_string: logger.info("Owned message: " + target_node.owned_string) - return ActionResult(reward=float(target_node_data.value) if was_previously_owned_at is None else 0.0, + return ActionResult(reward=float(target_node_data.value) if last_owned_at is None else 0.0, outcome=model.LateralMove()) def _check_service_running_and_authorized(self, diff --git a/cyberbattle/simulation/actions_test.py b/cyberbattle/simulation/actions_test.py index e1ed7ec9..e28c156f 100644 --- a/cyberbattle/simulation/actions_test.py +++ b/cyberbattle/simulation/actions_test.py @@ -276,7 +276,7 @@ def test_exploit_remote_vulnerability(actions_on_simple_environment: Fixture) -> # test a valid and functional one. result = actions_on_simple_environment.exploit_remote_vulnerability('a', 'c', "RDPBF") assert isinstance(result.outcome, model.LateralMove) - assert result.reward >= node.value + assert result.reward < node.value def test_exploit_local_vulnerability(actions_on_simple_environment: Fixture) -> None: diff --git a/cyberbattle/simulation/commandcontrol_test.py b/cyberbattle/simulation/commandcontrol_test.py index 844e1cce..8b5341f8 100644 --- a/cyberbattle/simulation/commandcontrol_test.py +++ b/cyberbattle/simulation/commandcontrol_test.py @@ -72,6 +72,6 @@ def test_toyctf() -> None: reward = command.total_reward() print('Total reward ' + str(reward)) - assert reward == 289.0 + assert reward == 389.0 assert github is not None pass