Client credential flow to interact with API secured with Azure AD

[rupesh3020]

Hi Team,

I am aware that as new PR is merged, we will be able to add requests headers to Http request created by HttpDispatcher.

We need to send lineage to a endpoint which is secured by Azure AD. Can't we make a functionality to get access token using client credential flow and send that in a header?

I think that would be great.

[wajda]

Sure. Propose a PR and we'll discuss it further.

[rupesh3020]

Hi @wajda , For the same functionality, i need to update the Authorization header whenever Spline makes a post request.
I want to insert this function call in restclient class for it. Will it serve my purpose or should i put this in Httpdispatcher class?

For us it might happen that post is done after 1 hr , in that case i want to regenerate token and add to header.

def getAccessToken(clientId: String, clientSecret: String, tenantId: String): String = {
  val tokenEndpoint = s"https://login.microsoftonline.com/$tenantId/oauth2/v2.0/token"
  val response = Http(tokenEndpoint)
    .postForm(Seq(
      "client_id" -> clientId,
      "client_secret" -> clientSecret,
      "grant_type" -> "client_credentials",
      "scope" -> "somescope/.default"
    ))
    .asString

  val responseJson = response.body
  val accessToken = responseJson

  accessToken
}

[rupesh3020]

I have done these changes. but getting 401 issue

/*
 * Copyright 2020 ABSA Group Limited
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.abnamro.dmp.mdc.splinecustomhttpdispatcher

import org.apache.spark.internal.Logging
import scalaj.http.HttpOptions.HttpOption
import scalaj.http.{BaseHttp, HttpOptions}
import za.co.absa.commons.lang.extensions.AnyExtension._
import za.co.absa.spline.harvester.dispatcher.SplineHeaders

import scala.concurrent.duration.Duration

trait AzureRestClient {
  def endpoint(url: String): AzureRestEndpoint
}

object AzureRestClient extends Logging {
  /**
   * @param baseHttp    HTTP client
   * @param baseURL     REST endpoint base URL
   * @param connTimeout timeout for establishing TCP connection
   * @param readTimeout timeout for each individual TCP packet (in already established connection)
   */
  def apply(
    baseHttp: BaseHttp,
    baseURL: String,
    connTimeout: Duration,
    readTimeout: Duration,
    disableSslValidation: Boolean,
    headers: Map[String, String]
  ): AzureRestClient = {

    logDebug(s"baseURL = $baseURL")
    logDebug(s"connTimeout = $connTimeout")
    logDebug(s"readTimeout = $readTimeout")
    logDebug(s"disableSslValidation = $disableSslValidation")
    logWarning(s"headers = $headers")

    val maybeDisableSslValidationOption: Option[HttpOption] =
      if (disableSslValidation) {
        logWarning(s"SSL validation is DISABLED -- not recommended for production!")
        Some(HttpOptions.allowUnsafeSSL)
      } else {
        None
      }

    //noinspection ConvertExpressionToSAM
    new AzureRestClient {
      override def endpoint(resource: String): AzureRestEndpoint = new AzureRestEndpoint(
        baseHttp(s"$baseURL/$resource")
          .option(HttpOptions.connTimeout(connTimeout.toMillis.toInt))
          .option(HttpOptions.readTimeout(readTimeout.toMillis.toInt))
          .having(maybeDisableSslValidationOption)(_ option _)
          .header(SplineHeaders.Timeout, readTimeout.toMillis.toString)
          .headers(headers)
          .compress(true),headers("clientId"),headers("clientSecret"),headers("tenantId"),headers("scope"))
    }
  }
}

/*
 * Copyright 2020 ABSA Group Limited
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.abnamro.dmp.mdc.splinecustomhttpdispatcher
import org.apache.spark.internal.Logging

import org.apache.http.HttpHeaders
import scalaj.http.{HttpRequest, HttpResponse}
import za.co.absa.commons.lang.ARM.using
import za.co.absa.spline.harvester.dispatcher.httpdispatcher.HttpConstants.Encoding
import org.abnamro.dmp.mdc.splinecustomhttpdispatcher.AzureRestEndpoint._

import java.io.ByteArrayOutputStream
import java.util.zip.GZIPOutputStream
import javax.ws.rs.HttpMethod
import scalaj.http._

class AzureRestEndpoint(val request: HttpRequest,val clientId: String, val clientSecret: String, val scope:String, val tenantId: String) extends Logging {

  def getAccessToken(clientId: String, clientSecret: String, tenantId: String,scope: String): String = {
    val tokenEndpoint = s"https://login.microsoftonline.com/$tenantId/oauth2/v2.0/token"
    val response = Http(tokenEndpoint)
      .postForm(Seq(
        "client_id" -> clientId,
        "client_secret" -> clientSecret,
        "grant_type" -> "client_credentials",
        "scope" -> scope.concat("/.default")
      )).options(HttpOptions.allowUnsafeSSL,HttpOptions.readTimeout(5000))
      .asString

    val responseJson = response.body
    import scala.util.matching.Regex
    val pattern: Regex = """"access_token"\s*:\s*"([^"]+)"""".r

    val accessToken = "Bearer ".concat(pattern.findFirstMatchIn(responseJson).map(_.group(1)).getOrElse(""))
    logWarning(accessToken)
    accessToken
  }
  
  def head(): HttpResponse[String] = request.header("Authorization", getAccessToken(clientId, clientSecret, tenantId, scope)).method(HttpMethod.HEAD).asString

  def post(data: String, contentType: String, enableRequestCompression: Boolean): HttpResponse[String] = {
    val jsonRequest = request
      .header(HttpHeaders.CONTENT_TYPE, contentType)
      .header("Authorization", getAccessToken(clientId, clientSecret, tenantId, scope))

    if (enableRequestCompression && data.length > GzipCompressionLengthThreshold) {
      jsonRequest
        .header(HttpHeaders.CONTENT_ENCODING, Encoding.GZIP)
        .postData(gzipContent(data.getBytes("UTF-8")))
        .asString
    } else {
      jsonRequest
        .postData(data)
        .asString
    }
  }
}

object AzureRestEndpoint {
  private val GzipCompressionLengthThreshold: Int = 2048

  private def gzipContent(bytes: Array[Byte]): Array[Byte] = {
    val byteStream = new ByteArrayOutputStream(bytes.length)
    using(new GZIPOutputStream(byteStream))(_.write(bytes))
    byteStream.toByteArray
  }
}

[wajda]

Yeah, it's not easy to implement that use-case in HttpLineageDispatcher. We need to make it more flexible.
But your code looks about right, I don't see any evident mistake.
Maybe verify if you get your access token correctly. Is your Regexp works correctly?
I'm not 100% sure as I haven't worked with Azure, but according to RFC6749 that access_token could contain any printable ASCII symbol (0x20-0x7E) including double quote (") which your regexp would not pick. I would rather parse JSON and get the access token out of it normally as a string property. But that's just my guess.

[wajda]

Created feature request #600

[rupesh3020]

Solved @wajda , i made a small mistake using positional arguments

[rupesh3020]

/*
 * Copyright 2020 ABSA Group Limited
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.abnamro.dmp.mdc.splinecustomhttpdispatcher
import org.apache.spark.internal.Logging

import org.apache.http.HttpHeaders
import scalaj.http.{HttpRequest, HttpResponse}
import za.co.absa.commons.lang.ARM.using
import za.co.absa.spline.harvester.dispatcher.httpdispatcher.HttpConstants.Encoding
import org.abnamro.dmp.mdc.splinecustomhttpdispatcher.AzureRestEndpoint._

import java.io.ByteArrayOutputStream
import java.util.zip.GZIPOutputStream
import javax.ws.rs.HttpMethod
import scalaj.http._

class AzureRestEndpoint(val request: HttpRequest,val clientId: String, val clientSecret: String, val scope:String, val tenantId: String) extends Logging {

  def getAccessToken(clientId: String, clientSecret: String, tenantId: String,scope: String): String = {
  val tokenEndpoint = s"https://login.microsoftonline.com/$tenantId/oauth2/v2.0/token"
  val response = Http(tokenEndpoint)
    .postForm(Seq(
      "client_id" -> clientId,
      "client_secret" -> clientSecret,
      "grant_type" -> "client_credentials",
      "scope" -> scope.concat("/.default")
    )).options(HttpOptions.allowUnsafeSSL,HttpOptions.readTimeout(5000))
    .asString

  val responseJson = response.body
  import scala.util.matching.Regex
  val pattern: Regex = """"access_token"\s*:\s*"([^"]+)"""".r

  val accessToken = "Bearer ".concat(pattern.findFirstMatchIn(responseJson).map(_.group(1)).getOrElse(""))
  logWarning(accessToken)
  accessToken
}
  
  def head(): HttpResponse[String] = request.header("Authorization", getAccessToken(clientId, clientSecret, tenantId, scope)).method(HttpMethod.HEAD).asString

  def post(data: String, contentType: String, enableRequestCompression: Boolean): HttpResponse[String] = {
    val jsonRequest = request
      .header(HttpHeaders.CONTENT_TYPE, contentType)
      .header("Authorization", getAccessToken(clientId, clientSecret, tenantId, scope))

    if (enableRequestCompression && data.length > GzipCompressionLengthThreshold) {
      jsonRequest
        .header(HttpHeaders.CONTENT_ENCODING, Encoding.GZIP)
        .postData(gzipContent(data.getBytes("UTF-8")))
        .asString
    } else {
      jsonRequest
        .postData(data)
        .asString
    }
  }
}

object AzureRestEndpoint {
  private val GzipCompressionLengthThreshold: Int = 2048

  private def gzipContent(bytes: Array[Byte]): Array[Byte] = {
    val byteStream = new ByteArrayOutputStream(bytes.length)
    using(new GZIPOutputStream(byteStream))(_.write(bytes))
    byteStream.toByteArray
  }
}

[rupesh3020]

I have made a pull request for this. Kindly check

[wajda]

Awesome! Thanks a lot! We will check it early next week and will contact you back.

[rupesh3020]

@wajda I have resolved all the comments in the PR can you please check once. I am not able to test the jar , can you please help in building it?

[wajda]

Builds are triggered automatically, currently they are all failing.
Do you have access to build logs?

[rupesh3020]

@wajda It is passing now

[wajda]

Spline Agent 1.1.0 has been released.
I changed the Http authentication config just a little bit for clarity and ease for supporting more auth types in the future.
@rupesh3020 can you please check if it still works for you as intended?
The new config is this:

spline.lineageDispatcher.http.authentication.type=OAUTH
spline.lineageDispatcher.http.authentication.grantType=client_credentials
spline.lineageDispatcher.http.authentication.clientId=<client_id>
spline.lineageDispatcher.http.authentication.clientSecret=<secret>
spline.lineageDispatcher.http.authentication.scope=<scope>
spline.lineageDispatcher.http.authentication.tokenUrl=<token_url>

[rupesh3020]

@wajda yes it works perfectly, out team is using it