
package org.apache.giraph.comm.netty;

import org.apache.commons.net.util.Base64;
import org.apache.hadoop.classification.InterfaceStability;





import org.apache.hadoop.mapreduce.security.token.JobTokenIdentifier;
import org.apache.hadoop.mapreduce.security.token.JobTokenSecretManager;
import org.apache.hadoop.security.SaslRpcServer;
import org.apache.log4j.Logger;

import javax.security.auth.callback.Callback;
import javax.security.auth.callback.CallbackHandler;
import javax.security.auth.callback.NameCallback;
import javax.security.auth.callback.PasswordCallback;
import javax.security.auth.callback.UnsupportedCallbackException;
import javax.security.sasl.AuthorizeCallback;
import javax.security.sasl.RealmCallback;
import javax.security.sasl.Sasl;
import javax.security.sasl.SaslException;
import javax.security.sasl.SaslServer;
import java.io.IOException;

/**
 * Encapsulates SASL server logic for Giraph BSP worker servers.
 */
public class SaslNettyServer extends SaslRpcServer {
  /** Logger */
  public static final Logger LOG = Logger.getLogger(SaslNettyServer.class);

  /**
   * Actual SASL work done by this object from javax.security.sasl.
   * Initialized below in constructor.
   */
  private SaslServer saslServer;

  /**
   * Constructor
   *
   * @param secretManager supplied by SaslServerHandler.
   */
  public SaslNettyServer(JobTokenSecretManager secretManager) {
    if (LOG.isDebugEnabled()) {
      LOG.debug("SaslNettyServer: Secret manager is: " + secretManager);
    }









    try {
      SaslDigestCallbackHandler ch =
          new SaslNettyServer.SaslDigestCallbackHandler(secretManager);
      saslServer = Sasl.createSaslServer(SaslNettyServer.AuthMethod.DIGEST
          .getMechanismName(), null, SaslRpcServer.SASL_DEFAULT_REALM,
          SaslRpcServer.SASL_PROPS, ch);
    } catch (SaslException e) {
      LOG.error("SaslNettyServer: Could not create SaslServer: " + e);
    }
  }

  public boolean isComplete() {
    return saslServer.isComplete();
  }

  public String getUserName() {
    return saslServer.getAuthorizationID();
  }

  /**
   * Used by SaslTokenMessage::processToken() to respond to server SASL tokens.
   *
   * @param token Server's SASL token
   * @return token to send back to the server.
   */
  public byte[] response(byte[] token) {
    try {
      if (LOG.isDebugEnabled()) {
        LOG.debug("response: Responding to input token of length: " +
            token.length);
      }
      byte[] retval = saslServer.evaluateResponse(token);
      if (LOG.isDebugEnabled()) {
        LOG.debug("response: Response token length: " + retval.length);
      }
      return retval;
    } catch (SaslException e) {
      LOG.error("response: Failed to evaluate client token of length: " +
          token.length + " : " + e);
      return null;
    }
  }

  /**
   * Encode a byte[] identifier as a Base64-encoded string.
   *
   * @param identifier identifier to encode
   * @return Base64-encoded string
   */
  static String encodeIdentifier(byte[] identifier) {
    return new String(Base64.encodeBase64(identifier));
  }

  /**
   * Encode a password as a base64-encoded char[] array.
   * @param password as a byte array.
   * @return password as a char array.
   */
  static char[] encodePassword(byte[] password) {
    return new String(Base64.encodeBase64(password)).toCharArray();
  }

  /** CallbackHandler for SASL DIGEST-MD5 mechanism */
  @InterfaceStability.Evolving
  public static class SaslDigestCallbackHandler implements CallbackHandler {
    /** Used to authenticate the clients */
    private JobTokenSecretManager secretManager;

    /**
     * Constructor
     *
     * @param secretManager used to authenticate clients
     */
    public SaslDigestCallbackHandler(
        JobTokenSecretManager secretManager) {
      if (LOG.isDebugEnabled()) {
        LOG.debug("SaslDigestCallback: Creating SaslDigestCallback handler " +
            "with secret manager: " + secretManager);
      }
      this.secretManager = secretManager;
    }

    /** {@inheritDoc} */
    @Override
    public void handle(Callback[] callbacks) throws IOException,
        UnsupportedCallbackException {
      NameCallback nc = null;
      PasswordCallback pc = null;
      AuthorizeCallback ac = null;
      for (Callback callback : callbacks) {
        if (callback instanceof AuthorizeCallback) {
          ac = (AuthorizeCallback) callback;
        } else if (callback instanceof NameCallback) {
          nc = (NameCallback) callback;
        } else if (callback instanceof PasswordCallback) {
          pc = (PasswordCallback) callback;
        } else if (callback instanceof RealmCallback) {
          continue; // realm is ignored
        } else {
          throw new UnsupportedCallbackException(callback,
              "handle: Unrecognized SASL DIGEST-MD5 Callback");
        }
      }
      if (pc != null) {
        JobTokenIdentifier tokenIdentifier = getIdentifier(nc.getDefaultName(),
            secretManager);
        char[] password =
          encodePassword(secretManager.retrievePassword(tokenIdentifier));

        if (LOG.isDebugEnabled()) {
          LOG.debug("handle: SASL server DIGEST-MD5 callback: setting " +
              "password for client: " + tokenIdentifier.getUser());
        }
        pc.setPassword(password);
      }
      if (ac != null) {
        String authid = ac.getAuthenticationID();
        String authzid = ac.getAuthorizationID();
        if (authid.equals(authzid)) {
          ac.setAuthorized(true);
        } else {
          ac.setAuthorized(false);
        }
        if (ac.isAuthorized()) {
          if (LOG.isDebugEnabled()) {
            String username =
              getIdentifier(authzid, secretManager).getUser().getUserName();
            if (LOG.isDebugEnabled()) {
              LOG.debug("handle: SASL server DIGEST-MD5 callback: setting " +
                  "canonicalized client ID: " + username);
            }
          }
          ac.setAuthorizedID(authzid);
        }
      }
    }
  }
}
