Azure Resource Graph

import json
import os

import pytest
from AzureResourceGraph import AzureResourceGraphClient, list_operations_command, pagination, query_resources_command

client = AzureResourceGraphClient(
    base_url="url",
    tenant_id="tenant",
    auth_id="auth_id",
    enc_key="enc_key",
    app_name="APP_NAME",
    verify="verify",
    proxy="proxy",
    self_deployed="self_deployed",
    ok_codes=(1, 2),
    server="server",
    certificate_thumbprint="",
    private_key="",
)


""" HELPER FUNCTIONS """


def util_load_json(path):
    full_path = os.path.join(os.path.dirname(__file__), path)
    with open(file=full_path, encoding="utf-8") as f:
        return json.loads(f.read())


def get_azure_access_token_mock() -> dict:
    """
    Mock Azure access token object.

    Returns:
        dict: Azure access token mock.
    """
    return {
        "access_token": "my-access-token",
        "expires_in": 3595,
        "refresh_token": "my-refresh-token",
    }


"""TESTS"""


@pytest.mark.parametrize(
    "function_name, data_file_name, args, additional_args",
    [
        ("list_operations", "test_list_operations_output", None, None),
        (
            "query_resources",
            "test_resources_query_output",
            {"query": "dummy"},
            {"paging_options": "dummy", "subscriptions": "dummy", "management_groups": "dummy"},
        ),
    ],
)
def test_client_api_calls(mocker, function_name, data_file_name, args, additional_args):
    data = util_load_json(f"test_data/{data_file_name}.json")
    mock_http_request = mocker.patch.object(client.ms_client, "http_request")
    mock_http_request.return_value = data
    if args and not additional_args:
        result = getattr(client, function_name)(args)
    elif args and additional_args:
        result = getattr(client, function_name)(args, **additional_args)
    else:
        result = getattr(client, function_name)()
    assert result == data


@pytest.mark.parametrize(
    "limit, page_size, page_number, expected_number_of_operations, expected_first_name",
    [
        (1, None, None, 1, "Microsoft.ResourceGraph/operations/read"),
        (3, 2, 2, 2, "Microsoft.ResourceGraph/resourcesHistory/read"),
    ],
)
def test_successful_list_operations_command(
    mocker, limit, page_size, page_number, expected_number_of_operations, expected_first_name
):
    operations_data = util_load_json("test_data/test_list_operations_output.json")
    mocker.patch.object(client, "list_operations", return_value=operations_data)
    args: dict = {"limit": limit, "page": page_number, "page_size": page_size}
    command_results = list_operations_command(client, args)
    assert command_results.outputs[0]["Name"] == expected_first_name
    assert len(command_results.outputs) == expected_number_of_operations


@pytest.mark.parametrize(
    "limit, page_size, page_number, expected_error_message",
    [
        (None, None, 2, 'Please enter a value for "page_size" when using "page".'),
        (None, 2, None, 'Please enter a value for "page" when using "page_size".'),
    ],
)
def test_failure_list_operations_command(mocker, limit, page_size, page_number, expected_error_message):
    try:
        operations_data = util_load_json("test_data/test_list_operations_output.json")
        mocker.patch.object(client, "list_operations", return_value=operations_data)
        args: dict = {"limit": limit, "page": page_number, "page_size": page_size}
        list_operations_command(client, args)
    except Exception as e:
        assert e.message == expected_error_message


@pytest.mark.parametrize(
    "query, data_file_name, expected_output",
    [("Resources | project id, name, type, location, tags | limit 3", "test_resources_query_output", "test-ssh-nsg")],
)
def test_query_resources_command_output(mocker, query, data_file_name, expected_output):
    query_data = util_load_json(f"test_data/{data_file_name}.json")
    args: dict = {"query": query, "limit": 1}
    mocker.patch.object(client, "query_resources", return_value=query_data)
    command_results = query_resources_command(client, args=args)
    assert command_results.outputs[0]["name"] == expected_output


@pytest.mark.parametrize(
    "query, data_file_name, extra_args, expected_length, expected_name",
    [
        (
            "Resources | project id, name, type, location, tags",
            "test_resources_query_paging_output",
            {"page": 1, "page_size": 3},
            3,
            "test-ssh-nsg-2",
        )
    ],
)
def test_query_resources_command_paging(mocker, query, data_file_name, extra_args, expected_length, expected_name):
    query_data = util_load_json(f"test_data/{data_file_name}.json")
    args: dict = {"query": query}
    args.update(extra_args)
    mocker.patch.object(client, "query_resources", return_value=query_data)
    command_results = query_resources_command(client, args=args)
    assert len(command_results.outputs) == expected_length
    assert command_results.outputs[0]["name"] == expected_name


def test_pagination_helper():
    fake_response = list(range(1, 101))
    results = pagination(fake_response, page_size=3, page_number=5)
    assert len(results) == 3


def test_test_module_command(mocker) -> None:
    """
    Scenario: run test module when managed identities client id provided.
    Given:
     - User has provided managed identities client oid.
    When:
     - test-module called.
    Then:
     - Ensure the output are as expected
    """
    import AzureResourceGraph
    import demistomock as demisto
    from AzureResourceGraph import MicrosoftClient, main

    params = {
        "auth_id": "test_client_id",
        "use_managed_identities": "True",
        "cred_token": {"password": "test"},
        "client_credentials": True,
        "host": "https://management.azure.com",
        "tenant_id": "1abc234d-12a3-12a3-12a3-1234abcde123",
        "cred_auth_id": {"password": "test_api"},
        "unsecure": False,
        "proxy": False,
        "private_key": "test-key",
        "self_deployed": True,
        "enc_key": "test",
    }

    # operations_data = util_load_json('./test_data/test_list_operations_output.json')
    mocker.patch.object(demisto, "params", return_value=params)
    mocker.patch.object(demisto, "args", return_value={})
    mocker.patch.object(demisto, "command", return_value="test-module")
    mocker.patch.object(MicrosoftClient, "http_request", return_value=get_azure_access_token_mock())
    mocker.patch.object(AzureResourceGraph, "return_results")

    main()

    assert "ok" in AzureResourceGraph.return_results.call_args[0][0]


class TestGovAccountToggle:
    """Tests for the Gov Account checkbox feature."""

    GOV_SCOPE = "https://management.usgovcloudapi.net/.default"
    GOV_SERVER = "https://management.usgovcloudapi.net"
    GOV_BASE_URL = "https://management.usgovcloudapi.net/providers/Microsoft.ResourceGraph"
    NORMAL_SCOPE = "https://management.azure.com/.default"
    NORMAL_SERVER = "https://management.azure.com"
    NORMAL_BASE_URL = "https://management.azure.com/providers/Microsoft.ResourceGraph"

    def test_client_with_gov_account_enabled(self, mocker):
        """
        Given:
            - is_gov=True is passed to AzureResourceGraphClient.
        When:
            - The client is initialized.
        Then:
            - MicrosoftClient is created with the gov scope, AZURE_US_GCC_HIGH_CLOUD azure_cloud,
              and the gov base_url.
            - The client server is set to the gov management URL.
        """
        from AzureResourceGraph import AZURE_US_GCC_HIGH_CLOUD, AzureResourceGraphClient, MicrosoftClient

        mock_ms_client = mocker.patch.object(MicrosoftClient, "__init__", return_value=None)

        gov_client = AzureResourceGraphClient(
            base_url=self.NORMAL_BASE_URL,
            tenant_id="tenant",
            auth_id="auth_id",
            enc_key="enc_key",
            app_name="APP_NAME",
            verify=False,
            proxy=False,
            self_deployed=True,
            ok_codes=(200,),
            server=self.NORMAL_SERVER,
            certificate_thumbprint="",
            private_key="",
            is_gov=True,
        )

        call_kwargs = mock_ms_client.call_args[1]
        assert call_kwargs["scope"] == self.GOV_SCOPE
        assert call_kwargs["azure_cloud"] == AZURE_US_GCC_HIGH_CLOUD
        assert call_kwargs["base_url"] == self.GOV_BASE_URL
        assert gov_client.server == self.GOV_SERVER

    def test_client_with_gov_account_disabled(self, mocker):
        """
        Given:
            - is_gov=False (default) is passed to AzureResourceGraphClient.
        When:
            - The client is initialized.
        Then:
            - MicrosoftClient is created with the standard scope, no azure_cloud parameter,
              and the standard base_url.
            - The client server is set to the standard management URL.
        """
        from AzureResourceGraph import AzureResourceGraphClient, MicrosoftClient

        mock_ms_client = mocker.patch.object(MicrosoftClient, "__init__", return_value=None)

        normal_client = AzureResourceGraphClient(
            base_url=self.NORMAL_BASE_URL,
            tenant_id="tenant",
            auth_id="auth_id",
            enc_key="enc_key",
            app_name="APP_NAME",
            verify=False,
            proxy=False,
            self_deployed=True,
            ok_codes=(200,),
            server=self.NORMAL_SERVER,
            certificate_thumbprint="",
            private_key="",
            is_gov=False,
        )

        call_kwargs = mock_ms_client.call_args[1]
        assert call_kwargs["scope"] == self.NORMAL_SCOPE
        assert "azure_cloud" not in call_kwargs
        assert call_kwargs["base_url"] == self.NORMAL_BASE_URL
        assert normal_client.server == self.NORMAL_SERVER

    def test_client_default_is_not_gov(self, mocker):
        """
        Given:
            - is_gov is not passed to AzureResourceGraphClient (uses default).
        When:
            - The client is initialized.
        Then:
            - MicrosoftClient is created with the standard scope (same as is_gov=False).
        """
        from AzureResourceGraph import AzureResourceGraphClient, MicrosoftClient

        mock_ms_client = mocker.patch.object(MicrosoftClient, "__init__", return_value=None)

        default_client = AzureResourceGraphClient(
            base_url=self.NORMAL_BASE_URL,
            tenant_id="tenant",
            auth_id="auth_id",
            enc_key="enc_key",
            app_name="APP_NAME",
            verify=False,
            proxy=False,
            self_deployed=True,
            ok_codes=(200,),
            server=self.NORMAL_SERVER,
            certificate_thumbprint="",
            private_key="",
        )

        call_kwargs = mock_ms_client.call_args[1]
        assert call_kwargs["scope"] == self.NORMAL_SCOPE
        assert "azure_cloud" not in call_kwargs
        assert call_kwargs["base_url"] == self.NORMAL_BASE_URL
        assert default_client.server == self.NORMAL_SERVER

    def test_main_passes_gov_account_param(self, mocker):
        """
        Given:
            - The gov_account parameter is set to True in demisto.params().
        When:
            - main() is called with test-module command.
        Then:
            - AzureResourceGraphClient is initialized with is_gov=True and the gov URLs are used.
        """
        import AzureResourceGraph
        import demistomock as demisto
        from AzureResourceGraph import AZURE_US_GCC_HIGH_CLOUD, MicrosoftClient, main

        params = {
            "auth_id": "test_client_id",
            "cred_token": {"password": "test"},
            "host": "https://management.azure.com",
            "tenant_id": "1abc234d-12a3-12a3-12a3-1234abcde123",
            "cred_auth_id": {"password": "test_api"},
            "unsecure": False,
            "proxy": False,
            "private_key": "test-key",
            "self_deployed": True,
            "enc_key": "test",
            "gov_account": True,
        }

        mocker.patch.object(demisto, "params", return_value=params)
        mocker.patch.object(demisto, "args", return_value={})
        mocker.patch.object(demisto, "command", return_value="test-module")
        mock_ms_init = mocker.patch.object(MicrosoftClient, "__init__", return_value=None)
        mocker.patch.object(MicrosoftClient, "http_request", return_value=get_azure_access_token_mock())
        mocker.patch.object(AzureResourceGraph, "return_results")

        main()

        ms_call_kwargs = mock_ms_init.call_args[1]
        assert ms_call_kwargs["scope"] == self.GOV_SCOPE
        assert ms_call_kwargs["azure_cloud"] == AZURE_US_GCC_HIGH_CLOUD
        assert ms_call_kwargs["base_url"] == self.GOV_BASE_URL
        assert "ok" in AzureResourceGraph.return_results.call_args[0][0]

    def test_main_without_gov_account_param(self, mocker):
        """
        Given:
            - The gov_account parameter is not set (defaults to False) in demisto.params().
        When:
            - main() is called with test-module command.
        Then:
            - AzureResourceGraphClient is initialized with is_gov=False and the standard URLs are used.
        """
        import AzureResourceGraph
        import demistomock as demisto
        from AzureResourceGraph import MicrosoftClient, main

        params = {
            "auth_id": "test_client_id",
            "cred_token": {"password": "test"},
            "host": "https://management.azure.com",
            "tenant_id": "1abc234d-12a3-12a3-12a3-1234abcde123",
            "cred_auth_id": {"password": "test_api"},
            "unsecure": False,
            "proxy": False,
            "private_key": "test-key",
            "self_deployed": True,
            "enc_key": "test",
        }

        mocker.patch.object(demisto, "params", return_value=params)
        mocker.patch.object(demisto, "args", return_value={})
        mocker.patch.object(demisto, "command", return_value="test-module")
        mock_ms_init = mocker.patch.object(MicrosoftClient, "__init__", return_value=None)
        mocker.patch.object(MicrosoftClient, "http_request", return_value=get_azure_access_token_mock())
        mocker.patch.object(AzureResourceGraph, "return_results")

        main()

        ms_call_kwargs = mock_ms_init.call_args[1]
        assert ms_call_kwargs["scope"] == self.NORMAL_SCOPE
        assert "azure_cloud" not in ms_call_kwargs
        assert ms_call_kwargs["base_url"] == self.NORMAL_BASE_URL
        assert "ok" in AzureResourceGraph.return_results.call_args[0][0]