GravityZone

import json
import os
from typing import cast
from unittest.mock import patch
from freezegun import freeze_time
import pytest

from CommonServerPython import CommandResults, EntryType


SERVER_URL = "https://localhost"


def util_load_json(path):
    with open(path, encoding="utf-8") as f:
        return json.loads(f.read())


def get_client():
    from GravityZone import Client

    return Client(url=SERVER_URL, verify=False, api_key="fake")


def load_api_mocked_data(requests_mock, command_name):
    command_renamed = command_name.replace("-", "_")
    function_name = f"{command_renamed}_command"

    api_file_path = os.path.join(os.path.dirname(__file__), "test_data", f"{function_name}_api.json")
    data = util_load_json(api_file_path)

    # Group responses by URL and method
    url_responses = {}
    for item in data:
        url = f"{SERVER_URL}{item['request']['url']}"
        method = "get" if item["request"].get("get", False) else "post"
        url_responses.setdefault((method, url), []).append(item["response"])

    # Register each URL/method with a callback to return responses sequentially
    for (method, url), responses in url_responses.items():
        response_iter = iter(responses)

        def get_callback(request, context, response_iter=response_iter, responses=responses):
            try:
                resp = next(response_iter)
            except StopIteration:
                resp = responses[-1]
            if "content" in resp:
                context.status_code = resp.get("status_code", 200)
                for k, v in resp.get("headers", {}).items():
                    context.headers[k] = v
                return resp["content"].encode()
            return resp

        def post_callback(request, context, response_iter=response_iter, responses=responses):
            try:
                return next(response_iter)
            except StopIteration:
                return responses[-1]

        if method == "post":
            requests_mock.post(url, json=post_callback)
        else:
            requests_mock.get(url, content=get_callback)


def assert_command_mocked_data(command_name, command_response, polling_func=None, client=None):
    test_data_dir = os.path.join(os.path.dirname(__file__), "test_data")

    command_renamed = command_name.replace("-", "_")
    function_name = f"{command_renamed}_command"

    if polling_func is not None and client is not None:
        while getattr(command_response, "scheduled_command", None):
            command_response = polling_func(command_response.scheduled_command._args, client)

    if isinstance(command_response, list):
        command_response = command_response[0]

    context = command_response.to_context()

    raw_response = context.get("Contents")
    human_readable = context.get("HumanReadable")
    outputs = context.get("EntryContext")

    def assert_file_content(actual, expected_paths, error_msg):
        if not any(os.path.exists(p) for p in expected_paths):
            raise FileNotFoundError(f"No expected file found for {function_name} in paths: {expected_paths}")
        for expected_path in expected_paths:
            if not os.path.exists(expected_path):
                continue
            with open(expected_path, encoding="utf-8") as f:
                expected = f.read() if expected_path.endswith((".md", ".txt")) else json.load(f)
                assert actual == expected, f"{error_msg} for {function_name}"

    assert_file_content(
        raw_response,
        [
            os.path.join(test_data_dir, f"{function_name}_raw_response.json"),
            os.path.join(test_data_dir, f"{function_name}_raw_response.txt"),
        ],
        "Raw response does not match",
    )
    assert_file_content(
        human_readable, [os.path.join(test_data_dir, f"{function_name}_human_readable.md")], "Human readable does not match"
    )
    assert_file_content(
        outputs, [os.path.join(test_data_dir, f"{function_name}_entry_context.json")], "Entry context does not match"
    )


@patch("GravityZone.demisto")
def test_fetch_incidents_command(mock_demisto, requests_mock):
    """
    Given
            All relevant arguments for the command that is executed
    When
            Calling fetch-incidents command
    Then
            Make sure the fetched incidents values are as expected.
    """

    # Prepare
    from GravityZone import fetch_incidents_command

    mock_demisto.command.return_value = "fetch-incidents"
    mock_demisto.args.return_value = {}
    mock_demisto.getLastRun.return_value = {}
    mock_demisto.integrationInstance.return_value = "GravityZone"
    mock_demisto.params.return_value = {"max_fetch": 2, "first_fetch": "3 days"}
    load_api_mocked_data(requests_mock, "fetch-incidents")
    client = get_client()
    fetch_incidents_command(client, {})

    assert mock_demisto.incidents.call_count == 1
    assert len(mock_demisto.incidents.call_args[0][0]) == 2

    incidents = mock_demisto.incidents.call_args[0][0]
    for incident in incidents:
        assert "name" in incident
        assert "occurred" in incident
        assert "rawJSON" in incident

    incident_names = [incident.get("name") for incident in incidents]
    assert incident_names == ["GravityZone Incident #477", "GravityZone Incident #475"]
    assert mock_demisto.setLastRun.call_count == 1


@patch("GravityZone.demisto")
def test_get_remote_data_command(mock_demisto, requests_mock):
    """
    Given
            All relevant arguments for the command that is executed
    When
            Calling get-remote-data command
    Then
            Make sure the mirrored incident values are as expected.
    """

    # Prepare
    from GravityZone import get_remote_data_command

    mock_demisto.command.return_value = "get-remote-data"
    mock_demisto.params.return_value = {"mirror_direction": "Both"}
    mock_demisto.integrationInstance.return_value = "GravityZone"
    load_api_mocked_data(requests_mock, "get-remote-data")
    client = get_client()
    result = get_remote_data_command(client, {"id": "6940ee975f5c8b75247c3f38", "lastUpdate": "2024-01-01T00:00:00Z"})
    assert result is not None, "get_remote_data_command returned None"
    local_data = result.extract_for_local() if hasattr(result, "extract_for_local") else None
    assert local_data is not None, "extract_for_local() returned None"
    assert isinstance(local_data, list), "extract_for_local() did not return a list"
    assert len(local_data) > 0, "extract_for_local() returned an empty list"
    assert local_data[0].get("id") == "6940ee975f5c8b75247c3f38"


@patch("GravityZone.demisto")
def test_get_modified_remote_data_command(mock_demisto, requests_mock):
    """
    Given
            All relevant arguments for the command that is executed
    When
            Calling get-modified-remote-data command
    Then
            Make sure the mirrored incident values are as expected.
    """
    # Prepare
    from GravityZone import get_modified_remote_data_command

    mock_demisto.command.return_value = "get-modified-remote-data"
    mock_demisto.params.return_value = {}
    load_api_mocked_data(requests_mock, "get-modified-remote-data")
    client = get_client()
    result = get_modified_remote_data_command(client, {"lastUpdate": "2026-01-01T00:00:00Z"})
    assert result.to_entry().get("Contents") == ["6940ee975f5c8b75247c3f38", "6940eeadaa87b3d57af088aa"]


@freeze_time("2025-12-16 13:00:00 UTC")
@patch("GravityZone.demisto")
def test_get_mapping_fields_command(mock_demisto, mocker, requests_mock):
    """
    Given
            All relevant arguments for the command that is executed
    When
            Calling get-mapping-fields command
    Then
            Make sure the mappings values are as expected.
    """
    # Prepare
    from GravityZone import get_mapping_fields_command

    mock_demisto.command.return_value = "get-mapping-fields"
    mock_demisto.params.return_value = {}
    client = get_client()
    response = get_mapping_fields_command(client, {})
    schemes = response.extract_mapping()
    assert len(schemes) == 2

    xdr_scheme = next((s for s in schemes if s == "GravityZone XDR"), None)
    edr_scheme = next((s for s in schemes if s == "GravityZone EDR"), None)

    assert xdr_scheme is not None
    assert edr_scheme is not None


@freeze_time("2025-12-16 13:00:00 UTC")
@patch("GravityZone.demisto")
def test_test_module(mock_demisto, requests_mock):
    """
    Given
            All relevant arguments for the command that is executed
    When
            Calling test-module command
    Then
            Make sure the command works.
    """

    # Prepare
    from GravityZone import test_module

    mock_demisto.command.return_value = "test-module"
    mock_demisto.params.return_value = {
        "first_fetch": "3 days",
    }
    mock_demisto.args.return_value = {}
    load_api_mocked_data(requests_mock, "test-module")
    client = get_client()

    # Execute command
    result = test_module(client=client, args={})

    # Assert
    assert result == "ok"


@patch("GravityZone.demisto")
def test_gz_endpoint_list_command(mock_demisto, requests_mock):
    """
    Given
            All relevant arguments for the command that is executed
    When
            Calling gz-endpoint-list command
    Then
            Make sure the outputs, outputs_prefix and outputs_key_field values are as expected.
    """

    # Prepare
    from GravityZone import gz_endpoint_list_command

    mock_demisto.command.return_value = "gz-endpoint-list"
    mock_demisto.params.return_value = {}
    load_api_mocked_data(requests_mock, "gz-endpoint-list")
    client = get_client()

    # Execute command
    command_response = gz_endpoint_list_command(client=client, args={})

    # Assert command response
    assert_command_mocked_data("gz-endpoint-list", command_response)


@patch("GravityZone.demisto")
@pytest.mark.parametrize(
    "valid_name",
    [
        "ABC",
    ],
)
def test_gz_endpoint_list_command_valid_name_values(mock_demisto, requests_mock, valid_name):
    """
    Given
            A valid `name` argument value
    When
            Calling gz-endpoint-list command
    Then
            Make sure command execution succeeds.
    """

    # Prepare
    from GravityZone import gz_endpoint_list_command

    mock_demisto.command.return_value = "gz-endpoint-list"
    mock_demisto.params.return_value = {}
    load_api_mocked_data(requests_mock, "gz-endpoint-list-name")
    client = get_client()

    # Execute
    command_response = gz_endpoint_list_command(client=client, args={"name": valid_name})

    # Assert
    assert command_response is not None
    assert requests_mock.called

    # Assert command response
    assert_command_mocked_data("gz-endpoint-list-name", command_response)


@patch("GravityZone.demisto")
@pytest.mark.parametrize(
    ("invalid_name", "expected_error"),
    [
        ("ab", r"The 'name' argument must be at least 3 characters long\."),
        ("ac*b", r"If using a wildcard, the 'name' argument must start with '\*'\."),
        ("*ac*b*d", r"The 'name' argument can only contain one wildcard '\*' character\."),
    ],
)
def test_gz_endpoint_list_command_invalid_name_values(mock_demisto, requests_mock, invalid_name, expected_error):
    """
    Given
            An invalid `name` argument value
    When
            Calling gz-endpoint-list command
    Then
            Make sure command validation fails.
    """

    # Prepare
    from GravityZone import gz_endpoint_list_command

    mock_demisto.command.return_value = "gz-endpoint-list"
    mock_demisto.params.return_value = {}
    client = get_client()

    # Execute + Assert
    with pytest.raises(Exception, match=expected_error):
        gz_endpoint_list_command(client=client, args={"name": invalid_name})

    # Invalid input should fail before API calls
    assert not requests_mock.called


@patch("GravityZone.demisto")
@pytest.mark.parametrize(
    "node_id",
    [
        "COMPANY_ID",
    ],
)
def test_gz_endpoint_list_command_with_node_id_skips_get_my_company(mock_demisto, requests_mock, node_id):
    """
    Given
            A `node_id` argument value
    When
            Calling gz-endpoint-list command
    Then
            Make sure get_my_company is not called.
    """

    # Prepare
    from GravityZone import gz_endpoint_list_command

    mock_demisto.command.return_value = "gz-endpoint-list"
    mock_demisto.params.return_value = {}
    load_api_mocked_data(requests_mock, "gz-endpoint-list-node")
    client = get_client()

    with patch.object(client, "get_my_company", side_effect=AssertionError("get_my_company should not be called")) as mocked:
        # Execute
        command_response = gz_endpoint_list_command(client=client, args={"node_id": node_id})

        # Assert
        assert command_response is not None
        mocked.assert_not_called()
        assert requests_mock.called

        # Assert command response
        assert_command_mocked_data("gz-endpoint-list-node", command_response)


@patch("GravityZone.demisto")
def test_gz_endpoint_get_command(mock_demisto, requests_mock):
    """
    Given
            All relevant arguments for the command that is executed
    When
            Calling gz-endpoint-get command
    Then
            Make sure the outputs, outputs_prefix and outputs_key_field values are as expected.
    """
    # Prepare
    from GravityZone import gz_endpoint_get_command

    mock_demisto.command.return_value = "gz-endpoint-get"
    mock_demisto.params.return_value = {}
    load_api_mocked_data(requests_mock, "gz-endpoint-get")
    client = get_client()

    # Execute command
    command_response = gz_endpoint_get_command(client=client, args={"id": "ENDPOINT_ID"})

    # Assert command response
    assert_command_mocked_data("gz-endpoint-get", command_response)


@patch("GravityZone.demisto")
def test_gz_endpoint_users_loggedin_command(mock_demisto, requests_mock):
    """
    Given
            All relevant arguments for the command that is executed
    When
            Calling gz-endpoint-users-loggedin command
    Then
            Make sure the outputs, outputs_prefix and outputs_key_field values are as expected.
    """

    # Prepare
    from GravityZone import gz_endpoint_users_loggedin_command, gz_poll_endpoint_users_loggedin_status_command

    mock_demisto.command.return_value = "gz-endpoint-users-loggedin"
    mock_demisto.params.return_value = {}
    load_api_mocked_data(requests_mock, "gz-endpoint-users-loggedin")
    client = get_client()

    # Execute command
    command_response = gz_endpoint_users_loggedin_command(client=client, args={"id": "ENDPOINT_ID"})

    # Assert command response
    assert_command_mocked_data(
        "gz-endpoint-users-loggedin",
        command_response,
        polling_func=gz_poll_endpoint_users_loggedin_status_command,
        client=client,
    )


@patch("GravityZone.demisto")
def test_gz_endpoint_create_memory_dump_command(mock_demisto, requests_mock):
    """
    Given
            All relevant arguments for the command that is executed
    When
            Calling gz-endpoint-create-memory-dump command
    Then
            Make sure the outputs, outputs_prefix and outputs_key_field values are as expected.
    """

    # Prepare
    from GravityZone import gz_endpoint_create_memory_dump_command, gz_endpoint_memory_dump_status_command

    mock_demisto.command.return_value = "gz-endpoint-create-memory-dump"
    mock_demisto.params.return_value = {}
    load_api_mocked_data(requests_mock, "gz-endpoint-create-memory-dump")
    client = get_client()

    # Execute command
    command_response = gz_endpoint_create_memory_dump_command(
        client=client,
        args={"id": "ENDPOINT_ID", "path": "C:\\dumps\\", "password": "ComplexPass123!"},
    )

    # Assert command response
    assert_command_mocked_data(
        "gz-endpoint-create-memory-dump",
        command_response,
        polling_func=lambda poll_args, poll_client: gz_endpoint_memory_dump_status_command(poll_client, poll_args),
        client=client,
    )


@patch("GravityZone.demisto")
def test_gz_endpoint_create_memory_dump_command_no_polling(mock_demisto, requests_mock):
    from GravityZone import gz_endpoint_create_memory_dump_command
    from CommonServerPython import CommandResults

    mock_demisto.command.return_value = "gz-endpoint-create-memory-dump"
    mock_demisto.params.return_value = {}
    load_api_mocked_data(requests_mock, "gz-endpoint-create-memory-dump")
    client = get_client()

    command_response = cast(
        CommandResults,
        gz_endpoint_create_memory_dump_command(
            client=client,
            args={"id": "ENDPOINT_ID", "path": "C:\\dumps\\", "password": "ComplexPass123!", "polling": "false"},
        ),
    )

    assert getattr(command_response, "scheduled_command", None) is None
    context = command_response.to_context()
    output = context["EntryContext"]["GravityZone.MemoryDump(val.EndpointID && val.EndpointID == obj.EndpointID)"]

    assert output["TaskID"] == "TASK_ID"
    assert output["TaskType"] == "CreateMemoryDump"
    assert output["Status"] == "Pending"
    assert output["EndpointID"] == "ENDPOINT_ID"
    assert output["EndDate"] == ""
    assert output["ErrorCode"] == ""
    assert output["Error"] == ""
    assert output["DownloadURL"] == ""


@patch("GravityZone.check_endpoint_memory_dump_status")
def test_gz_endpoint_memory_dump_status_command_defaults_polling_true(mock_check_status):
    from GravityZone import gz_endpoint_memory_dump_status_command
    from GravityZone import Client

    client = cast(Client, object())

    gz_endpoint_memory_dump_status_command(client, {"task_id": "TASK_ID", "endpoint_id": "ENDPOINT_ID"})

    mock_check_status.assert_called_once_with(
        {"task_id": "TASK_ID", "endpoint_id": "ENDPOINT_ID", "polling": True},
        client,
    )


@patch("GravityZone.demisto")
def test_gz_endpoint_memory_dump_status_command(mock_demisto, requests_mock):
    from GravityZone import gz_endpoint_memory_dump_status_command

    mock_demisto.command.return_value = "gz-endpoint-memory-dump-status"
    mock_demisto.params.return_value = {}
    load_api_mocked_data(requests_mock, "gz-endpoint-memory-dump-status")
    client = get_client()

    command_response = gz_endpoint_memory_dump_status_command(
        client,
        {"task_id": "TASK_ID", "endpoint_id": "ENDPOINT_ID"},
    )

    assert_command_mocked_data(
        "gz-endpoint-memory-dump-status",
        command_response,
        polling_func=lambda poll_args, poll_client: gz_endpoint_memory_dump_status_command(poll_client, poll_args),
        client=client,
    )


def test_check_endpoint_users_loggedin_status_returns_error_on_unexpected_task_type():
    from GravityZone import check_endpoint_users_loggedin_status

    class MockClient:
        def get_task_status(self, _task_id):
            return {"status": 3, "type": 27, "subtasks": []}

    result = check_endpoint_users_loggedin_status(
        {"task_id": "TASK_ID", "endpoint_id": "ENDPOINT_ID"},
        MockClient(),
    )

    assert isinstance(result, CommandResults)
    assert result.readable_output is not None
    assert "unexpected type" in result.readable_output
    assert "GetActiveSessions" in result.readable_output


def test_check_endpoint_users_loggedin_status_returns_error_when_task_type_is_missing():
    from GravityZone import check_endpoint_users_loggedin_status

    class MockClient:
        def get_task_status(self, _task_id):
            return {"status": 3, "subtasks": []}

    result = check_endpoint_users_loggedin_status(
        {"task_id": "TASK_ID", "endpoint_id": "ENDPOINT_ID"},
        MockClient(),
    )

    assert isinstance(result, CommandResults)
    assert result.readable_output is not None
    assert "does not have a type" in result.readable_output
    assert "gz-poll-endpoint-users-loggedin-status" in result.readable_output


def test_check_endpoint_memory_dump_status_returns_error_on_unexpected_task_type():
    from GravityZone import check_endpoint_memory_dump_status

    class MockClient:
        def get_task_status(self, _task_id):
            return {"status": 3, "type": 26, "subtasks": []}

    result = check_endpoint_memory_dump_status(
        {"task_id": "TASK_ID", "endpoint_id": "ENDPOINT_ID"},
        MockClient(),
    )

    assert isinstance(result, CommandResults)
    assert result.readable_output is not None
    assert "unexpected type" in result.readable_output
    assert "CreateMemoryDump" in result.readable_output


def test_check_endpoint_memory_dump_status_returns_error_when_task_type_is_missing():
    from GravityZone import check_endpoint_memory_dump_status

    class MockClient:
        def get_task_status(self, _task_id):
            return {"status": 3, "subtasks": []}

    result = check_endpoint_memory_dump_status(
        {"task_id": "TASK_ID", "endpoint_id": "ENDPOINT_ID"},
        MockClient(),
    )

    assert isinstance(result, CommandResults)
    assert result.readable_output is not None
    assert "does not have a type" in result.readable_output
    assert "gz-endpoint-memory-dump-status" in result.readable_output


def test_extract_active_sessions_from_task_handles_missing_optional_fields():
    from GravityZone import _extract_active_sessions_from_task

    task_output = {
        "status": 3,
        "subtasks": [
            {
                "endpointId": "endpoint-1",
                "endpointName": "host-1",
                "status": 3,
                "result": [
                    {
                        "connection": {
                            "started": "2021-05-19T10:37:56Z",
                            "type": "local",
                        },
                        "user": {
                            "displayName": "user@example.com",
                        },
                    }
                ],
            }
        ],
    }

    sessions = _extract_active_sessions_from_task(task_output)

    assert len(sessions) == 1
    assert sessions[0]["Username"] == "user@example.com"
    assert sessions[0]["ConnectionType"] == "local"
    assert sessions[0]["StartTime"] == "2021-05-19T10:37:56Z"
    assert "UserSID" not in sessions[0]
    assert "DomainSID" not in sessions[0]
    assert "OrganizationalUnitDN" not in sessions[0]
    assert "MemberOfSIDs" not in sessions[0]


def test_extract_endpoint_summary_from_task_returns_hostname_for_matching_endpoint_id():
    from GravityZone import _extract_endpoint_summary_from_task

    task_output = {
        "subtasks": [
            {
                "endpointId": "endpoint-2",
                "endpointName": "host-2",
            },
            {
                "endpointId": "endpoint-1",
                "endpointName": "host-1-updated",
            },
        ]
    }

    endpoint_id, hostname = _extract_endpoint_summary_from_task(task_output, "endpoint-1")

    assert endpoint_id == "endpoint-1"
    assert hostname == "host-1-updated"


def test_build_users_loggedin_results_outputs_endpoint_scoped_context():
    from GravityZone import _build_users_loggedin_results

    task_output = {
        "status": 3,
        "subtasks": [
            {
                "endpointId": "endpoint-1",
                "endpointName": "host-1-updated",
                "status": 3,
                "result": [],
            }
        ],
    }

    result = _build_users_loggedin_results(task_output, "endpoint-1")

    assert result.outputs_prefix == "GravityZone.Endpoint"
    assert result.outputs_key_field == "ID"
    assert result.outputs == {
        "ID": "endpoint-1",
        "Hostname": "host-1-updated",
        "ActiveSessions": [],
    }


def test_extract_memory_dump_summary_uses_matching_endpoint_subtask_and_download_url():
    from GravityZone import _extract_memory_dump_summary

    task_output = {
        "subtasks": [
            {
                "endpointId": "endpoint-other",
                "status": 3,
                "downloadURL": "https://example.com/other",
            },
            {
                "endpointId": "endpoint-1",
                "status": 3,
                "downloadURL": "https://example.com/memory-dump",
            },
        ]
    }

    endpoint_id, endpoint_hostname, subtask, download_url = _extract_memory_dump_summary(task_output, "endpoint-1")
    assert endpoint_id == "endpoint-1"
    assert endpoint_hostname == ""
    assert subtask == {
        "endpointId": "endpoint-1",
        "status": 3,
        "downloadURL": "https://example.com/memory-dump",
    }
    assert download_url == "https://example.com/memory-dump"


def test_extract_memory_dump_summary_ignores_non_processed_subtask_for_download_url():
    from GravityZone import _extract_memory_dump_summary

    task_output = {
        "subtasks": [
            {
                "endpointId": "endpoint-1",
                "status": 2,
                "downloadURL": "https://example.com/should-not-be-used",
            }
        ]
    }

    endpoint_id, endpoint_hostname, subtask, download_url = _extract_memory_dump_summary(task_output, "endpoint-1")
    assert endpoint_id == "endpoint-1"
    assert endpoint_hostname == ""
    assert subtask == {
        "endpointId": "endpoint-1",
        "status": 2,
        "downloadURL": "https://example.com/should-not-be-used",
    }
    assert download_url == ""


def test_extract_memory_dump_summary_returns_none_when_endpoint_not_found():
    from GravityZone import _extract_memory_dump_summary

    task_output = {
        "subtasks": [
            {
                "endpointId": "endpoint-1",
                "status": 3,
                "downloadURL": "https://example.com/memory-dump",
            }
        ]
    }

    endpoint_id, endpoint_hostname, subtask, download_url = _extract_memory_dump_summary(task_output, "endpoint-missing")
    assert endpoint_id == "endpoint-missing"
    assert endpoint_hostname == ""
    assert subtask is None
    assert download_url == ""


def test_extract_memory_dump_task_id_from_string_result():
    from GravityZone import _extract_memory_dump_task_id

    assert _extract_memory_dump_task_id("TASK_ID") == "TASK_ID"


def test_extract_memory_dump_task_id_raises_on_invalid_result():
    from GravityZone import _extract_memory_dump_task_id

    with pytest.raises(Exception, match="createMemoryDumpTask response is missing task ID"):
        _extract_memory_dump_task_id({})


def test_build_memory_dump_results_outputs_download_url_on_success():
    from GravityZone import _build_memory_dump_results

    task_output = {
        "status": 3,
        "subtasks": [
            {
                "endpointId": "endpoint-1",
                "endpointName": "host-1",
                "status": 3,
                "startDate": "2026-05-25T10:00:00",
                "endDate": "2026-05-25T10:01:00",
                "errorCode": "Success",
                "errorMessage": "Success",
                "downloadURL": "https://example.com/memory-dump",
            }
        ],
    }

    result = _build_memory_dump_results(task_output, "TASK_ID", "endpoint-1")

    assert result.outputs_prefix == "GravityZone.MemoryDump"
    assert result.outputs_key_field == "EndpointID"
    assert result.outputs == {
        "TaskID": "TASK_ID",
        "TaskType": "CreateMemoryDump",
        "Status": "Processed",
        "EndpointID": "endpoint-1",
        "Hostname": "host-1",
        "StartDate": "2026-05-25T10:00:00Z",
        "EndDate": "2026-05-25T10:01:00Z",
        "ErrorCode": "Success",
        "Error": "Success",
        "DownloadURL": "https://example.com/memory-dump",
    }


def test_build_memory_dump_results_omits_download_url_when_unavailable():
    from GravityZone import _build_memory_dump_results

    task_output = {
        "status": 1,
        "subtasks": [
            {
                "endpointId": "endpoint-1",
                "endpointName": "host-1",
                "status": 1,
                "startDate": "2026-05-25T10:00:00",
            }
        ],
    }

    result = _build_memory_dump_results(task_output, "TASK_ID", "endpoint-1")

    assert result.outputs == {
        "TaskID": "TASK_ID",
        "TaskType": "CreateMemoryDump",
        "Status": "Pending",
        "EndpointID": "endpoint-1",
        "Hostname": "host-1",
        "StartDate": "2026-05-25T10:00:00Z",
        "EndDate": "",
        "ErrorCode": "",
        "Error": "",
        "DownloadURL": "",
    }


def test_build_memory_dump_results_returns_error_when_endpoint_not_found():
    from GravityZone import _build_memory_dump_results

    task_output = {
        "status": 3,
        "subtasks": [
            {
                "endpointId": "endpoint-1",
                "endpointName": "host-1",
                "status": 3,
                "startDate": "2026-05-25T10:00:00",
                "endDate": "2026-05-25T10:01:00",
            }
        ],
    }

    result = _build_memory_dump_results(task_output, "TASK_ID", "endpoint-missing")

    assert result.entry_type == EntryType.ERROR
    assert result.readable_output == "Invalid Endpoint ID. Available endpoint IDs for task 'TASK_ID': ['endpoint-1']"


def test_build_memory_dump_results_returns_error_when_task_has_no_results():
    from GravityZone import _build_memory_dump_results

    task_output = {
        "status": 3,
        "subtasks": [],
    }

    result = _build_memory_dump_results(task_output, "TASK_ID", "")

    assert result.entry_type == EntryType.ERROR
    assert result.readable_output == "Task 'TASK_ID' has no memory dump results."


def test_generate_processed_task_command_result_maps_memory_dump_task_type():
    from GravityZone import generate_processed_task_command_result

    task_output = {
        "type": 27,
        "subtasks": [
            {
                "endpointId": "endpoint-1",
                "endpointName": "host-1",
                "startDate": "2026-05-26T10:00:00",
                "endDate": "2026-05-26T10:01:00",
            }
        ],
    }

    result = generate_processed_task_command_result("parent-task-id", task_output, {})

    assert result.outputs_prefix == "GravityZone.Command.CreateMemoryDump"


@patch("GravityZone.demisto")
def test_gz_endpoint_isolate_command(mock_demisto, requests_mock):
    """
    Given
            All relevant arguments for the command that is executed
    When
            Calling gz-endpoint-isolate command
    Then
            Make sure the outputs, outputs_prefix and outputs_key_field values are as expected.
    """

    # Prepare
    from GravityZone import gz_endpoint_isolate_command, gz_poll_task_status_command

    mock_demisto.command.return_value = "gz-endpoint-isolate"
    mock_demisto.params.return_value = {}
    load_api_mocked_data(requests_mock, "gz-endpoint-isolate")
    client = get_client()

    # Execute command
    command_response = gz_endpoint_isolate_command(client=client, args={"id": "ENDPOINT_ID"})

    # Assert command response
    assert_command_mocked_data("gz-endpoint-isolate", command_response, polling_func=gz_poll_task_status_command, client=client)


@patch("GravityZone.demisto")
def test_gz_endpoint_deisolate_command(mock_demisto, requests_mock):
    """
    Given
            All relevant arguments for the command that is executed
    When
            Calling gz-endpoint-deisolate command
    Then
            Make sure the outputs, outputs_prefix and outputs_key_field values are as expected.
    """

    # Prepare
    from GravityZone import gz_endpoint_deisolate_command, gz_poll_task_status_command

    mock_demisto.command.return_value = "gz-endpoint-deisolate"
    mock_demisto.params.return_value = {}
    load_api_mocked_data(requests_mock, "gz-endpoint-deisolate")
    client = get_client()

    # Execute command
    command_response = gz_endpoint_deisolate_command(client=client, args={"id": "ENDPOINT_ID"})

    # Assert command response
    assert_command_mocked_data("gz-endpoint-deisolate", command_response, polling_func=gz_poll_task_status_command, client=client)


@patch("GravityZone.demisto")
def test_gz_endpoint_kill_process_command(mock_demisto, requests_mock):
    """
    Given
            All relevant arguments for the command that is executed
    When
            Calling gz-endpoint-kill-process command
    Then
            Make sure the outputs, outputs_prefix and outputs_key_field values are as expected.
    """

    # Prepare
    from GravityZone import (
        gz_endpoint_kill_process_command,
        gz_poll_task_status_command,
    )

    mock_demisto.command.return_value = "gz-endpoint-kill-process"
    mock_demisto.params.return_value = {}
    load_api_mocked_data(requests_mock, "gz-endpoint-kill-process")
    client = get_client()

    # Execute command
    command_response = gz_endpoint_kill_process_command(
        client=client,
        args={
            "id": "ENDPOINT_ID",
            "pid": "5876",
        },
    )

    # Assert command response
    assert_command_mocked_data(
        "gz-endpoint-kill-process", command_response, polling_func=gz_poll_task_status_command, client=client
    )


@patch("GravityZone.demisto")
def test_gz_endpoint_run_command_command(mock_demisto, requests_mock):
    """
    Given
            All relevant arguments for the command that is executed
    When
            Calling gz-endpoint-run-command command
    Then
            Make sure the outputs, outputs_prefix and outputs_key_field values are as expected.
    """

    # Prepare
    from GravityZone import (
        gz_endpoint_run_command_command,
        gz_poll_investigation_activity_status_command,
    )

    mock_demisto.command.return_value = "gz-endpoint-run-command"
    mock_demisto.params.return_value = {}
    load_api_mocked_data(requests_mock, "gz-endpoint-run-command")
    client = get_client()

    # Execute command
    command_response = gz_endpoint_run_command_command(
        client=client,
        args={
            "id": "6942a43afe8d4e463ca5c197",
            "command": "ls -1",
        },
    )

    # Assert command response
    assert_command_mocked_data(
        "gz-endpoint-run-command", command_response, polling_func=gz_poll_investigation_activity_status_command, client=client
    )


@patch("GravityZone.demisto")
def test_gz_endpoint_get_process_tree_by_hash(mock_demisto, requests_mock):
    """
    Given
            All relevant arguments for the command that is executed
    When
            Calling gz-endpoint-get-process-tree-by-hash command
    Then
            Make sure the outputs, outputs_prefix and outputs_key_field values are as expected.
    """

    # Prepare
    from GravityZone import (
        gz_endpoint_get_process_tree_by_hash_command,
        gz_poll_live_search_status_command,
    )

    mock_demisto.command.return_value = "gz-endpoint-get-process-tree-by-hash"
    mock_demisto.params.return_value = {}
    load_api_mocked_data(requests_mock, "gz-endpoint-get-process-tree-by-hash")
    client = get_client()

    # Execute command
    command_response = gz_endpoint_get_process_tree_by_hash_command(
        client=client,
        args={
            "id": "ENDPOINT_ID",
            "process_hash": "PROCESS_HASH",
        },
    )

    # Assert command response
    assert_command_mocked_data(
        "gz-endpoint-get-process-tree-by-hash",
        command_response,
        polling_func=gz_poll_live_search_status_command,
        client=client,
    )


@patch("GravityZone.demisto")
def test_gz_endpoint_list_by_running_process_hash_command(mock_demisto, requests_mock):
    """
    Given
            All relevant arguments for the command that is executed
    When
            Calling gz-endpoint-list-by-running-process-hash command
    Then
            Make sure the outputs, outputs_prefix and outputs_key_field values are as expected.
    """

    # Prepare
    from GravityZone import (
        gz_endpoint_list_by_running_process_hash_command,
        gz_poll_live_search_status_command,
    )

    mock_demisto.command.return_value = "gz-endpoint-list-by-running-process-hash"
    mock_demisto.params.return_value = {}
    load_api_mocked_data(requests_mock, "gz-endpoint-list-by-running-process-hash")
    client = get_client()

    # Execute command
    command_response = gz_endpoint_list_by_running_process_hash_command(
        client=client,
        args={
            "process_hash": "PROCESS_HASH",
        },
    )

    # Assert command response
    assert_command_mocked_data(
        "gz-endpoint-list-by-running-process-hash",
        command_response,
        polling_func=gz_poll_live_search_status_command,
        client=client,
    )


@patch("GravityZone.demisto")
def test_gz_poll_task_status_command(mock_demisto, requests_mock):
    """
    Given
            All relevant arguments for the command that is executed
    When
            Calling gz-poll-task-status command
    Then
            Make sure the outputs, outputs_prefix and outputs_key_field values are as expected.
    """

    # Prepare
    from GravityZone import gz_poll_task_status_command

    mock_demisto.command.return_value = "gz-poll-task-status"
    mock_demisto.params.return_value = {}
    load_api_mocked_data(requests_mock, "gz-poll-task-status")
    client = get_client()

    # Execute command
    result = gz_poll_task_status_command(
        client=client,
        args={
            "task_id": "6941944218d1fb8aea04019b",
            "metadata": '{"targetId":"ENDPOINT_ID","processId":10252}',
        },
    )

    # Assert command response
    assert_command_mocked_data("gz-poll-task-status", result)


@patch("GravityZone.demisto")
def test_gz_poll_live_search_status_command(mock_demisto, requests_mock):
    """
    Given
            All relevant arguments for the command that is executed
    When
            Calling gz-poll-live-search-status command
    Then
            Make sure the outputs, outputs_prefix and outputs_key_field values are as expected.
    """

    # Prepare
    from GravityZone import gz_poll_live_search_status_command

    mock_demisto.command.return_value = "gz-poll-live-search-status"
    mock_demisto.params.return_value = {}
    load_api_mocked_data(requests_mock, "gz-poll-live-search-status")
    client = get_client()

    # Execute command
    result = gz_poll_live_search_status_command(
        client=client,
        args={
            "task_id": "694266e8b349a87d10079634",
            "search_type": "fake",
            "metadata": '{"process_hash": "process_hash"}',
        },
    )

    # Assert command response
    assert_command_mocked_data("gz-poll-live-search-status", result)


@patch("GravityZone.demisto")
def test_gz_poll_investigation_activity_status_command(mock_demisto, mocker, requests_mock):
    """
    Given
            All relevant arguments for the command that is executed
    When
            Calling gz-poll-investigation-activity-status command
    Then
            Make sure the outputs, outputs_prefix and outputs_key_field values are as expected.
    """

    # Prepare
    from GravityZone import gz_poll_investigation_activity_status_command

    mock_demisto.command.return_value = "gz-poll-investigation-activity-status"
    mock_demisto.params.return_value = {}
    mocker.patch(
        "GravityZone.fileResult",
        return_value={
            "Contents": "",
            "ContentsFormat": "text",
            "Type": 3,
            "File": "downloaded_file.txt",
            "FileID": "abc",
        },
    )
    load_api_mocked_data(requests_mock, "gz-poll-investigation-activity-status")
    client = get_client()

    # Execute command
    result = gz_poll_investigation_activity_status_command(
        client=client,
        args={
            "target_id": "6942a43afe8d4e463ca5c197",
            "activity_id": "6943db6115457d24450196d3",
            "metadata": '{"activityType":2,"command":"ls -1"}',
        },
    )

    # Assert command response
    assert_command_mocked_data("gz-poll-investigation-activity-status", result)


@pytest.mark.parametrize(
    "args, api_mock_name",
    [
        ({}, "gz-incident-list"),
        ({"endpoint_id": "ENDPOINT_ID"}, "gz-incident-list-2"),
    ],
)
@patch("GravityZone.demisto")
def test_gz_incident_list_command(mock_demisto, requests_mock, args, api_mock_name):
    """
    Given
            All relevant arguments for the command that is executed
    When
            Calling gz-incident-list command
    Then
            Make sure the outputs, outputs_prefix and outputs_key_field values are as expected.
    """

    # Prepare
    from GravityZone import gz_incident_list_command

    mock_demisto.command.return_value = "gz-incident-list"
    mock_demisto.params.return_value = {}
    load_api_mocked_data(requests_mock, api_mock_name)
    client = get_client()

    # Execute command
    command_response = gz_incident_list_command(client=client, args=args)

    # Assert command response
    assert_command_mocked_data(api_mock_name, command_response)


@pytest.mark.parametrize(
    "incident_id",
    [
        ("INCIDENT_ID_1"),
    ],
)
@patch("GravityZone.demisto")
def test_gz_incident_get_edr_command(mock_demisto, requests_mock, incident_id):
    """
    Given
            All relevant arguments for the command that is executed
    When
            Calling gz-incident-get command
    Then
            Make sure the outputs, outputs_prefix and outputs_key_field values are as expected.
    """

    # Prepare
    from GravityZone import gz_incident_get_command

    mock_demisto.command.return_value = "gz-incident-get"
    mock_demisto.params.return_value = {}
    load_api_mocked_data(requests_mock, "gz-incident-get-edr")
    client = get_client()

    # Execute command
    command_response = gz_incident_get_command(client=client, args={"id": incident_id})

    # Assert command response
    assert_command_mocked_data("gz-incident-get-edr", command_response)


@pytest.mark.parametrize(
    "incident_id",
    [
        ("INCIDENT_ID_1"),
    ],
)
@patch("GravityZone.demisto")
def test_gz_incident_get_xdr_command(mock_demisto, requests_mock, incident_id):
    """
    Given
            All relevant arguments for the command that is executed
    When
            Calling gz-incident-get command
    Then
            Make sure the outputs, outputs_prefix and outputs_key_field values are as expected.
    """

    # Prepare
    from GravityZone import gz_incident_get_command

    mock_demisto.command.return_value = "gz-incident-get"
    mock_demisto.params.return_value = {}
    load_api_mocked_data(requests_mock, "gz-incident-get-xdr")
    client = get_client()

    # Execute command
    command_response = gz_incident_get_command(client=client, args={"id": incident_id})

    # Assert command response
    assert_command_mocked_data("gz-incident-get-xdr", command_response)


@pytest.mark.parametrize(
    "incident_id, note",
    [
        ("INCIDENT_ID_1", "This is a test note from unit test."),
    ],
)
@patch("GravityZone.demisto")
def test_gz_incident_add_note_command(mock_demisto, requests_mock, incident_id, note):
    """
    Given
            All relevant arguments for the command that is executed
    When
            Calling gz-incident-add-note command
    Then
            Make sure the outputs, outputs_prefix and outputs_key_field values are as expected.
    """

    # Prepare
    from GravityZone import gz_incident_add_note_command

    mock_demisto.command.return_value = "gz-incident-add-note"
    mock_demisto.params.return_value = {}
    load_api_mocked_data(requests_mock, "gz-incident-add-note")
    client = get_client()

    # Execute command
    command_response = gz_incident_add_note_command(
        client=client,
        args={
            "id": incident_id,
            "note": note,
        },
    )

    # Assert command response
    assert_command_mocked_data("gz-incident-add-note", command_response)


@pytest.mark.parametrize(
    "incident_id, status_param, api_mock_name",
    [
        ("INCIDENT_ID_1", "PENDING", "gz-incident-change-status-pending"),
        ("INCIDENT_ID_1", "ACTIVE", "gz-incident-change-status-active"),
        ("INCIDENT_ID_1", "DONE", "gz-incident-change-status-done"),
    ],
)
@patch("GravityZone.demisto")
def test_gz_incident_change_status_command(mock_demisto, requests_mock, incident_id, status_param, api_mock_name):
    """
    Given
            All relevant arguments for the command that is executed
    When
            Calling gz-incident-change-status command
    Then
            Make sure the outputs, outputs_prefix and outputs_key_field values are as expected.
    """

    # Prepare
    from GravityZone import gz_incident_change_status_command

    mock_demisto.command.return_value = "gz-incident-change-status"
    mock_demisto.params.return_value = {}
    load_api_mocked_data(requests_mock, api_mock_name)
    client = get_client()

    # Execute command
    command_response = gz_incident_change_status_command(
        client=client,
        args={
            "id": incident_id,
            "status": status_param,
        },
    )

    # Assert command response
    assert_command_mocked_data(api_mock_name, command_response)


@pytest.mark.parametrize(
    "incident_id, status_param, note",
    [
        ("INCIDENT_ID_1", 0, "incident pending"),
        ("INCIDENT_ID_1", 1, "incident active"),
        ("INCIDENT_ID_1", 2, "incident done"),
    ],
)
@patch("GravityZone.demisto")
def test_update_remote_system_command(mock_demisto, requests_mock, incident_id, status_param, note):
    """
    Given
            All relevant arguments for the command that is executed
    When
            Calling update-remote-system command
    Then
            Make sure the the return value is as expected.
    """

    # Prepare
    from GravityZone import update_remote_system_command

    mock_demisto.command.return_value = "update-remote-system"
    mock_demisto.params.return_value = {}
    load_api_mocked_data(requests_mock, "update-remote-system")
    client = get_client()

    # Execute command
    command_response = update_remote_system_command(
        client=client,
        args={
            "remoteId": incident_id,
            "data": {},
            "entries": [],
            "status": status_param,
            "incidentChanged": True,
            "delta": {"closeNotes": note},
        },
    )

    # Assert command response
    assert command_response == incident_id
    assert requests_mock.called


@pytest.mark.parametrize(
    "endpoint_id, mock_data_suffix",
    [("6942a43afe8d4e463ca5c197", ""), ("ENDPOINT_ID", "-failed")],
)
@patch("GravityZone.demisto")
def test_gz_endpoint_download_file_command(mock_demisto, mocker, requests_mock, endpoint_id, mock_data_suffix):
    """
    Given
            All relevant arguments for the command that is executed
    When
            Calling gz-endpoint-download-file command
    Then
            Make sure the outputs, outputs_prefix and outputs_key_field values are as expected.
    """

    # Prepare
    from GravityZone import (
        gz_endpoint_download_file_command,
        gz_poll_investigation_activity_status_command,
    )

    mock_demisto.command.return_value = "gz-endpoint-download-file"
    mock_demisto.params.return_value = {}
    mocker.patch(
        "GravityZone.fileResult",
        return_value={
            "Contents": "",
            "ContentsFormat": "text",
            "Type": 3,
            "File": "downloaded_file.txt",
            "FileID": "abc",
        },
    )
    load_api_mocked_data(requests_mock, f"gz-endpoint-download-file{mock_data_suffix}")
    client = get_client()

    # Execute command
    command_response = gz_endpoint_download_file_command(
        client=client,
        args={
            "id": endpoint_id,
            "remote_file": "/home/ENDPOINT_NAME/test.txt",
            "output_file": "downloaded_file.txt",
        },
    )

    # Assert command response
    assert_command_mocked_data(
        f"gz-endpoint-download-file{mock_data_suffix}",
        command_response,
        polling_func=gz_poll_investigation_activity_status_command,
        client=client,
    )


@pytest.mark.parametrize(
    "endpoint_id",
    [("6942a43afe8d4e463ca5c197")],
)
@patch("GravityZone.demisto")
def test_gz_endpoint_download_investigation_package_command(mock_demisto, mocker, requests_mock, endpoint_id):
    """
    Given
            All relevant arguments for the command that is executed
    When
            Calling gz-endpoint-download-investigation-package command
    Then
            Make sure the outputs, outputs_prefix and outputs_key_field values are as expected.
    """

    # Prepare
    from GravityZone import (
        gz_endpoint_download_investigation_package_command,
        gz_poll_investigation_activity_status_command,
    )

    mock_demisto.command.return_value = "gz-endpoint-download-investigation-package"
    mock_demisto.params.return_value = {}
    mocker.patch(
        "GravityZone.fileResult",
        return_value={
            "Contents": "",
            "ContentsFormat": "text",
            "Type": 3,
            "File": "investigation_package.zip",
            "FileID": "abc",
        },
    )
    load_api_mocked_data(requests_mock, "gz-endpoint-download-investigation-package")
    client = get_client()

    # Execute command
    command_response = gz_endpoint_download_investigation_package_command(
        client=client,
        args={"id": endpoint_id, "output_file": "investigation_package.zip"},
    )

    # Assert command response
    assert_command_mocked_data(
        "gz-endpoint-download-investigation-package",
        command_response,
        polling_func=gz_poll_investigation_activity_status_command,
        client=client,
    )


@pytest.mark.parametrize(
    "endpoint_id",
    [("6942a43afe8d4e463ca5c197")],
)
@patch("GravityZone.demisto")
@patch("GravityZone.FileManagement.get_file")
def test_gz_endpoint_upload_file_command(mock_get_file, mock_demisto, requests_mock, endpoint_id):
    """
    Given
            All relevant arguments for the command that is executed
    When
            Calling gz-endpoint-upload-file command
    Then
            Make sure the outputs, outputs_prefix and outputs_key_field values are as expected.
    """

    # Prepare
    from GravityZone import (
        gz_endpoint_upload_file_command,
        gz_poll_task_status_command,
    )

    mock_demisto.command.return_value = "gz-endpoint-upload-file"
    mock_get_file.return_value = ("file_to_upload.txt", b"file_content")
    mock_demisto.params.return_value = {}
    load_api_mocked_data(requests_mock, "gz-endpoint-upload-file")
    client = get_client()

    # Execute command
    command_response = gz_endpoint_upload_file_command(
        client=client,
        args={"id": endpoint_id, "remote_location": "/home/ENDPOINT_NAME/", "entry_id": "abc"},
    )

    # Assert command response
    assert_command_mocked_data(
        "gz-endpoint-upload-file", command_response, polling_func=gz_poll_task_status_command, client=client
    )