/*
    * 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.hadoop.hbase.security;
  
  import java.io.ByteArrayInputStream;
  import java.io.DataInputStream;
  import java.io.IOException;
  import java.util.Locale;
  import java.util.Map;
  
  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 org.apache.commons.logging.Log;
  import org.apache.commons.logging.LogFactory;
  import org.apache.hadoop.conf.Configuration;
  import org.apache.hadoop.hbase.classification.InterfaceAudience;
  import org.apache.hadoop.hbase.ipc.RpcServer;
  import org.apache.hadoop.hbase.security.SaslUtil.QualityOfProtection;
  import org.apache.hadoop.security.UserGroupInformation;
  import org.apache.hadoop.security.token.SecretManager;
  import org.apache.hadoop.security.token.SecretManager.InvalidToken;
  import org.apache.hadoop.security.token.TokenIdentifier;
  
  /**
   * A utility class for dealing with SASL on RPC server
   */
  @InterfaceAudience.Private
  public class HBaseSaslRpcServer {
    private static final Log LOG = LogFactory.getLog(HBaseSaslRpcServer.class);
  
    private static Map<String, String> saslProps = null;
  
    public static void init(Configuration conf) {
      saslProps = SaslUtil.initSaslProperties(conf.get("hbase.rpc.protection",
            QualityOfProtection.AUTHENTICATION.name().toLowerCase(Locale.ROOT)));
    }
  
    public static Map<String, String> getSaslProps() {
      return saslProps;
    }
  
    public static <T extends TokenIdentifier> T getIdentifier(String id,
        SecretManager<T> secretManager) throws InvalidToken {
      byte[] tokenId = SaslUtil.decodeIdentifier(id);
      T tokenIdentifier = secretManager.createIdentifier();
      try {
        tokenIdentifier.readFields(new DataInputStream(new ByteArrayInputStream(
            tokenId)));
      } catch (IOException e) {
        throw (InvalidToken) new InvalidToken(
            "Can't de-serialize tokenIdentifier").initCause(e);
      }
      return tokenIdentifier;
    }
  
  
    /** CallbackHandler for SASL DIGEST-MD5 mechanism */
    public static class SaslDigestCallbackHandler implements CallbackHandler {
      private SecretManager<TokenIdentifier> secretManager;
      private RpcServer.Connection connection;
  
      public SaslDigestCallbackHandler(
          SecretManager<TokenIdentifier> secretManager,
          RpcServer.Connection connection) {
        this.secretManager = secretManager;
        this.connection = connection;
      }
  
      private char[] getPassword(TokenIdentifier tokenid) throws InvalidToken {
        return SaslUtil.encodePassword(secretManager.retrievePassword(tokenid));
      }
  
      /** {@inheritDoc} */
      @Override
      public void handle(Callback[] callbacks) throws InvalidToken,
          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,
               "Unrecognized SASL DIGEST-MD5 Callback");
         }
       }
       if (pc != null) {
         TokenIdentifier tokenIdentifier = getIdentifier(nc.getDefaultName(), secretManager);
         char[] password = getPassword(tokenIdentifier);
         UserGroupInformation user = null;
         user = tokenIdentifier.getUser(); // may throw exception
         connection.attemptingUser = user;
         if (LOG.isTraceEnabled()) {
           LOG.trace("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.isTraceEnabled()) {
             String username =
               getIdentifier(authzid, secretManager).getUser().getUserName();
             LOG.trace("SASL server DIGEST-MD5 callback: setting "
                 + "canonicalized client ID: " + username);
           }
           ac.setAuthorizedID(authzid);
         }
       }
     }
   }
 
   /** CallbackHandler for SASL GSSAPI Kerberos mechanism */
   public static class SaslGssCallbackHandler implements CallbackHandler {
 
     /** {@inheritDoc} */
     @Override
     public void handle(Callback[] callbacks) throws
         UnsupportedCallbackException {
       AuthorizeCallback ac = null;
       for (Callback callback : callbacks) {
         if (callback instanceof AuthorizeCallback) {
           ac = (AuthorizeCallback) callback;
         } else {
           throw new UnsupportedCallbackException(callback,
               "Unrecognized SASL GSSAPI Callback");
         }
       }
       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())
             LOG.debug("SASL server GSSAPI callback: setting "
                 + "canonicalized client ID: " + authzid);
           ac.setAuthorizedID(authzid);
         }
       }
     }
   }
 }