/*
 * 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.digital.envelope.factory;

import java.math.BigInteger;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.interfaces.RSAPrivateCrtKey;
import java.security.interfaces.RSAPublicKey;
import java.util.Arrays;
import java.util.Base64;
import java.util.List;

import javax.crypto.SecretKey;
import javax.crypto.spec.SecretKeySpec;

import com.scytl.cryptolib.api.exceptions.GeneralCryptoLibException;
import com.scytl.cryptolib.asymmetric.cipher.factory.AsymmetricCipherFactory;
import com.scytl.cryptolib.asymmetric.cipher.factory.CryptoAsymmetricCipher;
import com.scytl.cryptolib.digital.envelope.bean.SecretKeyEncryption;
import com.scytl.cryptolib.digital.envelope.configuration.DigitalEnvelopePolicy;
import com.scytl.cryptolib.digital.envelope.cryptoapi.CryptoAPIDigitalEnvelope;
import com.scytl.cryptolib.symmetric.cipher.factory.CryptoSymmetricCipher;
import com.scytl.cryptolib.symmetric.cipher.factory.SymmetricCipherFactory;
import com.scytl.cryptolib.symmetric.mac.factory.CryptoMac;
import com.scytl.cryptolib.symmetric.mac.factory.MacFactory;

/**
 * Defines a digital envelope opener.
 * <P>
 * Instances of this class are immutable.
 */
public class CryptoDigitalEnvelopeOpener {

    private final int _encryptionSecretKeyByteLength;

    private final int _macSecretKeyByteLength;

    private final String _secretKeyAlgorithm;

    private final CryptoMac _macGenerator;

    private final CryptoSymmetricCipher _symmetricCipher;

    private final CryptoAsymmetricCipher _asymmetricCipher;

    /**
     * Creates an instance of a digital envelope, using the given cryptographic
     * policy.
     * 
     * @param digitalEnvelopePolicy
     *            the digital envelope cryptographic policy.
     */
    CryptoDigitalEnvelopeOpener(
            final DigitalEnvelopePolicy digitalEnvelopePolicy) {

        _encryptionSecretKeyByteLength = digitalEnvelopePolicy
            .getSecretKeyAlgorithmAndSpec().getKeyLength() / Byte.SIZE;
        _macSecretKeyByteLength =
            digitalEnvelopePolicy.getHmacSecretKeyAlgorithmAndSpec()
                .getKeyLengthInBits() / Byte.SIZE;
        _secretKeyAlgorithm = digitalEnvelopePolicy
            .getSecretKeyAlgorithmAndSpec().getAlgorithm();

        _macGenerator = new MacFactory(digitalEnvelopePolicy).create();

        _symmetricCipher =
            new SymmetricCipherFactory(digitalEnvelopePolicy).create();

        _asymmetricCipher =
            new AsymmetricCipherFactory(digitalEnvelopePolicy).create();
    }

    /**
     * Opens a digital envelope and retrieves its data.
     * 
     * @param digitalEnvelope
     *            the digital envelope, as an object of type
     *            {@link CryptoAPIDigitalEnvelope}.
     * @param privateKey
     *            the private key, as an object of type
     *            {@link java.security.PrivateKey}, corresponding to one of the
     *            public keys used to create the digital envelope.
     * @return the data contained in the digital envelope.
     * @throws GeneralCryptoLibException
     *             if the input data is invalid, the digital envelope cannot be
     *             opened or if the data that it contains is invalid.
     */
    public byte[] open(final CryptoAPIDigitalEnvelope digitalEnvelope,
            final PrivateKey privateKey) throws GeneralCryptoLibException {

        // Retrieve data from digital envelope.
        byte[] encryptedData = digitalEnvelope.getEncryptedData();
        byte[] mac = digitalEnvelope.getMac();
        List<SecretKeyEncryption> secretKeyEncryptions =
            digitalEnvelope.getSecretKeyEncryptions();

        // Retrieve asymmetric encryption of bit-wise concatenation of
        // encryption and MAC secret keys that can be decrypted with private
        // key provided as input.
        byte[] encryptedSecretKeyConcatenation =
            getEncryptedSecretKeyConcatenation(secretKeyEncryptions,
                privateKey, mac);

        // Asymmetrically decrypt bit-wise concatenation of encryption and MAC
        // secret keys and extract secret keys from result.
        byte[] secretKeyConcatenation = _asymmetricCipher
            .decrypt(privateKey, encryptedSecretKeyConcatenation);
        byte[] encryptionSecretKeyBytes = Arrays.copyOfRange(
            secretKeyConcatenation, 0, _encryptionSecretKeyByteLength);
        SecretKey encryptionSecretKey = new SecretKeySpec(
            encryptionSecretKeyBytes, _secretKeyAlgorithm);
        byte[] macSecretKeyBytes = Arrays.copyOfRange(
            secretKeyConcatenation, _encryptionSecretKeyByteLength,
            _encryptionSecretKeyByteLength + _macSecretKeyByteLength);
        SecretKey macSecretKey =
            new SecretKeySpec(macSecretKeyBytes, _secretKeyAlgorithm);

        // Check integrity of digital envelope.
        byte[] encryptedDataBase64 =
            Base64.getEncoder().encode(encryptedData);
        if (!_macGenerator.verify(macSecretKey, mac,
            encryptedDataBase64)) {
            throw new GeneralCryptoLibException(
                "The integrity of the digital envelope could not be verified.");
        }

        return _symmetricCipher.decrypt(encryptionSecretKey,
            encryptedData);
    }

    /**
     * Retrieves the asymmetric encryption of the digital envelope secret key
     * concatenation that can be decrypted with the private key provided as
     * input.
     * 
     * @param secretKeyEncryptions
     *            the asymmetric encryptions of the secret key concatenation,
     *            each created with a different public key provided as input to
     *            the digital envelope generator.
     * @param privateKey
     *            the private key, of type {@link java.security.PrivateKey}, to
     *            be used for decryption of the encrypted secret key
     *            concatenation.
     * @throws GeneralCryptoLibException
     *             if no asymmetric encryption of the secret key concatenation
     *             could be found that could be decrypted using the provided
     *             private key.
     */
    private byte[] getEncryptedSecretKeyConcatenation(
            final List<SecretKeyEncryption> secretKeyEncryptions,
            final PrivateKey privateKey, final byte[] mac)
            throws GeneralCryptoLibException {

        for (SecretKeyEncryption secretKeyEncryption : secretKeyEncryptions) {
            if (keysMatch(secretKeyEncryption.getPublicKey(), privateKey,
                mac)) {
                return secretKeyEncryption
                    .getEncryptedSecretKeyConcatenation();
            }
        }

        throw new GeneralCryptoLibException(
            "Private key provided to open digital envelope does not correspond to a public key used to generate it.");
    }

    private boolean keysMatch(final PublicKey publicKey,
            final PrivateKey privateKey, final byte[] mac)
            throws GeneralCryptoLibException {

        String algorithm = publicKey.getAlgorithm();
        if ("RSA".equals(algorithm)) {
            return rsaKeysMatch(publicKey, privateKey);
        } else {
            try {
                byte[] encryptedMac =
                    _asymmetricCipher.encrypt(publicKey, mac);
                byte[] decryptedMac =
                    _asymmetricCipher.decrypt(privateKey, encryptedMac);

                return Arrays.equals(decryptedMac, mac);
            } catch (GeneralCryptoLibException e) {
                return false;
            }
        }
    }

    private boolean rsaKeysMatch(final PublicKey publicKey,
            final PrivateKey privateKey) {

        RSAPrivateCrtKey rsaPrivateCrtKey = (RSAPrivateCrtKey) privateKey;
        BigInteger modulus = rsaPrivateCrtKey.getModulus();
        BigInteger publicExponent = rsaPrivateCrtKey.getPublicExponent();

        RSAPublicKey rsaPublicKey = (RSAPublicKey) publicKey;

        return rsaPublicKey.getModulus().equals(modulus)
            && rsaPublicKey.getPublicExponent().equals(publicExponent);
    }
}
