/*
 * Copyright 2018 Scytl Secure Electronic Voting SA
 *
 * All rights reserved
 *
 * See our extended copyright notice in *file 'Copyright.txt' which is part of this source code package
 */
package com.scytl.cryptolib.elgamal.encrytionparams;

import java.math.BigInteger;

import org.bouncycastle.util.encoders.Hex;

import com.scytl.cryptolib.api.exceptions.GeneralCryptoLibException;
import com.scytl.cryptolib.primitives.messagedigest.factory.CryptoMessageDigest;
import com.scytl.cryptolib.primitives.securerandom.factory.CryptoRandomBytes;
import com.scytl.cryptolib.securerandom.cryptoapi.CryptoAPIRandomBytes;
import com.scytl.math.BigIntegers;

/**
 * An encryption parameter generator, which complies with FIPS186.
 * <P>
 * This class is a modified version of the class
 * {@link org.bouncycastle.crypto.generators.DSAParametersGenerator}. We had to
 * make this version as the original did not permit the generation of both 'g'
 * and 'f' parameters for the same p and q parameters (it forced the generation
 * of a new p and q parameters each time a new generator was created).
 * <P>
 * The principle changes that were made to the original class are the following:
 * <ul>
 * <li>Accept a {@link ElGamalEncryptionParametersGenerationParameters} as a
 * constructor input.</li>
 * <li>The method {@link #generateParameters_FIPS186_3()} is used to only
 * generate the parameters p, q, seed and counter (unlike the original method,
 * this version does not return a generator). These outputs are returned within
 * an instance of {@link EncryptionParametersPrimeOutputs}.</li>
 * <li>The method {@link #generateParameters_FIPS186_3()} uses a
 * {@link CryptoRandomBytes} as the source of random bytes.</li>
 * <li>The method
 * {@link #calculateGenerator_FIPS186_3_Verifiable (org.bouncycastle.crypto.Digest, java.math.BigInteger, java.math.BigInteger, byte[], int)}
 * has been made public, and can be used for obtaining generators. It can be
 * called as many times as is required, and a set of parameters (p, q and seed)
 * can be passed to it.</li>
 * <li>The message digest (hash) that is used has been changed to a
 * {@link CryptoMessageDigest}.</li>
 * <li>Unneeded methods have been deleted.</li>
 * </ul>
 * <P>
 * Note: If in future the class
 * {@link org.bouncycastle.crypto.generators.DSAParametersGenerator} is updated
 * to provide the functionality that is required then we could consider
 * replacing this class.
 */
public final class EncryptionParameterGenerator {

    private final CryptoMessageDigest _digest;

    private final int _L;

    private final int _N;

    private final int _certainty;

    private final CryptoAPIRandomBytes _cryptoRandomBytes;

    private static final BigInteger ZERO = BigInteger.valueOf(0);

    private static final BigInteger ONE = BigInteger.valueOf(1);

    private static final BigInteger TWO = BigInteger.valueOf(2);

    /**
     * Creates and initialize the key generator for DSA 2 with the specified
     * {@code params}.
     * 
     * @param params
     *            generation parameters.
     */
    public EncryptionParameterGenerator(
            final ElGamalEncryptionParametersGenerationParameters params) {

        _L = params.getL();
        _N = params.getN();
        _certainty = params.getCertainty();
        _digest = params.getDigester();
        _cryptoRandomBytes = params.getCryptoRandomBytes();

        if ((_L < 1024 || _L > 3072) || _L % 1024 != 0) {
            throw new IllegalArgumentException(
                "L values must be between 1024 and 3072 and a multiple of 1024");
        } else if (_L == 1024 && _N != 160) {
            throw new IllegalArgumentException(
                "N must be 160 for L = 1024");
        } else if (_L == 2048 && (_N != 224 && _N != 256)) {
            throw new IllegalArgumentException(
                "N must be 224 or 256 for L = 2048");
        } else if (_L == 3072 && _N != 256) {
            throw new IllegalArgumentException(
                "N must be 256 for L = 3072");
        }

        if (_digest.getDigestLength() * Byte.SIZE < _N) {
            throw new IllegalStateException(
                "Digest output size too small for value of N");
        }
    }

    /**
     * Generates suitable parameters for DSA, in line with <i>FIPS 186-3 A.1
     * Generation of the FFC Primes p and q</i>.
     * 
     * @return parameters for DSA.
     * @throws GeneralCryptoLibException
     *             if the defined in constructor's params N is not positive.
     */
    public EncryptionParametersPrimeOutputs generateParameters_FIPS186_3()
            throws GeneralCryptoLibException {

        // A.1.1.2 Generation of the Probable Primes p and q Using an Approved
        // Hash Function
        CryptoMessageDigest d = _digest;
        int outlen = d.getDigestLength() * Byte.SIZE;

        // 1. Check that the (L, N) pair is in the list of acceptable (L, N
        // pairs) (see Section 4.2). If the pair is not in the list, then return
        // INVALID.

        // 2. If (seedlen < N), then return INVALID.
        int seedlen = _N;
        byte[] seed = new byte[seedlen / Byte.SIZE];

        // 3. n = ceiling(L / outlen) - 1.
        int n = (_L - 1) / outlen;

        // 4. b = L - 1 - (n * outlen).
        int b = (_L - 1) % outlen;

        byte[] output;
        for (;;) {
            // 5. Get an arbitrary sequence of seedlen bits as the
            // domain_parameter_seed.
            seed = _cryptoRandomBytes.nextRandom(seed.length);

            // 6. U = Hash (domain_parameter_seed) mod 2^(N–1).
            output = d.generate(seed);

            BigInteger U =
                new BigInteger(1, output).mod(ONE.shiftLeft(_N - 1));

            // 7. q = 2^(N–1) + U + 1 – ( U mod 2).
            BigInteger q =
                ONE.shiftLeft(_N - 1).add(U).add(ONE).subtract(U.mod(TWO));

            // 8. Test whether or not q is prime as specified in Appendix C.3.
            if (!q.isProbablePrime(_certainty)) {
                // 9. If q is not a prime, then go to step 5.
                continue;
            }

            // 10. offset = 1.
            // Note: 'offset' value managed incrementally
            byte[] offset = seed.clone();

            // 11. For counter = 0 to (4L – 1) do
            int counterLimit = 4 * _L;
            for (int counter = 0; counter < counterLimit; ++counter) {
                // 11.1 For j = 0 to n do
                // Vj = Hash ((domain_parameter_seed + offset + j) mod
                // 2^seedlen).
                // 11.2 W = V0 + (V1 ∗ 2^outlen) + ... + (V^(n–1) ∗ 2^((n–1) ∗
                // outlen)) + ((Vn mod 2^b) ∗ 2^(n ∗ outlen)).
                BigInteger W = ZERO;
                for (int j = 0, exp = 0; j <= n; ++j, exp += outlen) {
                    inc(offset);
                    output = d.generate(offset);

                    BigInteger Vj = new BigInteger(1, output);
                    if (j == n) {
                        Vj = Vj.mod(ONE.shiftLeft(b));
                    }

                    W = W.add(Vj.shiftLeft(exp));
                }

                // 11.3 X = W + 2^(L–1). Comment: 0 ≤ W < 2L–1; hence, 2L–1 ≤ X
                // < 2L.
                BigInteger X = W.add(ONE.shiftLeft(_L - 1));

                // 11.4 c = X mod 2q.
                BigInteger c = X.mod(q.shiftLeft(1));

                // 11.5 p = X - (c - 1). Comment: p ≡ 1 (mod 2q).
                BigInteger p = X.subtract(c.subtract(ONE));

                // 11.6 If (p < 2^(L - 1)), then go to step 11.9
                if (p.bitLength() != _L) {
                    continue;
                }

                // 11.7 Test whether or not p is prime as specified in Appendix
                // C.3.
                if (p.isProbablePrime(_certainty)) {
                    // 11.8 If p is determined to be prime, then return VALID
                    // and the values of p, q and
                    // (optionally) the values of domain_parameter_seed and
                    // counter.

                    // =======================================================================
                    // The following section of code is slightly different from
                    // the Bouncy Castle implementation. At this point we return
                    // the p, q, seed, but we DO NOT obtain any
                    // generators at this point. Generators are later obtained
                    // using the method
                    // calculateGenerator_FIPS186_3_Verifiable(). Also, we don't
                    // return the counter variable, as we are not interested in
                    // it.
                    return new EncryptionParametersPrimeOutputs(p, q,
                        seed);
                    // =======================================================================
                }

                // 11.9 offset = offset + n + 1. Comment: Increment offset;
                // then, as part of
                // the loop in step 11, increment counter; if
                // counter < 4L, repeat steps 11.1 through 11.8.
                // Note: 'offset' value already incremented in inner loop
            }
            // 12. Go to step 5.
        }
    }

    /**
     * Obtain a generator for the group defined by the received parameters.
     * <P>
     * This is a copy of the private method from
     * {@link org.bouncycastle.crypto.generators.DSAParametersGenerator}. The
     * only changes that have been made are that the message digest that is used
     * has been changed, and the method has been made made public.
     * 
     * @param d
     *            the hash to be used.
     * @param p
     *            the p parameter to be used.
     * @param q
     *            the q parameter to be used.
     * @param seed
     *            the seed to be used. Should be the same seed as was generated
     *            when the p and q parameters were generated.
     * @param index
     *            the index (use) of the generator to be created.
     * @return the obtained generator.
     */
    public BigInteger calculateGenerator_FIPS186_3_Verifiable(
            final CryptoMessageDigest d, final BigInteger p,
            final BigInteger q, final byte[] seed, final int index) {
        // A.2.3 Verifiable Canonical Generation of the Generator g
        BigInteger e = p.subtract(ONE).divide(q);
        byte[] ggen = Hex.decode("6767656E");

        // 7. U = domain_parameter_seed || "ggen" || index || count.
        byte[] U = new byte[seed.length + ggen.length + 1 + 2];
        System.arraycopy(seed, 0, U, 0, seed.length);
        System.arraycopy(ggen, 0, U, seed.length, ggen.length);
        U[U.length - 3] = (byte) index;

        byte[] w;
        for (int count = 1; count < (1 << 16); ++count) {
            inc(U);
            w = d.generate(U);

            BigInteger W = new BigInteger(1, w);
            BigInteger g = BigIntegers.modPow(W, e, p);
            if (g.compareTo(TWO) >= 0) {
                return g;
            }
        }

        return null;
    }

    private void inc(final byte[] buf) {
        for (int i = buf.length - 1; i >= 0; --i) {
            byte b = (byte) ((buf[i] + 1) & 0xff);
            buf[i] = b;

            if (b != 0) {
                break;
            }
        }
    }
}
