feat(azure): add test_connection method (#4615)

Co-authored-by: Pepe Fagoaga <pepe@prowler.com>

prowler/providers/azure/azure_provider.py

@@ -1,8 +1,10 @@
 import asyncio
 import sys
+from argparse import ArgumentTypeError
 from os import getenv
 
 import requests
+from azure.core.exceptions import HttpResponseError
 from azure.identity import DefaultAzureCredential, InteractiveBrowserCredential
 from azure.mgmt.subscription import SubscriptionClient
 from colorama import Fore, Style
@@ -14,6 +16,7 @@ from prowler.config.config import (
 )
 from prowler.lib.logger import logger
 from prowler.lib.utils.utils import print_boxes
+from prowler.providers.azure.lib.arguments.arguments import validate_azure_region
 from prowler.providers.azure.lib.mutelist.mutelist import AzureMutelist
 from prowler.providers.azure.lib.regions.regions import get_regions_config
 from prowler.providers.azure.models import (
@@ -21,11 +24,47 @@ from prowler.providers.azure.models import (
     AzureOutputOptions,
     AzureRegionConfig,
 )
-from prowler.providers.common.models import Audit_Metadata
+from prowler.providers.common.models import Audit_Metadata, Connection
 from prowler.providers.common.provider import Provider
 
 
 class AzureProvider(Provider):
+    """
+    Represents an Azure provider.
+
+    This class provides functionality to interact with the Azure cloud provider.
+    It handles authentication, region configuration, and provides access to various properties and methods
+    related to the Azure provider.
+
+    Attributes:
+        _type (str): The type of the provider, which is set to "azure".
+        _session (DefaultAzureCredential): The session object associated with the Azure provider.
+        _identity (AzureIdentityInfo): The identity information for the Azure provider.
+        _audit_config (dict): The audit configuration for the Azure provider.
+        _region_config (AzureRegionConfig): The region configuration for the Azure provider.
+        _locations (dict): A dictionary containing the available locations for the Azure provider.
+        _output_options (AzureOutputOptions): The output options for the Azure provider.
+        _mutelist (AzureMutelist): The mutelist object associated with the Azure provider.
+        audit_metadata (Audit_Metadata): The audit metadata for the Azure provider.
+
+    Methods:
+        __init__(self, arguments): Initializes the AzureProvider object.
+        identity(self): Returns the identity of the Azure provider.
+        type(self): Returns the type of the Azure provider.
+        session(self): Returns the session object associated with the Azure provider.
+        region_config(self): Returns the region configuration for the Azure provider.
+        locations(self): Returns a list of available locations for the Azure provider.
+        audit_config(self): Returns the audit configuration for the Azure provider.
+        fixer_config(self): Returns the fixer configuration.
+        output_options(self, options: tuple): Sets the output options for the Azure provider.
+        mutelist(self) -> AzureMutelist: Returns the mutelist object associated with the Azure provider.
+        get_output_mapping(self): Returns a dictionary that maps output keys to their corresponding values.
+        validate_arguments(cls, az_cli_auth, sp_env_auth, browser_auth, managed_identity_auth, tenant_id): Validates the authentication arguments for the Azure provider.
+        setup_region_config(cls, region): Sets up the region configuration for the Azure provider.
+        print_credentials(self): Prints the Azure credentials information.
+        setup_session(cls, az_cli_auth, sp_env_auth, browser_auth, managed_identity_auth, tenant_id, region_config): Set up the Azure session with the specified authentication method.
+    """
+
     _type: str = "azure"
     _session: DefaultAzureCredential
     _identity: AzureIdentityInfo
@@ -45,23 +84,34 @@ class AzureProvider(Provider):
         az_cli_auth = arguments.az_cli_auth
         sp_env_auth = arguments.sp_env_auth
         browser_auth = arguments.browser_auth
-        managed_entity_auth = arguments.managed_identity_auth
+        managed_identity_auth = arguments.managed_identity_auth
         tenant_id = arguments.tenant_id
 
+        # Validate the authentication arguments
+        self.validate_arguments(
+            az_cli_auth, sp_env_auth, browser_auth, managed_identity_auth, tenant_id
+        )
+
         logger.info("Checking if region is different than default one")
         region = arguments.azure_region
-        self.validate_arguments(
-            az_cli_auth, sp_env_auth, browser_auth, managed_entity_auth, tenant_id
-        )
         self._region_config = self.setup_region_config(region)
+
+        # Set up the Azure session
         self._session = self.setup_session(
-            az_cli_auth, sp_env_auth, browser_auth, managed_entity_auth, tenant_id
+            az_cli_auth,
+            sp_env_auth,
+            browser_auth,
+            managed_identity_auth,
+            tenant_id,
+            self._region_config,
         )
+
+        # Set up the identity
         self._identity = self.setup_identity(
             az_cli_auth,
             sp_env_auth,
             browser_auth,
-            managed_entity_auth,
+            managed_identity_auth,
             subscription_ids,
         )
 
@@ -79,38 +129,56 @@ class AzureProvider(Provider):
 
     @property
     def identity(self):
+        """Returns the identity of the Azure provider."""
         return self._identity
 
     @property
     def type(self):
+        """Returns the type of the Azure provider."""
         return self._type
 
     @property
     def session(self):
+        """Returns the session object associated with the Azure provider."""
         return self._session
 
     @property
     def region_config(self):
+        """Returns the region configuration for the Azure provider."""
         return self._region_config
 
     @property
     def locations(self):
+        """Returns a list of available locations for the Azure provider."""
         return self._locations
 
     @property
     def audit_config(self):
+        """Returns the audit configuration for the Azure provider."""
         return self._audit_config
 
     @property
     def fixer_config(self):
+        """Returns the fixer configuration."""
         return self._fixer_config
 
     @property
     def output_options(self):
+        """Returns the output options for the Azure provider."""
         return self._output_options
 
     @output_options.setter
     def output_options(self, options: tuple):
+        """Set output options for the Azure provider.
+
+        Sets the output options for the Azure provider using the provided arguments and bulk checks metadata.
+
+        Args:
+            options (tuple): A tuple containing the arguments and bulk checks metadata.
+
+        Returns:
+            None
+        """
         arguments, bulk_checks_metadata = options
         self._output_options = AzureOutputOptions(
             arguments, bulk_checks_metadata, self._identity
@@ -118,9 +186,7 @@ class AzureProvider(Provider):
 
     @property
     def mutelist(self) -> AzureMutelist:
-        """
-        mutelist method returns the provider's mutelist.
-        """
+        """Mutelist object associated with this Azure provider."""
         return self._mutelist
 
     @mutelist.setter
@@ -136,6 +202,7 @@ class AzureProvider(Provider):
 
     @property
     def get_output_mapping(self):
+        """Dictionary that maps output keys to their corresponding values."""
         return {
             # identity_type: identity_id
             # "auth_method": "identity.profile",
@@ -155,14 +222,34 @@ class AzureProvider(Provider):
 
     # TODO: this should be moved to the argparse, if not we need to enforce it from the Provider
     # previously was using the AzureException
+    @staticmethod
     def validate_arguments(
-        self, az_cli_auth, sp_env_auth, browser_auth, managed_entity_auth, tenant_id
+        az_cli_auth: bool,
+        sp_env_auth: bool,
+        browser_auth: bool,
+        managed_identity_auth: bool,
+        tenant_id: str,
     ):
+        """
+        Validates the authentication arguments for the Azure provider.
+
+        Args:
+            az_cli_auth (bool): Flag indicating whether AZ CLI authentication is enabled.
+            sp_env_auth (bool): Flag indicating whether Service Principal environment authentication is enabled.
+            browser_auth (bool): Flag indicating whether browser authentication is enabled.
+            managed_identity_auth (bool): Flag indicating whether managed identity authentication is enabled.
+            tenant_id (str): The Azure Tenant ID.
+
+        Raises:
+            SystemExit: If none of the authentication methods are set.
+            SystemExit: If browser authentication is enabled but the tenant ID is not provided.
+            SystemExit: If tenant ID is provided but browser authentication is not enabled.
+        """
         if (
             not az_cli_auth
             and not sp_env_auth
             and not browser_auth
-            and not managed_entity_auth
+            and not managed_identity_auth
         ):
             raise SystemExit(
                 "Azure provider requires at least one authentication method set: [--az-cli-auth | --sp-env-auth | --browser-auth | --managed-identity-auth]"
@@ -177,16 +264,51 @@ class AzureProvider(Provider):
                 "Azure Tenant ID (--tenant-id) is required only for browser authentication mode"
             )
 
-    def setup_region_config(self, region):
-        config = get_regions_config(region)
-        return AzureRegionConfig(
-            name=region,
-            authority=config["authority"],
-            base_url=config["base_url"],
-            credential_scopes=config["credential_scopes"],
-        )
+    @staticmethod
+    def setup_region_config(region):
+        """
+        Sets up the region configuration for the Azure provider.
+
+        Args:
+            region (str): The name of the region.
+
+        Returns:
+            AzureRegionConfig: The region configuration object.
+
+        """
+        try:
+            validate_azure_region(region)
+            config = get_regions_config(region)
+
+            return AzureRegionConfig(
+                name=region,
+                authority=config["authority"],
+                base_url=config["base_url"],
+                credential_scopes=config["credential_scopes"],
+            )
+        except ArgumentTypeError as validation_error:
+            logger.error(
+                f"{validation_error.__class__.__name__}[{validation_error.__traceback__.tb_lineno}]: {validation_error}"
+            )
+            raise validation_error
+        except Exception as validation_error:
+            logger.error(
+                f"{validation_error.__class__.__name__}[{validation_error.__traceback__.tb_lineno}]: {validation_error}"
+            )
+            raise validation_error
 
     def print_credentials(self):
+        """Azure credentials information.
+
+        This method prints the Azure Tenant Domain, Azure Tenant ID, Azure Region,
+        Azure Subscriptions, Azure Identity Type, and Azure Identity ID.
+
+        Args:
+            None
+
+        Returns:
+            None
+        """
         printed_subscriptions = []
         for key, value in self._identity.subscriptions.items():
             intermediate = key + ": " + value
@@ -203,20 +325,46 @@ class AzureProvider(Provider):
         print_boxes(report_lines, report_title)
 
     # TODO: setup_session or setup_credentials?
+    # This should be setup_credentials, since it is setting up the credentials for the provider
+    @staticmethod
     def setup_session(
-        self, az_cli_auth, sp_env_auth, browser_auth, managed_entity_auth, tenant_id
+        az_cli_auth: bool,
+        sp_env_auth: bool,
+        browser_auth: bool,
+        managed_identity_auth: bool,
+        tenant_id: str,
+        region_config: AzureRegionConfig,
     ):
+        """Returns the Azure credentials object.
+
+        Set up the Azure session with the specified authentication method.
+
+        Args:
+            az_cli_auth (bool): Flag indicating whether to use Azure CLI authentication.
+            sp_env_auth (bool): Flag indicating whether to use Service Principal authentication with environment variables.
+            browser_auth (bool): Flag indicating whether to use interactive browser authentication.
+            managed_identity_auth (bool): Flag indicating whether to use managed identity authentication.
+            tenant_id (str): The Azure Active Directory tenant ID.
+            region_config (AzureRegionConfig): The region configuration object.
+
+        Returns:
+            credentials: The Azure credentials object.
+
+        Raises:
+            Exception: If failed to retrieve Azure credentials.
+
+        """
         # Browser auth creds cannot be set with DefaultAzureCredentials()
         if not browser_auth:
             if sp_env_auth:
-                self.__check_service_principal_creds_env_vars__()
+                AzureProvider.check_service_principal_creds_env_vars()
             try:
                 # Since the input vars come as True when it is wanted to be used, we need to inverse it since
                 # DefaultAzureCredential sets the auth method excluding the others
                 credentials = DefaultAzureCredential(
                     exclude_environment_credential=not sp_env_auth,
                     exclude_cli_credential=not az_cli_auth,
-                    exclude_managed_identity_credential=not managed_entity_auth,
+                    exclude_managed_identity_credential=not managed_identity_auth,
                     # Azure Auth using Visual Studio is not supported
                     exclude_visual_studio_code_credential=True,
                     # Azure Auth using Shared Token Cache is not supported
@@ -224,7 +372,7 @@ class AzureProvider(Provider):
                     # Azure Auth using PowerShell is not supported
                     exclude_powershell_credential=True,
                     # set Authority of a Microsoft Entra endpoint
-                    authority=self.region_config.authority,
+                    authority=region_config.authority,
                 )
             except Exception as error:
                 logger.critical("Failed to retrieve azure credentials")
@@ -244,14 +392,113 @@ class AzureProvider(Provider):
 
         return credentials
 
-    def __check_service_principal_creds_env_vars__(self):
+    @staticmethod
+    def test_connection(
+        az_cli_auth=False,
+        sp_env_auth=False,
+        browser_auth=False,
+        managed_identity_auth=False,
+        tenant_id=None,
+        region="AzureCloud",
+        raise_on_exception=True,
+    ) -> Connection:
+        """Test connection to Azure subscription.
+
+        Test the connection to an Azure subscription using the provided credentials.
+
+        Args:
+            az_cli_auth (bool): Flag indicating if Azure CLI authentication is used.
+            sp_env_auth (bool): Flag indicating if Service Principal environment authentication is used.
+            browser_auth (bool): Flag indicating if browser authentication is used.
+            managed_identity_auth (bool): Flag indicating if managed entity authentication is used.
+            tenant_id (str): The Azure Active Directory tenant ID.
+            region (str): The Azure region.
+            raise_on_exception (bool): Flag indicating whether to raise an exception if the connection fails.
+
+        Returns:
+            bool: True if the connection is successful, False otherwise.
+
+        Raises:
+            Exception: If failed to test the connection to Azure subscription.
+
+        Examples:
+            >>> AzureProvider.test_connection(az_cli_auth=True)
+            True
+        """
+        try:
+            AzureProvider.validate_arguments(
+                az_cli_auth, sp_env_auth, browser_auth, managed_identity_auth, tenant_id
+            )
+            region_config = AzureProvider.setup_region_config(region)
+            # Set up the Azure session
+            credentials = AzureProvider.setup_session(
+                az_cli_auth,
+                sp_env_auth,
+                browser_auth,
+                managed_identity_auth,
+                tenant_id,
+                region_config,
+            )
+            # Create a SubscriptionClient
+            subscription_client = SubscriptionClient(credentials)
+
+            # Get info from the first subscription
+            subscription = next(subscription_client.subscriptions.list())
+
+            logger.info(f"Connected to Azure subscription: {subscription.display_name}")
+            return Connection(is_connected=True)
+
+        except HttpResponseError as credentials_error:
+            logger.error(
+                f"{credentials_error.__class__.__name__}[{credentials_error.__traceback__.tb_lineno}]: {credentials_error}"
+            )
+            if raise_on_exception:
+                raise credentials_error
+            return Connection(error=credentials_error)
+
+        except ArgumentTypeError as validation_error:
+            logger.error(
+                f"{validation_error.__class__.__name__}[{validation_error.__traceback__.tb_lineno}]: {validation_error}"
+            )
+            if raise_on_exception:
+                raise validation_error
+            return Connection(error=validation_error)
+
+        except SystemExit as exit_error:
+            logger.error(
+                f"{exit_error.__class__.__name__}[{exit_error.__traceback__.tb_lineno}]: {exit_error}"
+            )
+            if raise_on_exception:
+                raise exit_error
+            return Connection(error=exit_error)
+
+        except Exception as error:
+            logger.critical(
+                f"{error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}"
+            )
+            if raise_on_exception:
+                raise error
+            return Connection(error=error)
+
+    @staticmethod
+    def check_service_principal_creds_env_vars():
+        """
+        Checks the presence of required environment variables for service principal authentication against Azure.
+
+        This method checks for the presence of the following environment variables:
+        - AZURE_CLIENT_ID: Azure client ID
+        - AZURE_TENANT_ID: Azure tenant ID
+        - AZURE_CLIENT_SECRET: Azure client secret
+
+        If any of the environment variables is missing, it logs a critical error and exits the program.
+        """
         logger.info(
             "Azure provider: checking service principal environment variables  ..."
         )
         for env_var in ["AZURE_CLIENT_ID", "AZURE_TENANT_ID", "AZURE_CLIENT_SECRET"]:
             if not getenv(env_var):
                 logger.critical(
-                    f"Azure provider: Missing environment variable {env_var} needed to autenticate against Azure"
+                    f"Azure provider: Missing environment variable {env_var} needed to authenticate against Azure"
                 )
                 sys.exit(1)
 
@@ -260,9 +507,22 @@ class AzureProvider(Provider):
         az_cli_auth,
         sp_env_auth,
         browser_auth,
-        managed_entity_auth,
+        managed_identity_auth,
         subscription_ids,
     ):
+        """
+        Sets up the identity for the Azure provider.
+
+        Args:
+            az_cli_auth (bool): Flag indicating if Azure CLI authentication is used.
+            sp_env_auth (bool): Flag indicating if Service Principal environment authentication is used.
+            browser_auth (bool): Flag indicating if browser authentication is used.
+            managed_identity_auth (bool): Flag indicating if managed entity authentication is used.
+            subscription_ids (list): List of subscription IDs.
+
+        Returns:
+            AzureIdentityInfo: An instance of AzureIdentityInfo containing the identity information.
+        """
         credentials = self.session
         # TODO: fill this object with real values not default and set to none
         identity = AzureIdentityInfo()
@@ -319,7 +579,7 @@ class AzureProvider(Provider):
             asyncio.get_event_loop().run_until_complete(get_azure_identity())
 
         # Managed identities only can be assigned resource, resource group and subscription scope permissions
-        elif managed_entity_auth:
+        elif managed_identity_auth:
             identity.identity_id = "Default Managed Identity ID"
             identity.identity_type = "Managed Identity"
             # Pending extracting info from managed identity
@@ -373,6 +633,16 @@ class AzureProvider(Provider):
         return identity
 
     def get_locations(self, credentials) -> dict[str, list[str]]:
+        """
+        Retrieves the locations available for each subscription using the provided credentials.
+
+        Args:
+            credentials: The credentials object used to authenticate the request.
+
+        Returns:
+            A dictionary containing the locations available for each subscription. The dictionary
+            has subscription display names as keys and lists of location names as values.
+        """
         locations = None
         if credentials:
             locations = {}

prowler/providers/azure/lib/arguments/arguments.py

@@ -12,12 +12,12 @@ def init_parser(self):
     azure_auth_modes_group.add_argument(
         "--az-cli-auth",
         action="store_true",
-        help="Use Azure cli credentials to log in against azure",
+        help="Use Azure CLI credentials to log in against Azure",
     )
     azure_auth_modes_group.add_argument(
         "--sp-env-auth",
         action="store_true",
-        help="Use service principal env variables authentication to log in against azure",
+        help="Use Service Principal environment variables authentication to log in against Azure",
     )
     azure_auth_modes_group.add_argument(
         "--browser-auth",
@@ -27,7 +27,7 @@ def init_parser(self):
     azure_auth_modes_group.add_argument(
         "--managed-identity-auth",
         action="store_true",
-        help="Use managed identity authentication to log in against azure ",
+        help="Use managed identity authentication to log in against Azure ",
     )
     # Subscriptions
     azure_subscriptions_subparser = azure_parser.add_argument_group("Subscriptions")

tests/providers/azure/azure_provider_test.py

@@ -1,11 +1,15 @@
 from argparse import Namespace
 from datetime import datetime
 from os import rmdir
+from unittest.mock import patch
+from uuid import uuid4
 
 import pytest
+from azure.core.credentials import AccessToken
+from azure.core.exceptions import ClientAuthenticationError, HttpResponseError
 from azure.identity import DefaultAzureCredential
 from freezegun import freeze_time
-from mock import patch
+from mock import MagicMock
 
 from prowler.config.config import (
     default_config_file_path,
@@ -17,6 +21,7 @@ from prowler.providers.azure.models import (
     AzureOutputOptions,
     AzureRegionConfig,
 )
+from prowler.providers.common.models import Connection
 
 
 class TestAzureProvider:
@@ -225,3 +230,100 @@ class TestAzureProvider:
                 # TODO: move this to a fixtures file
                 rmdir(f"{arguments.output_directory}/compliance")
                 rmdir(arguments.output_directory)
+
+    def test_test_connection_browser_auth(self):
+        with patch(
+            "prowler.providers.azure.azure_provider.DefaultAzureCredential"
+        ) as mock_default_credential, patch(
+            "prowler.providers.azure.azure_provider.AzureProvider.setup_session"
+        ) as mock_setup_session, patch(
+            "prowler.providers.azure.azure_provider.SubscriptionClient"
+        ) as mock_resource_client:
+
+            # Mock the return value of DefaultAzureCredential
+            mock_credentials = MagicMock()
+            mock_credentials.get_token.return_value = AccessToken(
+                token="fake_token", expires_on=9999999999
+            )
+            mock_default_credential.return_value = mock_credentials
+
+            # Mock setup_session to return a mocked session object
+            mock_session = MagicMock()
+            mock_setup_session.return_value = mock_session
+
+            # Mock ResourceManagementClient to avoid real API calls
+            mock_client = MagicMock()
+            mock_resource_client.return_value = mock_client
+
+            test_connection = AzureProvider.test_connection(
+                browser_auth=True,
+                tenant_id=str(uuid4()),
+                region="AzureCloud",
+                raise_on_exception=False,
+            )
+
+            assert isinstance(test_connection, Connection)
+            assert test_connection.is_connected
+            assert test_connection.error is None
+
+    def test_test_connection_with_ClientAuthenticationError(self):
+        with pytest.raises(ClientAuthenticationError) as exception:
+            tenant_id = str(uuid4())
+            AzureProvider.test_connection(
+                browser_auth=True,
+                tenant_id=tenant_id,
+                region="AzureCloud",
+            )
+
+        assert exception.type == ClientAuthenticationError
+        assert (
+            exception.value.args[0]
+            == f"Authentication failed: Unable to get authority configuration for https://login.microsoftonline.com/{tenant_id}. Authority would typically be in a format of https://login.microsoftonline.com/your_tenant or https://tenant_name.ciamlogin.com or https://tenant_name.b2clogin.com/tenant.onmicrosoft.com/policy.  Also please double check your tenant name or GUID is correct."
+        )
+
+    def test_test_connection_without_any_method(self):
+        with pytest.raises(SystemExit) as exception:
+            AzureProvider.test_connection()
+
+        assert exception.type == SystemExit
+        assert (
+            exception.value.args[0]
+            == "Azure provider requires at least one authentication method set: [--az-cli-auth | --sp-env-auth | --browser-auth | --managed-identity-auth]"
+        )
+
+    def test_test_connection_with_httpresponseerror(self):
+        with patch(
+            "prowler.providers.azure.azure_provider.AzureProvider.get_locations",
+            return_value={},
+        ), patch(
+            "prowler.providers.azure.azure_provider.AzureProvider.setup_session"
+        ) as mock_setup_session:
+
+            mock_setup_session.side_effect = HttpResponseError(
+                "Simulated HttpResponseError"
+            )
+
+            with pytest.raises(HttpResponseError) as exception:
+                AzureProvider.test_connection(
+                    az_cli_auth=True,
+                    raise_on_exception=True,
+                )
+
+            assert exception.type == HttpResponseError
+            assert exception.value.args[0] == "Simulated HttpResponseError"
+
+    def test_test_connection_with_exception(self):
+        with patch(
+            "prowler.providers.azure.azure_provider.AzureProvider.setup_session"
+        ) as mock_setup_session:
+
+            mock_setup_session.side_effect = Exception("Simulated Exception")
+
+            with pytest.raises(Exception) as exception:
+                AzureProvider.test_connection(
+                    sp_env_auth=True,
+                    raise_on_exception=True,
+                )
+
+            assert exception.type == Exception
+            assert exception.value.args[0] == "Simulated Exception"