import pytest from ForcepointSecurityManagementCenter import ( create_iplist_command, update_iplist_command, list_iplist_command, delete_iplist_command, create_host_command, list_host_command, delete_host_command, update_host_command, create_domain_command, list_domain_command, delete_domain_command, list_policy_template_command, list_firewall_policy_command, create_firewall_policy_command, delete_firewall_policy_command, update_rule_command, list_engine_command, list_rule_command, delete_rule_command, create_rule_command, engine_refresh_command, IPList, Host, DomainName, FirewallPolicy, DemistoException, ) from smc.api.exceptions import ElementNotFound, TaskRunFailed from smc.base.collection import CollectionManager class mock_Rule: def __init__(self): self.name = "name" self.tag = "tag" self.sources = mock_GenericRuleEntity() self.destinations = mock_GenericRuleEntity() self.action = mock_GenericRuleEntity() self.services = mock_GenericRuleEntity() self.comment = "comment" def all(self): return [self] def create(self, name, sources, destinations, action, comment): return self def update(self, **kwargs): return self def save( self, ): return self class mock_IPList: def __init__(self, name: str, iplist: list, comment: str): self.name = name self.iplist = iplist self.comment = comment class mock_Host: def __init__(self, name: str, address: str, ipv6_address: str, secondary: str, comment: str): self.name = name self.address = address self.ipv6_address = ipv6_address self.secondary = secondary self.comment = comment class mock_Domain: def __init__(self, name: str, comment: str): self.name = name self.comment = comment class mock_FirewallRule: def __init__(self): self.fw_ipv4_access_rules = mock_Rule() self.fw_ipv6_access_rules = mock_Rule() class mock_GenericRuleEntity: def __init__(self): self.name = "name" self.action = "ALLOW" def all(self): return [self] def filter(self): return self def add_many(self, elements): return self class mock_Policy: def __init__(self, name: str, comment: str): self.name = name self.comment = comment self.fw_ipv4_access_rules = mock_Rule() self.fw_ipv6_access_rules = mock_Rule() def open(self): return def save(self): return class mock_Task: def __init__(self, href: str, in_progress: bool = True): self.href = href self.in_progress = in_progress def update_status(self): return self class mock_Poller: def __init__(self, task: mock_Task): self.task = task class mock_Engine: def __init__(self, name: str, comment: str = ""): self.name = name self.comment = comment self.pending_changes = mock_Changes() def objects(self): return mock_GenericRuleEntity() def refresh(self): return mock_Poller(mock_Task(href=self.name)) class mock_Changes: def __init__(self): self.element_name = "element_name" self.modifier = "modifier" self.changed_on = "changed_on" def all(self): return self def approve_all(self): return def test_create_iplist_command(mocker): """ Given: - demisto args When: - Calling function create_address_command Then: - Ensure the results holds the expected data """ args = {"name": "name", "addresses": "1.1.1.1,8.8.8.8", "comment": "comment"} mocker.patch.object(IPList, "create", return_value=mock_IPList(name="name", iplist=["1.1.1.1", "8.8.8.8"], comment="comment")) response = create_iplist_command(args) assert response.readable_output == "IP List name was created successfully." assert response.outputs.get("Name") == "name" @pytest.mark.parametrize("is_override,returned_iplist", [(True, ["1.2.3.4"]), (False, ["1.1.1.1", "1.2.3.4"])]) def test_update_iplist_command(mocker, is_override, returned_iplist): """ Given: - demisto args Case 1: overwriting the existing ip list Case 2: appending to the existing list When: - Calling function update_address_command Then: - Ensure the results holds the expected data """ args = {"name": "name", "addresses": "1.2.3.4", "is_override": is_override} ip_list = mock_IPList(name="name", iplist=returned_iplist, comment="new_comment") mocker.patch.object(CollectionManager, "filter", return_value=[ip_list]) mocker.patch.object(IPList, "update_or_create", return_value=ip_list) response = update_iplist_command(args) assert response.readable_output == "IP List name was updated successfully." assert response.outputs.get("Addresses") == returned_iplist @pytest.mark.parametrize("args,returned_results", [({"name": "name"}, 1), ({"limit": "2"}, 2), ({"all_results": "True"}, 3)]) def test_list_iplist_command(mocker, args, returned_results): """ Given: - demisto args: Case 1: stating a specific IPList name Case 2: getting 2 results Case 3: getting all of the results (3 results) When: - Calling function list_address_command Then: - Ensure the results holds the expected data and the correct number of results """ ip_list = mock_IPList(name="name", iplist=["1.1.1.1"], comment="new_comment") mocker.patch.object(CollectionManager, "filter", return_value=[ip_list]) mocker.patch.object(CollectionManager, "limit", return_value=[ip_list, ip_list]) mocker.patch.object(CollectionManager, "all", return_value=[ip_list, ip_list, ip_list]) response = list_iplist_command(args) assert "IP Lists" in response.readable_output assert len(response.outputs) == returned_results def mock_delete(): raise ElementNotFound def test_delete_address_command(mocker): """ Given: - demisto args When: - Calling function delete_address_command Then: - Ensure the results holds the expected data in case of an ElementNotFound exception """ mocker.patch.object(IPList, "delete", side_effect=mock_delete) response = delete_iplist_command({"name": "name"}) assert response.readable_output == "IP List name was not found." def test_create_host_command(mocker): """ Given: - demisto args When: - Calling function create_host_command Then: - Ensure the results holds the expected data """ args = {"name": "name", "address": "1.1.1.1", "comment": "comment"} mocker.patch.object( Host, "create", return_value=mock_Host(name="name", address="1.1.1.1", ipv6_address="", secondary="", comment="comment") ) response = create_host_command(args) assert response.readable_output == "Host name was created successfully." assert response.outputs.get("Name") == "name" @pytest.mark.parametrize("args,returned_results", [({"name": "name"}, 1), ({"limit": "2"}, 2), ({"all_results": "True"}, 3)]) def test_list_host_command(mocker, args, returned_results): """ Given: - demisto args: Case 1: stating a specific IPList name Case 2: getting 2 results Case 3: getting all of the results (3 results) When: - Calling function list_host_command Then: - Ensure the results holds the expected data and the correct number of results """ host = mock_Host(name="name", address="1.1.1.1", ipv6_address="", secondary="", comment="comment") mocker.patch.object(CollectionManager, "filter", return_value=[host]) mocker.patch.object(CollectionManager, "limit", return_value=[host, host]) mocker.patch.object(CollectionManager, "all", return_value=[host, host, host]) response = list_host_command(args) assert "Hosts:" in response.readable_output assert len(response.outputs) == returned_results @pytest.mark.parametrize("is_override,returned_host", [(True, ["1.2.3.4"]), (False, ["1.1.1.1", "1.2.3.4"])]) def test_update_host_command(mocker, is_override, returned_host): """ Given: - demisto args Case 1: overwriting the existing host Case 2: appending to the existing host When: - Calling function update_host_command Then: - Ensure the results holds the expected data """ args = {"name": "name", "address": "1.2.3.4", "comment": "new_comment", "is_override": is_override} host = mock_Host(name="name", address="1.1.1.1", ipv6_address="", secondary=returned_host, comment="comment") mocker.patch.object(CollectionManager, "filter", return_value=[host]) mocker.patch.object(Host, "update_or_create", return_value=host) response = update_host_command(args) assert response.readable_output == "Host name was updated successfully." assert response.outputs.get("Secondary_address") == returned_host def mock_return_error(error_str: str): raise Exception(error_str) def test_update_host_with_host_not_found(mocker): """ Given: - name of host When: - Calling function update_host_command Then: - Ensure the exception is raised with the correct data when a host to update was not found """ args = { "name": "name", } mocker.patch("ForcepointSecurityManagementCenter.return_error", side_effect=mock_return_error) with pytest.raises(Exception) as e: mocker.patch.object(CollectionManager, "filter", return_value=[]) update_host_command(args) assert str(e) == "Host name was not found" def test_update_iplist_with_iplist_not_found(mocker): """ Given: - name of host When: - Calling function update_host_command Then: - Ensure the exception is raised with the correct data when a host to update was not found """ args = { "name": "name", } mocker.patch("ForcepointSecurityManagementCenter.return_error", side_effect=mock_return_error) with pytest.raises(Exception) as e: mocker.patch.object(CollectionManager, "filter", return_value=[]) update_iplist_command(args) assert str(e) == "IP List name was not found" def test_delete_host_command(mocker): """ Given: - demisto args When: - Calling function delete_host_command Then: - Ensure the results holds the expected data in case of an ElementNotFound exception """ mocker.patch.object(Host, "delete", side_effect=mock_delete) response = delete_host_command({"name": "name"}) assert response.readable_output == "Host name was not found." def test_create_domain_command(mocker): """ Given: - demisto args When: - Calling function create_domain_command Then: - Ensure the results holds the expected data """ args = {"name": "name", "comment": "comment"} mocker.patch.object(DomainName, "create", return_value=mock_Domain(name="name", comment="comment")) response = create_domain_command(args) assert response.readable_output == "Domain name was created successfully." assert response.outputs.get("Name") == "name" @pytest.mark.parametrize("args,returned_results", [({"name": "name"}, 1), ({"limit": "2"}, 2), ({"all_results": "True"}, 3)]) def test_list_domain_command(mocker, args, returned_results): """ Given: - demisto args: Case 1: stating a specific Domain name Case 2: getting 2 results Case 3: getting all of the results (3 results) When: - Calling function list_domain_command Then: - Ensure the results holds the expected data and the correct number of results """ domain = mock_Domain(name="name", comment="comment") mocker.patch.object(CollectionManager, "filter", return_value=[domain]) mocker.patch.object(CollectionManager, "limit", return_value=[domain, domain]) mocker.patch.object(CollectionManager, "all", return_value=[domain, domain, domain]) response = list_domain_command(args) assert "Domains:" in response.readable_output assert len(response.outputs) == returned_results def test_delete_domain_command(mocker): """ Given: - demisto args When: - Calling function delete_domain_command Then: - Ensure the results holds the expected data in case of an ElementNotFound exception """ mocker.patch.object(DomainName, "delete", side_effect=mock_delete) response = delete_domain_command({"name": "name"}) assert response.readable_output == "Domain name was not found." @pytest.mark.parametrize("args,returned_results", [({"limit": "2"}, 2), ({"all_results": "True"}, 3)]) def test_list_policy_template_command(mocker, args, returned_results): """ Given: - demisto args: Case 1: stating a specific Domain name Case 2: getting 2 results Case 3: getting all of the results (3 results) When: - Calling function list_domain_command Then: - Ensure the results holds the expected data and the correct number of results """ policy = mock_Policy(name="name", comment="comment") mocker.patch.object(CollectionManager, "limit", return_value=[policy, policy]) mocker.patch.object(CollectionManager, "all", return_value=[policy, policy, policy]) response = list_policy_template_command(args) assert "Policy template:" in response.readable_output assert len(response.outputs) == returned_results @pytest.mark.parametrize("args,returned_results", [({"limit": "2"}, 2), ({"all_results": "True"}, 3)]) def test_list_firewall_policy_command(mocker, args, returned_results): """ Given: - demisto args: Case 1: stating a specific Domain name Case 2: getting 2 results Case 3: getting all of the results (3 results) When: - Calling function list_domain_command Then: - Ensure the results holds the expected data and the correct number of results """ policy = mock_Policy(name="name", comment="comment") mocker.patch.object(CollectionManager, "limit", return_value=[policy, policy]) mocker.patch.object(CollectionManager, "all", return_value=[policy, policy, policy]) response = list_firewall_policy_command(args) assert "Firewall policies:" in response.readable_output assert len(response.outputs) == returned_results def test_create_firewall_policy_command(mocker): """ Given: - demisto args When: - Calling function create_domain_command Then: - Ensure the results holds the expected data """ args = {"name": "name", "comment": "comment"} mocker.patch.object(FirewallPolicy, "create", return_value=mock_Policy(name="name", comment="comment")) response = create_firewall_policy_command(args) assert response.readable_output == "Firewall policy name was created successfully." assert response.outputs.get("Name") == "name" def test_firewall_policy_delete_command(mocker): """ Given: - demisto args When: - Calling function delete_firewall_policy_command Then: - Ensure the results holds the expected data in case of an ElementNotFound exception """ mocker.patch.object(FirewallPolicy, "delete", side_effect=mock_delete) response = delete_firewall_policy_command({"name": "name"}) assert response.readable_output == "Firewall policy name was not found." @pytest.mark.parametrize("args,returned_results", [({"limit": "2"}, 2), ({"all_results": "True"}, 3)]) def test_list_engine_command(mocker, args, returned_results): """ Given: - demisto args: Case 1: stating a specific Domain name Case 2: getting 2 results Case 3: getting all of the results (3 results) When: - Calling function list_domain_command Then: - Ensure the results holds the expected data and the correct number of results """ engine = mock_Engine(name="name", comment="comment") mocker.patch.object(CollectionManager, "limit", return_value=[engine, engine]) mocker.patch.object(CollectionManager, "all", return_value=[engine, engine, engine]) response = list_engine_command(args) assert "Engines:" in response.readable_output assert len(response.outputs) == returned_results def test_list_rule_command(mocker): """ Given: - demisto args: When: - Calling function list_rule_command Then: - Ensure the results holds the expected data and the correct number of results """ policy = mock_Policy("name", "comment") mocker.patch("ForcepointSecurityManagementCenter.FirewallPolicy", return_value=policy) response = list_rule_command({"policy_name": "name"}) assert "Rules:" in response.readable_output assert len(response.outputs) == 2 def test_delete_rule_command_no_policy(mocker): """ Given: - demisto args When: - Calling function delete_rule_command Then: - Ensure the results holds the expected data in case of an ElementNotFound exception """ mocker.patch.object(CollectionManager, "filter", return_value=[]) args = {"policy_name": "name", "rule_id": "id"} response = delete_rule_command(args) assert "Firewall policy name was not found." in response.readable_output def test_delete_rule_command_no_rule(mocker): """ Given: - demisto args When: - Calling function delete_rule_command Then: - Ensure the results holds the expected data in case of an ElementNotFound exception """ rule = mock_Rule() policy = mock_Policy("name", "comment") mocker.patch.object(CollectionManager, "filter", return_value=[policy]) mocker.patch("ForcepointSecurityManagementCenter.get_policy_rules", return_value=[rule]) args = {"policy_name": "name", "rule_name": "test", "ip_version": "ip_version"} with pytest.raises(DemistoException, match="Rule with name test was not found in policy name."): delete_rule_command(args) def test_create_rule_command(mocker): """ Given: - demisto args: When: - Calling function create_rule_command Then: - Ensure the results holds the expected data """ mocker.patch.object(CollectionManager, "filter", return_value=[mock_Domain("name", "comment")]) args = { "policy_name": "name", "rule_name": "name", "ip_version": "V4", "source_ip_list": ["source_ip_list"], "source_host": ["source_host"], "source_domain": ["source_domain"], "destination_ip_list": ["destination_ip_list"], "destination_host": ["destination_host"], "destination_domain": ["destination_domain"], "action": "ALLOW", "comment": "test_comment", } policy = mock_Policy("name", "comment") mocker.patch("ForcepointSecurityManagementCenter.FirewallPolicy", return_value=policy) response = create_rule_command(args) assert "The rule name to the policy name was created successfully." in response.readable_output def test_create_rule_command_no_sources_or_destinations(mocker): """ Given: - demisto args: When: - Calling function create_rule_command Then: - Ensure the results holds the expected data """ mocker.patch.object(CollectionManager, "filter", return_value=[mock_Domain("name", "comment")]) args = {"policy_name": "name", "rule_name": "name", "ip_version": "V4", "action": "ALLOW", "comment": "test_comment"} policy = mock_Policy("name", "comment") mocker.patch("ForcepointSecurityManagementCenter.FirewallPolicy", return_value=policy) response = create_rule_command(args) assert "No sources or destinations were provided, provide at least one." in response.readable_output def test_update_rule_command(mocker): """ Given: - demisto args: When: - Calling function create_rule_command Then: - Ensure the results holds the expected data """ mocker.patch.object(CollectionManager, "filter", return_value=[mock_Domain("name", "comment")]) args = { "policy_name": "name", "rule_name": "name", "ip_version": "V4", "source_ip_list": ["source_ip_list"], "source_host": ["source_host"], "source_domain": ["source_domain"], "destination_ip_list": ["destination_ip_list"], "destination_host": ["destination_host"], "destination_domain": ["destination_domain"], "comment": "test_comment", } policy = mock_Policy("name", "comment") mocker.patch("ForcepointSecurityManagementCenter.FirewallPolicy", return_value=policy) response = update_rule_command(args) assert "The rule name to the policy name was updated successfully." in response.readable_output def test_engine_refresh_command_initiation(mocker): """ Given: - engine_name argument When: - Calling engine_refresh_command for the first time (no task_ids) Then: - Ensure refresh is initiated, task_ids are returned in args_for_next_run, and continue_to_poll is True """ args = {"engine_name": ["engine1", "engine2"]} mocker.patch("ForcepointSecurityManagementCenter.Engine", side_effect=lambda name: mock_Engine(name)) mocker.patch("ForcepointSecurityManagementCenter.Task", side_effect=lambda d: mock_Task(d.get("follower"))) result = engine_refresh_command(args) assert result.scheduled_command is not None assert result.scheduled_command._command is not None assert result.scheduled_command._command == "forcepoint-smc-engine-refresh" assert result.scheduled_command._args is not None assert "task_ids" in result.scheduled_command._args assert result.scheduled_command._args["task_ids"] == ["engine1", "engine2"] def test_engine_refresh_command_polling_in_progress(mocker): """ Given: - task_ids and engine_name arguments When: - Calling engine_refresh_command while tasks are still in progress Then: - Ensure continue_to_poll is True """ args = {"engine_name": ["engine1"], "task_ids": ["engine1"]} mocker.patch("ForcepointSecurityManagementCenter.Task", side_effect=lambda d: mock_Task(d.get("follower"), in_progress=True)) result = engine_refresh_command(args) assert result.scheduled_command is not None assert result.scheduled_command._command is not None assert result.scheduled_command._command == "forcepoint-smc-engine-refresh" assert result.scheduled_command._args is not None assert "task_ids" in result.scheduled_command._args assert result.scheduled_command._args["task_ids"] == ["engine1"] def test_engine_refresh_command_polling_complete(mocker): """ Given: - task_ids and engine_name arguments When: - Calling engine_refresh_command when all tasks are finished Then: - Ensure continue_to_poll is False and success message is returned """ args = {"engine_name": ["engine1", "engine2"], "task_ids": ["engine1", "engine2"]} mocker.patch("ForcepointSecurityManagementCenter.Task", side_effect=lambda d: mock_Task(d.get("follower"), in_progress=False)) result = engine_refresh_command(args) assert result.scheduled_command is None assert "refreshed successfully" in result.readable_output def test_engine_refresh_command_failure(mocker): """ Given: - engine_name argument When: - engine.refresh() raises TaskRunFailed Then: - Ensure DemistoException is raised with the correct engine names """ args = {"engine_name": ["engine1", "engine2"]} def mock_engine_side_effect(name): engine = mock_Engine(name) if name == "engine2": mocker.patch.object(engine, "refresh", side_effect=TaskRunFailed("Locked")) return engine mocker.patch("ForcepointSecurityManagementCenter.Engine", side_effect=mock_engine_side_effect) with pytest.raises(DemistoException, match=r"Failed to initiate refresh tasks for engines \(engine2\)"): engine_refresh_command(args)