/*
 * 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.asymmetric.keypair.openssl;

import static java.text.MessageFormat.format;
import static java.util.Arrays.fill;

import java.math.BigInteger;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidParameterException;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGeneratorSpi;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.security.interfaces.RSAPrivateCrtKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.AlgorithmParameterSpec;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.RSAKeyGenParameterSpec;
import java.security.spec.RSAPublicKeySpec;
import java.util.HashMap;
import java.util.Map;

import com.scytl.cryptolib.api.exceptions.CryptoLibException;

/**
 * Implementation of {@link KeyPairGeneratorSpi} for RSA.
 */
class RSAKeyPairGeneratorSpi extends KeyPairGeneratorSpi {
    private final OpenSSL openSSL;

    private final KeyFactory factory;

    private int keysize = 1024;

    private BigInteger publicExponent = RSAKeyGenParameterSpec.F4;

    /**
     * Constructor. For internal use only.
     *
     * @param openSSL
     * @param factory
     */
    RSAKeyPairGeneratorSpi(final OpenSSL openSSL,
            final KeyFactory factory) {
        this.openSSL = openSSL;
        this.factory = factory;
    }

    /**
     * Creates a new instance.
     *
     * @param openSSL
     * @return the instance
     * @throws NoSuchAlgorithmException
     *             RSA is not supported.
     */
    public static RSAKeyPairGeneratorSpi newInstance(final OpenSSL openSSL)
            throws NoSuchAlgorithmException {
        KeyFactory factory = KeyFactory.getInstance("RSA");
        return new RSAKeyPairGeneratorSpi(openSSL, factory);
    }

    @Override
    public KeyPair generateKeyPair() throws CryptoLibException {
        RSAPrivateCrtKey privateKey;
        byte[] pem = generatePrivateKeyPEM();
        try {
            privateKey = decodePrivateKey(pem);
        } finally {
            destroyPrivateKeyPEM(pem);
        }
        RSAPublicKey publicKey = extractPublicKey(privateKey);
        return new KeyPair(publicKey, privateKey);
    }

    @Override
    public void initialize(final AlgorithmParameterSpec spec,
            final SecureRandom random)
            throws InvalidAlgorithmParameterException {
        if (!(spec instanceof RSAKeyGenParameterSpec)) {
            throw new InvalidAlgorithmParameterException(
                format("Invalid algorithm parameter spec ''{0}''", spec));
        }
        RSAKeyGenParameterSpec rsaSpec = (RSAKeyGenParameterSpec) spec;
        if (rsaSpec.getKeysize() <= 0) {
            throw new InvalidAlgorithmParameterException(
                format("Invalid keysize ''{0}''", rsaSpec.getKeysize()));
        }
        if (rsaSpec.getPublicExponent() == null
            || rsaSpec.getPublicExponent().signum() <= 0) {
            throw new InvalidAlgorithmParameterException(
                format("Invalid public exponent ''{0}''",
                    rsaSpec.getPublicExponent()));
        }
        keysize = rsaSpec.getKeysize();
        publicExponent = rsaSpec.getPublicExponent();
    }

    @Override
    public void initialize(final int keysize, final SecureRandom random) {
        if (keysize <= 0) {
            throw new InvalidParameterException(
                "Key size is not positive.");
        }
        this.keysize = keysize;
    }

    private RSAPrivateCrtKey decodePrivateKey(final byte[] pem)
            throws CryptoLibException {
        RSAPrivateCrtKey privateKey;
        PEMEncodedKeySpec spec = PEMEncodedKeySpec.newInstance(pem);
        try {
            privateKey = (RSAPrivateCrtKey) factory.generatePrivate(spec);
        } catch (InvalidKeySpecException e) {
            throw new CryptoLibException(
                "Failed to decode the private key.", e);
        } finally {
            spec.destroy();
        }
        return privateKey;
    }

    private void destroyPrivateKeyPEM(final byte[] pem) {
        fill(pem, (byte) 0);
    }

    private RSAPublicKey extractPublicKey(
            final RSAPrivateCrtKey privateKey) throws CryptoLibException {
        BigInteger modulus = privateKey.getModulus();
        BigInteger pubExponent = privateKey.getPublicExponent();
        RSAPublicKeySpec spec =
            new RSAPublicKeySpec(modulus, pubExponent);
        RSAPublicKey publicKey;
        try {
            publicKey = (RSAPublicKey) factory.generatePublic(spec);
        } catch (InvalidKeySpecException e) {
            throw new CryptoLibException("Failed to extract public key.",
                e);
        }
        return publicKey;
    }

    private byte[] generatePrivateKeyPEM() throws CryptoLibException {
        byte[] pem;
        Map<String, String> options = new HashMap<>(2);
        options.put("rsa_keygen_bits", Integer.toString(keysize));
        options.put("rsa_keygen_pubexp", publicExponent.toString());
        try {
            pem = openSSL.genpkey("RSA", options);
        } catch (OpenSSLException e) {
            throw new CryptoLibException(
                "Failed to generate private key PEM.", e);
        }
        return pem;
    }
}
