Azure Databricks: Machine to Machine Authentication using SPN with Certificate for Azure Databricks Tasks.

Introduction:

Azure Databricks offers a robust platform for big data analytics and machine learning tasks. When performing various operations within Databricks, such as accessing resources or executing tasks, authentication is crucial. Traditionally, Personal Access Tokens (PAT) have been used for authentication, but leveraging Service Principal Name (SPN) with certificate authentication provides an alternative that offers enhanced security and flexibility.

Why SPN with Certificate? Using SPN with certificate authentication for Azure Databricks tasks offers several advantages over PAT tokens as well as over SPN with secrets.

Why SPN with certificates rather than PAT Tokens/SPN with secrets?

1. Enhanced Security: Certificates provide a more secure way to authenticate, reducing the risk associated with token-based authentication.
2. Long-term Authentication: Unlike tokens, certificates do not expire frequently, reducing the need for frequent renewal and maintenance.
3. No User Dependency: SPN with certificate authentication enables automation and non-interactive workflows, making it ideal for scheduled tasks and automated pipelines.
4. Granular Access Control: SPN authentication allows for fine-grained access control through Azure Active Directory (AAD), ensuring that only authorized entities can access Databricks resources.
5. No track of PAT Tokens and you may end up deleting after your use case. Even if any user leaves Organization/Project we have no visibility whether PAT token’s are in use or not.

Prerequisites:

Before anyone starts the actual implementation, ensure you have these initial prerequisites in place,

1. Include “2ff814a6–3304–4ab8–85cb-cd0e6f879c1d” a unique Databricks Identity, typically found under Enterprise App.

You can find more info about Databricks enterprise app: https://learn.microsoft.com/en-us/azure/databricks/administration-guide/access-control/conditional-access

2. Verify if your Azure Databricks resides within a Managed Resource group, then navigate to the corresponding link.

check if there is any Managed Identity exists, if yes copy its ClientID and add it to Azure Keyvault.

This is my Azure Keyvault screen looks after adding both keys.

3. Basic understanding of Python Programming.

4. Make sure you add your SPN Name under Service Principals.

5. Next you need to add this Service Principal under Admins group.

Implementation Steps:

Let’s now proceed with the actual implementation steps:

1. Begin by installing the necessary libraries. You can install them on your compute. Refer below screenshot.

for more details : https://learn.microsoft.com/en-us/azure/databricks/libraries/cluster-libraries

2. Ensure all essential libraries are added to facilitate interaction with Azure services, secret management, and SSL certificate cryptographic operations. For that I have added these parameters.

from azure.identity import DefaultAzureCredential
from azure.keyvault.secrets import SecretClient
from cryptography.hazmat.primitives.serialization import load_pem_private_key
from cryptography.hazmat.backends import default_backend
import msal

3. Convert the .pfx file to a valid .pem format, there are lot of articles available to convert however there is a specific approach for Databricks, which requires a properly formatted PKCS12 data as it looks for unicode object and it needs to deserialize with valid format/content for PKCS12 data.

For that I have used these 3 commands to convert .pfx to .pem so that Databricks can read.

openssl pkcs12 -in cert.pfx -out cert_export.pem -nodes -password pass:"certpass"
openssl x509 -in cert_export.pem >> temp.pem
openssl pkcs8 -topk8 -nocrypt -in cert_export.pem >> temp.pem

Here is the reference link: https://github.com/MicrosoftDocs/azure-docs/issues/23558

4. Define the Azure Keyvault keys. After obtaining a valid .pem file, add the key to Azure Keyvault along with its associated password. Ensure to provide the SPN certification Application ID, Azure Tenant ID, and SPN certificate’s thumbprint.

key_vault_name = "nonprodkeyvault"
pem_secret_name = "1649-spn-pem"  
pfx_password_secret = "1649-pfx-secret"  
app_id = "1649-spn-appid"
tenant_id = "1649-spn-tenantid"
certificate_thumbprint = "1649-spn-thumbprint"

5. Proceed to read the keys from AzureKeyVault.

credential = DefaultAzureCredential()
secret_client = SecretClient(vault_url=f"https://{key_vault_name}.vault.azure.net/", credential=credential)

pem_secret = secret_client.get_secret(pem_secret_name)
pem_bytes = pem_secret.value

pfx_password_secret = secret_client.get_secret(pfx_password_secret)
pfx_password = pfx_password_secret.value

app_id_secret = secret_client.get_secret(app_id)
app_id = app_id_secret.value

tenant_id_secret = secret_client.get_secret(tenant_id)
tenant_id = tenant_id_secret.value

certificate_thumbprint_secret = secret_client.get_secret(certificate_thumbprint)
certificate_thumbprint = certificate_thumbprint_secret.value

6. Read the .pem file in the correct byte format.

if isinstance(pem_bytes, str):
    pem_bytes = pem_bytes.encode()

private_key = load_pem_private_key(pem_bytes, password=None, backend=default_backend())

7. Instantiate the SPN certificate through the constructor method, initializing the AzureSPNWithCertificate class with all necessary information required for authentication and cryptographic operations involving the SPN’s certificate.

def __init__(
            self,
            tenant_id: str,
            spn_app_id: str,
            certificate_thumbprint: str,
            pfx_password: str = None,
            pem_bytes: bytes = None):
        """Instantiate an AzureSPNWithCertificate class.
        
        :param tenant_id: Azure tenant id
        :param spn_app_id: SPN application (client) id
        :param certificate_thumbprint: SPN's certificate thumbprint
        :param pfx_password: password used for the certificate
        :param pem_bytes: certificate file as bytes.
        """
        self.tenant_id = tenant_id
        self.spn_app_id = spn_app_id
        self.pem_bytes = pem_bytes
        self.pfx_password = pfx_password
        self.certificate_thumbprint = certificate_thumbprint
        self.private_key_bytes = None

8. Finally, generate a new bearer token with Databricks scope.

 def generate_ad_token(self) -> str:
        """Generate a short lived AD token for the SPN."""
        authority = f"https://login.microsoftonline.com/{self.tenant_id}"
        azure_dbx_scope = ["2ff814a6-3304-4ab8-85cb-cd0e6f879c1d/.default"]
        
        app = msal.ConfidentialClientApplication(
            self.spn_app_id,
            authority=authority,
            client_credential={
                "thumbprint": self.certificate_thumbprint,
                "private_key": self.pem_bytes,
                "password": self.pfx_password,
            }
        )

Output:

Testing:

Now that we have the Bearer token, let’s proceed with testing to ensure its usability. Now you can execute various tasks within Azure Databricks, such as running notebooks, managing clusters, or accessing data sources

For this purpose, I’ll check to list our clusters.

import requests

# Endpoint for listing clusters
clusters_endpoint = "https://adb-xxxxxxxxxxx.14.azuredatabricks.net/api/2.0/clusters/list"

# Set the token in the header
headers = {
    "Authorization": f"Bearer {new_token}"
}

# Make the API request
response = requests.get(clusters_endpoint, headers=headers)

# Check if the request was successful
if response.status_code == 200:
    clusters = response.json()["clusters"]
    print("List of clusters:")
    for cluster in clusters:
        print(cluster["cluster_name"])
else:
    print("Error:", response.text)

Here is the result(s).

Meanwhile I have tested the same from Postman as well.

Code:

Here is the full script which I have used it for my work.

from azure.identity import DefaultAzureCredential
from azure.keyvault.secrets import SecretClient
from cryptography.hazmat.primitives.serialization import load_pem_private_key
from cryptography.hazmat.backends import default_backend
import msal

# Azure Key Vault details
# https://github.com/MicrosoftDocs/azure-docs/issues/23558
#openssl pkcs12 -in cert.pfx -out cert_export.pem -nodes -password pass:"certpass"
#openssl x509 -in cert_export.pem >> temp.pem
#openssl pkcs8 -topk8 -nocrypt -in cert_export.pem >> temp.pem

key_vault_name = "nonprodkeyvault"
pem_secret_name = "1649-spn-pem"  
pfx_password_secret = "1649-pfx-secret"  
app_id = "1649-spn-appid"
tenant_id = "1649-spn-tenantid"
certificate_thumbprint = "1649-spn-thumbprint"

credential = DefaultAzureCredential()
secret_client = SecretClient(vault_url=f"https://{key_vault_name}.vault.azure.net/", credential=credential)

pem_secret = secret_client.get_secret(pem_secret_name)
pem_bytes = pem_secret.value

pfx_password_secret = secret_client.get_secret(pfx_password_secret)
pfx_password = pfx_password_secret.value

app_id_secret = secret_client.get_secret(app_id)
app_id = app_id_secret.value

tenant_id_secret = secret_client.get_secret(tenant_id)
tenant_id = tenant_id_secret.value

certificate_thumbprint_secret = secret_client.get_secret(certificate_thumbprint)
certificate_thumbprint = certificate_thumbprint_secret.value


if isinstance(pem_bytes, str):
    pem_bytes = pem_bytes.encode()

private_key = load_pem_private_key(pem_bytes, password=None, backend=default_backend())



class AzureSPNWithCertificateNewCode:
    """SPN with certificate."""
    
    def __init__(
            self,
            tenant_id: str,
            spn_app_id: str,
            certificate_thumbprint: str,
            pfx_password: str = None,
            pem_bytes: bytes = None):
        """Instantiate an AzureSPNWithCertificate class.
        
        :param tenant_id: Azure tenant id
        :param spn_app_id: SPN application (client) id
        :param certificate_thumbprint: SPN's certificate thumbprint
        :param pfx_password: password used for the certificate
        :param pem_bytes: certificate file as bytes.
        """
        self.tenant_id = tenant_id
        self.spn_app_id = spn_app_id
        self.pem_bytes = pem_bytes
        self.pfx_password = pfx_password
        self.certificate_thumbprint = certificate_thumbprint
        self.private_key_bytes = None

    def generate_ad_token(self) -> str:
        """Generate a short lived AD token for the SPN."""
        authority = f"https://login.microsoftonline.com/{self.tenant_id}"
        azure_dbx_scope = ["2ff814a6-3304-4ab8-85cb-cd0e6f879c1d/.default"]
        
        app = msal.ConfidentialClientApplication(
            self.spn_app_id,
            authority=authority,
            client_credential={
                "thumbprint": self.certificate_thumbprint,
                "private_key": self.pem_bytes,
                "password": self.pfx_password,
            }
        )


        result = app.acquire_token_for_client(scopes=azure_dbx_scope)
        new_token = result["access_token"]
        return new_token


new_spn = AzureSPNWithCertificateNewCode(tenant_id, app_id, certificate_thumbprint, pfx_password, pem_bytes)
new_token = new_spn.generate_ad_token()
print("New token:", new_token)


# For Testing Purpose
import requests
clusters_endpoint = "https://adb-xxxxxxxxxxxxxxx.14.azuredatabricks.net/api/2.0/clusters/list"

headers = {
    "Authorization": f"Bearer {new_token}"
}


response = requests.get(clusters_endpoint, headers=headers)


if response.status_code == 200:
    clusters = response.json()["clusters"]
    print("List of clusters:")
    for cluster in clusters:
        print(cluster["cluster_name"])
else:
    print("Error:", response.text)

# If you want to check & Decode the token to see if it has right scope/tenant details etc
import jwt
decoded_token = jwt.decode(new_token, options={"verify_signature": False})

print(decoded_token)

Conclusion:

Meanwhile some of you guys might have wondered I could have leveraged DBFS to store .pfx file, however Databricks/any Organization policy recommends against storing any sensitive or production code or data in the DBFS root.

By leveraging Service Principal Name (SPN) with certificate authentication for Azure Databricks tasks, organizations can enhance security, streamline automation, and reduce dependency on user interactions while using with PAT token. This approach provides a robust and scalable solution for integrating Databricks into automated workflows, enabling efficient data processing, analysis, and machine learning operations.

Happy reading, feel free to provide your comments if any.