/*
 * 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.ambari.server.serveraction.kerberos;

import junit.framework.Assert;
import org.apache.commons.codec.binary.Base64;
import org.apache.directory.server.kerberos.shared.keytab.Keytab;
import org.apache.directory.server.kerberos.shared.keytab.KeytabEntry;
import org.apache.directory.shared.kerberos.codec.types.EncryptionType;
import org.easymock.EasyMockSupport;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TemporaryFolder;

import java.io.File;
import java.io.FileInputStream;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

public abstract class KerberosOperationHandlerTest extends EasyMockSupport {

  @Rule
  public TemporaryFolder folder = new TemporaryFolder();

  @Test
  public void testCreateSecurePassword() throws Exception {

    KerberosOperationHandler handler1 = createHandler();
    KerberosOperationHandler handler2 = createHandler();

    String password1 = handler1.createSecurePassword();
    Assert.assertNotNull(password1);
    Assert.assertEquals(KerberosOperationHandler.SECURE_PASSWORD_LENGTH, password1.length());

    String password2 = handler2.createSecurePassword();
    Assert.assertNotNull(password2);
    Assert.assertEquals(KerberosOperationHandler.SECURE_PASSWORD_LENGTH, password2.length());

    // Make sure the passwords are different... if they are the same, that indicated the random
    // number generators are generating using the same pattern and that is not secure.
    Assert.assertFalse((password1.equals(password2)));
  }

  @Test
  public void testCreateSecurePasswordWithSize() throws Exception {
    KerberosOperationHandler handler = createHandler();

    String password;

    password = handler.createSecurePassword(10);
    Assert.assertNotNull(password);
    Assert.assertEquals(10, password.length());

    password = handler.createSecurePassword(0);
    Assert.assertNotNull(password);
    Assert.assertEquals(KerberosOperationHandler.SECURE_PASSWORD_LENGTH, password.length());

    password = handler.createSecurePassword(-20);
    Assert.assertNotNull(password);
    Assert.assertEquals(KerberosOperationHandler.SECURE_PASSWORD_LENGTH, password.length());
  }

  @Test
  public void testCreateKeytabFileOneAtATime() throws Exception {
    KerberosOperationHandler handler = createHandler();
    File file = folder.newFile();
    final String principal1 = "principal1@REALM.COM";
    final String principal2 = "principal2@REALM.COM";
    int count;

    Assert.assertTrue(handler.createKeytabFile(principal1, handler.createSecurePassword(), 0, file));

    Keytab keytab = Keytab.read(file);
    Assert.assertNotNull(keytab);

    List<KeytabEntry> entries = keytab.getEntries();
    Assert.assertNotNull(entries);
    Assert.assertFalse(entries.isEmpty());

    count = entries.size();

    for (KeytabEntry entry : entries) {
      Assert.assertEquals(principal1, entry.getPrincipalName());
    }

    Assert.assertTrue(handler.createKeytabFile(principal2, handler.createSecurePassword(), 0, file));

    keytab = Keytab.read(file);
    Assert.assertNotNull(keytab);

    entries = keytab.getEntries();
    Assert.assertNotNull(entries);
    Assert.assertFalse(entries.isEmpty());

    Assert.assertEquals(count * 2, entries.size());
  }

  @Test
  public void testEnsureKeytabFileContainsNoDuplicates() throws Exception {
    KerberosOperationHandler handler = createHandler();
    File file = folder.newFile();
    final String principal1 = "principal1@REALM.COM";
    final String principal2 = "principal2@REALM.COM";
    Set<String> seenEntries = new HashSet<String>();

    Assert.assertTrue(handler.createKeytabFile(principal1, handler.createSecurePassword(), 0, file));
    Assert.assertTrue(handler.createKeytabFile(principal2, handler.createSecurePassword(), 0, file));

    // Attempt to add duplicate entries
    Assert.assertTrue(handler.createKeytabFile(principal2, handler.createSecurePassword(), 0, file));

    Keytab keytab = Keytab.read(file);
    Assert.assertNotNull(keytab);

    List<KeytabEntry> entries = keytab.getEntries();
    Assert.assertNotNull(entries);
    Assert.assertFalse(entries.isEmpty());

    for (KeytabEntry entry : entries) {
      String seenEntry = String.format("%s|%s", entry.getPrincipalName(), entry.getKey().getKeyType().toString());
      Assert.assertFalse(seenEntries.contains(seenEntry));
      seenEntries.add(seenEntry);
    }
  }

  @Test
  public void testCreateKeytabFileExceptions() throws Exception {
    KerberosOperationHandler handler = createHandler();
    File file = folder.newFile();
    final String principal1 = "principal1@REALM.COM";

    try {
      handler.createKeytabFile(null, handler.createSecurePassword(), 0, file);
      Assert.fail("KerberosOperationException not thrown with null principal");
    } catch (Throwable t) {
      Assert.assertEquals(KerberosOperationException.class, t.getClass());
    }

    try {
      handler.createKeytabFile(principal1, null, null, file);
      Assert.fail("KerberosOperationException not thrown with null password");
    } catch (Throwable t) {
      Assert.assertEquals(KerberosOperationException.class, t.getClass());
    }

    try {
      handler.createKeytabFile(principal1, handler.createSecurePassword(), 0, null);
      Assert.fail("KerberosOperationException not thrown with null file");
    } catch (Throwable t) {
      Assert.assertEquals(KerberosOperationException.class, t.getClass());
    }
  }

  @Test
  public void testCreateKeytabFileFromBase64EncodedData() throws Exception {
    KerberosOperationHandler handler = createHandler();
    File file = folder.newFile();
    final String principal = "principal@REALM.COM";

    Assert.assertTrue(handler.createKeytabFile(principal, handler.createSecurePassword(), 0, file));

    FileInputStream fis = new FileInputStream(file);
    byte[] data = new byte[(int) file.length()];

    Assert.assertEquals(data.length, fis.read(data));
    fis.close();

    File f = handler.createKeytabFile(Base64.encodeBase64String(data));

    try {
      Keytab keytab = Keytab.read(f);
      Assert.assertNotNull(keytab);

      List<KeytabEntry> entries = keytab.getEntries();
      Assert.assertNotNull(entries);
      Assert.assertFalse(entries.isEmpty());

      for (KeytabEntry entry : entries) {
        Assert.assertEquals(principal, entry.getPrincipalName());
      }
    } finally {
      if (!f.delete()) {
        f.deleteOnExit();
      }
    }
  }

  @Test
  public void testMergeKeytabs() throws KerberosOperationException {
    KerberosOperationHandler handler = createHandler();

    Keytab keytab1 = handler.createKeytab("principal@EXAMPLE.COM", "password", 1);
    Keytab keytab2 = handler.createKeytab("principal@EXAMPLE.COM", "password1", 1);
    Keytab keytab3 = handler.createKeytab("principal1@EXAMPLE.COM", "password", 4);

    Keytab merged;

    merged = handler.mergeKeytabs(keytab1, keytab2);
    Assert.assertEquals(keytab1.getEntries().size(), merged.getEntries().size());

    merged = handler.mergeKeytabs(keytab1, keytab3);
    Assert.assertEquals(keytab1.getEntries().size() + keytab3.getEntries().size(), merged.getEntries().size());

    merged = handler.mergeKeytabs(keytab2, keytab3);
    Assert.assertEquals(keytab2.getEntries().size() + keytab3.getEntries().size(), merged.getEntries().size());

    merged = handler.mergeKeytabs(keytab2, merged);
    Assert.assertEquals(keytab2.getEntries().size() + keytab3.getEntries().size(), merged.getEntries().size());
  }

  @Test
  public void testTranslateEncryptionTypes() throws Exception {
    KerberosOperationHandler handler = createHandler();

    Assert.assertEquals(
        new HashSet<EncryptionType>() {{
          add(EncryptionType.AES256_CTS_HMAC_SHA1_96);
          add(EncryptionType.AES128_CTS_HMAC_SHA1_96);
          add(EncryptionType.DES3_CBC_SHA1_KD);
          add(EncryptionType.DES_CBC_MD5);
          add(EncryptionType.DES_CBC_MD4);
          add(EncryptionType.DES_CBC_CRC);
          add(EncryptionType.UNKNOWN);
        }},
        handler.translateEncryptionTypes("aes256-cts-hmac-sha1-96\n aes128-cts-hmac-sha1-96\tdes3-cbc-sha1 arcfour-hmac-md5 " +
            "camellia256-cts-cmac camellia128-cts-cmac des-cbc-crc des-cbc-md5 des-cbc-md4", "\\s+")
    );

    Assert.assertEquals(
        new HashSet<EncryptionType>() {{
          add(EncryptionType.AES256_CTS_HMAC_SHA1_96);
          add(EncryptionType.AES128_CTS_HMAC_SHA1_96);
        }},
        handler.translateEncryptionTypes("aes", " ")
    );

    Assert.assertEquals(
        new HashSet<EncryptionType>() {{
          add(EncryptionType.AES256_CTS_HMAC_SHA1_96);
        }},
        handler.translateEncryptionTypes("aes-256", " ")
    );

    Assert.assertEquals(
        new HashSet<EncryptionType>() {{
          add(EncryptionType.DES3_CBC_SHA1_KD);
        }},
        handler.translateEncryptionTypes("des3", " ")
    );
  }

  @Test
  public void testEscapeCharacters() throws KerberosOperationException {
    KerberosOperationHandler handler = createHandler();

    HashSet<Character> specialCharacters = new HashSet<Character>() {
      {
        add('/');
        add(',');
        add('\\');
        add('#');
        add('+');
        add('<');
        add('>');
        add(';');
        add('"');
        add('=');
        add(' ');
      }
    };

    Assert.assertEquals("\\/\\,\\\\\\#\\+\\<\\>\\;\\\"\\=\\ ", handler.escapeCharacters("/,\\#+<>;\"= ", specialCharacters, '\\'));
    Assert.assertNull(handler.escapeCharacters(null, specialCharacters, '\\'));
    Assert.assertEquals("", handler.escapeCharacters("", specialCharacters, '\\'));
    Assert.assertEquals("nothing_special_here", handler.escapeCharacters("nothing_special_here", specialCharacters, '\\'));
    Assert.assertEquals("\\/\\,\\\\\\#\\+\\<\\>\\;\\\"\\=\\ ", handler.escapeCharacters("/,\\#+<>;\"= ", specialCharacters, '\\'));

    Assert.assertEquals("nothing<>special#here!", handler.escapeCharacters("nothing<>special#here!", null, '\\'));
    Assert.assertEquals("nothing<>special#here!", handler.escapeCharacters("nothing<>special#here!", Collections.<Character>emptySet(), '\\'));
    Assert.assertEquals("nothing<>special#here!", handler.escapeCharacters("nothing<>special#here!", Collections.singleton('?'), '\\'));
    Assert.assertEquals("\\A's are special!", handler.escapeCharacters("A's are special!", Collections.singleton('A'), '\\'));
  }

  @Test(expected = KerberosAdminAuthenticationException.class)
  public void testAdminCredentialsNullPrincipal() throws KerberosOperationException {
    KerberosOperationHandler handler = createHandler();

    KerberosCredential credentials = new KerberosCredential(null, "password", null);
    handler.setAdministratorCredentials(credentials);
  }

  @Test(expected = KerberosAdminAuthenticationException.class)
  public void testAdminCredentialsEmptyPrincipal() throws KerberosOperationException {
    KerberosOperationHandler handler = createHandler();

    KerberosCredential credentials = new KerberosCredential("", "password", null);
    handler.setAdministratorCredentials(credentials);
  }

  @Test(expected = KerberosAdminAuthenticationException.class)
  public void testAdminCredentialsNullCredential() throws KerberosOperationException {
    KerberosOperationHandler handler = createHandler();

    KerberosCredential credentials = new KerberosCredential("principal", null, null);
    handler.setAdministratorCredentials(credentials);
  }

  @Test(expected = KerberosAdminAuthenticationException.class)
  public void testAdminCredentialsEmptyCredential1() throws KerberosOperationException {
    KerberosOperationHandler handler = createHandler();

    KerberosCredential credentials = new KerberosCredential("principal", "", null);
    handler.setAdministratorCredentials(credentials);
  }

  @Test(expected = KerberosAdminAuthenticationException.class)
  public void testAdminCredentialsEmptyCredential2() throws KerberosOperationException {
    KerberosOperationHandler handler = createHandler();

    KerberosCredential credentials = new KerberosCredential("principal", null, "");
    handler.setAdministratorCredentials(credentials);
  }

  private KerberosOperationHandler createHandler() throws KerberosOperationException {
    KerberosOperationHandler handler = new KerberosOperationHandler() {

      @Override
      public void open(KerberosCredential administratorCredentials, String defaultRealm, Map<String, String> kerberosConfiguration) throws KerberosOperationException {
        setAdministratorCredentials(administratorCredentials);
        setDefaultRealm(defaultRealm);
      }

      @Override
      public void close() throws KerberosOperationException {

      }

      @Override
      public boolean principalExists(String principal) throws KerberosOperationException {
        return false;
      }

      @Override
      public Integer createPrincipal(String principal, String password, boolean service) throws KerberosOperationException {
        return 0;
      }

      @Override
      public Integer setPrincipalPassword(String principal, String password) throws KerberosOperationException {
        return 0;
      }

      @Override
      public boolean removePrincipal(String principal) throws KerberosOperationException {
        return false;
      }
    };

    handler.open(new KerberosCredential("admin/admin", "hadoop", null), "EXAMPLE.COM", null);
    return handler;
  }
}
