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,34 +52,48 @@ class AzureCredentials(BaseCredentials): ...@@ -48,34 +52,48 @@ 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._client_id is None: if self._azure_paas_podidentity_isEnabled == "true":
self._populate_ad_credentials() try:
if self._tenant_id is None: print("MSI Token generation")
logger.error('TenantId is not set properly') headers = {
raise ValueError("TenantId is not set properly") 'Metadata': 'true'
if self._resource_id is None: }
logger.error('ResourceId is not set properly') url = 'http://169.254.169.254/metadata/identity/oauth2/token?api-version=2018-02-01&resource=https%3A%2F%2Fmanagement.azure.com%2F'
raise ValueError("ResourceId is not set properly") response = requests.request("GET", url, headers=headers)
if self._client_id is None: data_msi = json.loads(response.text)
logger.error('Please pass client Id to generate token') token = data_msi["access_token"]
raise ValueError("Please pass client Id to generate token") return token
if self._client_secret is None: except Exception as e:
logger.error('Please pass client secret to generate token') logger.error(e)
raise ValueError("Please pass client secret to generate token") raise e
else:
try: if self._client_id is None:
authority_host_uri = 'https://login.microsoftonline.com' self._populate_ad_credentials()
authority_uri = authority_host_uri + '/' + self._tenant_id if self._tenant_id is None:
scopes = [self._resource_id + '/.default'] logger.error('TenantId is not set properly')
app = msal.ConfidentialClientApplication(client_id = self._client_id, raise ValueError("TenantId is not set properly")
authority = authority_uri, if self._resource_id is None:
client_credential = self._client_secret) logger.error('ResourceId is not set properly')
result = app.acquire_token_for_client(scopes=scopes) raise ValueError("ResourceId is not set properly")
return result.get('access_token') if self._client_id is None:
except Exception as e: logger.error('Please pass client Id to generate token')
logger.error(e) raise ValueError("Please pass client Id to generate token")
raise e if self._client_secret is None:
logger.error('Please pass client secret to generate token')
raise ValueError("Please pass client secret to generate token")
try:
authority_host_uri = 'https://login.microsoftonline.com'
authority_uri = authority_host_uri + '/' + self._tenant_id
scopes = [self._resource_id + '/.default']
app = msal.ConfidentialClientApplication(client_id = self._client_id,
authority = authority_uri,
client_credential = self._client_secret)
result = app.acquire_token_for_client(scopes=scopes)
return result.get('access_token')
except Exception as e:
logger.error(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:
......
Markdown is supported
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