diff --git a/prometrix/auth.py b/prometrix/auth.py index 4919f9d..d5a3150 100644 --- a/prometrix/auth.py +++ b/prometrix/auth.py @@ -16,7 +16,9 @@ def azure_authorization(cls, config: PrometheusConfig) -> bool: if not isinstance(config, AzurePrometheusConfig): return False return (config.azure_client_id != "" and config.azure_tenant_id != "") and ( - config.azure_client_secret != "" or config.azure_use_managed_id != "" + config.azure_client_secret != "" or # Service Principal Auth + config.azure_use_managed_id != "" or # Managed Identity Auth + config.azure_use_workload_id != "" # Workload Identity Auth ) @classmethod @@ -48,15 +50,33 @@ def _get_azure_metadata_endpoint(cls, config: PrometheusConfig): @no_type_check @classmethod def _post_azure_token_endpoint(cls, config: PrometheusConfig): - return requests.post( - url=config.azure_token_endpoint, - headers={"Content-Type": "application/x-www-form-urlencoded"}, - data={ + # Try Azure Workload Identity + with open("/var/run/secrets/azure/tokens/azure-identity-token", "r") as token_file: + token = token_file.read() + data = { + "grant_type": "client_credentials", + "client_assertion_type": "urn:ietf:params:oauth:client-assertion-type:jwt-bearer", + "client_assertion": token, + "client_id": config.azure_client_id, + "scope": f"{config.azure_resource}/.default", + } + # Fallback to Azure Service Principal + if not token: + if config.azure_use_workload_id: + return { + "ok": False, + "reason": f"Could not open token file from {token_file}", + } + data = { "grant_type": "client_credentials", "client_id": config.azure_client_id, "client_secret": config.azure_client_secret, "resource": config.azure_resource, - }, + } + return requests.post( + url=config.azure_token_endpoint, + headers={"Content-Type": "application/x-www-form-urlencoded"}, + data=data, ) @classmethod @@ -67,7 +87,7 @@ def request_new_token(cls, config: PrometheusConfig) -> bool: try: if config.azure_use_managed_id: res = cls._get_azure_metadata_endpoint(config) - else: + else: # Service Principal and Workload Identity res = cls._post_azure_token_endpoint(config) except Exception: logging.exception( diff --git a/prometrix/models/prometheus_config.py b/prometrix/models/prometheus_config.py index d5e6793..9b40492 100644 --- a/prometrix/models/prometheus_config.py +++ b/prometrix/models/prometheus_config.py @@ -64,6 +64,7 @@ class AzurePrometheusConfig(PrometheusConfig): azure_metadata_endpoint: str azure_token_endpoint: str azure_use_managed_id: Optional[str] + azure_use_workload_id: Optional[str] azure_client_id: Optional[str] azure_tenant_id: Optional[str] azure_client_secret: Optional[str]