Commit ae14d37d authored by harshit aggarwal's avatar harshit aggarwal
Browse files

Merge branch 'users/preetisingh/spnManifestIngestion' into 'master'

Users/preetisingh/spn manifest ingestion

See merge request !44
parents c34f2678 68d67504
Pipeline #67544 passed with stages
in 1 minute and 46 seconds
......@@ -22,6 +22,9 @@ import msal
import os
from azure.keyvault import secrets
from azure import identity
import requests
import json
logger = logging.getLogger(__name__)
RETRIES = 3
......@@ -37,6 +40,7 @@ class AzureCredentials(BaseCredentials):
self._client_secret = None
self._tenant_id = None
self._resource_id = None
self._azure_paas_podidentity_isEnabled= os.getenv("AIRFLOW_VAR_AZURE_ENABLE_MSI")
def _populate_ad_credentials(self) -> None:
uri = os.getenv("AIRFLOW_VAR_KEYVAULT_URI")
......@@ -48,34 +52,48 @@ class AzureCredentials(BaseCredentials):
self._resource_id = client.get_secret("aad-client-id").value
def _generate_token(self) -> str:
if self._client_id is None:
self._populate_ad_credentials()
if self._tenant_id is None:
logger.error('TenantId is not set properly')
raise ValueError("TenantId is not set properly")
if self._resource_id is None:
logger.error('ResourceId is not set properly')
raise ValueError("ResourceId is not set properly")
if self._client_id is None:
logger.error('Please pass client Id to generate token')
raise ValueError("Please pass client Id to generate token")
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
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:
self._populate_ad_credentials()
if self._tenant_id is None:
logger.error('TenantId is not set properly')
raise ValueError("TenantId is not set properly")
if self._resource_id is None:
logger.error('ResourceId is not set properly')
raise ValueError("ResourceId is not set properly")
if self._client_id is None:
logger.error('Please pass client Id to generate token')
raise ValueError("Please pass client Id to generate token")
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))
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