/**
  * Licensed to the Apache Software Foundation (ASF) under one or more
  * contributor license agreements.  See the NOTICE file distributed with
  * this work for additional information regarding copyright ownership.
  * The ASF licenses this file to You 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 kafka.api

import java.security.AccessController
import java.util.Properties

import javax.security.auth.callback._
import javax.security.auth.Subject
import javax.security.auth.login.AppConfigurationEntry

import scala.collection.Seq
import kafka.server.KafkaConfig
import kafka.utils.TestUtils
import kafka.utils.JaasTestUtils._
import org.apache.kafka.common.config.SaslConfigs
import org.apache.kafka.common.config.internals.BrokerSecurityConfigs
import org.apache.kafka.common.network.Mode
import org.apache.kafka.common.security.auth._
import org.apache.kafka.common.security.authenticator.DefaultKafkaPrincipalBuilder
import org.apache.kafka.common.security.plain.PlainAuthenticateCallback
import org.junit.jupiter.api.Assertions.assertTrue
import org.junit.jupiter.api.Test

object SaslPlainSslEndToEndAuthorizationTest {

  class TestPrincipalBuilder extends DefaultKafkaPrincipalBuilder(null, null) {

    override def build(context: AuthenticationContext): KafkaPrincipal = {
      val saslContext = context.asInstanceOf[SaslAuthenticationContext]

      // Verify that peer principal can be obtained from the SSLSession provided in the context
      // since we have enabled TLS mutual authentication for the listener
      val sslPrincipal = saslContext.sslSession.get.getPeerPrincipal.getName
      assertTrue(sslPrincipal.endsWith(s"CN=${TestUtils.SslCertificateCn}"), s"Unexpected SSL principal $sslPrincipal")

      saslContext.server.getAuthorizationID match {
        case KafkaPlainAdmin =>
          new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "admin")
        case KafkaPlainUser =>
          new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "user")
        case _ =>
          KafkaPrincipal.ANONYMOUS
      }
    }
  }

  object Credentials {
    val allUsers = Map(KafkaPlainUser -> "user1-password",
      KafkaPlainUser2 -> KafkaPlainPassword2,
      KafkaPlainAdmin -> "broker-password")
  }

  class TestServerCallbackHandler extends AuthenticateCallbackHandler {
    def configure(configs: java.util.Map[String, _], saslMechanism: String, jaasConfigEntries: java.util.List[AppConfigurationEntry]): Unit = {}
    def handle(callbacks: Array[Callback]): Unit = {
      var username: String = null
      for (callback <- callbacks) {
        if (callback.isInstanceOf[NameCallback])
          username = callback.asInstanceOf[NameCallback].getDefaultName
        else if (callback.isInstanceOf[PlainAuthenticateCallback]) {
          val plainCallback = callback.asInstanceOf[PlainAuthenticateCallback]
          plainCallback.authenticated(Credentials.allUsers(username) == new String(plainCallback.password))
        } else
          throw new UnsupportedCallbackException(callback)
      }
    }
    def close(): Unit = {}
  }

  class TestClientCallbackHandler extends AuthenticateCallbackHandler {
    def configure(configs: java.util.Map[String, _], saslMechanism: String, jaasConfigEntries: java.util.List[AppConfigurationEntry]): Unit = {}
    def handle(callbacks: Array[Callback]): Unit = {
      val subject = Subject.getSubject(AccessController.getContext())
      val username = subject.getPublicCredentials(classOf[String]).iterator().next()
      for (callback <- callbacks) {
        if (callback.isInstanceOf[NameCallback])
          callback.asInstanceOf[NameCallback].setName(username)
        else if (callback.isInstanceOf[PasswordCallback]) {
          if (username == KafkaPlainUser || username == KafkaPlainAdmin)
            callback.asInstanceOf[PasswordCallback].setPassword(Credentials.allUsers(username).toCharArray)
        } else
          throw new UnsupportedCallbackException(callback)
      }
    }
    def close(): Unit = {}
  }
}


// This test uses SASL callback handler overrides for server connections of Kafka broker
// and client connections of Kafka producers and consumers. Client connections from Kafka brokers
// used for inter-broker communication also use custom callback handlers. The second client used in
// the multi-user test SaslEndToEndAuthorizationTest#testTwoConsumersWithDifferentSaslCredentials uses
// static JAAS configuration with default callback handlers to test those code paths as well.
class SaslPlainSslEndToEndAuthorizationTest extends SaslEndToEndAuthorizationTest {
  import SaslPlainSslEndToEndAuthorizationTest._

  this.serverConfig.setProperty(s"${listenerName.configPrefix}${KafkaConfig.SslClientAuthProp}", "required")
  this.serverConfig.setProperty(BrokerSecurityConfigs.PRINCIPAL_BUILDER_CLASS_CONFIG, classOf[TestPrincipalBuilder].getName)
  this.serverConfig.put(KafkaConfig.SaslClientCallbackHandlerClassProp, classOf[TestClientCallbackHandler].getName)
  val mechanismPrefix = listenerName.saslMechanismConfigPrefix("PLAIN")
  this.serverConfig.put(s"$mechanismPrefix${KafkaConfig.SaslServerCallbackHandlerClassProp}", classOf[TestServerCallbackHandler].getName)
  this.producerConfig.put(SaslConfigs.SASL_CLIENT_CALLBACK_HANDLER_CLASS, classOf[TestClientCallbackHandler].getName)
  this.consumerConfig.put(SaslConfigs.SASL_CLIENT_CALLBACK_HANDLER_CLASS, classOf[TestClientCallbackHandler].getName)
  this.adminClientConfig.put(SaslConfigs.SASL_CLIENT_CALLBACK_HANDLER_CLASS, classOf[TestClientCallbackHandler].getName)
  private val plainLogin = s"org.apache.kafka.common.security.plain.PlainLoginModule username=$KafkaPlainUser required;"
  this.producerConfig.put(SaslConfigs.SASL_JAAS_CONFIG, plainLogin)
  this.consumerConfig.put(SaslConfigs.SASL_JAAS_CONFIG, plainLogin)
  this.adminClientConfig.put(SaslConfigs.SASL_JAAS_CONFIG, plainLogin)

  override protected def kafkaClientSaslMechanism = "PLAIN"
  override protected def kafkaServerSaslMechanisms = List("PLAIN")

  override val clientPrincipal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "user")
  override val kafkaPrincipal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "admin")

  override def jaasSections(kafkaServerSaslMechanisms: Seq[String],
                            kafkaClientSaslMechanism: Option[String],
                            mode: SaslSetupMode,
                            kafkaServerEntryName: String): Seq[JaasSection] = {
    val brokerLogin = PlainLoginModule(KafkaPlainAdmin, "") // Password provided by callback handler
    val clientLogin = PlainLoginModule(KafkaPlainUser2, KafkaPlainPassword2)
    Seq(JaasSection(kafkaServerEntryName, Seq(brokerLogin)),
      JaasSection(KafkaClientContextName, Seq(clientLogin))) ++ zkSections
  }

  // Generate SSL certificates for clients since we are enabling TLS mutual authentication
  // in this test for the SASL_SSL listener.
  override def clientSecurityProps(certAlias: String): Properties = {
    TestUtils.securityConfigs(Mode.CLIENT, securityProtocol, trustStoreFile, certAlias, TestUtils.SslCertificateCn,
      clientSaslProperties, needsClientCert = Some(true))
  }

  /**
   * Checks that secure paths created by broker and acl paths created by AclCommand
   * have expected ACLs.
   */
  @Test
  def testAcls(): Unit = {
    TestUtils.verifySecureZkAcls(zkClient, 1)
  }
}
