RemoveFileWrapper

import CommonServerPython
import pytest
import RemoveFileWrapper

test_data = [
    CommonServerPython.CommandRunner.Result(
        command="endpoint",
        args={"id": "device1,device2", "using-brand": "CrowdstrikeFalcon"},
        brand="CrowdstrikeFalcon",
        instance="CrowdstrikeFalcon_instance_1",
        result={
            "errors": [],
            "resources": [
                {"device_id": "device1", "platform_name": "Windows"},
                {"device_id": "device2", "platform_name": "Linux"},
            ],
        },
    )
]


def test_get_crowdstrike_os_to_id(mocker):
    """
    Given:
        A list of devices ids which are on crowdstrike
    When:
        Getting the os for each device
    Then:
        Return a mapping between an OS and a set of device-ids what are on the OS.

    """
    from RemoveFileWrapper import get_crowdstrike_os_to_id

    mocker.patch.object(RemoveFileWrapper.CommandRunner, "execute_commands", return_value=(test_data, []))
    os_to_id = get_crowdstrike_os_to_id(["device2", "device1"])
    assert os_to_id == {"Windows": {"device1"}, "Linux": {"device2"}}


def test_create_command_executors(mocker):
    """
    Given:
        the action to perform (allow or block)
    When:
        Calling `create_command_wrappers` to get all the command wrappers for the script.
    Then:
        Ensure the right commands wrappers are being returned.
    """
    from RemoveFileWrapper import create_commands, demisto

    device_ids = ["device1", "device2", "device3", "device4"]
    file_path = "filepath"
    file_hash = "filehash"
    mock_incident_id = 123
    mocker.patch.object(demisto, "incident", return_value={"id": mock_incident_id})
    mocker.patch.object(
        RemoveFileWrapper,
        "get_crowdstrike_os_to_id",
        return_value={"Windows": {"device1", "device2"}, "Linux": {"device3"}, "Mac": {"device4"}},
    )
    commands = create_commands(device_ids, file_path, file_hash)
    assert len(commands) == 3
    for command in commands:
        command_names = set(command.commands)
        if "xdr-run-script-delete-file" in command_names:
            assert command_names == {"xdr-run-script-delete-file"}
            assert command.args_lst == [{"endpoint_ids": ",".join(device_ids), "file_path": file_path}]
        if "cs-falcon-rtr-remove-file" in command_names:
            assert len(command.commands) == 3
            assert len(command.args_lst) == 3
            assert command_names == {"cs-falcon-rtr-remove-file"}
            assert set(command.args_lst[0].get("host_ids", "").split(",")) == {"device1", "device2"}
            assert command.args_lst[1:] == [
                {"host_ids": "device3", "file_path": file_path, "os": "Linux"},
                {
                    "host_ids": "device4",
                    "file_path": file_path,
                    "os": "Mac",
                },
            ]
        if "microsoft-ato-stop-and-quarantine-file" in command_names:
            assert command.commands == ["microsoft-ato-stop-and-quarantine-file"]
            assert command.args_lst == [
                {
                    "machine_id": "device1,device2,device3,device4",
                    "file_hash": "filehash",
                    "comment": f"Action was taken by Cortex XSOAR - incident #{mock_incident_id}",
                }
            ]


@pytest.mark.parametrize("approved", ("yes", "no", False))
def test_approve_action(mocker, approved: bool):
    """
    Given   a value for the `approve_action` argument
    When    calling main()
    Then    make sure an error is raised only when it should (when approved = 'no')
    """
    from RemoveFileWrapper import demisto, main

    mocker.patch.object(RemoveFileWrapper, "run_remove_file", return_value=None)
    mocker.patch.object(demisto, "error", return_value=None)
    return_error = mocker.patch.object(RemoveFileWrapper, "return_error", return_value=None)
    mocker.patch.object(
        demisto, "args", return_value={"approve_action": approved, "device_ids": ["device1"], "file_path": "file_path"}
    )
    main()
    assert return_error.call_count == int(approved != "yes")


@pytest.mark.parametrize("file_path", ("", "path"))
@pytest.mark.parametrize("file_hash", ("", "hash"))
def test_created_command_count(mocker, file_path: str, file_hash: str):
    """
    Given   the file_path and file_hash arguments
    When    calling create_commands
    Then    make sure the number of command created is correct
    """
    from RemoveFileWrapper import create_commands, demisto

    mocker.patch.object(RemoveFileWrapper, "run_remove_file", return_value=None)
    mocker.patch.object(demisto, "error", return_value=None)
    mocker.patch.object(
        demisto, "args", return_value={"approve_action": "yes", "device_ids": ["device1"], "file_path": "file_path"}
    )

    path_based_commands = {"xdr", "falcon"}
    hash_based_commands = {"microsoft_atp"}
    expected_commands = (path_based_commands if file_path else set()) | (hash_based_commands if file_hash else set())

    commands = create_commands(["id"], file_path, file_hash)
    assert len(commands) == len(expected_commands)