Commit 68d67504 authored by preeti singh[Microsoft]'s avatar preeti singh[Microsoft] Committed by harshit aggarwal
Browse files

Users/preetisingh/spn manifest ingestion

parent c34f2678
...@@ -22,6 +22,9 @@ import msal ...@@ -22,6 +22,9 @@ import msal
import os import os
from azure.keyvault import secrets from azure.keyvault import secrets
from azure import identity from azure import identity
import requests
import json
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
RETRIES = 3 RETRIES = 3
...@@ -37,6 +40,7 @@ class AzureCredentials(BaseCredentials): ...@@ -37,6 +40,7 @@ class AzureCredentials(BaseCredentials):
self._client_secret = None self._client_secret = None
self._tenant_id = None self._tenant_id = None
self._resource_id = None self._resource_id = None
self._azure_paas_podidentity_isEnabled= os.getenv("AIRFLOW_VAR_AZURE_ENABLE_MSI")
def _populate_ad_credentials(self) -> None: def _populate_ad_credentials(self) -> None:
uri = os.getenv("AIRFLOW_VAR_KEYVAULT_URI") uri = os.getenv("AIRFLOW_VAR_KEYVAULT_URI")
...@@ -48,6 +52,21 @@ class AzureCredentials(BaseCredentials): ...@@ -48,6 +52,21 @@ class AzureCredentials(BaseCredentials):
self._resource_id = client.get_secret("aad-client-id").value self._resource_id = client.get_secret("aad-client-id").value
def _generate_token(self) -> str: def _generate_token(self) -> str:
if self._azure_paas_podidentity_isEnabled == "true":
try:
print("MSI Token generation")
headers = {
'Metadata': 'true'
}
url = 'http://169.254.169.254/metadata/identity/oauth2/token?api-version=2018-02-01&resource=https%3A%2F%2Fmanagement.azure.com%2F'
response = requests.request("GET", url, headers=headers)
data_msi = json.loads(response.text)
token = data_msi["access_token"]
return token
except Exception as e:
logger.error(e)
raise e
else:
if self._client_id is None: if self._client_id is None:
self._populate_ad_credentials() self._populate_ad_credentials()
if self._tenant_id is None: if self._tenant_id is None:
...@@ -76,7 +95,6 @@ class AzureCredentials(BaseCredentials): ...@@ -76,7 +95,6 @@ class AzureCredentials(BaseCredentials):
logger.error(e) logger.error(e)
raise e raise e
@retry(stop=stop_after_attempt(RETRIES)) @retry(stop=stop_after_attempt(RETRIES))
def refresh_token(self) -> str: def refresh_token(self) -> str:
"""Refresh token. """Refresh token.
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment