diff --git a/caracara/modules/custom_ioa/custom_ioa.py b/caracara/modules/custom_ioa/custom_ioa.py index 6e59240..0ce37ba 100644 --- a/caracara/modules/custom_ioa/custom_ioa.py +++ b/caracara/modules/custom_ioa/custom_ioa.py @@ -96,7 +96,7 @@ def create_rule_group( rule.group = new_group # Update the rules - new_group = self._create_update_delete_rules(new_group, comment=comment) + new_group = self._update_create_delete_rules(new_group, comment=comment) return new_group @@ -138,7 +138,7 @@ def update_rule_group( new_group.rules_to_delete = group.rules_to_delete # Update the rules - new_group = self._create_update_delete_rules(new_group, comment=comment) + new_group = self._update_create_delete_rules(new_group, comment=comment) return new_group @@ -167,7 +167,7 @@ def delete_rule_groups( ids_to_delete.append(rule_group) instr(self.custom_ioa_api.delete_rule_groups)(ids=ids_to_delete, comment=comment) - def _create_update_delete_rules(self, group: IoaRuleGroup, comment: str) -> IoaRuleGroup: + def _update_create_delete_rules(self, group: IoaRuleGroup, comment: str) -> IoaRuleGroup: existing_rules = [] to_be_created = [] for rule in group.rules: @@ -176,25 +176,13 @@ def _create_update_delete_rules(self, group: IoaRuleGroup, comment: str) -> IoaR else: to_be_created.append(rule) - # Create the new rules - new_rules = [] - for rule in to_be_created: - resp = instr(self.custom_ioa_api.create_rule)(body=rule.dump_create(comment=comment)) - raw_rule = resp["body"]["resources"][0] - new_rule = CustomIoaRule.from_data_dict( - raw_rule, - rule_type=self._get_rule_types_cached()[raw_rule["ruletype_id"]], - ) - new_rule.rulegroup_id = group.id_ - new_rules.append(new_rule) - # Update the existing rules, if there are any if len(existing_rules) > 0: response = instr(self.custom_ioa_api.update_rules)( body={ "comment": comment, "rule_updates": [rule.dump_update() for rule in existing_rules], - "rulegroup_version": group.version + 1, + "rulegroup_version": group.version, "rulegroup_id": group.id_, } ) @@ -204,9 +192,22 @@ def _create_update_delete_rules(self, group: IoaRuleGroup, comment: str) -> IoaR rule_types = self._get_rule_types_cached() new_group = IoaRuleGroup.from_data_dict(data_dict=raw_group, rule_type_map=rule_types) else: - group.rules = new_rules new_group = group + # Create the new rules + new_rules = [] + for rule in to_be_created: + resp = instr(self.custom_ioa_api.create_rule)(body=rule.dump_create(comment=comment)) + raw_rule = resp["body"]["resources"][0] + new_rule = CustomIoaRule.from_data_dict( + raw_rule, + rule_type=self._get_rule_types_cached()[raw_rule["ruletype_id"]], + ) + new_rule.rulegroup_id = group.id_ + new_rules.append(new_rule) + + new_group.rules.extend(new_rules) + # Delete rules queued for deletion, if any if len(group.rules_to_delete) > 0: ids_to_delete = [rule.instance_id for rule in group.rules_to_delete] @@ -215,6 +216,8 @@ def _create_update_delete_rules(self, group: IoaRuleGroup, comment: str) -> IoaR ) # If successful (i.e. no exceptions raised), clear the deletion queue group.rules_to_delete = [] + + new_group.version += len(new_group.rules) + bool(group.rules_to_delete) return new_group @filter_string diff --git a/caracara/modules/custom_ioa/rules.py b/caracara/modules/custom_ioa/rules.py index 5314e48..881636a 100644 --- a/caracara/modules/custom_ioa/rules.py +++ b/caracara/modules/custom_ioa/rules.py @@ -262,7 +262,7 @@ def dump_rules_update(self, comment: str) -> dict: return { "comment": comment, "rule_updates": [rule.dump_update(group=self) for rule in self.rules], - "rulegroup_version": self.version + 1, + "rulegroup_version": self.version, "rulegroup_id": self.id_, } diff --git a/tests/unit_tests/test_custom_ioas.py b/tests/unit_tests/test_custom_ioas.py index 647496c..1921c5c 100644 --- a/tests/unit_tests/test_custom_ioas.py +++ b/tests/unit_tests/test_custom_ioas.py @@ -575,13 +575,14 @@ def mock_update_rule_group(body): raw_group["description"] = body["description"] raw_group["enabled"] = body["enabled"] raw_group["comment"] = body["comment"] + raw_group["version"] += 1 return {"body": {"resources": [raw_group]}} custom_ioa_api.update_rule_group.side_effect = mock_update_rule_group def mock_update_rules(body): assert body["rulegroup_id"] == raw_group["id"] - assert body["rulegroup_version"] == raw_group["version"] + 1 + assert body["rulegroup_version"] == raw_group["version"] raw_group["version"] = body["rulegroup_version"] for raw_rule_update in body["rule_updates"]: matching_rules = [ @@ -647,8 +648,8 @@ def mock_create_rule(body): # Assert falconpy called correctly # This consists of # - A rule group update - # - A rule deletion # - A rule update + # - A rule deletion # - A rule creation custom_ioa_api.update_rule_group.assert_called_once_with( body={ @@ -660,11 +661,6 @@ def mock_create_rule(body): "comment": "test update comment", } ) - custom_ioa_api.delete_rules.assert_called_once_with( - rule_group_id="test_group_01", - ids=["test_rule_01"], - comment="test update comment", - ) custom_ioa_api.update_rules.assert_called_once_with( body={ "rulegroup_id": "test_group_01", @@ -695,5 +691,10 @@ def mock_create_rule(body): "comment": "test update comment", } ) + custom_ioa_api.delete_rules.assert_called_once_with( + rule_group_id="test_group_01", + ids=["test_rule_01"], + comment="test update comment", + ) # Assert new group is as expected - assert new_group.version == group.version + 1 + assert new_group.version == group.version + 4