/*

 * Copyright (c) 2016, 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 sun.security.provider.EntropySource;

import sun.security.provider.MoreDrbgParameters;

import javax.crypto.Cipher;

import java.io.BufferedReader;

import java.io.ByteArrayOutputStream;

import java.io.File;

import java.io.InputStream;

import java.io.InputStreamReader;

import java.io.PrintStream;

import java.lang.*;

import java.security.NoSuchAlgorithmException;

import java.security.SecureRandom;

import java.security.DrbgParameters;

import java.util.ArrayDeque;

import java.util.Arrays;

import java.util.Queue;

import java.util.stream.Stream;

import java.util.zip.ZipEntry;

import java.util.zip.ZipFile;

import java.util.zip.ZipInputStream;

import static java.security.DrbgParameters.Capability.*;

/**

 * The Known-output DRBG test. The test vector can be obtained from

 * http://csrc.nist.gov/groups/STM/cavp/documents/drbg/drbgtestvectors.zip.

 *

 * Manually run this test with

 *

 *   java DrbgCavp drbgtestvectors.zip

 *

 */

public class DrbgCavp {

    // the current nonce

    private static byte[] nonce;

    // A buffer to store test materials for the current call and

    // can be printed out of an error occurs.

    private static ByteArrayOutputStream bout = new ByteArrayOutputStream();

    // Save err for restoring

    private static PrintStream err = System.err;

    private static final int AES_LIMIT;

    static {

        try {

            AES_LIMIT = Cipher.getMaxAllowedKeyLength("AES");

        } catch (Exception e) {

            // should not happen

            throw new AssertionError("Cannot detect AES");

        }

    }

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

        if (args.length != 1) {

            System.out.println("Usage: java DrbgCavp drbgtestvectors.zip");

            return;

        }

        File tv = new File(args[0]);

        EntropySource es = new TestEntropySource();

        System.setErr(new PrintStream(bout));

        // The testsuite is a zip file containing more zip files for different

        // working modes. Each internal zip file contains test materials for

        // different mechanisms.

        try (ZipFile zf = new ZipFile(tv)) {

            String[] modes = {"no_reseed", "pr_false", "pr_true"};

            for (String mode : modes) {

                try (ZipInputStream zis = new ZipInputStream(zf.getInputStream(

                        zf.getEntry("drbgvectors_" + mode + ".zip")))) {

                    while (true) {

                        ZipEntry ze = zis.getNextEntry();

                        if (ze == null) {

                            break;

                        }

                        String fname = ze.getName();

                        if (fname.equals("Hash_DRBG.txt")

                                || fname.equals("HMAC_DRBG.txt")

                                || fname.equals("CTR_DRBG.txt")) {

                            String algorithm

                                    = fname.substring(0, fname.length() - 4);

                            test(mode, algorithm, es, zis);

                        }

                    }

                }

            }

        } finally {

            System.setErr(err);

        }

    }

    /**

     * A special entropy source you can set entropy input at will.

     */

    private static class TestEntropySource implements EntropySource {

        private static Queue<byte[]> data = new ArrayDeque<>();

        @Override

        public byte[] getEntropy(int minEntropy, int minLength,

                                 int maxLength, boolean pr) {

            byte[] result = data.poll();

            if (result == null

                    || result.length < minLength

                    || result.length > maxLength) {

                throw new RuntimeException("Invalid entropy: " +

                        "need [" + minLength + ", " + maxLength + "], " +

                        (result == null ? "none" : "has " + result.length));

            }

            return result;

        }

        private static void setEntropy(byte[] input) {

            data.offer(input);

        }

        private static void clearEntropy() {

            data.clear();

        }

    }

    /**

     * The test.

     *

     * // Algorithm line, might contain usedf flag

     * [AES-128 use df]

     * // Ignored, use mode argument

     * [PredictionResistance = True]

     * // Ignored, just read EntropyInput

     * [EntropyInputLen = 128]

     * // Ignored, just read Nonce

     * [NonceLen = 64]

     * // Ignored, just read PersonalizationString

     * [PersonalizationStringLen = 128]

     * // Ignored, just read AdditionalInput

     * [AdditionalInputLen = 128]

     * // Used to allocate buffer for nextBytes() call

     * [ReturnedBitsLen = 512]

     *

     * // A sign we can ignore old unused entropy input

     * COUNT = 0

     *

     * // Instantiate

     * EntropyInput = 92898f...

     * Nonce = c2a4d9...

     * PersonalizationString = ea65ee...  // Enough to call getInstance()

     *

     * // Reseed

     * EntropyInputReseed = bfd503...

     * AdditionalInputReseed = 009e0b... // Enough to call reseed()

     *

     * // Generation

     * AdditionalInput = 1a40fa....  // Enough to call nextBytes() for PR off

     * EntropyInputPR = 20728a...  // Enough to call nextBytes() for PR on

     * ReturnedBits = 5a3539...  // Compare this to last nextBytes() output

     *

     * @param mode one of "no_reseed", "pr_false", "pr_true"

     * @param mech one of "Hash_DRBG", "HMAC_DRBG", "CTR_DRBG"

     * @param es our own entropy source

     * @param is test material

     */

    private static void test(String mode, String mech, EntropySource es,

            InputStream is) throws Exception {

        SecureRandom hd = null;

        // Expected output length in bits as in [ReturnedBitsLen]

        int outLen = 0;

        // DRBG algorithm as in the algorithm line

        String algorithm = null;

        // When CTR_DRBG uses a derivation function as in the algorithm line

        boolean usedf = false;

        // Additional input as in "AdditionalInput"

        byte[] additional = null;

        // Random bits generated

        byte[] output = null;

        // Prediction resistance flag, determined by mode

        boolean isPr = false;

        StringBuilder sb = new StringBuilder();

        int lineno = 0;

        System.out.println(mode + "/" + mech);

        try (Stream<String> lines =

                     new BufferedReader(new InputStreamReader(is)).lines()) {

            for (String s: (Iterable<String>) lines::iterator) {

                lineno++;

                err.print(hd == null ? '-' : '*');

                Line l = new Line(s);

                if (l.key.contains("no df") || l.key.contains("use df") ||

                        l.key.startsWith("SHA-")) {

                    sb = new StringBuilder();

                    bout.reset();

                }

                sb.append(String.format(

                        "%9s %4s %5d %s\n", mode, mech, lineno, s));

                switch (l.key) {

                    case "3KeyTDEA no df":

                    case "AES-128 no df":

                    case "AES-192 no df":

                    case "AES-256 no df":

                    case "3KeyTDEA use df":

                    case "AES-128 use df":

                    case "AES-192 use df":

                    case "AES-256 use df":

                        algorithm = l.key.split(" ")[0];

                        usedf = l.key.contains("use df");

                        break;

                    case "ReturnedBitsLen":

                        outLen = l.vint();

                        output = new byte[outLen / 8];

                        break;

                    case "EntropyInput":

                        TestEntropySource.setEntropy(l.vdata());

                        break;

                    case "Nonce":

                        nonce = l.vdata();

                        break;

                    case "COUNT":

                        // Remove unused entropy (say, when AES-256 is skipped)

                        TestEntropySource.clearEntropy();

                        break;

                    case "PersonalizationString":

                        try {

                            isPr = mode.equals("pr_true");

                            byte[] ps = null;

                            if (l.vdata().length != 0) {

                                ps = l.vdata();

                            }

                            // MoreDrbgParameters must be used because we

                            // want to set entropy input and nonce. Since

                            // it can also set mechanism, algorithm and usedf,

                            // we don't need to touch securerandom.drbg.config.

                            hd = SecureRandom.getInstance("DRBG",

                                    new MoreDrbgParameters(es, mech, algorithm,

                                            nonce, usedf,

                                            DrbgParameters.instantiation(

                                                    -1,

                                                    isPr ? PR_AND_RESEED

                                                            : RESEED_ONLY,

                                                    ps)),

                                    "SUN");

                        } catch (NoSuchAlgorithmException iae) {

                            // AES-256 might not be available. This is OK.

                            if ((algorithm.equals("AES-192")

                                    || algorithm.equals("AES-256"))

                                    && AES_LIMIT == 128) {

                                hd = null;

                            } else {

                                throw iae;

                            }

                        }

                        break;

                    case "EntropyInputReseed":

                        TestEntropySource.setEntropy(l.vdata());

                        break;

                    case "AdditionalInputReseed":

                        if (l.vdata().length == 0) {

                            additional = null;

                        } else {

                            additional = l.vdata();

                        }

                        if (hd != null) {

                            if (additional == null) {

                                hd.reseed();

                            } else {

                                hd.reseed(DrbgParameters.reseed(

                                        isPr, additional));

                            }

                        }

                        break;

                    case "EntropyInputPR":

                        if (l.vdata().length != 0) {

                            TestEntropySource.setEntropy(l.vdata());

                        }

                        if (mode.equals("pr_true")) {

                            if (hd != null) {

                                if (additional == null) {

                                    hd.nextBytes(output);

                                } else {

                                    hd.nextBytes(output,

                                            DrbgParameters.nextBytes(

                                                    -1, isPr, additional));

                                }

                            }

                        }

                        break;

                    case "AdditionalInput":

                        if (l.vdata().length == 0) {

                            additional = null;

                        } else {

                            additional = l.vdata();

                        }

                        if (!mode.equals("pr_true")) {

                            if (hd != null) {

                                if (additional == null) {

                                    hd.nextBytes(output);

                                } else {

                                    hd.nextBytes(output,

                                            DrbgParameters.nextBytes(

                                                    -1, isPr, additional));

                                }

                            }

                        }

                        break;

                    case "ReturnedBits":

                        if (hd != null) {

                            if (!Arrays.equals(output, l.vdata())) {

                                throw new Exception("\nExpected: " +

                                        l.value + "\n  Actual: " + hex(output));

                            }

                        }

                        break;

                    default:

                        // Algorithm line for Hash_DRBG and HMAC_DRBG

                        if (l.key.startsWith("SHA-")) {

                            algorithm = l.key;

                        }

                }

            }

            err.println();

        } catch (Exception e) {

            err.println();

            err.println(sb.toString());

            err.println(bout.toString());

            throw e;

        }

    }

    /**

     * Parse a line from test material.

     *

     * Brackets are removed. Key and value separated.

     */

    static class Line {

        final String key;

        final String value;

        Line(String s) {

            s = s.trim();

            if (s.length() >= 2) {

                if (s.charAt(0) == '[') {

                    s = s.substring(1, s.length() - 1);

                }

            }

            if (s.indexOf('=') < 0) {

                key = s;

                value = null;

            } else {

                key = s.substring(0, s.indexOf('=')).trim();

                value = s.substring(s.indexOf('=') + 1).trim();

            }

        }

        int vint() {

            return Integer.parseInt(value);

        }

        byte[] vdata() {

            return xeh(value);

        }

    }

    // Bytes to HEX

    private static String hex(byte[] in) {

        StringBuilder sb = new StringBuilder();

        for (byte b: in) {

            sb.append(String.format("%02x", b&0xff));

        }

        return sb.toString();

    }

    // HEX to bytes

    private static byte[] xeh(String in) {

        in = in.replaceAll(" ", "");

        int len = in.length() / 2;

        byte[] out = new byte[len];

        for (int i = 0; i < len; i++) {

            out[i] = (byte) Integer.parseInt(

                    in.substring(i * 2, i * 2 + 2), 16);

        }

        return out;

    }

}