import pytest
import demistomock as demisto # noqa: F401
from QuarantineFile import (
QuarantineException,
QuarantineOrchestrator,
Brands,
EndpointBrandMapper,
QuarantineResult,
handler_factory,
XDRHandler,
MDEHandler,
Command,
main,
)
SHA_256_HASH = "sha256sha256sha256sha256sha256sha256sha256sha256sha256sha256sha2"
SHA_1_HASH = "sha1sha1sha1sha1sha1sha1sha1sha1sha1sha1"
# Pytest fixture to patch demisto functions
@pytest.fixture(autouse=True)
def mock_demisto(mocker):
"""
This fixture automatically mocks all required demisto functions for each test.
"""
mocker.patch.object(demisto, "error")
mocker.patch.object(demisto, "results")
mocker.patch.object(
demisto,
"getModules",
return_value={
"Cortex XDR - IR": {"state": "disabled", "brand": Brands.CORTEX_XDR_IR},
"Cortex Core - IR": {"state": "active", "brand": Brands.CORTEX_CORE_IR},
},
)
def _get_orchestrator(args: dict) -> QuarantineOrchestrator:
"""Helper function to create a QuarantineOrchestrator instance for testing."""
return QuarantineOrchestrator(args)
class TestCommand:
def test_get_entry_context_only_returns_populated_entries(self):
"""
Given: A raw response from Cortex Core - IR.
When: get_entry_contexts is called.
Then: Ensure only the populated entries are returned.
"""
# Arrange
raw_response = [
{
"Type": 1,
"EntryContext": {},
},
{
"Type": 1,
"EntryContext": {
"EndpointData(val.Brand && val.Brand == obj.Brand && val.ID && val.ID == "
"obj.ID && val.Hostname && val.Hostname == obj.Hostname)": [
{
"Brand": "Cortex Core - IR",
"Hostname": "FT",
"ID": "EP1_ID",
"IPAddress": ["192.0.2.1"],
"IsIsolated": "No",
"Message": "Command successful",
"Status": "Online",
},
{
"Brand": "Cortex Core - IR",
"Hostname": "FT",
"ID": "EP2_ID",
"IPAddress": ["192.0.2.1"],
"IsIsolated": "No",
"Message": "Command successful",
"Status": "Online",
},
]
},
},
{
"Type": 16,
"EntryContext": None,
},
]
# Act
entry_contexts = Command.get_entry_contexts(raw_response)
# Assert
assert len(entry_contexts) == 1
assert entry_contexts[0] == raw_response[1]["EntryContext"]
def test_get_entry_context_doesnt_return_error_entry_context(self):
"""
Given: A raw response from a command.
When: get_entry_contexts is called with a raw response containing an error entry context.
Then: Ensure only the non-error entries are returned.
"""
# Arrange
raw_response = [
{
"Type": 1,
"EntryContext": {
"EndpointData(val.Brand && val.Brand == obj.Brand && val.ID && val.ID == "
"obj.ID && val.Hostname && val.Hostname == obj.Hostname)": [
{
"Brand": "Cortex Core - IR",
"Hostname": "FT",
"ID": "EP1_ID",
"IPAddress": ["192.0.2.1"],
"IsIsolated": "No",
"Message": "Command successful",
"Status": "Online",
},
]
},
},
{
"Type": 4, # Type 4 means it's an error
"EntryContext": {"Error": []},
},
]
# Act
entry_contexts = Command.get_entry_contexts(raw_response)
# Assert
assert len(entry_contexts) == 1
assert entry_contexts[0] == raw_response[0]["EntryContext"]
def test_get_entry_context_object_containing_key_returns_first_entry_with_key(self):
"""
Given: A raw response from a command and a key.
When: get_entry_context_object_containing_key is called.
Then: Ensure the first EntryContext object containing the key is returned.
"""
# Arrange
raw_response = [
{
"Type": 1,
"EntryContext": {
"EndpointData(val.ID && val.ID == obj.ID)": [
{"ID": "EP1_ID", "Status": "Online"},
{"ID": "EP2_ID", "Status": "Online"},
]
},
},
{
"Type": 1,
"EntryContext": {
"AnotherKey(val.ID && val.ID == obj.ID)": [
{"ID": "blabla", "Status": "Online"},
{"ID": "anotherblabla", "Status": "Online"},
]
},
},
]
# Act
entry_context = Command.get_entry_context_object_containing_key(raw_response, "ID")
# Assert
assert entry_context == raw_response[0]["EntryContext"]["EndpointData(val.ID && val.ID == obj.ID)"]
def test_get_entry_context_object_containing_key_returns_none_if_no_entry_with_key(self):
"""
Given: A raw response from a command and a key that doesn't exist.
When: get_entry_context_object_containing_key is called.
Then: Ensure None is returned.
"""
# Arrange
raw_response = [
{
"Type": 1,
"EntryContext": {
"EndpointData(val.ID && val.ID == obj.ID)": [
{"ID": "EP1_ID", "Status": "Online"},
{"ID": "EP2_ID", "Status": "Online"},
]
},
}
]
# Act
entry_context = Command.get_entry_context_object_containing_key(raw_response, "blabla")
# Assert
assert entry_context is None
class TestBrands:
def test_get_values_returns_all_values(self):
"""
When: Executing Brands.values()
Then: Ensure all values are returned.
"""
assert Brands.values() == ["Cortex XDR - IR", "Cortex Core - IR", "Microsoft Defender Advanced Threat Protection"]
def test_normalize_returns_normalized_brand_name_when_alias_given(self):
"""
Given: An alias of a brand
When: normalize() is called
Then: The brand name is returned
"""
assert Brands.normalize("Microsoft Defender ATP") == "Microsoft Defender Advanced Threat Protection"
def test_normalize_returns_itself_when_not_given_alias(self):
"""
Given: A non alias of a brand
When: normalize() is called
Then: The given brand name is returned
"""
assert Brands.normalize("Something") == "Something"
class TestEndpointBrandMapper:
"""
Unit tests for the EndpointBrandMapper class, which handles endpoint discovery and grouping.
"""
def _get_mapper(self, args: dict, orchestrator=None) -> EndpointBrandMapper:
"""Helper to create an EndpointBrandMapper instance for testing."""
if orchestrator is None:
orchestrator = _get_orchestrator(args)
return EndpointBrandMapper(args, orchestrator)
def test_group_by_brand_end_to_end(self, mocker):
"""
Given:
- A list of endpoint IDs containing online, offline, and not-found endpoints.
When:
- The group_by_brand method is called.
Then:
- Ensure online endpoints are grouped correctly by brand.
- Ensure failure results are created for offline and not-found endpoints.
"""
# Arrange
args = {
"endpoint_id": "xdr-online-1,core-online-1,offline-1,not-found-1,xdr-online-2,ep-without-brand",
"file_hash": "some_hash",
"file_path": "/path",
}
mock_response = [
{
"Type": 1,
"EntryContext": {
"EndpointData(val.ID && val.ID == obj.ID)": [
{
"ID": "xdr-online-1",
"Brand": Brands.CORTEX_XDR_IR,
"Status": "Online",
"Message": "Command successful",
},
{
"ID": "core-online-1",
"Brand": Brands.CORTEX_CORE_IR,
"Status": "Online",
"Message": "Command successful",
},
{"ID": "offline-1", "Brand": Brands.CORTEX_XDR_IR, "Status": "Offline", "Message": "Command successful"},
{
"ID": "not-found-1",
"Brand": Brands.CORTEX_XDR_IR,
"Status": "Unknown",
"Message": "Command failed - no endpoint found",
},
{
"ID": "xdr-online-2",
"Brand": Brands.CORTEX_XDR_IR,
"Status": "Online",
"Message": "Command successful",
},
]
},
}
]
mocker.patch.object(demisto, "executeCommand", return_value=mock_response)
mapper = self._get_mapper(args)
# Act
grouped_endpoints = mapper.group_by_brand()
# Assert Grouping
expected_groups = {Brands.CORTEX_XDR_IR: ["xdr-online-1", "xdr-online-2"], Brands.CORTEX_CORE_IR: ["core-online-1"]}
assert grouped_endpoints == expected_groups
# Assert Initial Failure Results
assert len(mapper.initial_results) == 3
offline_result = next((r for r in mapper.initial_results if r.EndpointID == "offline-1"), None)
assert offline_result is not None
assert offline_result.Status == QuarantineResult.Statuses.FAILED
assert "Endpoint status is 'Offline'" in offline_result.Message
not_found_result = next((r for r in mapper.initial_results if r.EndpointID == "not-found-1"), None)
assert not_found_result is not None
assert not_found_result.Status == QuarantineResult.Statuses.FAILED
assert not_found_result.Message == QuarantineResult.Messages.FAILED_WITH_REASON.format(
reason="Command failed - no endpoint found"
)
ep_without_brand_result = next((r for r in mapper.initial_results if r.EndpointID == "ep-without-brand"), None)
assert ep_without_brand_result is not None
assert ep_without_brand_result.Status == QuarantineResult.Statuses.FAILED
assert ep_without_brand_result.Message == QuarantineResult.Messages.ENDPOINT_NOT_FOUND
def test_all_endpoints_offline_or_not_found_dont_raise_exception(self, mocker):
"""
Given:
- A get-endpoint-data response where no endpoints are 'Online'.
When:
- The group_by_brand method is called.
Then:
- Ensure the grouped_endpoints result is empty.
- Ensure failure results are created for all endpoints.
"""
# Arrange
args = {"endpoint_id": "offline-1,not-found-1"}
mock_response = [
{
"Type": 1,
"EntryContext": {
"EndpointData(val.ID && val.ID == obj.ID)": [
{"ID": "offline-1", "Brand": Brands.CORTEX_XDR_IR, "Status": "Offline", "Message": "Command successful"},
{
"ID": "not-found-1",
"Brand": Brands.CORTEX_XDR_IR,
"Status": "Unknown",
"Message": "Command failed - no endpoint found",
},
]
},
}
]
mocker.patch.object(demisto, "executeCommand", return_value=mock_response)
mapper = self._get_mapper(args)
# Act
grouped_endpoints = mapper.group_by_brand()
# Assert
assert not grouped_endpoints # Should be an empty dict
assert len(mapper.initial_results) == 2
def test_get_endpoint_data_command_fails(self, mocker):
"""
Given:
- The underlying 'get-endpoint-data' command will raise an exception.
When:
- The group_by_brand method is called.
Then:
- Ensure the exception is propagated upwards.
"""
# Arrange
args = {"endpoint_id": "any-id"}
mocker.patch.object(demisto, "executeCommand", side_effect=QuarantineException("API limit reached"))
mapper = self._get_mapper(args)
# Act & Assert
with pytest.raises(QuarantineException) as e:
mapper.group_by_brand()
assert "API limit reached" in str(e.value)
def test_handles_duplicate_api_entries_for_same_endpoint(self, mocker):
"""
Given:
- The get-endpoint-data response contains two entries for the same endpoint ID,
one 'Online' and one 'Offline'.
When:
- The group_by_brand method is called.
Then:
- Ensure the endpoint is correctly identified as 'Online'.
- Ensure no failure result is created for the 'Offline' duplicate.
"""
# Arrange
args = {"endpoint_id": "duplicate-id"}
mock_response = [
{
"Type": 1,
"EntryContext": {
"EndpointData(val.ID && val.ID == obj.ID)": [
# The script should prioritize the 'Online' entry
{
"ID": "duplicate-id",
"Brand": Brands.CORTEX_XDR_IR,
"Status": "Online",
"Message": "Command successful",
},
{
"ID": "duplicate-id",
"Brand": Brands.CORTEX_XDR_IR,
"Status": "Offline",
"Message": "Command successful",
},
]
},
}
]
mocker.patch.object(demisto, "executeCommand", return_value=mock_response)
mapper = self._get_mapper(args)
# Act
grouped_endpoints = mapper.group_by_brand()
# Assert
assert grouped_endpoints == {Brands.CORTEX_XDR_IR: ["duplicate-id"]}
assert not mapper.initial_results # No failure results should be created
def test_empty_api_response_raises_exception(self, mocker):
"""
Given:
- The 'get-endpoint-data' command returns a response with an empty data list.
When:
- The group_by_brand method is called.
Then:
- Ensure a QuarantineException is raised.
"""
# Arrange
args = {"endpoint_id": "any-id"}
mock_response = [{"Type": 1, "Contents": {}, "EntryContext": {"EndpointData(val.ID && val.ID == obj.ID)": []}}]
mocker.patch.object(demisto, "executeCommand", return_value=mock_response)
mapper = self._get_mapper(args)
# Act & Assert
with pytest.raises(QuarantineException) as e:
mapper.group_by_brand()
assert "Could not retrieve endpoint data" in str(e.value)
def test_verbose_mode_populates_verbose_results(self, mocker):
"""
Given:
- The 'verbose' argument is set to True.
When:
- The group_by_brand method is called, which executes a command.
Then:
- Ensure the orchestrator's verbose_results list is populated.
"""
# Arrange
args = {
"endpoint_id": "any-id",
"verbose": True, # Enable verbose mode
}
mock_response = [
{
"Type": 1,
"HumanReadable": "This is a verbose message.",
"EntryContext": {"EndpointData(val.ID && val.ID == obj.ID)": [{"ID": "any-id", "Status": "Online"}]},
}
]
mocker.patch.object(demisto, "executeCommand", return_value=mock_response)
orchestrator = _get_orchestrator(args)
mapper = self._get_mapper(args, orchestrator)
# Act
mapper.group_by_brand()
# Assert
assert len(orchestrator.verbose_results) == 1
assert orchestrator.verbose_results[0].readable_output == "This is a verbose message."
class TestBrandFactory:
def test_handler_factory_returns_correct_handler(self):
"""
Given: Brand "Cortex XDR - IR"
When: handler_factory is called
Then: The XDRBrandHandler is called
"""
orchestrator = _get_orchestrator({"endpoint_id": "any-id"})
brand = Brands.CORTEX_XDR_IR
handler = handler_factory(brand, orchestrator)
assert isinstance(handler, XDRHandler)
brand = Brands.CORTEX_CORE_IR
handler = handler_factory(brand, orchestrator)
assert isinstance(handler, XDRHandler)
def test_handler_factory_raises_exception_for_invalid_brand(self):
"""
Given: An unknown brand
When: handler_factory is called
Then: An exception is raised indicating the brand is unknown.
"""
orchestrator = _get_orchestrator({"endpoint_id": "any-id"})
with pytest.raises(QuarantineException) as e:
handler_factory("invalid-brand", orchestrator)
assert "No handler available for brand: invalid-brand" in str(e.value)
class TestXDRHandler:
def test_constructor_sets_correct_properties(self):
"""
Given: EndpointID, FileHash and FilePath
When: XDRHandler constructor is called
Then: It properly sets its properties.
"""
args = {"endpoint_id": "id1", "file_hash": SHA_256_HASH, "file_path": "/path"}
orchestrator = _get_orchestrator(args)
handler = XDRHandler(Brands.CORTEX_XDR_IR, orchestrator)
assert handler.command_prefix == "xdr"
assert handler.quarantine_command == "xdr-file-quarantine"
handler = XDRHandler(Brands.CORTEX_CORE_IR, orchestrator)
assert handler.command_prefix == "core"
assert handler.quarantine_command == "core-quarantine-files"
class TestPreProcessing:
def test_validate_args_raises_exception_for_missing_file_path(self):
"""
Given: Demisto args missing a file path
When: XDRHandler validate_args is called
Then: Raises error indicating missing file_path param.
"""
args = {"endpoint_id": "id1", "file_hash": SHA_256_HASH}
orchestrator = _get_orchestrator(args)
handler = XDRHandler(Brands.CORTEX_XDR_IR, orchestrator)
with pytest.raises(QuarantineException) as e:
handler.validate_args(args)
assert "The 'file_path' argument is required for brand Cortex XDR - IR." in str(e.value)
class TestInitialQuarantine:
"""Tests the quarantine kickoff flow of the XDRHandler."""
def test_initiate_quarantine_calls_expected_xdr_command(self, mocker):
"""
Given: args for quarantining multiple endpoints
When: XDRHandler initiate_quarantine is called
Then: Calls the expected XDR command
"""
args = {"endpoint_id": ["id1", "id2"], "file_hash": "sha256", "file_path": "/path"}
orchestrator = _get_orchestrator(args)
handler = XDRHandler(Brands.CORTEX_XDR_IR, orchestrator)
# mock the Command class execute() method, and check that it was called with the expected arguments
mock_command_instance = mocker.Mock()
mock_command_instance.execute.return_value = (
[{"Metadata": {"pollingCommand": "xdr-get-quarantine-status", "pollingArgs": {"action_id": "123"}}}],
[],
)
mock_command_class = mocker.patch("QuarantineFile.Command", return_value=mock_command_instance)
job = handler.initiate_quarantine(args)
# Assert that the Command class was instantiated correctly
expected_command_args = {
"endpoint_id_list": ["id1", "id2"],
"file_hash": "sha256",
"file_path": "/path",
"timeout_in_seconds": 300, # the default timeout
}
mock_command_class.assert_called_once_with(
name="xdr-file-quarantine", args=expected_command_args, brand=Brands.CORTEX_XDR_IR
)
# Assert that the returned job object is correct
expected_job = {
"brand": Brands.CORTEX_XDR_IR,
"poll_command": "xdr-get-quarantine-status",
"poll_args": {"action_id": "123"},
"finalize_args": {"file_hash": "sha256", "file_path": "/path"},
}
assert job == expected_job
def test_initiate_quarantine_calls_expected_xdr_command_with_timeout(self, mocker):
"""
Given: args for quarantining multiple endpoints with a timeout
When: XDRHandler initiate_quarantine is called
Then: Calls the expected XDR command with the given timeout
"""
args = {"endpoint_id": ["id1", "id2"], "file_hash": SHA_256_HASH, "file_path": "/path", "timeout": 123}
orchestrator = _get_orchestrator(args)
handler = XDRHandler(Brands.CORTEX_XDR_IR, orchestrator)
# mock the Command class execute() method, and check that it was called with the expected arguments
mock_command_instance = mocker.Mock()
mock_command_instance.execute.return_value = (
[{"Metadata": {"pollingCommand": "xdr-get-quarantine-status", "pollingArgs": {"action_id": "123"}}}],
[],
)
mock_command_class = mocker.patch("QuarantineFile.Command", return_value=mock_command_instance)
job = handler.initiate_quarantine(args)
# Assert that the Command class was instantiated correctly
expected_command_args = {
"endpoint_id_list": ["id1", "id2"],
"file_hash": SHA_256_HASH,
"file_path": "/path",
"timeout_in_seconds": 123,
}
mock_command_class.assert_called_once_with(
name="xdr-file-quarantine", args=expected_command_args, brand=Brands.CORTEX_XDR_IR
)
# Assert that the returned job object is correct
expected_job = {
"brand": Brands.CORTEX_XDR_IR,
"poll_command": "xdr-get-quarantine-status",
"poll_args": {"action_id": "123"},
"finalize_args": {"file_hash": SHA_256_HASH, "file_path": "/path"},
}
assert job == expected_job
def test_initiate_quarantine_adds_verbose_results_when_requested(self, mocker):
"""
Given: Args for quarantine with multiple endpoints and verbose enabled
When: XDRHandler initiate_quarantine is called
Then: Adds verbose results to the orchestrator
"""
args = {"endpoint_id": ["id1", "id2"], "file_hash": SHA_256_HASH, "file_path": "/path", "verbose": True}
orchestrator = _get_orchestrator(args)
handler = XDRHandler(Brands.CORTEX_XDR_IR, orchestrator)
mock_command_instance = mocker.Mock()
mock_command_instance.execute.return_value = (
[{"Metadata": {"pollingCommand": "xdr-get-quarantine-status", "pollingArgs": {"action_id": "123"}}}],
[
{
"Type": 1,
"HumanReadable": "This is a verbose message.",
"EntryContext": {"EndpointData(val.ID && val.ID == obj.ID)": [{"ID": "any-id", "Status": "Online"}]},
}
],
)
mocker.patch("QuarantineFile.Command", return_value=mock_command_instance)
assert orchestrator.verbose_results == []
assert orchestrator.verbose
handler.initiate_quarantine(args)
assert orchestrator.verbose_results == [
{
"Type": 1,
"HumanReadable": "This is a verbose message.",
"EntryContext": {"EndpointData(val.ID && val.ID == obj.ID)": [{"ID": "any-id", "Status": "Online"}]},
}
]
class TestFinalizeQuarantine:
"""Tests the finalization flow of the XDRHandler."""
@pytest.fixture
def setup_finalize(self, mocker):
"""A fixture to set up a handler and job object for finalize tests."""
args = {"file_hash": "hash123", "file_path": "/path/test.txt"}
orchestrator = _get_orchestrator(args)
handler = XDRHandler(Brands.CORTEX_CORE_IR, orchestrator)
job = {
"brand": "Cortex Core - IR",
"poll_command": "core-quarantine-files",
"poll_args": {
"action_id": [6],
"endpoint_id": "ep1",
"endpoint_id_list": ["ep1"],
"file_hash": "hash123",
"file_path": "/path/test.txt",
"integration_context_brand": "Core",
"integration_name": "Cortex Core - IR",
"interval_in_seconds": 60,
"timeout_in_seconds": "300",
},
"finalize_args": {"file_hash": "hash123", "file_path": "/path/test.txt"},
}
mocker.patch.object(handler, "_execute_quarantine_status_command")
return handler, job
def test_process_final_endpoint_status_receives_successfully_quarantined(self, setup_finalize):
"""
Given:
- The polling action completes successfully.
- The final status check confirms the file is quarantined.
When:
- _process_final_endpoint_status is called.
Then:
- Ensure a 'Success' result is returned.
"""
# Arrange
handler, job_data = setup_finalize
handler._execute_quarantine_status_command.return_value = {"status": True}
# Act
final_result = handler._process_final_endpoint_status(
{"action_id": 123, "endpoint_id": "ep1", "status": "COMPLETED_SUCCESSFULLY"}
)
# Assert
assert handler._execute_quarantine_status_command.call_args[0][0] == "ep1"
assert handler._execute_quarantine_status_command.call_args[0][1] == "hash123"
assert handler._execute_quarantine_status_command.call_args[0][2] == "/path/test.txt"
assert final_result.Status == QuarantineResult.Statuses.SUCCESS
assert final_result.Message == QuarantineResult.Messages.SUCCESS
assert final_result.EndpointID == "ep1"
assert final_result.FileHash == "hash123"
assert final_result.FilePath == "/path/test.txt"
assert final_result.Brand == Brands.CORTEX_CORE_IR
def test_process_final_endpoint_status_receives_unsuccessfully_quarantined(self, setup_finalize, mocker):
"""
Given:
- The polling action completes successfully.
- The final status check returns that the file is not quarantined.
When:
- _process_final_endpoint_status is called.
Then:
- Ensure a 'Failed' result is returned.
"""
# Arrange
handler, job_data = setup_finalize
handler._execute_quarantine_status_command.return_value = {"status": True}
# Act
final_result = handler._process_final_endpoint_status(
{"action_id": 123, "endpoint_id": "ep1", "status": "FAILED", "error_description": "Error from xdr agent"},
)
# Assert that _execute_quarantine_status_command was not called
handler._execute_quarantine_status_command.assert_not_called()
assert final_result.Status == QuarantineResult.Statuses.FAILED
assert final_result.Message == QuarantineResult.Messages.FAILED_WITH_REASON.format(reason="Error from xdr agent")
assert final_result.EndpointID == "ep1"
assert final_result.FileHash == "hash123"
assert final_result.FilePath == "/path/test.txt"
assert final_result.Brand == Brands.CORTEX_CORE_IR
def test_finalize_all_eps_success(self, setup_finalize, mocker):
"""
Given:
- The polling action completes successfully.
- The final status check confirms the file is quarantined.
When:
- finalize is called.
Then:
- Ensure a 'Success' result is returned.
"""
# Arrange
handler, job = setup_finalize
mocker.patch(
"QuarantineFile.Command.get_entry_context_object_containing_key",
return_value=[
{"action_id": 123, "endpoint_id": "ep1", "status": "COMPLETED_SUCCESSFULLY"},
{"action_id": 123, "endpoint_id": "ep2", "status": "COMPLETED_SUCCESSFULLY"},
],
)
handler._execute_quarantine_status_command.return_value = {"status": True}
# Act
final_results = handler.finalize([])
# Assert
assert len(final_results) == 2
result = final_results[0]
assert result.Status == QuarantineResult.Statuses.SUCCESS
assert result.Message == QuarantineResult.Messages.SUCCESS
assert result.EndpointID == "ep1"
result = final_results[1]
assert result.Status == QuarantineResult.Statuses.SUCCESS
assert result.Message == QuarantineResult.Messages.SUCCESS
assert result.EndpointID == "ep2"
def test_finalize_all_some_eps_failed(self, setup_finalize, mocker):
"""
Given:
- The polling action completes successfully.
- The final status check returns some eps are not quarantined.
When:
- finalize is called.
Then:
- Ensure a the failed results are returned.
"""
# Arrange
handler, job = setup_finalize
mocker.patch(
"QuarantineFile.Command.get_entry_context_object_containing_key",
return_value=[
{"action_id": 123, "endpoint_id": "ep1", "status": "COMPLETED_SUCCESSFULLY"},
{"action_id": 123, "endpoint_id": "ep2", "status": "Failed", "error_description": "Error from xdr agent"},
],
)
handler._execute_quarantine_status_command.return_value = {"status": True}
# Act
final_results = handler.finalize([])
# Assert
assert len(final_results) == 2
result = final_results[0]
assert result.Status == QuarantineResult.Statuses.SUCCESS
assert result.Message == QuarantineResult.Messages.SUCCESS
assert result.EndpointID == "ep1"
result = final_results[1]
assert result.Status == QuarantineResult.Statuses.FAILED
assert result.Message == QuarantineResult.Messages.FAILED_WITH_REASON.format(reason="Error from xdr agent")
assert result.EndpointID == "ep2"
def test_finalize_returns_failed_when_unexpected_exception(self, setup_finalize, mocker):
"""
Given:
- The polling action completes successfully.
When:
- finalize is called.
- During the final status check, an unexpected exception is raised.
Then:
- Ensure a 'Success' result is returned.
"""
# Arrange
handler, job = setup_finalize
mocker.patch(
"QuarantineFile.Command.get_entry_context_object_containing_key",
return_value=[{"action_id": 123, "endpoint_id": "ep1", "status": "COMPLETED_SUCCESSFULLY"}],
)
mocker.patch.object(handler, "_process_final_endpoint_status", side_effect=Exception("some error"))
# Act
result = handler.finalize([])
# Assert
assert len(result) == 1
result = result[0]
assert result.Status == QuarantineResult.Statuses.FAILED
assert result.Message == QuarantineResult.Messages.GENERAL_FAILURE
assert result.EndpointID == "ep1"
assert result.FileHash == "hash123"
assert result.FilePath == "/path/test.txt"
assert result.Brand == Brands.CORTEX_CORE_IR
class TestXDRHandlerBuiltinPath:
"""Tests for the Cortex platform Builtin path of the XDRHandler.
On the platform the Core-IR brand runs the Builtin `quarantineFile` command and
reports its brand as "Builtin". The quarantine confirmation is produced by the
Builtin command itself (asked for via the hidden `verify_quarantine` arg) and
surfaced under the `Core.QuarantineFiles.status` context path, so no identity is
forwarded from caller-supplied args into any RBAC-enforced status command.
The Builtin path additionally requires a minimum platform version; these tests
satisfy that gate by default via the autouse fixture below.
"""
@pytest.fixture(autouse=True)
def _mock_builtin_version(self, mocker):
"""Satisfy the minimum-version gate on `use_builtin` for the Builtin-path tests.
The version-gate behavior itself is covered by
test_constructor_keeps_legacy_command_for_old_version_on_platform.
"""
mocker.patch("QuarantineFile.is_demisto_version_ge", return_value=True)
def test_constructor_uses_builtin_command_for_core_on_platform(self, mocker):
"""
Given: The Core-IR brand while running on the Cortex platform.
When: The XDRHandler is constructed.
Then: use_builtin is True and the Builtin quarantine command is selected.
"""
mocker.patch("QuarantineFile.is_platform", return_value=True)
args = {"endpoint_id": "id1", "file_hash": SHA_256_HASH, "file_path": "/path"}
handler = XDRHandler(Brands.CORTEX_CORE_IR, _get_orchestrator(args))
assert handler.use_builtin is True
assert handler.quarantine_command == XDRHandler.BUILTIN_QUARANTINE_COMMAND
def test_constructor_keeps_legacy_command_for_xdr_on_platform(self, mocker):
"""
Given: The XDR-IR brand while running on the Cortex platform.
When: The XDRHandler is constructed.
Then: use_builtin is False (Builtin path is Core-only) and the legacy command is used.
"""
mocker.patch("QuarantineFile.is_platform", return_value=True)
args = {"endpoint_id": "id1", "file_hash": SHA_256_HASH, "file_path": "/path"}
handler = XDRHandler(Brands.CORTEX_XDR_IR, _get_orchestrator(args))
assert handler.use_builtin is False
assert handler.quarantine_command == "xdr-file-quarantine"
def test_constructor_keeps_legacy_command_for_old_version_on_platform(self, mocker):
"""
Given: The Core-IR brand on the platform, but the platform version is below the
minimum required for the Builtin quarantineFile command.
When: The XDRHandler is constructed.
Then: use_builtin is False and the legacy Core command is used (version gate).
"""
mocker.patch("QuarantineFile.is_platform", return_value=True)
# Override the autouse version mock: simulate an older platform version.
mocker.patch("QuarantineFile.is_demisto_version_ge", return_value=False)
args = {"endpoint_id": "id1", "file_hash": SHA_256_HASH, "file_path": "/path"}
handler = XDRHandler(Brands.CORTEX_CORE_IR, _get_orchestrator(args))
assert handler.use_builtin is False
assert handler.quarantine_command == "core-quarantine-files"
def test_display_brand_is_builtin_on_platform(self, mocker):
"""
Given: The Core-IR brand running on the platform (Builtin path).
When: display_brand is read.
Then: It returns "Builtin".
"""
mocker.patch("QuarantineFile.is_platform", return_value=True)
args = {"endpoint_id": "id1", "file_hash": SHA_256_HASH, "file_path": "/path"}
handler = XDRHandler(Brands.CORTEX_CORE_IR, _get_orchestrator(args))
assert handler.display_brand == "Cortex Builtin"
def test_display_brand_is_real_brand_off_platform(self):
"""
Given: The Core-IR brand running off-platform (legacy path).
When: display_brand is read.
Then: It returns the real brand name.
"""
args = {"endpoint_id": "id1", "file_hash": SHA_256_HASH, "file_path": "/path"}
handler = XDRHandler(Brands.CORTEX_CORE_IR, _get_orchestrator(args))
assert handler.use_builtin is False
assert handler.display_brand == Brands.CORTEX_CORE_IR
def test_initiate_quarantine_uses_endpoint_ids_plural_on_platform(self, mocker):
"""
Given: The Core-IR brand on the platform (Builtin path).
When: initiate_quarantine is called.
Then: The Builtin `quarantineFile` command is called with the plural `endpoint_ids` arg.
"""
mocker.patch("QuarantineFile.is_platform", return_value=True)
args = {"endpoint_id": ["id1", "id2"], "file_hash": SHA_256_HASH, "file_path": "/path"}
handler = XDRHandler(Brands.CORTEX_CORE_IR, _get_orchestrator(args))
mock_command_instance = mocker.Mock()
mock_command_instance.execute.return_value = (
[{"Metadata": {"pollingCommand": "quarantineFile", "pollingArgs": {"action_id": "123"}}}],
[],
)
mock_command_class = mocker.patch("QuarantineFile.Command", return_value=mock_command_instance)
handler.initiate_quarantine(args)
called_kwargs = mock_command_class.call_args.kwargs
assert called_kwargs["name"] == XDRHandler.BUILTIN_QUARANTINE_COMMAND
assert called_kwargs["args"]["endpoint_ids"] == ["id1", "id2"]
assert "endpoint_id_list" not in called_kwargs["args"]
def test_initiate_quarantine_requests_verification_on_platform(self, mocker):
"""
Given: The Core-IR brand on the platform (Builtin path).
When: initiate_quarantine is called.
Then: The Builtin `quarantineFile` command is asked to self-verify via the hidden
`verify_quarantine=true` arg (so the RBAC-gated status call runs inside the
command's own trusted poll context, not from caller-supplied identity args).
"""
mocker.patch("QuarantineFile.is_platform", return_value=True)
args = {"endpoint_id": ["id1"], "file_hash": SHA_256_HASH, "file_path": "/path"}
handler = XDRHandler(Brands.CORTEX_CORE_IR, _get_orchestrator(args))
mock_command_instance = mocker.Mock()
mock_command_instance.execute.return_value = (
[{"Metadata": {"pollingCommand": "quarantineFile", "pollingArgs": {"action_id": "123"}}}],
[],
)
mock_command_class = mocker.patch("QuarantineFile.Command", return_value=mock_command_instance)
handler.initiate_quarantine(args)
called_kwargs = mock_command_class.call_args.kwargs
assert called_kwargs["name"] == XDRHandler.BUILTIN_QUARANTINE_COMMAND
assert called_kwargs["args"]["verify_quarantine"] == "true"
def test_finalize_uses_builtin_confirmation_context_on_platform(self, mocker):
"""
Given: The Core-IR brand on the platform, a successful action, and a Builtin
quarantine confirmation surfaced under `Core.QuarantineFiles.status`.
When: finalize processes the results.
Then: The confirmation is read from context (no legacy status command is executed)
and the endpoint is reported SUCCESS.
"""
mocker.patch("QuarantineFile.is_platform", return_value=True)
args = {"file_hash": SHA_256_HASH, "file_path": "/path/test.txt"}
handler = XDRHandler(Brands.CORTEX_CORE_IR, _get_orchestrator(args))
# get_entry_context_object_containing_key is called twice in finalize:
# first for the action-status list, then for the QuarantineFiles confirmation.
mocker.patch(
"QuarantineFile.Command.get_entry_context_object_containing_key",
side_effect=[
[{"ActionID": 123, "EndpointID": "ep1", "Status": "COMPLETED_SUCCESSFULLY"}],
[{"status": True, "endpointId": "ep1", "filePath": "/path/test.txt"}],
],
)
legacy_status = mocker.patch.object(handler, "_execute_quarantine_status_command")
final_results = handler.finalize([])
legacy_status.assert_not_called()
assert len(final_results) == 1
assert final_results[0].Status == QuarantineResult.Statuses.SUCCESS
assert final_results[0].EndpointID == "ep1"
def test_finalize_reports_failed_when_builtin_confirmation_is_false(self, mocker):
"""
Given: The Core-IR brand on the platform, a successful action, but the Builtin
confirmation reports the file is NOT quarantined (e.g. already quarantined).
When: finalize processes the results.
Then: The endpoint is reported FAILED (guards action-status false positives).
"""
mocker.patch("QuarantineFile.is_platform", return_value=True)
args = {"file_hash": SHA_256_HASH, "file_path": "/path/test.txt"}
handler = XDRHandler(Brands.CORTEX_CORE_IR, _get_orchestrator(args))
mocker.patch(
"QuarantineFile.Command.get_entry_context_object_containing_key",
side_effect=[
[{"ActionID": 123, "EndpointID": "ep1", "Status": "COMPLETED_SUCCESSFULLY"}],
[{"status": False, "endpointId": "ep1", "filePath": "/path/test.txt"}],
],
)
final_results = handler.finalize([])
assert len(final_results) == 1
assert final_results[0].Status == QuarantineResult.Statuses.FAILED
def test_status_command_does_not_forward_identity_off_platform(self, mocker):
"""
Given: The Core-IR brand off-platform (legacy path).
When: _execute_quarantine_status_command is called.
Then: The legacy status command is used with only file coordinates - no identity args.
"""
args = {"endpoint_id": "ep1", "file_hash": SHA_256_HASH, "file_path": "/path"}
handler = XDRHandler(Brands.CORTEX_CORE_IR, _get_orchestrator(args))
mock_command_instance = mocker.Mock()
mock_command_instance.execute.return_value = ([], [])
mock_command_class = mocker.patch("QuarantineFile.Command", return_value=mock_command_instance)
mocker.patch("QuarantineFile.Command.get_entry_contexts", return_value=[])
handler._execute_quarantine_status_command("ep1", SHA_256_HASH, "/path")
called_kwargs = mock_command_class.call_args.kwargs
assert called_kwargs["name"] == "core-get-quarantine-status"
assert called_kwargs["args"] == {"endpoint_id": "ep1", "file_hash": SHA_256_HASH, "file_path": "/path"}
assert "user_id" not in called_kwargs["args"]
def test_finalize_reports_builtin_brand_on_platform(self, mocker):
"""
Given: The Core-IR brand on the platform with a successful action.
When: finalize processes the results.
Then: The resulting Brand column reads "Builtin".
"""
mocker.patch("QuarantineFile.is_platform", return_value=True)
args = {"file_hash": "hash123", "file_path": "/path/test.txt"}
handler = XDRHandler(Brands.CORTEX_CORE_IR, _get_orchestrator(args))
mocker.patch(
"QuarantineFile.Command.get_entry_context_object_containing_key",
return_value=[{"ActionID": 123, "EndpointID": "ep1", "Status": "COMPLETED_SUCCESSFULLY"}],
)
mocker.patch.object(handler, "_execute_quarantine_status_command", return_value={"status": True})
final_results = handler.finalize([])
assert len(final_results) == 1
assert final_results[0].Status == QuarantineResult.Statuses.SUCCESS
assert final_results[0].Brand == "Cortex Builtin"
class TestMDEHandler:
def test_constructor_sets_correct_properties(self):
"""
Given: EndpointID and FileHash
When: XDRHandler constructor is called
Then: It properly sets its properties.
"""
args = {"endpoint_id": "id1", "file_hash": SHA_1_HASH}
orchestrator = _get_orchestrator(args)
handler = MDEHandler(orchestrator)
assert handler.brand == "Microsoft Defender Advanced Threat Protection"
class TestInitialQuarantine:
"""Tests the quarantine kickoff flow of the MDEHandler."""
def test_initiate_quarantine_calls_expected_command_with_timeout(self, mocker):
"""
Given: args for quarantining multiple endpoints with a timeout
When: MDEHandler initiate_quarantine is called
Then: Calls the expected MDE command with the given timeout
"""
args = {"endpoint_id": ["id1", "id2"], "file_hash": SHA_1_HASH, "timeout": 123}
orchestrator = _get_orchestrator(args)
handler = MDEHandler(orchestrator)
# mock the Command class execute() method, and check that it was called with the expected arguments
# mock_execute = mocker.patch("QuarantineFile.Command.execute")
mock_response = [
{
"Type": 1,
"HumanReadable": "Quarantine operations are still in progress...",
"EntryContext": {
"MicrosoftATP.MachineAction(val.ID && val.ID == obj.ID)": [
{
"Commands": None,
"ComputerDNSName": None,
"CreationDateTimeUtc": None,
"ID": None,
"LastUpdateTimeUtc": None,
"MachineID": "id1",
"RelatedFileInfo": {
"FileIdentifier": "sha1sha1",
"FileIdentifierType": None,
},
"Requestor": None,
"RequestorComment": "Quarantine file hash: sha1sha1",
"Scope": None,
"Status": "Failed",
"Type": None,
},
{
"Commands": None,
"ComputerDNSName": None,
"CreationDateTimeUtc": "2025-09-30T10:51:12.8554335Z",
"ID": "111111",
"LastUpdateTimeUtc": "2025-09-30T10:51:12.8554346Z",
"MachineID": "id2",
"RelatedFileInfo": {
"FileIdentifier": None,
"FileIdentifierType": None,
},
"Requestor": "Cortex XSOAR - Microsoft Defender ATP - Dev",
"RequestorComment": "Quarantine file hash: sha1sha1",
"Scope": None,
"Status": "Pending",
"Type": "StopAndQuarantineFile",
},
]
},
"Metadata": {
"pollingCommand": "microsoft-atp-stop-and-quarantine-file",
"pollingArgs": {
"action_ids": ["111111"],
"machine_id": ["22222", "33333"],
"file_hash": "sha1sha1",
"timeout_in_seconds": "300",
},
},
"Contents": {},
},
]
mock_execute = mocker.patch.object(demisto, "executeCommand", return_value=mock_response)
job = handler.initiate_quarantine(args)
# Assert that the Command class was instantiated correctly
expected_command_args = {
"comment": "Quarantine file hash: sha1sha1sha1sha1sha1sha1sha1sha1sha1sha1",
"file_hash": "sha1sha1sha1sha1sha1sha1sha1sha1sha1sha1",
"machine_id": ["id1", "id2"],
"polling": True,
"timeout_in_seconds": 123,
}
mock_execute.assert_called_once_with("microsoft-atp-stop-and-quarantine-file", expected_command_args)
# Assert that the returned job object is correct
expected_job = {
"brand": Brands.MDE,
"poll_command": "microsoft-atp-stop-and-quarantine-file",
"poll_args": {
"action_ids": ["111111"],
"machine_id": ["22222", "33333"],
"file_hash": "sha1sha1",
"timeout_in_seconds": "300",
},
}
assert job == expected_job
def test_initiate_quarantine_adds_verbose_results_when_requested(self, mocker):
"""
Given: Args for quarantine with multiple endpoints and verbose enabled
When: MDEHandler initiate_quarantine is called
Then: Adds verbose results to the orchestrator
"""
args = {"endpoint_id": ["id1", "id2"], "file_hash": SHA_1_HASH, "timeout": 123, "verbose": True}
orchestrator = _get_orchestrator(args)
handler = XDRHandler(Brands.CORTEX_XDR_IR, orchestrator)
mock_command_instance = mocker.Mock()
mock_command_instance.execute.return_value = (
[
{
"Metadata": {
"pollingCommand": "microsoft-atp-stop-and-quarantine-file",
"pollingArgs": {
"action_ids": ["111111"],
"machine_id": ["22222", "33333"],
"file_hash": "sha1sha1",
"timeout_in_seconds": "300",
},
}
}
],
[
{
"Type": 1,
"HumanReadable": "This is a verbose message.",
"EntryContext": {"EndpointData(val.ID && val.ID == obj.ID)": [{"ID": "any-id", "Status": "Online"}]},
}
],
)
mocker.patch("QuarantineFile.Command", return_value=mock_command_instance)
assert orchestrator.verbose_results == []
assert orchestrator.verbose
handler.initiate_quarantine(args)
assert orchestrator.verbose_results == [
{
"Type": 1,
"HumanReadable": "This is a verbose message.",
"EntryContext": {"EndpointData(val.ID && val.ID == obj.ID)": [{"ID": "any-id", "Status": "Online"}]},
}
]
class TestFinalizeQuarantine:
"""Tests the finalization flow of the MDEHandler."""
@pytest.fixture
def setup_finalize(self):
"""A fixture to set up a handler and job object for finalize tests."""
args = {"file_hash": "hash123"}
orchestrator = _get_orchestrator(args)
handler = MDEHandler(orchestrator)
last_poll_response = [
{
"Type": 1,
"EntryContext": {
"MicrosoftATP.MachineAction(val.ID && val.ID == obj.ID)": [
{
"Commands": None,
"ComputerDNSName": "win10",
"CreationDateTimeUtc": "2025-09-04T16:30:40.190142Z",
"ID": "a6e0718b-0267-461a-9a0b-dbd7284d0bde",
"LastUpdateTimeUtc": "2025-09-04T16:30:55.3588394Z",
"MachineID": "123",
"RelatedFileInfo": {
"FileIdentifier": "hash123",
"FileIdentifierType": "Sha1",
},
"Requestor": "Cortex XSOAR - Microsoft Defender ATP",
"RequestorComment": "Quarantine file hash: hash123",
"Scope": None,
"Status": "Succeeded",
"Type": "StopAndQuarantineFile",
}
]
},
},
{
"Type": 16,
"EntryContext": None,
},
]
return handler, last_poll_response
def test_finalize_returns_expected_final_results(self, setup_finalize):
handler, last_poll_response = setup_finalize
final_results = handler.finalize(last_poll_response)
assert len(final_results) == 1
result = final_results[0]
assert result.EndpointID == "123"
assert result.Brand == "Microsoft Defender Advanced Threat Protection"
assert result.FileHash == "hash123"
assert result.Message == "File successfully quarantined."
assert result.Status == "Succeeded"
class TestQuarantineOrchestrator:
class TestArgumentValidationSanitization:
"""
Unit tests for the _sanitize_and_validate_args method of the QuarantineOrchestrator.
This approach directly tests the validation logic in isolation.
"""
class TestEndpointId:
"""Tests specifically for the 'endpoint_id' script argument."""
def test_missing_endpoint_id_raises_exception(self):
"""
Given: Args without 'endpoint_id'.
When: _sanitize_and_validate_args is called.
Then: Ensure QuarantineException is raised.
"""
args = {"file_hash": "sha256_hash", "file_path": "/path"}
orchestrator = QuarantineOrchestrator(args)
with pytest.raises(QuarantineException) as e:
orchestrator._sanitize_and_validate_args()
assert "Missing required argument" in str(e.value)
assert QuarantineOrchestrator.ENDPOINT_IDS_ARG in str(e.value)
def test_duplicate_endpoint_ids_are_deduplicated(self, mocker):
"""
Given: Args with duplicate endpoint IDs.
When: _sanitize_and_validate_args is called.
Then: Ensure the 'endpoint_id' list in orchestrator.args is deduplicated.
"""
args = {
"endpoint_id": "id1,id2,id1,id3,id2",
"file_hash": SHA_256_HASH,
"file_path": "/path",
}
mocker.patch.object(
demisto,
"getModules",
return_value={
"Cortex XDR - IR": {"state": "disabled", "brand": Brands.CORTEX_XDR_IR},
"Cortex Core - IR": {"state": "active", "brand": Brands.CORTEX_CORE_IR},
},
)
orchestrator = _get_orchestrator(args)
orchestrator._sanitize_and_validate_args()
actual_ids_list = orchestrator.args[QuarantineOrchestrator.ENDPOINT_IDS_ARG]
actual_ids_list.sort() # sort to have same order for comparison
expected_sorted_list = ["id1", "id2", "id3"]
assert actual_ids_list == expected_sorted_list
# --- Brands Tests ---
class TestBrands:
"""Tests specifically for the 'brands' script argument."""
def test_no_brands_provided_defaults_to_all_active_brands(self, mocker):
"""
Given: No 'brands' argument is provided.
When: _sanitize_and_validate_args is called.
Then: Ensure orchestrator.args['brands'] contains all active brands.
"""
args = {
"endpoint_id": "id1",
"file_hash": SHA_256_HASH,
"file_path": "/path",
}
mocker.patch.object(
demisto,
"getModules",
return_value={
"Cortex XDR - IR": {"state": "disabled", "brand": Brands.CORTEX_XDR_IR},
"Cortex Core - IR": {"state": "active", "brand": Brands.CORTEX_CORE_IR},
},
)
orchestrator = QuarantineOrchestrator(args)
orchestrator._sanitize_and_validate_args()
expected_brands = [Brands.CORTEX_CORE_IR]
actual_brands = orchestrator.args[QuarantineOrchestrator.BRANDS_ARG]
assert actual_brands == expected_brands
def test_invalid_brand_name_raises_exception(self, mocker):
"""
Given: An invalid brand name is provided in 'brands'.
When: _sanitize_and_validate_args is called.
Then: Ensure QuarantineException is raised.
"""
args = {
"endpoint_id": "id1",
"file_hash": SHA_256_HASH,
"file_path": "/path",
"brands": "Invalid Brand Name",
}
mocker.patch.object(
demisto,
"getModules",
return_value={
"Cortex XDR - IR": {"state": "disabled", "brand": Brands.CORTEX_XDR_IR},
"Cortex Core - IR": {"state": "active", "brand": Brands.CORTEX_CORE_IR},
},
)
orchestrator = QuarantineOrchestrator(args)
with pytest.raises(QuarantineException) as e:
orchestrator._sanitize_and_validate_args()
assert "Invalid brand" in str(e.value)
def test_disabled_brand_is_filtered_out(self, mocker):
"""
Given: A mix of active and disabled brands are provided.
When: _sanitize_and_validate_args is called.
Then: Ensure only the active brand remains in the arguments.
"""
# Arrange: Mock getModules to have one disabled brand
mocker.patch.object(
demisto,
"getModules",
return_value={
"Cortex XDR - IR": {"state": "active", "brand": Brands.CORTEX_XDR_IR},
"Cortex Core - IR": {"state": "disabled", "brand": Brands.CORTEX_CORE_IR}, # This one is disabled
},
)
args = {
"endpoint_id": "id1",
"file_hash": SHA_256_HASH,
"file_path": "/path",
"brands": f"{Brands.CORTEX_CORE_IR},{Brands.CORTEX_XDR_IR}",
}
orchestrator = QuarantineOrchestrator(args)
orchestrator._sanitize_and_validate_args()
# Assert that only the active brand is left
assert orchestrator.args[QuarantineOrchestrator.BRANDS_ARG] == [Brands.CORTEX_XDR_IR]
def test_no_active_brands_raises_exception(self, mocker):
"""
Given: All integrations for the target brands are disabled.
When: _sanitize_and_validate_args is called.
Then: Ensure DemistoException is raised.
"""
# Arrange: Mock getModules to have all relevant brands disabled
mocker.patch.object(
demisto,
"getModules",
return_value={
"Cortex XDR - IR": {"state": "disabled", "brand": Brands.CORTEX_XDR_IR},
"Cortex Core - IR": {"state": "disabled", "brand": Brands.CORTEX_CORE_IR},
},
)
args = {
"endpoint_id": "id1",
"file_hash": SHA_256_HASH,
"file_path": "/path",
}
orchestrator = QuarantineOrchestrator(args)
with pytest.raises(QuarantineException) as e:
orchestrator._sanitize_and_validate_args()
assert (
"have an enabled "
"integration instance. Ensure valid integration IDs are specified, and that "
"the integrations are enabled."
) in str(e.value)
class TestFileHash:
"""Tests specifically for the 'file_hash' script argument."""
def test_missing_file_hash_raises_exception(self):
"""
Given: Args without 'file_hash'.
When: _sanitize_and_validate_args is called.
Then: Ensure DemistoException is raised.
"""
args = {"endpoint_id": "id1", "file_path": "/path"}
orchestrator = QuarantineOrchestrator(args)
with pytest.raises(QuarantineException) as e:
orchestrator._sanitize_and_validate_args()
assert "Missing required argument" in str(e.value)
assert QuarantineOrchestrator.FILE_HASH_ARG in str(e.value)
@pytest.mark.parametrize(
"unsupported_hash",
[
"md5md5md5md5md5md5md5md5md5md5md", # md5
],
)
def test_unsupported_hash_type_raises_exception(self, unsupported_hash):
"""
Given: An unsupported hash type (MD5, SHA1).
When: _sanitize_and_validate_args is called.
Then: Ensure DemistoException is raised with 'Unsupported hash type'.
"""
args = {"endpoint_id": "id1", "file_hash": unsupported_hash, "file_path": "/path"}
orchestrator = QuarantineOrchestrator(args)
with pytest.raises(QuarantineException) as e:
orchestrator._sanitize_and_validate_args()
assert "Unsupported hash type" in str(e.value)
def test_hash_type_not_matching_active_brands_raises_exception(self, mocker):
"""
Given: A hash type that is not supported by the active brands.
When: _sanitize_and_validate_args is called.
Then: Ensure DemistoException is raised with 'Unsupported hash type'.
"""
mocker.patch.object(
demisto,
"getModules",
return_value={
"Cortex XDR - IR": {"state": "active", "brand": Brands.CORTEX_XDR_IR},
"Cortex Core - IR": {"state": "active", "brand": Brands.CORTEX_CORE_IR},
},
)
args = {"endpoint_id": "id1", "file_hash": SHA_1_HASH, "file_path": "/path"}
orchestrator = QuarantineOrchestrator(args)
with pytest.raises(QuarantineException) as e:
orchestrator._sanitize_and_validate_args()
assert str(e.value) == (
"Could not find enabled integrations for the requested hash type.\n"
"For hash_type SHA1 please use one of the following brands: Microsoft "
"Defender Advanced Threat Protection"
)
def test_brands_to_run_will_only_be_those_active_and_matching_hash_type(self, mocker):
"""
Given: A hash type that is supported by 1 active brand.
When: _sanitize_and_validate_args is called.
Then: brands_to_run is only of the active and matching brands of the hash type
"""
mocker.patch.object(
demisto,
"getModules",
return_value={
"Cortex XDR - IR": {"state": "active", "brand": Brands.CORTEX_XDR_IR},
"Cortex Core - IR": {"state": "disabled", "brand": Brands.CORTEX_CORE_IR},
"Microsoft Defender Advanced Threat Protection": {"state": "disabled", "brand": Brands.MDE},
},
)
args = {"endpoint_id": "id1", "file_hash": SHA_256_HASH, "file_path": "/path"}
orchestrator = QuarantineOrchestrator(args)
orchestrator._sanitize_and_validate_args()
assert orchestrator.args["brands"] == [Brands.CORTEX_XDR_IR]
class TestConstructor:
def test_constructor_sets_args(self):
"""
Given: Args for quarantine.
When: QuarantineOrchestrator is instantiated.
Then: Args are properly set
"""
args = {
"endpoint_id": "id1",
"file_hash": "sha256",
"file_path": "/path",
"brands": f"{Brands.CORTEX_CORE_IR},{Brands.CORTEX_XDR_IR}",
}
orchestrator = QuarantineOrchestrator(args)
assert orchestrator.args == args
assert not orchestrator.verbose
def test_constructor_sets_pending_and_completed_jobs(self, mocker):
"""
Given: Args for quarantine.
When: QuarantineOrchestrator is instantiated and there are pending jobs from initial poll
Then: Args are properly set, load pending jobs and completed results from the given args
"""
completed_result = QuarantineResult.create(
endpoint_id="id2",
status=QuarantineResult.Statuses.FAILED,
message=QuarantineResult.Messages.ENDPOINT_OFFLINE,
brand=Brands.MDE,
script_args={"file_hash": "sha256", "file_path": "/path"},
)
pending_job = [
{
"brand": "Microsoft Defender Advanced Threat Protection",
"poll_args": {
"action_ids": ["111"],
"comment": "Quarantine file hash: SHA1",
"file_hash": "SHA1",
"hide_polling_output": True,
"machine_id": ["id1"],
"polling": "true",
"timeout_in_seconds": "300",
},
"poll_command": "microsoft-atp-stop-and-quarantine-file",
}
]
args = {
"endpoint_id": "id1,is2",
"file_hash": "sha1",
"file_path": "/path",
"brands": f"{Brands.CORTEX_CORE_IR},{Brands.CORTEX_XDR_IR}",
"verbose": True,
"pending_jobs": pending_job,
"completed_results": QuarantineResult.to_context_entry([completed_result]),
}
orchestrator = QuarantineOrchestrator(args)
assert orchestrator.args == args
assert orchestrator.verbose
assert orchestrator.pending_jobs == pending_job
assert orchestrator.completed_results == [completed_result]
class TestRun:
def test_run_first_run_handles_mixed_results_and_schedules_poll(self, mocker):
"""
Given:
- A first run with a mix of endpoints (online, offline).
When:
- The orchestrator's run() method is called.
Then:
- Ensure a polling job is created only for endpoints needing action.
- Ensure initial results from both the mapper and handler are collected.
- Ensure both pending jobs and completed results are saved in args for next poll.
"""
# Arrange
args = {
"endpoint_id": "ep1,ep2,ep3,offline-ep",
"file_hash": "sha256sha256sha256sha256sha256sha256sha256sha256sha256sha256sha2",
"file_path": "/path",
}
# Mock EndpointBrandMapper to find some endpoints and fail others
mock_mapper_instance = mocker.Mock()
mock_mapper_instance.group_by_brand.return_value = {Brands.CORTEX_CORE_IR: ["ep1", "ep2", "ep3"]}
offline_result = QuarantineResult.create("offline-ep", "Failed", "Offline", "Unknown", args)
mock_mapper_instance.initial_results = [offline_result]
mocker.patch("QuarantineFile.EndpointBrandMapper", return_value=mock_mapper_instance)
mock_handler_instance = mocker.Mock()
mock_handler_instance.initiate_quarantine.return_value = {
"brand": Brands.CORTEX_CORE_IR,
"poll_command": "some-poll-cmd",
}
mocker.patch("QuarantineFile.handler_factory", return_value=mock_handler_instance)
orchestrator = _get_orchestrator(args)
# Act
result = orchestrator.run()
# Assert Polling is Scheduled with completed/pending jobs
assert result.continue_to_poll is True
# The result from your command
actual_args = result.args_for_next_run
# 1. Pop the unpredictable list from the actual results
actual_endpoint_ids = actual_args.pop("endpoint_id")
# Define your expected results, also without the endpoint_id list
expected_args = {
"brands": [Brands.CORTEX_CORE_IR],
"completed_results": [
{
"Brand": "Unknown",
"EndpointID": "offline-ep",
"FileHash": "sha256sha256sha256sha256sha256sha256sha256sha256sha256sha256sha2",
"FilePath": "/path",
"Message": "Offline",
"Status": "Failed",
}
],
# 'endpoint_id' key is removed
"file_hash": "sha256sha256sha256sha256sha256sha256sha256sha256sha256sha256sha2",
"file_path": "/path",
"pending_jobs": [{"brand": Brands.CORTEX_CORE_IR, "poll_command": "some-poll-cmd"}],
}
# 2. Assert that the rest of the dictionary is an exact match
assert actual_args == expected_args
# 3. Assert that the lists have the same content by comparing them as sets
expected_endpoint_ids = ["ep2", "ep1", "ep3", "offline-ep"]
assert set(actual_endpoint_ids) == set(expected_endpoint_ids)
assert result.response.readable_output == "Quarantine operations are still in progress..."
def test_run_first_run_handles_arg_validation_error(self, mocker):
"""
Given:
- A first run with a mix of endpoints (online, offline, already quarantined).
When:
- The orchestrator's run() method is called with invalid arguments.
Then:
- Ensure get_final_results is called.
- Ensure a polling job is not scheduled.
- Ensure the correct context data is returned
"""
# Arrange
args = {
"endpoint_id": "ep1,ep2,ep3,offline-ep",
"file_hash": "sha256sha256sha256sha256sha256sha256sha256sha256sha256sha256sha2",
"file_path": "/path",
}
mock_demisto_results = mocker.patch.object(demisto, "results")
mock_verify_and_get_valid_brands = mocker.patch.object(QuarantineOrchestrator, "_verify_and_get_valid_brands")
mock_verify_and_get_valid_brands.side_effect = QuarantineException("Invalid argument")
# spy on but dont mock _get_final_results because we need it to actually run
mock_get_final_results = mocker.spy(QuarantineOrchestrator, "_get_final_results")
mock_initiate_job = mocker.patch.object(QuarantineOrchestrator, "_initiate_jobs")
# Act
orchestrator = _get_orchestrator(args)
result = orchestrator.run()
# Assert get_final_results is called
mock_get_final_results.assert_called_once()
# Assert initiate_job is not called
mock_initiate_job.assert_not_called()
mock_demisto_results.assert_called_once_with(
{"Contents": "Invalid argument", "ContentsFormat": "text", "EntryContext": {}, "Type": 4},
)
assert len(result.response.outputs) == 4
assert result.response.outputs[0]["FilePath"] == "/path"
assert result.response.outputs[0]["FileHash"] == "sha256sha256sha256sha256sha256sha256sha256sha256sha256sha256sha2"
assert result.response.outputs[0]["Status"] == "Failed"
assert result.response.outputs[0]["Message"] == "Failed to quarantine file. Invalid argument"
assert result.response.outputs[0]["Brand"] == "Unknown"
assert len(orchestrator.completed_results) == 4
assert not orchestrator.pending_jobs # The list should now be empty
def test_run_polling_run_job_still_polling(self, mocker):
"""
Given:
- A polling run with a pending job in the given args for the poll run.
- The polling command indicates the action is still in progress.
When:
- The orchestrator's run() method is called.
Then:
- Ensure the method returns a PollResult to continue polling.
"""
# Arrange
pending_job = {"brand": Brands.CORTEX_CORE_IR, "poll_command": "some-poll-cmd", "poll_args": {"action_id": "123"}}
args = {"file_hash": "hash123", "file_path": "/path", "pending_jobs": [pending_job]}
# Mock the polling command to return 'polling: True'
polling_response = [{"Type": 1, "Contents": {}, "Metadata": {"polling": True, "pollingArgs": {"action_id": "456"}}}]
mocker.patch.object(demisto, "executeCommand", return_value=polling_response)
orchestrator = _get_orchestrator(args)
# Act
result = orchestrator.run()
# Assert continuing to poll because response stated so.
assert result.continue_to_poll is True
def test_run_polling_run_job_finishes(self, mocker):
"""
Given:
- A polling run with a pending job in the args.
- The polling command indicates the action is complete.
When:
- The orchestrator's run() method is called.
Then:
- Ensure the handler's finalize method is called.
- Ensure the pending jobs list becomes empty.
- Ensure the method returns a final result (continue_to_poll=False).
"""
# Arrange
pending_job = {"brand": Brands.CORTEX_CORE_IR, "poll_command": "some-poll-cmd", "poll_args": {"action_id": "123"}}
args = {"endpoint_id": "ep1", "file_hash": "hash123", "file_path": "/path", "pending_jobs": [pending_job]}
# Mock the polling command to return 'polling: False'
polling_response = [{"Type": 1, "Contents": {}, "Metadata": {"polling": False}}]
mocker.patch.object(demisto, "executeCommand", return_value=polling_response)
# Mock the handler's finalize method
mock_handler_instance = mocker.Mock()
final_result = QuarantineResult.create("ep1", "Success", "File quarantined", Brands.CORTEX_CORE_IR, args)
mock_handler_instance.finalize.return_value = [final_result]
mocker.patch("QuarantineFile.handler_factory", return_value=mock_handler_instance)
orchestrator = _get_orchestrator(args)
# Act
result = orchestrator.run()
# Assert
assert result.continue_to_poll is False
assert result.response.outputs_prefix == "QuarantineFile"
assert result.response.outputs_key_field == ["EndpointID", "FilePath", "FileHash"]
assert len(result.response.outputs) == 1
assert result.response.outputs[0]["EndpointID"] == "ep1"
assert result.response.outputs[0]["FilePath"] == "/path"
assert result.response.outputs[0]["FileHash"] == "hash123"
assert result.response.outputs[0]["Status"] == "Success"
assert result.response.outputs[0]["Message"] == "File quarantined"
assert result.response.outputs[0]["Brand"] == Brands.CORTEX_CORE_IR
# finalize receives the final polling response.
mock_handler_instance.finalize.assert_called_once_with(polling_response)
assert len(orchestrator.completed_results) == 1
assert orchestrator.completed_results[0].EndpointID == "ep1"
assert not orchestrator.pending_jobs # The list should now be empty
class TestScriptEntrypoints:
"""
Tests for the main script entry points, main() and quarantine_file_script().
"""
def test_main_function_success_path(self, mocker):
"""
Given:
- Script arguments provided via demisto.args().
When:
- The main() function is called.
Then:
- Ensure the polling entry point is called with polling enabled.
- Ensure return_results is called with the result.
"""
# Arrange
args = {"endpoint_id": "ep1", "file_hash": "hash123"}
mocker.patch.object(demisto, "args", return_value=args)
mock_return_results = mocker.patch("QuarantineFile.return_results")
mock_script_func = mocker.patch("QuarantineFile.quarantine_file_script", return_value="SUCCESS")
# Act
main()
# Assert
expected_args = args.copy()
expected_args["polling"] = True
mock_script_func.assert_called_once_with(expected_args)
mock_return_results.assert_called_once_with("SUCCESS")
if __name__ == "__main__":
pytest.main()