/**
 * 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.hdfs.protocol.datatransfer.sasl;

import static org.apache.hadoop.fs.CommonConfigurationKeysPublic.HADOOP_RPC_PROTECTION;
import static org.apache.hadoop.fs.CommonConfigurationKeysPublic.HADOOP_SECURITY_SASL_PROPS_RESOLVER_CLASS;
import static org.apache.hadoop.hdfs.DFSConfigKeys.DFS_DATA_TRANSFER_PROTECTION_KEY;
import static org.apache.hadoop.hdfs.DFSConfigKeys.DFS_DATA_TRANSFER_SASL_PROPS_RESOLVER_CLASS_KEY;
import static org.apache.hadoop.hdfs.DFSConfigKeys.DFS_ENCRYPT_DATA_TRANSFER_CIPHER_KEY_BITLENGTH_KEY;
import static org.apache.hadoop.hdfs.DFSConfigKeys.DFS_ENCRYPT_DATA_TRANSFER_CIPHER_KEY_BITLENGTH_DEFAULT;
import static org.apache.hadoop.hdfs.DFSConfigKeys.DFS_ENCRYPT_DATA_TRANSFER_CIPHER_SUITES_KEY;
import static org.apache.hadoop.hdfs.protocolPB.PBHelper.vintPrefixed;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.InetAddress;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.Set;
import javax.security.sasl.Sasl;

import org.apache.commons.codec.binary.Base64;
import org.apache.hadoop.classification.InterfaceAudience;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.crypto.CipherOption;
import org.apache.hadoop.crypto.CipherSuite;
import org.apache.hadoop.crypto.CryptoCodec;
import org.apache.hadoop.crypto.CryptoInputStream;
import org.apache.hadoop.crypto.CryptoOutputStream;
import org.apache.hadoop.hdfs.net.Peer;
import org.apache.hadoop.hdfs.protocol.datatransfer.IOStreamPair;
import org.apache.hadoop.hdfs.protocol.datatransfer.InvalidEncryptionKeyException;
import org.apache.hadoop.hdfs.protocol.proto.DataTransferProtos.DataTransferEncryptorMessageProto;
import org.apache.hadoop.hdfs.protocol.proto.DataTransferProtos.DataTransferEncryptorMessageProto.DataTransferEncryptorStatus;
import org.apache.hadoop.hdfs.protocol.proto.HdfsProtos.CipherOptionProto;
import org.apache.hadoop.hdfs.protocolPB.PBHelper;
import org.apache.hadoop.security.SaslPropertiesResolver;
import org.apache.hadoop.security.SaslRpcServer.QualityOfProtection;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.common.base.Charsets;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Maps;
import com.google.common.net.InetAddresses;
import com.google.protobuf.ByteString;

/**
 * Utility methods implementing SASL negotiation for DataTransferProtocol.
 */
@InterfaceAudience.Private
public final class DataTransferSaslUtil {

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

  /**
   * Delimiter for the three-part SASL username string.
   */
  public static final String NAME_DELIMITER = " ";

  /**
   * Sent by clients and validated by servers. We use a number that's unlikely
   * to ever be sent as the value of the DATA_TRANSFER_VERSION.
   */
  public static final int SASL_TRANSFER_MAGIC_NUMBER = 0xDEADBEEF;

  /**
   * Checks that SASL negotiation has completed for the given participant, and
   * the negotiated quality of protection is included in the given SASL
   * properties and therefore acceptable.
   *
   * @param sasl participant to check
   * @param saslProps properties of SASL negotiation
   * @throws IOException for any error
   */
  public static void checkSaslComplete(SaslParticipant sasl,
      Map<String, String> saslProps) throws IOException {
    if (!sasl.isComplete()) {
      throw new IOException("Failed to complete SASL handshake");
    }
    Set<String> requestedQop = ImmutableSet.copyOf(Arrays.asList(
      saslProps.get(Sasl.QOP).split(",")));
    String negotiatedQop = sasl.getNegotiatedQop();
    LOG.debug("Verifying QOP, requested QOP = {}, negotiated QOP = {}",
      requestedQop, negotiatedQop);
    if (!requestedQop.contains(negotiatedQop)) {
      throw new IOException(String.format("SASL handshake completed, but " +
        "channel does not have acceptable quality of protection, " +
        "requested = %s, negotiated = %s", requestedQop, negotiatedQop));
    }
  }
  
  /**
   * Check whether requested SASL Qop contains privacy.
   * 
   * @param saslProps properties of SASL negotiation
   * @return boolean true if privacy exists
   */
  public static boolean requestedQopContainsPrivacy(
      Map<String, String> saslProps) {
    Set<String> requestedQop = ImmutableSet.copyOf(Arrays.asList(
        saslProps.get(Sasl.QOP).split(",")));
    return requestedQop.contains("auth-conf");
  }

  /**
   * Creates SASL properties required for an encrypted SASL negotiation.
   *
   * @param encryptionAlgorithm to use for SASL negotation
   * @return properties of encrypted SASL negotiation
   */
  public static Map<String, String> createSaslPropertiesForEncryption(
      String encryptionAlgorithm) {
    Map<String, String> saslProps = Maps.newHashMapWithExpectedSize(3);
    saslProps.put(Sasl.QOP, QualityOfProtection.PRIVACY.getSaslQop());
    saslProps.put(Sasl.SERVER_AUTH, "true");
    saslProps.put("com.sun.security.sasl.digest.cipher", encryptionAlgorithm);
    return saslProps;
  }

  /**
   * For an encrypted SASL negotiation, encodes an encryption key to a SASL
   * password.
   *
   * @param encryptionKey to encode
   * @return key encoded as SASL password
   */
  public static char[] encryptionKeyToPassword(byte[] encryptionKey) {
    return new String(Base64.encodeBase64(encryptionKey, false), Charsets.UTF_8)
      .toCharArray();
  }

  /**
   * Returns InetAddress from peer.  The getRemoteAddressString has the form
   * [host][/ip-address]:port.  The host may be missing.  The IP address (and
   * preceding '/') may be missing.  The port preceded by ':' is always present.
   *
   * @param peer
   * @return InetAddress from peer
   */
  public static InetAddress getPeerAddress(Peer peer) {
    String remoteAddr = peer.getRemoteAddressString().split(":")[0];
    int slashIdx = remoteAddr.indexOf('/');
    return InetAddresses.forString(slashIdx != -1 ?
        remoteAddr.substring(slashIdx + 1, remoteAddr.length()) :
        remoteAddr);
  }

  /**
   * Creates a SaslPropertiesResolver from the given configuration.  This method
   * works by cloning the configuration, translating configuration properties
   * specific to DataTransferProtocol to what SaslPropertiesResolver expects,
   * and then delegating to SaslPropertiesResolver for initialization.  This
   * method returns null if SASL protection has not been configured for
   * DataTransferProtocol.
   *
   * @param conf configuration to read
   * @return SaslPropertiesResolver for DataTransferProtocol, or null if not
   *   configured
   */
  public static SaslPropertiesResolver getSaslPropertiesResolver(
      Configuration conf) {
    String qops = conf.get(DFS_DATA_TRANSFER_PROTECTION_KEY);
    if (qops == null || qops.isEmpty()) {
      LOG.debug("DataTransferProtocol not using SaslPropertiesResolver, no " +
        "QOP found in configuration for {}", DFS_DATA_TRANSFER_PROTECTION_KEY);
      return null;
    }
    Configuration saslPropsResolverConf = new Configuration(conf);
    saslPropsResolverConf.set(HADOOP_RPC_PROTECTION, qops);
    Class<? extends SaslPropertiesResolver> resolverClass = conf.getClass(
      HADOOP_SECURITY_SASL_PROPS_RESOLVER_CLASS,
      SaslPropertiesResolver.class, SaslPropertiesResolver.class);
    resolverClass = conf.getClass(DFS_DATA_TRANSFER_SASL_PROPS_RESOLVER_CLASS_KEY,
      resolverClass, SaslPropertiesResolver.class);
    saslPropsResolverConf.setClass(HADOOP_SECURITY_SASL_PROPS_RESOLVER_CLASS,
      resolverClass, SaslPropertiesResolver.class);
    SaslPropertiesResolver resolver = SaslPropertiesResolver.getInstance(
      saslPropsResolverConf);
    LOG.debug("DataTransferProtocol using SaslPropertiesResolver, configured " +
      "QOP {} = {}, configured class {} = {}", DFS_DATA_TRANSFER_PROTECTION_KEY, qops, 
      DFS_DATA_TRANSFER_SASL_PROPS_RESOLVER_CLASS_KEY, resolverClass);
    return resolver;
  }

  /**
   * Reads a SASL negotiation message.
   *
   * @param in stream to read
   * @return bytes of SASL negotiation messsage
   * @throws IOException for any error
   */
  public static byte[] readSaslMessage(InputStream in) throws IOException {
    DataTransferEncryptorMessageProto proto =
        DataTransferEncryptorMessageProto.parseFrom(vintPrefixed(in));
    if (proto.getStatus() == DataTransferEncryptorStatus.ERROR_UNKNOWN_KEY) {
      throw new InvalidEncryptionKeyException(proto.getMessage());
    } else if (proto.getStatus() == DataTransferEncryptorStatus.ERROR) {
      throw new IOException(proto.getMessage());
    } else {
      return proto.getPayload().toByteArray();
    }
  }
  
  /**
   * Reads a SASL negotiation message and negotiation cipher options. 
   * 
   * @param in stream to read
   * @param cipherOptions list to store negotiation cipher options
   * @return byte[] SASL negotiation message
   * @throws IOException for any error
   */
  public static byte[] readSaslMessageAndNegotiationCipherOptions(
      InputStream in, List<CipherOption> cipherOptions) throws IOException {
    DataTransferEncryptorMessageProto proto =
        DataTransferEncryptorMessageProto.parseFrom(vintPrefixed(in));
    if (proto.getStatus() == DataTransferEncryptorStatus.ERROR_UNKNOWN_KEY) {
      throw new InvalidEncryptionKeyException(proto.getMessage());
    } else if (proto.getStatus() == DataTransferEncryptorStatus.ERROR) {
      throw new IOException(proto.getMessage());
    } else {
      List<CipherOptionProto> optionProtos = proto.getCipherOptionList();
      if (optionProtos != null) {
        for (CipherOptionProto optionProto : optionProtos) {
          cipherOptions.add(PBHelper.convert(optionProto));
        }
      }
      return proto.getPayload().toByteArray();
    }
  }
  
  /**
   * Negotiate a cipher option which server supports.
   * 
   * @param conf the configuration
   * @param options the cipher options which client supports
   * @return CipherOption negotiated cipher option
   */
  public static CipherOption negotiateCipherOption(Configuration conf,
      List<CipherOption> options) throws IOException {
    // Negotiate cipher suites if configured.  Currently, the only supported
    // cipher suite is AES/CTR/NoPadding, but the protocol allows multiple
    // values for future expansion.
    String cipherSuites = conf.get(DFS_ENCRYPT_DATA_TRANSFER_CIPHER_SUITES_KEY);
    if (cipherSuites == null || cipherSuites.isEmpty()) {
      return null;
    }
    if (!cipherSuites.equals(CipherSuite.AES_CTR_NOPADDING.getName())) {
      throw new IOException(String.format("Invalid cipher suite, %s=%s",
          DFS_ENCRYPT_DATA_TRANSFER_CIPHER_SUITES_KEY, cipherSuites));
    }
    if (options != null) {
      for (CipherOption option : options) {
        CipherSuite suite = option.getCipherSuite();
        if (suite == CipherSuite.AES_CTR_NOPADDING) {
          int keyLen = conf.getInt(
              DFS_ENCRYPT_DATA_TRANSFER_CIPHER_KEY_BITLENGTH_KEY,
              DFS_ENCRYPT_DATA_TRANSFER_CIPHER_KEY_BITLENGTH_DEFAULT) / 8;
          CryptoCodec codec = CryptoCodec.getInstance(conf, suite);
          byte[] inKey = new byte[keyLen];
          byte[] inIv = new byte[suite.getAlgorithmBlockSize()];
          byte[] outKey = new byte[keyLen];
          byte[] outIv = new byte[suite.getAlgorithmBlockSize()];
          codec.generateSecureRandom(inKey);
          codec.generateSecureRandom(inIv);
          codec.generateSecureRandom(outKey);
          codec.generateSecureRandom(outIv);
          return new CipherOption(suite, inKey, inIv, outKey, outIv);
        }
      }
    }
    return null;
  }
  
  /**
   * Send SASL message and negotiated cipher option to client.
   * 
   * @param out stream to receive message
   * @param payload to send
   * @param option negotiated cipher option
   * @throws IOException for any error
   */
  public static void sendSaslMessageAndNegotiatedCipherOption(
      OutputStream out, byte[] payload, CipherOption option) 
          throws IOException {
    DataTransferEncryptorMessageProto.Builder builder =
        DataTransferEncryptorMessageProto.newBuilder();
    
    builder.setStatus(DataTransferEncryptorStatus.SUCCESS);
    if (payload != null) {
      builder.setPayload(ByteString.copyFrom(payload));
    }
    if (option != null) {
      builder.addCipherOption(PBHelper.convert(option));
    }
    
    DataTransferEncryptorMessageProto proto = builder.build();
    proto.writeDelimitedTo(out);
    out.flush();
  }
  
  /**
   * Create IOStreamPair of {@link org.apache.hadoop.crypto.CryptoInputStream}
   * and {@link org.apache.hadoop.crypto.CryptoOutputStream}
   * 
   * @param conf the configuration
   * @param cipherOption negotiated cipher option
   * @param out underlying output stream
   * @param in underlying input stream
   * @param isServer is server side
   * @return IOStreamPair the stream pair
   * @throws IOException for any error
   */
  public static IOStreamPair createStreamPair(Configuration conf,
      CipherOption cipherOption, OutputStream out, InputStream in, 
      boolean isServer) throws IOException {
    if (LOG.isDebugEnabled()) {
      LOG.debug("Creating IOStreamPair of CryptoInputStream and " +
          "CryptoOutputStream.");
    }
    CryptoCodec codec = CryptoCodec.getInstance(conf, 
        cipherOption.getCipherSuite());
    byte[] inKey = cipherOption.getInKey();
    byte[] inIv = cipherOption.getInIv();
    byte[] outKey = cipherOption.getOutKey();
    byte[] outIv = cipherOption.getOutIv();
    InputStream cIn = new CryptoInputStream(in, codec, 
        isServer ? inKey : outKey, isServer ? inIv : outIv);
    OutputStream cOut = new CryptoOutputStream(out, codec, 
        isServer ? outKey : inKey, isServer ? outIv : inIv);
    return new IOStreamPair(cIn, cOut);
  }

  /**
   * Sends a SASL negotiation message indicating an error.
   *
   * @param out stream to receive message
   * @param message to send
   * @throws IOException for any error
   */
  public static void sendGenericSaslErrorMessage(OutputStream out,
      String message) throws IOException {
    sendSaslMessage(out, DataTransferEncryptorStatus.ERROR, null, message);
  }

  /**
   * Sends a SASL negotiation message.
   *
   * @param out stream to receive message
   * @param payload to send
   * @throws IOException for any error
   */
  public static void sendSaslMessage(OutputStream out, byte[] payload)
      throws IOException {
    sendSaslMessage(out, DataTransferEncryptorStatus.SUCCESS, payload, null);
  }
  
  /**
   * Send a SASL negotiation message and negotiation cipher options to server.
   * 
   * @param out stream to receive message
   * @param payload to send
   * @param options cipher options to negotiate
   * @throws IOException for any error
   */
  public static void sendSaslMessageAndNegotiationCipherOptions(
      OutputStream out, byte[] payload, List<CipherOption> options)
          throws IOException {
    DataTransferEncryptorMessageProto.Builder builder =
        DataTransferEncryptorMessageProto.newBuilder();
    
    builder.setStatus(DataTransferEncryptorStatus.SUCCESS);
    if (payload != null) {
      builder.setPayload(ByteString.copyFrom(payload));
    }
    if (options != null) {
      builder.addAllCipherOption(PBHelper.convertCipherOptions(options));
    }
    
    DataTransferEncryptorMessageProto proto = builder.build();
    proto.writeDelimitedTo(out);
    out.flush();
  }
  
  /**
   * Read SASL message and negotiated cipher option from server.
   * 
   * @param in stream to read
   * @return SaslResponseWithNegotiatedCipherOption SASL message and 
   * negotiated cipher option
   * @throws IOException for any error
   */
  public static SaslResponseWithNegotiatedCipherOption
      readSaslMessageAndNegotiatedCipherOption(InputStream in)
          throws IOException {
    DataTransferEncryptorMessageProto proto =
        DataTransferEncryptorMessageProto.parseFrom(vintPrefixed(in));
    if (proto.getStatus() == DataTransferEncryptorStatus.ERROR_UNKNOWN_KEY) {
      throw new InvalidEncryptionKeyException(proto.getMessage());
    } else if (proto.getStatus() == DataTransferEncryptorStatus.ERROR) {
      throw new IOException(proto.getMessage());
    } else {
      byte[] response = proto.getPayload().toByteArray();
      List<CipherOption> options = PBHelper.convertCipherOptionProtos(
          proto.getCipherOptionList());
      CipherOption option = null;
      if (options != null && !options.isEmpty()) {
        option = options.get(0);
      }
      return new SaslResponseWithNegotiatedCipherOption(response, option);
    }
  }
  
  /**
   * Encrypt the key and iv of the negotiated cipher option.
   * 
   * @param option negotiated cipher option
   * @param sasl SASL participant representing server
   * @return CipherOption negotiated cipher option which contains the 
   * encrypted key and iv
   * @throws IOException for any error
   */
  public static CipherOption wrap(CipherOption option, SaslParticipant sasl) 
      throws IOException {
    if (option != null) {
      byte[] inKey = option.getInKey();
      if (inKey != null) {
        inKey = sasl.wrap(inKey, 0, inKey.length);
      }
      byte[] outKey = option.getOutKey();
      if (outKey != null) {
        outKey = sasl.wrap(outKey, 0, outKey.length);
      }
      return new CipherOption(option.getCipherSuite(), inKey, option.getInIv(),
          outKey, option.getOutIv());
    }
    
    return null;
  }
  
  /**
   * Decrypt the key and iv of the negotiated cipher option.
   * 
   * @param option negotiated cipher option
   * @param sasl SASL participant representing client
   * @return CipherOption negotiated cipher option which contains the 
   * decrypted key and iv
   * @throws IOException for any error
   */
  public static CipherOption unwrap(CipherOption option, SaslParticipant sasl)
      throws IOException {
    if (option != null) {
      byte[] inKey = option.getInKey();
      if (inKey != null) {
        inKey = sasl.unwrap(inKey, 0, inKey.length);
      }
      byte[] outKey = option.getOutKey();
      if (outKey != null) {
        outKey = sasl.unwrap(outKey, 0, outKey.length);
      }
      return new CipherOption(option.getCipherSuite(), inKey, option.getInIv(),
          outKey, option.getOutIv());
    }
    
    return null;
  }

  /**
   * Sends a SASL negotiation message.
   *
   * @param out stream to receive message
   * @param status negotiation status
   * @param payload to send
   * @param message to send
   * @throws IOException for any error
   */
  public static void sendSaslMessage(OutputStream out,
      DataTransferEncryptorStatus status, byte[] payload, String message)
          throws IOException {
    DataTransferEncryptorMessageProto.Builder builder =
        DataTransferEncryptorMessageProto.newBuilder();
    
    builder.setStatus(status);
    if (payload != null) {
      builder.setPayload(ByteString.copyFrom(payload));
    }
    if (message != null) {
      builder.setMessage(message);
    }
    
    DataTransferEncryptorMessageProto proto = builder.build();
    proto.writeDelimitedTo(out);
    out.flush();
  }

  /**
   * There is no reason to instantiate this class.
   */
  private DataTransferSaslUtil() {
  }
}
