/*
 * Copyright (c) 2012, 2022, Oracle and/or its affiliates. All rights reserved.
 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
 *
 * This code is free software; you can redistribute it and/or modify it
 * under the terms of the GNU General Public License version 2 only, as
 * published by the Free Software Foundation.
 *
 * This code is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
 * version 2 for more details (a copy is included in the LICENSE file that
 * accompanied this code).
 *
 * You should have received a copy of the GNU General Public License version
 * 2 along with this work; if not, write to the Free Software Foundation,
 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
 *
 * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
 * or visit www.oracle.com if you need additional information or have any
 * questions.
 */

import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.spec.InvalidKeySpecException;
import java.util.Arrays;
import java.util.Random;
import javax.crypto.SecretKey;
import javax.crypto.SecretKeyFactory;
import javax.crypto.interfaces.PBEKey;
import javax.crypto.spec.PBEKeySpec;
import javax.crypto.spec.SecretKeySpec;

/**
 * @test
 * @bug 8041781 7181214 8288050
 * @summary Verify if the SecretKeyFactory.translateKey() method works
 * @author Alexander Fomin
 * @run main PBKDF2Translate
 * @key randomness
 */
public class PBKDF2Translate {

    private static final String[] ALGO_TO_TEST = {
        "PBKDF2WithHmacSHA1",
        "PBKDF2WithHmacSHA224",
        "PBKDF2WithHmacSHA256",
        "PBKDF2WithHmacSHA384",
        "PBKDF2WithHmacSHA512",
        "PBKDF2WithHmacSHA512/224",
        "PBKDF2WithHmacSHA512/256",
    };

    private static final String PASS_PHRASE = "some hidden string";
    private static final int ITERATION_COUNT = 1000;
    private static final int KEY_SIZE = 128;

    private final String algoToTest;
    private final byte[] salt = new byte[8];

    public static void main(String[] args) throws Exception {

        boolean failed = false;

        for (String algo : ALGO_TO_TEST) {

            System.out.println("Testing " + algo + ":");
            PBKDF2Translate theTest = new PBKDF2Translate(algo);

            try {
                if (!theTest.testMyOwnSecretKey()
                        || !theTest.generateAndTranslateKey()
                        || !theTest.translateSpoiledKey()
                        || !theTest.testGeneralSecretKey()) {
                    // we don't want to set failed to false
                    failed = true;
                }
            } catch (InvalidKeyException | NoSuchAlgorithmException |
                    InvalidKeySpecException e) {
                e.printStackTrace(System.err);
                failed = true;
            }
        }

        if (failed) {
            throw new RuntimeException("One or more tests failed....");
        }
    }

    public PBKDF2Translate(String algoToTest) {
        this.algoToTest = algoToTest;
        new Random().nextBytes(this.salt);
    }

    /**
     * The test case scenario implemented in the method: - derive PBKDF2 key
     * using the given algorithm; - translate the key - check if the translated
     * and original keys have the same key value.
     *
     * @return true if the test case passed; false - otherwise.
     * @throws NoSuchAlgorithmException
     * @throws InvalidKeySpecException
     * @throws InvalidKeyException
     */
    public boolean generateAndTranslateKey() throws NoSuchAlgorithmException,
            InvalidKeySpecException, InvalidKeyException {
        // derive PBKDF2 key
        SecretKey key1 = getSecretKeyForPBKDF2(algoToTest);

        // translate key
        SecretKeyFactory skf = SecretKeyFactory.getInstance(algoToTest);
        SecretKey key2 = skf.translateKey(key1);

        // check if it still the same after translation
        if (!Arrays.equals(key1.getEncoded(), key2.getEncoded())) {
            System.err.println("generateAndTranslateKey test case failed: the "
                    + "key1 and key2 values in its primary encoding format are "
                    + "not the same for " + algoToTest + "algorithm.");
            return false;
        }

        return true;
    }

    /**
     * The test case scenario implemented in the method: - derive Key1 for the
     * given PBKDF2 algorithm - create my own secret Key2 as an instance of a
     * class implements PBEKey - translate Key2 - check if the key value of the
     * translated key and Key1 are the same.
     *
     * @return true if the test case passed; false - otherwise.
     * @throws NoSuchAlgorithmException
     * @throws InvalidKeySpecException
     * @throws InvalidKeyException
     */
    public boolean testMyOwnSecretKey()
            throws NoSuchAlgorithmException, InvalidKeySpecException,
            InvalidKeyException {
        SecretKey key1 = getSecretKeyForPBKDF2(algoToTest);
        SecretKey key2 = getMyOwnSecretKey();

        // Is it actually the same?
        if (!Arrays.equals(key1.getEncoded(), key2.getEncoded())) {
            System.err.println("We shouldn't be here. The key1 and key2 values "
                    + "in its primary encoding format have to be the same!");
            return false;
        }

        // Translate key
        SecretKeyFactory skf = SecretKeyFactory.getInstance(algoToTest);
        SecretKey key3 = skf.translateKey(key2);

        // Check if it still the same after translation
        if (!Arrays.equals(key1.getEncoded(), key3.getEncoded())) {
            System.err.println("testMyOwnSecretKey test case failed: the key1 "
                    + "and key3 values in its primary encoding format are not "
                    + "the same for " + algoToTest + "algorithm.");
            return false;
        }

        return true;
    }

    /**
     * The test case scenario implemented in the method: - create my own secret
     * Key2 as an instance of a class implements PBEKey - spoil the key (set
     * iteration count to 0, for example) - try to translate key -
     * InvalidKeyException is expected.
     *
     * @return true if InvalidKeyException occurred; false - otherwise.
     * @throws NoSuchAlgorithmException
     * @throws InvalidKeySpecException
     */
    public boolean translateSpoiledKey() throws NoSuchAlgorithmException,
            InvalidKeySpecException {
        // derive the key
        SecretKey key1 = getMyOwnSecretKey();

        // spoil the key
        ((MyPBKDF2SecretKey) key1).spoil();

        // translate key
        SecretKeyFactory skf = SecretKeyFactory.getInstance(algoToTest);
        try {
            SecretKey key2 = skf.translateKey(key1);
        } catch (InvalidKeyException ike) {
            // this is expected
            return true;
        }

        return false;
    }

    /**
     * The test case scenario implemented in the method: - create a general
     * secret key (does not implement PBEKey) - try calling
     * translate and getKeySpec methods and see if the expected
     * InvalidKeyException and InvalidKeySpecException is thrown.
     *
     * @return true if the expected Exception occurred; false - otherwise
     * @throws NoSuchAlgorithmException
     */
    public boolean testGeneralSecretKey() throws NoSuchAlgorithmException {
        SecretKey key = new SecretKeySpec("random#s".getBytes(), algoToTest);
        SecretKeyFactory skf = SecretKeyFactory.getInstance(algoToTest);
        try {
            skf.translateKey(key);
            System.out.println("Error: expected IKE not thrown");
            return false;
        } catch (InvalidKeyException e) {
            if (e.getMessage().indexOf("PBEKey") == -1) {
                System.out.println("Error: IKE message should " +
                        "indicate that PBEKey is required");
                return false;
            }
        }

        try {
            skf.getKeySpec(key, PBEKeySpec.class);
            System.out.println("Error: expected IKSE not thrown");
            return false;
        } catch (InvalidKeySpecException e) {
            if (e.getMessage().indexOf("PBEKey") == -1) {
                System.out.println("Error: IKSE message should " +
                        "indicate that PBEKey is required");
                return false;
            }
        }

        return true;
    }

    /**
     * Generate a PBKDF2 secret key using given algorithm.
     *
     * @param algoToDeriveKey PBKDF2 algorithm
     * @return PBKDF2 secret key
     * @throws NoSuchAlgorithmException
     * @throws InvalidKeySpecException
     */
    private SecretKey getSecretKeyForPBKDF2(String algoToDeriveKey)
            throws NoSuchAlgorithmException, InvalidKeySpecException {
        SecretKeyFactory skf = SecretKeyFactory.getInstance(algoToDeriveKey);

        PBEKeySpec spec = new PBEKeySpec(PASS_PHRASE.toCharArray(),
                this.salt, ITERATION_COUNT, KEY_SIZE);

        return skf.generateSecret(spec);
    }

    /**
     * Generate a secrete key as an instance of a class implements PBEKey.
     *
     * @return secrete key
     * @throws InvalidKeySpecException
     * @throws NoSuchAlgorithmException
     */
    private SecretKey getMyOwnSecretKey() throws InvalidKeySpecException,
            NoSuchAlgorithmException {
        return new MyPBKDF2SecretKey(PASS_PHRASE, this.algoToTest, this.salt,
                ITERATION_COUNT, KEY_SIZE);
    }
}

/**
 * An utility class to check the SecretKeyFactory.translateKey() method.
 */
class MyPBKDF2SecretKey implements PBEKey {

    private final byte[] key;
    private final byte[] salt;
    private final String algorithm;
    private final int keySize, keyLength;
    private int itereationCount;
    private final String pass;

    @Override
    public String getAlgorithm() {
        return algorithm;
    }

    @Override
    public String getFormat() {
        return "RAW";
    }

    @Override
    public byte[] getEncoded() {
        byte[] copy = new byte[keyLength];
        System.arraycopy(this.key, 0, copy, 0, keyLength);
        return copy;
    }

    /**
     * The key is generating by SecretKeyFactory and its value just copying in
     * the key field of MySecretKey class. So, this is real key derived using
     * the given algorithm.
     *
     * @param passPhrase some string intended to be a password
     * @param algo PBKDF2 algorithm
     * @param salt slat for PBKDF2
     * @param iterationCount iteration count
     * @param keySize key size in bits
     * @throws InvalidKeySpecException
     * @throws NoSuchAlgorithmException
     */
    public MyPBKDF2SecretKey(String passPhrase, String algo, byte[] salt,
            int iterationCount, int keySize)
            throws InvalidKeySpecException, NoSuchAlgorithmException {
        this.algorithm = algo;
        this.salt = salt;
        this.itereationCount = iterationCount;
        this.keySize = keySize;
        this.pass = passPhrase;

        PBEKeySpec spec = new PBEKeySpec(passPhrase.toCharArray(),
                this.salt, iterationCount, this.keySize);

        SecretKeyFactory keyFactory
                = SecretKeyFactory.getInstance(algo);

        SecretKey realKey = keyFactory.generateSecret(spec);

        this.keyLength = realKey.getEncoded().length;

        this.key = new byte[this.keyLength];
        System.arraycopy(realKey.getEncoded(), 0, this.key, 0,
                this.keyLength);
    }

    @Override
    public int getIterationCount() {
        return itereationCount;
    }

    @Override
    public byte[] getSalt() {
        return salt;
    }

    @Override
    public char[] getPassword() {
        return this.pass.toCharArray();
    }

    /**
     * Spoil the generated key (before translation) to cause an
     * InvalidKeyException
     */
    public void spoil() {
        this.itereationCount = -1;
    }

}