/*
 * 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.dse.driver.api.core.graph.GraphStatement;
import com.datastax.oss.driver.api.core.cql.Statement;
import com.datastax.oss.driver.shaded.guava.common.base.Charsets;
import com.datastax.oss.protocol.internal.util.collection.NullAllowingImmutableMap;
import edu.umd.cs.findbugs.annotations.NonNull;
import java.nio.ByteBuffer;
import java.util.Map;

public class ProxyAuthentication {
  private static final String PROXY_EXECUTE = "ProxyExecute";

  
Adds proxy authentication information to a CQL statement.
This allows executing a statement as another role than the one the session is currently
authenticated as.
Params:
userOrRole – the role to use for execution. If the statement was already configured with
    another role, it will get replaced by this one.
statement – the statement to modify.
See Also:
Setting
   *     up roles for applications (DSE 6.0 admin guide)
Returns:
a statement that will run the same CQL query as statement, but acting as the provided role. Note: with the driver's default implementations, this will always be a copy; but if you use a custom implementation, it might return the same instance (depending on the behavior of statement.setCustomPayload())./**
   * Adds proxy authentication information to a CQL statement.
   *
   * <p>This allows executing a statement as another role than the one the session is currently
   * authenticated as.
   *
   * @param userOrRole the role to use for execution. If the statement was already configured with
   *     another role, it will get replaced by this one.
   * @param statement the statement to modify.
   * @return a statement that will run the same CQL query as {@code statement}, but acting as the
   *     provided role. Note: with the driver's default implementations, this will always be a copy;
   *     but if you use a custom implementation, it might return the same instance (depending on the
   *     behavior of {@link Statement#setCustomPayload(Map) statement.setCustomPayload()}).
   * @see <a
   *     href="https://docs.datastax.com/en/dse/6.0/dse-admin/datastax_enterprise/security/secProxy.html">Setting
   *     up roles for applications (DSE 6.0 admin guide)</a>
   */
  @NonNull
  public static <StatementT extends Statement<StatementT>> StatementT executeAs(
      @NonNull String userOrRole, @NonNull StatementT statement) {
    return statement.setCustomPayload(
        addProxyExecuteEntry(statement.getCustomPayload(), userOrRole));
  }

  
Adds proxy authentication information to a graph statement.
See Also:
executeAs(String, Statement)/**
   * Adds proxy authentication information to a graph statement.
   *
   * @see #executeAs(String, Statement)
   */
  @NonNull
  public static <StatementT extends GraphStatement<StatementT>> StatementT executeAs(
      @NonNull String userOrRole, @NonNull StatementT statement) {
    return statement.setCustomPayload(
        addProxyExecuteEntry(statement.getCustomPayload(), userOrRole));
  }

  private static Map<String, ByteBuffer> addProxyExecuteEntry(
      Map<String, ByteBuffer> currentPayload, @NonNull String userOrRole) {
    NullAllowingImmutableMap.Builder<String, ByteBuffer> builder =
        NullAllowingImmutableMap.builder();
    builder.put(PROXY_EXECUTE, ByteBuffer.wrap(userOrRole.getBytes(Charsets.UTF_8)));
    if (!currentPayload.isEmpty()) {
      for (Map.Entry<String, ByteBuffer> entry : currentPayload.entrySet()) {
        String key = entry.getKey();
        if (!key.equals(PROXY_EXECUTE)) {
          builder.put(key, entry.getValue());
        }
      }
    }
    return builder.build();
  }
}