/*
 * Copyright DataStax, Inc.
 *
 * 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 com.datastax.dse.driver.api.core.auth;

import com.datastax.oss.driver.api.core.auth.AuthProvider;
import com.datastax.oss.driver.api.core.auth.AuthenticationException;
import com.datastax.oss.driver.api.core.auth.Authenticator;
import com.datastax.oss.driver.api.core.metadata.EndPoint;
import com.datastax.oss.driver.api.core.session.Session;
import com.datastax.oss.driver.shaded.guava.common.base.Charsets;
import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableMap;
import com.datastax.oss.protocol.internal.util.Bytes;
import edu.umd.cs.findbugs.annotations.NonNull;
import edu.umd.cs.findbugs.annotations.Nullable;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.security.PrivilegedActionException;
import java.security.PrivilegedExceptionAction;
import java.util.HashMap;
import java.util.Map;
import java.util.Objects;
import javax.security.auth.Subject;
import javax.security.auth.login.AppConfigurationEntry;
import javax.security.auth.login.Configuration;
import javax.security.auth.login.LoginContext;
import javax.security.auth.login.LoginException;
import javax.security.sasl.Sasl;
import javax.security.sasl.SaslClient;
import javax.security.sasl.SaslException;
import net.jcip.annotations.Immutable;
import net.jcip.annotations.NotThreadSafe;
import net.jcip.annotations.ThreadSafe;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

@ThreadSafe
public abstract class DseGssApiAuthProviderBase implements AuthProvider {

  
The default SASL service name used by this auth provider. /** The default SASL service name used by this auth provider. */
  public static final String DEFAULT_SASL_SERVICE_NAME = "dse";

  
The name of the system property to use to specify the SASL service name. /** The name of the system property to use to specify the SASL service name. */
  public static final String SASL_SERVICE_NAME_PROPERTY = "dse.sasl.service";

  
Legacy system property for SASL protocol name. Clients should migrate to
SASL_SERVICE_NAME_PROPERTY above.
/**
   * Legacy system property for SASL protocol name. Clients should migrate to
   * SASL_SERVICE_NAME_PROPERTY above.
   */
  private static final String LEGACY_SASL_PROTOCOL_PROPERTY = "dse.sasl.protocol";

  private static final Logger LOG = LoggerFactory.getLogger(DseGssApiAuthProviderBase.class);

  private final String logPrefix;

  
Params:
logPrefix – a string that will get prepended to the logs (this is used for discrimination when you have multiple driver instances executing in the same JVM). Config-based implementations fill this with Session.getName()./**
   * @param logPrefix a string that will get prepended to the logs (this is used for discrimination
   *     when you have multiple driver instances executing in the same JVM). Config-based
   *     implementations fill this with {@link Session#getName()}.
   */
  protected DseGssApiAuthProviderBase(@NonNull String logPrefix) {
    this.logPrefix = Objects.requireNonNull(logPrefix);
  }

  @NonNull
  protected abstract GssApiOptions getOptions(
      @NonNull EndPoint endPoint, @NonNull String serverAuthenticator);

  @NonNull
  @Override
  public Authenticator newAuthenticator(
      @NonNull EndPoint endPoint, @NonNull String serverAuthenticator)
      throws AuthenticationException {
    return new GssApiAuthenticator(
        getOptions(endPoint, serverAuthenticator), endPoint, serverAuthenticator);
  }

  @Override
  public void onMissingChallenge(@NonNull EndPoint endPoint) {
    LOG.warn(
        "[{}] {} did not send an authentication challenge; "
            + "This is suspicious because the driver expects authentication",
        logPrefix,
        endPoint);
  }

  @Override
  public void close() {
    // nothing to do
  }

  
The options to initialize a new authenticator.
Use builder() to create an instance. /**
   * The options to initialize a new authenticator.
   *
   * <p>Use {@link #builder()} to create an instance.
   */
  @Immutable
  public static class GssApiOptions {

    @NonNull
    public static Builder builder() {
      return new Builder();
    }

    private final Configuration loginConfiguration;
    private final Subject subject;
    private final String saslProtocol;
    private final String authorizationId;
    private final Map<String, String> saslProperties;

    private GssApiOptions(
        @Nullable Configuration loginConfiguration,
        @Nullable Subject subject,
        @Nullable String saslProtocol,
        @Nullable String authorizationId,
        @NonNull Map<String, String> saslProperties) {
      this.loginConfiguration = loginConfiguration;
      this.subject = subject;
      this.saslProtocol = saslProtocol;
      this.authorizationId = authorizationId;
      this.saslProperties = saslProperties;
    }

    @Nullable
    public Configuration getLoginConfiguration() {
      return loginConfiguration;
    }

    @Nullable
    public Subject getSubject() {
      return subject;
    }

    @Nullable
    public String getSaslProtocol() {
      return saslProtocol;
    }

    @Nullable
    public String getAuthorizationId() {
      return authorizationId;
    }

    @NonNull
    public Map<String, String> getSaslProperties() {
      return saslProperties;
    }

    @NotThreadSafe
    public static class Builder {

      private Configuration loginConfiguration;
      private Subject subject;
      private String saslProtocol;
      private String authorizationId;
      private final Map<String, String> saslProperties = new HashMap<>();

      public Builder() {
        saslProperties.put(Sasl.SERVER_AUTH, "true");
        saslProperties.put(Sasl.QOP, "auth");
      }

      
Sets a login configuration that will be used to create a LoginContext. 
You MUST call either a withLoginConfiguration method or withSubject(Subject); if both are called, the subject takes precedence, and the login configuration will be ignored. 
See Also:
withLoginConfiguration(Map<String,String>)/**
       * Sets a login configuration that will be used to create a {@link LoginContext}.
       *
       * <p>You MUST call either a withLoginConfiguration method or {@link #withSubject(Subject)};
       * if both are called, the subject takes precedence, and the login configuration will be
       * ignored.
       *
       * @see #withLoginConfiguration(Map)
       */
      @NonNull
      public Builder withLoginConfiguration(@Nullable Configuration loginConfiguration) {
        this.loginConfiguration = loginConfiguration;
        return this;
      }
      
Sets a login configuration that will be used to create a LoginContext. 
This is an alternative to withLoginConfiguration(Configuration), that builds the configuration from Krb5LoginModule with the given options. 
You MUST call either a withLoginConfiguration method or withSubject(Subject); if both are called, the subject takes precedence, and the login configuration will be ignored. /**
       * Sets a login configuration that will be used to create a {@link LoginContext}.
       *
       * <p>This is an alternative to {@link #withLoginConfiguration(Configuration)}, that builds
       * the configuration from {@code Krb5LoginModule} with the given options.
       *
       * <p>You MUST call either a withLoginConfiguration method or {@link #withSubject(Subject)};
       * if both are called, the subject takes precedence, and the login configuration will be
       * ignored.
       */
      @NonNull
      public Builder withLoginConfiguration(@Nullable Map<String, String> loginConfiguration) {
        this.loginConfiguration = fetchLoginConfiguration(loginConfiguration);
        return this;
      }

      
Sets a previously authenticated subject to reuse.
You MUST call either this method or withLoginConfiguration(Configuration); if both are called, the subject takes precedence, and the login configuration will be ignored. /**
       * Sets a previously authenticated subject to reuse.
       *
       * <p>You MUST call either this method or {@link #withLoginConfiguration(Configuration)}; if
       * both are called, the subject takes precedence, and the login configuration will be ignored.
       */
      @NonNull
      public Builder withSubject(@Nullable Subject subject) {
        this.subject = subject;
        return this;
      }

      
Sets the SASL protocol name to use; should match the username of the Kerberos service
principal used by the DSE server.
/**
       * Sets the SASL protocol name to use; should match the username of the Kerberos service
       * principal used by the DSE server.
       */
      @NonNull
      public Builder withSaslProtocol(@Nullable String saslProtocol) {
        this.saslProtocol = saslProtocol;
        return this;
      }

      
Sets the authorization ID (allows proxy authentication). /** Sets the authorization ID (allows proxy authentication). */
      @NonNull
      public Builder withAuthorizationId(@Nullable String authorizationId) {
        this.authorizationId = authorizationId;
        return this;
      }

      
Add a SASL property to use when creating the SASL client.
Note that this builder pre-initializes these two default properties:
javax.security.sasl.server.authentication = true
javax.security.sasl.qop = auth

/**
       * Add a SASL property to use when creating the SASL client.
       *
       * <p>Note that this builder pre-initializes these two default properties:
       *
       * <pre>
       * javax.security.sasl.server.authentication = true
       * javax.security.sasl.qop = auth
       * </pre>
       */
      @NonNull
      public Builder addSaslProperty(@NonNull String name, @NonNull String value) {
        this.saslProperties.put(Objects.requireNonNull(name), Objects.requireNonNull(value));
        return this;
      }

      @NonNull
      public GssApiOptions build() {
        return new GssApiOptions(
            loginConfiguration,
            subject,
            saslProtocol,
            authorizationId,
            ImmutableMap.copyOf(saslProperties));
      }

      public static Configuration fetchLoginConfiguration(Map<String, String> options) {
        return new Configuration() {

          @Override
          public AppConfigurationEntry[] getAppConfigurationEntry(String name) {
            return new AppConfigurationEntry[] {
              new AppConfigurationEntry(
                  "com.sun.security.auth.module.Krb5LoginModule",
                  AppConfigurationEntry.LoginModuleControlFlag.REQUIRED,
                  options)
            };
          }
        };
      }
    }
  }

  protected static class GssApiAuthenticator extends BaseDseAuthenticator {

    private static final ByteBuffer MECHANISM =
        ByteBuffer.wrap("GSSAPI".getBytes(Charsets.UTF_8)).asReadOnlyBuffer();
    private static final ByteBuffer SERVER_INITIAL_CHALLENGE =
        ByteBuffer.wrap("GSSAPI-START".getBytes(Charsets.UTF_8)).asReadOnlyBuffer();
    private static final ByteBuffer EMPTY_BYTE_ARRAY =
        ByteBuffer.wrap(new byte[0]).asReadOnlyBuffer();
    private static final String JAAS_CONFIG_ENTRY = "DseClient";
    private static final String[] SUPPORTED_MECHANISMS = new String[] {"GSSAPI"};

    private Subject subject;
    private SaslClient saslClient;
    private EndPoint endPoint;

    protected GssApiAuthenticator(
        GssApiOptions options, EndPoint endPoint, String serverAuthenticator) {
      super(serverAuthenticator);

      try {
        if (options.getSubject() != null) {
          this.subject = options.getSubject();
        } else {
          Configuration loginConfiguration = options.getLoginConfiguration();
          if (loginConfiguration == null) {
            throw new IllegalArgumentException("Must provide one of subject or loginConfiguration");
          }
          LoginContext login = new LoginContext(JAAS_CONFIG_ENTRY, null, null, loginConfiguration);
          login.login();
          this.subject = login.getSubject();
        }
        String protocol = options.getSaslProtocol();
        if (protocol == null) {
          protocol =
              System.getProperty(
                  SASL_SERVICE_NAME_PROPERTY,
                  System.getProperty(LEGACY_SASL_PROTOCOL_PROPERTY, DEFAULT_SASL_SERVICE_NAME));
        }
        this.saslClient =
            Sasl.createSaslClient(
                SUPPORTED_MECHANISMS,
                options.getAuthorizationId(),
                protocol,
                ((InetSocketAddress) endPoint.resolve()).getAddress().getCanonicalHostName(),
                options.getSaslProperties(),
                null);
      } catch (LoginException | SaslException e) {
        throw new AuthenticationException(endPoint, e.getMessage());
      }
      this.endPoint = endPoint;
    }

    @NonNull
    @Override
    protected ByteBuffer getMechanism() {
      return MECHANISM;
    }

    @NonNull
    @Override
    protected ByteBuffer getInitialServerChallenge() {
      return SERVER_INITIAL_CHALLENGE;
    }

    @Nullable
    @Override
    public ByteBuffer evaluateChallengeSync(@Nullable ByteBuffer challenge) {

      byte[] challengeBytes;
      if (SERVER_INITIAL_CHALLENGE.equals(challenge)) {
        if (!saslClient.hasInitialResponse()) {
          return EMPTY_BYTE_ARRAY;
        }
        challengeBytes = new byte[0];
      } else {
        // The native protocol spec says the incoming challenge can be null depending on the
        // implementation. But saslClient.evaluateChallenge clearly documents that the byte array
        // can't be null, which probably means that a SASL authenticator never sends back null.
        if (challenge == null) {
          throw new AuthenticationException(this.endPoint, "Unexpected null challenge from server");
        }
        challengeBytes = Bytes.getArray(challenge);
      }
      try {

        return ByteBuffer.wrap(
            Subject.doAs(
                subject,
                new PrivilegedExceptionAction<byte[]>() {
                  @Override
                  public byte[] run() throws SaslException {
                    return saslClient.evaluateChallenge(challengeBytes);
                  }
                }));
      } catch (PrivilegedActionException e) {
        throw new AuthenticationException(this.endPoint, e.getMessage(), e.getException());
      }
    }
  }
}