From eb67e381cf68c634406aeeaa87ed65ec52df8786 Mon Sep 17 00:00:00 2001
From: pedrooot <pedromarting3@gmail.com>
Date: Wed, 19 Mar 2025 10:17:34 +0100
Subject: [PATCH] feat(azure): change logic WIP

---
 prowler/providers/azure/azure_provider.py     | 22 +++++++++++++++----
 .../providers/azure/lib/regions/regions.py    | 20 ++++++++++-------
 prowler/providers/azure/models.py             |  1 +
 3 files changed, 31 insertions(+), 12 deletions(-)

diff --git a/prowler/providers/azure/azure_provider.py b/prowler/providers/azure/azure_provider.py
index 26f42c94d2..88e5015334 100644
--- a/prowler/providers/azure/azure_provider.py
+++ b/prowler/providers/azure/azure_provider.py
@@ -17,6 +17,9 @@ from azure.identity import (
 )
 from azure.mgmt.subscription import SubscriptionClient
 from colorama import Fore, Style
+from kiota_authentication_azure.azure_identity_authentication_provider import (
+    AzureIdentityAuthenticationProvider,
+)
 from msgraph import GraphServiceClient
 
 from prowler.config.config import (
@@ -407,6 +410,7 @@ class AzureProvider(Provider):
                 base_url=config["base_url"],
                 credential_scopes=config["credential_scopes"],
                 graph_credential_scopes=config["graph_credential_scopes"],
+                graph_base_url=config["graph_base_url"],
             )
         except ArgumentTypeError as validation_error:
             logger.error(
@@ -892,10 +896,20 @@ class AzureProvider(Provider):
                     logger.info(
                         "Trying to retrieve tenant domain from AAD to populate identity structure ..."
                     )
-                    client = GraphServiceClient(
-                        credentials=credentials,
-                        scopes=self.region_config.graph_credential_scopes,
-                    )
+                    if self.region_config.name == "AzureCloud":
+                        client = GraphServiceClient(
+                            credentials=credentials,
+                            scopes=self.region_config.graph_credential_scopes,
+                        )
+                    else:
+                        auth_provider = AzureIdentityAuthenticationProvider(
+                            self.credential,
+                            scopes=[self.region_config.graph_credential_scopes],
+                            allowed_hosts=[self.region_config.graph_base_url],
+                        )
+                        http_client = GraphClientFactory.create_with_default_middleware(
+                            host=self.region_config.graph_base_url
+                        )
 
                     domain_result = await client.domains.get()
                     if getattr(domain_result, "value"):
diff --git a/prowler/providers/azure/lib/regions/regions.py b/prowler/providers/azure/lib/regions/regions.py
index 07c1c91afa..1ae62c8fc4 100644
--- a/prowler/providers/azure/lib/regions/regions.py
+++ b/prowler/providers/azure/lib/regions/regions.py
@@ -1,14 +1,14 @@
 from azure.identity import AzureAuthorityHosts
 
 AZURE_GENERIC_CLOUD = "https://management.azure.com"
-AZURE_GRAPH_GLOBAL = "https://graph.microsoft.com/.default"
+AZURE_GRAPH_GLOBAL = "https://graph.microsoft.com"
 
 AZURE_US_GOV_CLOUD = "https://management.usgovcloudapi.net"
-AZURE_GRAPH_GOV_US_L4 = "https://graph.microsoft.us/.default"
-AZURE_GRAPH_GOV_US_L5 = "https://dod-graph.microsoft.us/.default"
+AZURE_GRAPH_GOV_US_L4 = "https://graph.microsoft.us"
+AZURE_GRAPH_GOV_US_L5 = "https://dod-graph.microsoft.us"
 
 AZURE_CHINA_CLOUD = "https://management.chinacloudapi.cn"
-AZURE_GRAPH_CHINA = "https://microsoftgraph.chinacloudapi.cn/.default"
+AZURE_GRAPH_CHINA = "https://microsoftgraph.chinacloudapi.cn"
 
 
 def get_regions_config(region):
@@ -17,25 +17,29 @@ def get_regions_config(region):
             "authority": AzureAuthorityHosts.AZURE_PUBLIC_CLOUD,
             "base_url": AZURE_GENERIC_CLOUD,
             "credential_scopes": [AZURE_GENERIC_CLOUD + "/.default"],
-            "graph_credential_scopes": [AZURE_GRAPH_GLOBAL],
+            "graph_credential_scopes": [AZURE_GRAPH_GLOBAL + "/.default"],
+            "graph_base_url": AZURE_GRAPH_GLOBAL,
         },
         "AzureChinaCloud": {
             "authority": AzureAuthorityHosts.AZURE_CHINA,
             "base_url": AZURE_CHINA_CLOUD,
             "credential_scopes": [AZURE_CHINA_CLOUD + "/.default"],
-            "graph_credential_scopes": [AZURE_GRAPH_CHINA],
+            "graph_credential_scopes": [AZURE_GRAPH_CHINA + "/.default"],
+            "graph_base_url": AZURE_GRAPH_CHINA,
         },
         "AzureUSGovernmentL4": {
             "authority": AzureAuthorityHosts.AZURE_GOVERNMENT,
             "base_url": AZURE_US_GOV_CLOUD,
             "credential_scopes": [AZURE_US_GOV_CLOUD + "/.default"],
-            "graph_credential_scopes": [AZURE_GRAPH_GOV_US_L4],
+            "graph_credential_scopes": [AZURE_GRAPH_GOV_US_L4 + "/.default"],
+            "graph_base_url": AZURE_GRAPH_GOV_US_L4,
         },
         "AzureUSGovernmentL5": {
             "authority": AzureAuthorityHosts.AZURE_GOVERNMENT,
             "base_url": AZURE_US_GOV_CLOUD,
             "credential_scopes": [AZURE_US_GOV_CLOUD + "/.default"],
-            "graph_credential_scopes": [AZURE_GRAPH_GOV_US_L5],
+            "graph_credential_scopes": [AZURE_GRAPH_GOV_US_L5 + "/.default"],
+            "graph_base_url": AZURE_GRAPH_GOV_US_L5,
         },
     }
     return allowed_regions[region]
diff --git a/prowler/providers/azure/models.py b/prowler/providers/azure/models.py
index c469238c83..1215550cee 100644
--- a/prowler/providers/azure/models.py
+++ b/prowler/providers/azure/models.py
@@ -19,6 +19,7 @@ class AzureRegionConfig(BaseModel):
     base_url: str = ""
     credential_scopes: list = []
     graph_credential_scopes: list = []
+    graph_base_url: str = ""
 
 
 class AzureSubscription(BaseModel):