/*
    * 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 org.apache.giraph.comm.netty;
  
  import org.apache.giraph.comm.requests.SaslTokenMessageRequest;
  import org.apache.hadoop.conf.Configuration;
  import org.apache.hadoop.mapred.JobConf;
  import org.apache.hadoop.mapreduce.security.TokenCache;
  import org.apache.hadoop.mapreduce.security.token.JobTokenIdentifier;
  import org.apache.hadoop.security.Credentials;
  /*if_not[STATIC_SASL_SYMBOL]*/
  import org.apache.hadoop.security.SaslPropertiesResolver;
  /*end[STATIC_SASL_SYMBOL]*/
  import org.apache.hadoop.security.SaslRpcServer;
  import org.apache.hadoop.security.SaslRpcServer.AuthMethod;
  import org.apache.hadoop.security.token.Token;
  import org.apache.hadoop.security.token.TokenIdentifier;
  import org.apache.hadoop.security.UserGroupInformation;
  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.RealmCallback;
  import javax.security.sasl.RealmChoiceCallback;
  import javax.security.sasl.Sasl;
  import javax.security.sasl.SaslClient;
  import javax.security.sasl.SaslException;
  import java.io.IOException;
  
  /**
   * Implements SASL logic for Giraph BSP worker clients.
   */
  public class SaslNettyClient {
    /** Class logger */
    public static final Logger LOG = Logger.getLogger(SaslNettyClient.class);
  
    /**
     * Used to synchronize client requests: client's work-related requests must
     * wait until SASL authentication completes.
     */
    private Object authenticated = new Object();
  
    /**
     * Used to respond to server's counterpart, SaslServer with SASL tokens
     * represented as byte arrays.
     */
    private SaslClient saslClient;
  
    /**
     * Create a SaslNettyClient for authentication with BSP servers.
     */
    public SaslNettyClient() {
      try {
        Token<? extends TokenIdentifier> token =
            createJobToken(new Configuration());
        if (LOG.isDebugEnabled()) {
          LOG.debug("SaslNettyClient: Creating SASL " +
              AuthMethod.DIGEST.getMechanismName() +
              " client to authenticate to service at " + token.getService());
        }
  /*if[STATIC_SASL_SYMBOL]
        saslClient =
            Sasl.createSaslClient(
                new String[] { AuthMethod.DIGEST.getMechanismName() }, null,
                null, SaslRpcServer.SASL_DEFAULT_REALM, SaslRpcServer.SASL_PROPS,
                new SaslClientCallbackHandler(token));
  else[STATIC_SASL_SYMBOL]*/
        SaslPropertiesResolver saslPropsResolver =
            SaslPropertiesResolver.getInstance(new Configuration());
        saslClient =
            Sasl.createSaslClient(
                new String[] { AuthMethod.DIGEST.getMechanismName() }, null,
                null, SaslRpcServer.SASL_DEFAULT_REALM,
                saslPropsResolver.getDefaultProperties(),
                new SaslClientCallbackHandler(token));
  /*end[STATIC_SASL_SYMBOL]*/
      } catch (IOException e) {
        LOG.error("SaslNettyClient: Could not obtain job token for Netty " +
            "Client to use to authenticate with a Netty Server.");
        saslClient = null;
     }
   }
 
   public Object getAuthenticated() {
     return authenticated;
   }
 
   /**
    * Obtain JobToken, which we'll use as a credential for SASL authentication
    * when connecting to other Giraph BSPWorkers.
    *
    * @param conf Configuration
    * @return a JobToken containing username and password so that client can
    * authenticate with a server.
    */
   private Token<JobTokenIdentifier> createJobToken(Configuration conf)
     throws IOException {
     String localJobTokenFile = System.getenv().get(
         UserGroupInformation.HADOOP_TOKEN_FILE_LOCATION);
     if (localJobTokenFile != null) {
       JobConf jobConf = new JobConf(conf);
       Credentials credentials =
           TokenCache.loadTokens(localJobTokenFile, jobConf);
       return TokenCache.getJobToken(credentials);
     } else {
       throw new IOException("createJobToken: Cannot obtain authentication " +
           "credentials for job: file: '" +
           UserGroupInformation.HADOOP_TOKEN_FILE_LOCATION + "' not found");
     }
   }
 
   /**
    * Used by authenticateOnChannel() to initiate SASL handshake with server.
    * @return SaslTokenMessageRequest message to be sent to server.
    * @throws IOException
    */
   public SaslTokenMessageRequest firstToken()
     throws IOException {
     byte[] saslToken = new byte[0];
     if (saslClient.hasInitialResponse()) {
       saslToken = saslClient.evaluateChallenge(saslToken);
     }
     SaslTokenMessageRequest saslTokenMessage =
         new SaslTokenMessageRequest();
     saslTokenMessage.setSaslToken(saslToken);
     return saslTokenMessage;
   }
 
   public boolean isComplete() {
     return saslClient.isComplete();
   }
 
   /**
    * Respond to server's SASL token.
    * @param saslTokenMessage contains server's SASL token
    * @return client's response SASL token
    */
   public byte[] saslResponse(SaslTokenMessageRequest saslTokenMessage) {
     try {
       byte[] retval =
           saslClient.evaluateChallenge(saslTokenMessage.getSaslToken());
       return retval;
     } catch (SaslException e) {
       LOG.error("saslResponse: Failed to respond to SASL server's token:", e);
       return null;
     }
   }
 
   /**
    * Implementation of javax.security.auth.callback.CallbackHandler
    * that works with Hadoop JobTokens.
    */
   private static class SaslClientCallbackHandler implements CallbackHandler {
     /** Generated username contained in JobToken */
     private final String userName;
     /** Generated password contained in JobToken */
     private final char[] userPassword;
 
     /**
      * Set private members using token.
      * @param token Hadoop JobToken.
      */
     public SaslClientCallbackHandler(Token<? extends TokenIdentifier> token) {
       this.userName = SaslNettyServer.encodeIdentifier(token.getIdentifier());
       this.userPassword = SaslNettyServer.encodePassword(token.getPassword());
     }
 
     /**
      * Implementation used to respond to SASL tokens from server.
      *
      * @param callbacks objects that indicate what credential information the
      *                  server's SaslServer requires from the client.
      * @throws UnsupportedCallbackException
      */
     public void handle(Callback[] callbacks)
       throws UnsupportedCallbackException {
       NameCallback nc = null;
       PasswordCallback pc = null;
       RealmCallback rc = null;
       for (Callback callback : callbacks) {
         if (callback instanceof RealmChoiceCallback) {
           continue;
         } else if (callback instanceof NameCallback) {
           nc = (NameCallback) callback;
         } else if (callback instanceof PasswordCallback) {
           pc = (PasswordCallback) callback;
         } else if (callback instanceof RealmCallback) {
           rc = (RealmCallback) callback;
         } else {
           throw new UnsupportedCallbackException(callback,
               "handle: Unrecognized SASL client callback");
         }
       }
       if (nc != null) {
         if (LOG.isDebugEnabled()) {
           LOG.debug("handle: SASL client callback: setting username: " +
               userName);
         }
         nc.setName(userName);
       }
       if (pc != null) {
         if (LOG.isDebugEnabled()) {
           LOG.debug("handle: SASL client callback: setting userPassword");
         }
         pc.setPassword(userPassword);
       }
       if (rc != null) {
         if (LOG.isDebugEnabled()) {
           LOG.debug("handle: SASL client callback: setting realm: " +
               rc.getDefaultText());
         }
         rc.setText(rc.getDefaultText());
       }
     }
   }
 }