/*
 * 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.config.DseDriverOption;
import com.datastax.dse.driver.internal.core.auth.DseGssApiAuthProvider;
import com.datastax.oss.driver.api.core.CqlSession;
import com.datastax.oss.driver.api.core.Version;
import com.datastax.oss.driver.api.core.config.DefaultDriverOption;
import com.datastax.oss.driver.api.testinfra.DseRequirement;
import com.datastax.oss.driver.api.testinfra.ccm.CcmBridge;
import com.datastax.oss.driver.api.testinfra.ccm.CustomCcmRule;
import com.datastax.oss.driver.api.testinfra.session.SessionUtils;
import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableMap;
import java.io.File;
import java.util.HashMap;
import java.util.Map;
import org.junit.AssumptionViolatedException;
import org.junit.rules.ExternalResource;
import org.junit.runner.Description;
import org.junit.runners.model.Statement;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * A testing rule that wraps the EmbeddedAds server, and ccmRule into one rule This is needed
 * because ccm needs to be aware of the kerberos server settings prior to it's initialization.
 */
public class EmbeddedAdsRule extends ExternalResource {

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

  public CustomCcmRule ccm;
  // Realm for the KDC.
  private final String realm = "DATASTAX.COM";
  private final String address = "127.0.0.1";

  private final EmbeddedAds adsServer =
      EmbeddedAds.builder().withKerberos().withRealm(realm).withAddress(address).build();

  // Principal for DSE service ( = kerberos_options.service_principal)
  private final String servicePrincipal = "dse/" + adsServer.getHostname() + "@" + realm;

  // A non-standard principal for DSE service, to test SASL protocol names
  private final String alternateServicePrincipal =
      "alternate/" + adsServer.getHostname() + "@" + realm;

  // Principal for the default cassandra user.
  private final String userPrincipal = "cassandra@" + realm;

  // Principal for a user that doesn't exist.
  private final String unknownPrincipal = "unknown@" + realm;

  // Keytabs to use for auth.
  private static File userKeytab;
  private static File unknownKeytab;
  private static File dseKeytab;
  private static File alternateKeytab;
  private static Map<String, File> customKeytabs = new HashMap<>();

  private boolean alternate = false;

  public EmbeddedAdsRule(boolean alternate) {
    this.alternate = alternate;
  }

  public EmbeddedAdsRule() {
    this(false);
  }

  @Override
  protected void before() {
    try {
      if (adsServer.isStarted()) {
        return;
      }
      // Start ldap/kdc server.
      adsServer.start();

      // Create users and keytabs for the DSE principal and cassandra user.
      dseKeytab = adsServer.addUserAndCreateKeytab("dse", "fakePasswordForTests", servicePrincipal);
      alternateKeytab =
          adsServer.addUserAndCreateKeytab(
              "alternate", "fakePasswordForTests", alternateServicePrincipal);
      userKeytab =
          adsServer.addUserAndCreateKeytab("cassandra", "fakePasswordForTests", userPrincipal);
      unknownKeytab = adsServer.createKeytab("unknown", "fakePasswordForTests", unknownPrincipal);

      String authenticationOptions =
          ""
              + "authentication_options:\n"
              + "  enabled: true\n"
              + "  default_scheme: kerberos\n"
              + "  other_schemes:\n"
              + "    - internal";

      if (alternate) {
        ccm =
            CustomCcmRule.builder()
                .withCassandraConfiguration(
                    "authorizer", "com.datastax.bdp.cassandra.auth.DseAuthorizer")
                .withCassandraConfiguration(
                    "authenticator", "com.datastax.bdp.cassandra.auth.DseAuthenticator")
                .withDseConfiguration("authorization_options.enabled", true)
                .withDseConfiguration(authenticationOptions)
                .withDseConfiguration("kerberos_options.qop", "auth-conf")
                .withDseConfiguration(
                    "kerberos_options.keytab", getAlternateKeytab().getAbsolutePath())
                .withDseConfiguration(
                    "kerberos_options.service_principal", "alternate/_HOST@" + getRealm())
                .withJvmArgs(
                    "-Dcassandra.superuser_setup_delay_ms=0",
                    "-Djava.security.krb5.conf=" + getAdsServer().getKrb5Conf().getAbsolutePath())
                .build();
      } else {
        ccm =
            CustomCcmRule.builder()
                .withCassandraConfiguration(
                    "authorizer", "com.datastax.bdp.cassandra.auth.DseAuthorizer")
                .withCassandraConfiguration(
                    "authenticator", "com.datastax.bdp.cassandra.auth.DseAuthenticator")
                .withDseConfiguration("authorization_options.enabled", true)
                .withDseConfiguration(authenticationOptions)
                .withDseConfiguration("kerberos_options.qop", "auth")
                .withDseConfiguration("kerberos_options.keytab", getDseKeytab().getAbsolutePath())
                .withDseConfiguration(
                    "kerberos_options.service_principal", "dse/_HOST@" + getRealm())
                .withJvmArgs(
                    "-Dcassandra.superuser_setup_delay_ms=0",
                    "-Djava.security.krb5.conf=" + getAdsServer().getKrb5Conf().getAbsolutePath())
                .build();
      }
      ccm.getCcmBridge().create();
      ccm.getCcmBridge().start();

    } catch (Exception e) {
      LOG.error("Unable to start ads server ", e);
    }
  }

  private Statement buildErrorStatement(
      Version requirement, Version actual, String description, boolean lessThan) {
    return new Statement() {

      @Override
      public void evaluate() {
        throw new AssumptionViolatedException(
            String.format(
                "Test requires %s %s %s but %s is configured.  Description: %s",
                lessThan ? "less than" : "at least", "DSE", requirement, actual, description));
      }
    };
  }

  @Override
  public Statement apply(Statement base, Description description) {
    DseRequirement dseRequirement = description.getAnnotation(DseRequirement.class);
    if (dseRequirement != null) {
      if (!CcmBridge.DSE_ENABLEMENT) {
        return new Statement() {
          @Override
          public void evaluate() {
            throw new AssumptionViolatedException("Test Requires DSE but C* is configured.");
          }
        };
      } else {
        Version dseVersion = CcmBridge.VERSION;
        if (!dseRequirement.min().isEmpty()) {
          Version minVersion = Version.parse(dseRequirement.min());
          if (minVersion.compareTo(dseVersion) > 0) {
            return buildErrorStatement(dseVersion, dseVersion, dseRequirement.description(), false);
          }
        }

        if (!dseRequirement.max().isEmpty()) {
          Version maxVersion = Version.parse(dseRequirement.max());

          if (maxVersion.compareTo(dseVersion) <= 0) {
            return buildErrorStatement(dseVersion, dseVersion, dseRequirement.description(), true);
          }
        }
      }
    }
    return super.apply(base, description);
  }

  @Override
  protected void after() {
    adsServer.stop();
    ccm.getCcmBridge().stop();
  }

  public CqlSession newKeyTabSession(String userPrincipal, String keytabPath) {
    return SessionUtils.newSession(
        getCcm(),
        SessionUtils.configLoaderBuilder()
            .withClass(DefaultDriverOption.AUTH_PROVIDER_CLASS, DseGssApiAuthProvider.class)
            .withStringMap(
                DseDriverOption.AUTH_PROVIDER_LOGIN_CONFIGURATION,
                ImmutableMap.of(
                    "principal",
                    userPrincipal,
                    "useKeyTab",
                    "true",
                    "refreshKrb5Config",
                    "true",
                    "keyTab",
                    keytabPath))
            .build());
  }

  public CqlSession newKeyTabSession(String userPrincipal, String keytabPath, String authId) {
    return SessionUtils.newSession(
        getCcm(),
        SessionUtils.configLoaderBuilder()
            .withClass(DefaultDriverOption.AUTH_PROVIDER_CLASS, DseGssApiAuthProvider.class)
            .withStringMap(
                DseDriverOption.AUTH_PROVIDER_LOGIN_CONFIGURATION,
                ImmutableMap.of(
                    "principal",
                    userPrincipal,
                    "useKeyTab",
                    "true",
                    "refreshKrb5Config",
                    "true",
                    "keyTab",
                    keytabPath))
            .withString(DseDriverOption.AUTH_PROVIDER_AUTHORIZATION_ID, authId)
            .build());
  }

  public CqlSession newKeyTabSession() {
    return newKeyTabSession(getUserPrincipal(), getUserKeytab().getAbsolutePath());
  }

  public CqlSession newTicketSession() {
    return SessionUtils.newSession(
        getCcm(),
        SessionUtils.configLoaderBuilder()
            .withClass(DefaultDriverOption.AUTH_PROVIDER_CLASS, DseGssApiAuthProvider.class)
            .withStringMap(
                DseDriverOption.AUTH_PROVIDER_LOGIN_CONFIGURATION,
                ImmutableMap.of(
                    "principal",
                    userPrincipal,
                    "useTicketCache",
                    "true",
                    "refreshKrb5Config",
                    "true",
                    "renewTGT",
                    "true"))
            .build());
  }

  public CustomCcmRule getCcm() {
    return ccm;
  }

  public String getRealm() {
    return realm;
  }

  public String getAddress() {
    return address;
  }

  public EmbeddedAds getAdsServer() {
    return adsServer;
  }

  public String getServicePrincipal() {
    return servicePrincipal;
  }

  public String getAlternateServicePrincipal() {
    return alternateServicePrincipal;
  }

  public String getUserPrincipal() {
    return userPrincipal;
  }

  public String getUnknownPrincipal() {
    return unknownPrincipal;
  }

  public File getUserKeytab() {
    return userKeytab;
  }

  public File getUnknownKeytab() {
    return unknownKeytab;
  }

  public File getDseKeytab() {
    return dseKeytab;
  }

  public File getAlternateKeytab() {
    return alternateKeytab;
  }

  public String addUserAndCreateKeyTab(String user, String password) {
    String principal = user + "@" + realm;
    try {
      File keytabFile = adsServer.addUserAndCreateKeytab(user, password, principal);
      customKeytabs.put(principal, keytabFile);
    } catch (Exception e) {
      LOG.error("Unable to add user and create keytab for " + user + " ", e);
    }
    return principal;
  }

  public File getKeytabForPrincipal(String prinicipal) {
    return customKeytabs.get(prinicipal);
  }
}
