/*
 * 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.symmetric.cipher.factory;

import static com.scytl.cryptolib.symmetric.cipher.constants.SymmetricCipherConstants.ALGORITHM_FIELD_INDEX;
import static com.scytl.cryptolib.symmetric.cipher.constants.SymmetricCipherConstants.CIPHER_SPECIFICATION_DELIMITER;
import static com.scytl.cryptolib.symmetric.cipher.constants.SymmetricCipherConstants.ENCRYPTION_OPERATING_MODE_FIELD_INDEX;

import java.security.GeneralSecurityException;
import java.util.Arrays;

import javax.crypto.Cipher;
import javax.crypto.SecretKey;
import javax.crypto.spec.GCMParameterSpec;
import javax.crypto.spec.IvParameterSpec;

import com.scytl.cryptolib.api.exceptions.CryptoLibException;
import com.scytl.cryptolib.api.exceptions.GeneralCryptoLibException;
import com.scytl.cryptolib.commons.binary.ByteArrayBuilder;
import com.scytl.cryptolib.commons.configuration.Provider;
import com.scytl.cryptolib.primitives.securerandom.configuration.ConfigSecureRandomAlgorithmAndProvider;
import com.scytl.cryptolib.primitives.securerandom.configuration.SecureRandomPolicy;
import com.scytl.cryptolib.primitives.securerandom.factory.CryptoRandomBytes;
import com.scytl.cryptolib.primitives.securerandom.factory.SecureRandomFactory;
import com.scytl.cryptolib.symmetric.cipher.configuration.EncryptionOperatingMode;
import com.scytl.cryptolib.symmetric.cipher.configuration.SymmetricCipherPolicy;

/**
 * A symmetric cipher.
 * <P>
 * Instances of this class are immutable.
 */
public class CryptoSymmetricCipher {

    private Cipher _cipher;

    private final CryptoRandomBytes _randomByteGenerator;

    private final String _algorithm;

    private final String _mode;

    private final int _initVectorByteLength;

    private final int _authTagBitLength;

    /**
     * Creates an instance of a symmetric cipher, using the specified policy.
     * 
     * @param symmetricCipherPolicy
     *            the policy for symmetric ciphers.
     * @throws GeneralCryptoLibException
     *             if the creation of the symmetric cipher instance fails.
     */
    CryptoSymmetricCipher(final SymmetricCipherPolicy symmetricCipherPolicy)
            throws GeneralCryptoLibException {

        // Create symmetric cipher instance.
        String algorithmAndModeAndPadding =
            symmetricCipherPolicy.getSymmetricCipherAlgorithmAndSpec()
                .getAlgorithmAndModeAndPadding();
        String provider =
            symmetricCipherPolicy.getSymmetricCipherAlgorithmAndSpec()
                .getProvider().getProviderName();
        try {
            if (Provider.DEFAULT == symmetricCipherPolicy
                .getSymmetricCipherAlgorithmAndSpec().getProvider()) {
                _cipher = Cipher.getInstance(algorithmAndModeAndPadding);
            } else {
                _cipher =
                    Cipher.getInstance(algorithmAndModeAndPadding,
                        provider);
            }
        } catch (GeneralSecurityException e) {
            throw new GeneralCryptoLibException(
                "Failed to create symmetric cipher in this environment. Attempted to use the provider: "
                    + provider
                    + ", and the algorithm/mode/padding: "
                    + algorithmAndModeAndPadding
                    + ". Error message was "
                    + e.getMessage(), e);
        }

        // Create instance of random number generator for creating
        // initialization vector.
        _randomByteGenerator =
            new SecureRandomFactory(
                buildSecureRandomPolicy(symmetricCipherPolicy
                    .getSecureRandomAlgorithmAndProvider()))
                .createByteRandom();

        // Retrieve cipher algorithm and encryption operating mode.
        String[] algorithmAndModeAndPaddingFields =
            algorithmAndModeAndPadding
                .split(CIPHER_SPECIFICATION_DELIMITER);
        _algorithm =
            algorithmAndModeAndPaddingFields[ALGORITHM_FIELD_INDEX];
        _mode =
            algorithmAndModeAndPaddingFields[ENCRYPTION_OPERATING_MODE_FIELD_INDEX];

        // Retrieve initialization vector byte length. If this property is set
        // to zero by policy, then retrieve it from cipher itself.
        int initVectorBitLengthFromPolicy =
            symmetricCipherPolicy.getSymmetricCipherAlgorithmAndSpec()
                .getInitVectorBitLength();
        if (initVectorBitLengthFromPolicy > 0) {
            _initVectorByteLength =
                initVectorBitLengthFromPolicy / Byte.SIZE;
        } else {
            _initVectorByteLength = _cipher.getBlockSize();
        }
        if (_initVectorByteLength <= 0) {
            throw new GeneralCryptoLibException(
                "Initialization vector byte length "
                    + _initVectorByteLength + " is invalid.");
        }

        // Retrieve authentication tag bit length.
        _authTagBitLength =
            symmetricCipherPolicy.getSymmetricCipherAlgorithmAndSpec()
                .getAuthTagBitLength();
    }

    /**
     * Uses a symmetric cipher to encrypt the given data using the given
     * {@link javax.crypto.SecretKey}.
     * 
     * @param key
     *            the {@link javax.crypto.SecretKey} to use.
     * @param data
     *            the data to be encrypted.
     * @return the bitwise concatenation of the initialization vector and
     *         encrypted data.
     * @throws GeneralCryptoLibException
     *             if the input validation or the cipher initialization fails.
     */
    public byte[] encrypt(final SecretKey key, final byte[] data)
            throws GeneralCryptoLibException {

        validateSecretKeyAlgorithm(key);
        
        ByteArrayBuilder initVectorAndEncryptedData =
            new ByteArrayBuilder();
        
        // Generate initialization vector.
        byte[] initVector =
            _randomByteGenerator.nextRandom(_initVectorByteLength);
        initVectorAndEncryptedData.append(initVector);

        // Initialize symmetric cipher for encryption.
        try {
            if (usingAuthenticationTag()) {
                if (_mode.equals(EncryptionOperatingMode.GCM.getMode())) {
                    _cipher
                        .init(Cipher.ENCRYPT_MODE, key,
                            new GCMParameterSpec(_authTagBitLength,
                                initVector));
                } else {
                    throw new CryptoLibException(
                        "Encryption operating mode '"
                            + _mode
                            + "' is not a recognized mode for use with authentication tag generation.");
                }
            } else {
                _cipher.init(Cipher.ENCRYPT_MODE, key,
                    new IvParameterSpec(initVector));
            }
        } catch (GeneralSecurityException e) {
            throw new GeneralCryptoLibException(
                "Could not initialize symmetric cipher for encryption.", e);
        }

        // Encrypt data.
        try {
            initVectorAndEncryptedData.append(_cipher.doFinal(data));
        } catch (GeneralSecurityException e) {
            throw new GeneralCryptoLibException(
                "Could not symmetrically encrypt data.", e);
        }

        return initVectorAndEncryptedData.build();
    }

    /**
     * Uses a symmetric cipher to decrypt the given data using the given
     * SecretKey.
     * 
     * @param key
     *            the {@link javax.crypto.SecretKey} to use.
     * @param data
     *            the bitwise concatenation of the initialization vector and the
     *            encrypted data.
     * @return the decrypted data.
     * @throws GeneralCryptoLibException
     *             if the input validation or the cipher initialization fails.
     */
    public byte[] decrypt(final SecretKey key, final byte[] data)
            throws GeneralCryptoLibException {

        validateSecretKeyAlgorithm(key);

        validateEncryptedData(data);

        // Retrieve initialization vector and encrypted data.
        byte[] initVector = getInitVector(data);
        byte[] encryptedData = getEncryptedData(data);

        // Initialize cipher for decryption.
        try {
            if (usingAuthenticationTag()) {
                if (_mode.equals(EncryptionOperatingMode.GCM.getMode())) {
                    _cipher
                        .init(Cipher.DECRYPT_MODE, key,
                            new GCMParameterSpec(_authTagBitLength,
                                initVector));
                } else {
                    throw new CryptoLibException(
                        "Decryption operating mode '"
                            + _mode
                            + "' is not a recognized mode for use with authentication tag generation.");
                }
            } else {
                _cipher.init(Cipher.DECRYPT_MODE, key,
                    new IvParameterSpec(initVector));
            }
        } catch (GeneralSecurityException e) {
            throw new GeneralCryptoLibException(
                "Could not initialize symmetric cipher for decryption.", e);
        }

        // Decrypt data.
        try {
            return _cipher.doFinal(encryptedData);
        } catch (GeneralSecurityException e) {
            throw new GeneralCryptoLibException(
                "Could not symmetrically decrypt data.", e);
        }
    }

    /**
     * Retrieves the initialization vector from the bitwise concatenation of the
     * initialization vector and the encrypted data.
     * 
     * @param data
     *            the bitwise concatenation of the initialization vector and the
     *            encrypted data.
     * @return the initialization vector.
     */
    public byte[] getInitVector(final byte[] data) {

        return Arrays.copyOfRange(data, 0, _initVectorByteLength);
    }

    /**
     * Retrieves the encrypted data from the bitwise concatenation of the
     * initialization vector and the encrypted data.
     * 
     * @param data
     *            the bitwise concatenation of the initialization vector and the
     *            encrypted data.
     * @return the encrypted data.
     */
    public byte[] getEncryptedData(final byte[] data) {

        return Arrays
            .copyOfRange(data, _initVectorByteLength, data.length);
    }

    /**
     * Constructs a SecureRandomPolicy, which is needed to create a
     * {@link com.scytl.cryptolib.primitives.securerandom.factory}.
     * 
     * @param configSecureRandomAlgorithmAndProvider
     *            the
     *            {@link com.scytl.cryptolib.primitives.securerandom.configuration.ConfigSecureRandomAlgorithmAndProvider}
     *            that should be returned when the method
     *            {@link com.scytl.cryptolib.primitives.securerandom.configuration.SecureRandomPolicy#getSecureRandomAlgorithmAndProvider()}
     *            is called.
     * @return a new SecureRandomPolicy.
     */
    private SecureRandomPolicy buildSecureRandomPolicy(
            final ConfigSecureRandomAlgorithmAndProvider configSecureRandomAlgorithmAndProvider) {

        return new SecureRandomPolicy() {

            @Override
            public ConfigSecureRandomAlgorithmAndProvider getSecureRandomAlgorithmAndProvider() {

                return configSecureRandomAlgorithmAndProvider;
            }
        };
    }

    /**
     * Validates the secret key algorithm against the symmetric cipher
     * algorithm.
     * 
     * @param key
     *            the {@link javax.crypto.SecretKey} for which to check the
     *            algorithm.
     * @throws GeneralCryptoLibException
     *             if the secret key algorithm does not match the symmetric
     *             cipher algorithm.
     */
    private void validateSecretKeyAlgorithm(final SecretKey key)
            throws GeneralCryptoLibException {

        String secretKeyAlgorithm = key.getAlgorithm();

        if (!secretKeyAlgorithm.equals(_algorithm)) {
            String errorMsg =
                "Secret key algorithm '" + secretKeyAlgorithm
                    + "' does not match symmetric cipher algorithm '"
                    + _algorithm + "'.";
            throw new GeneralCryptoLibException(errorMsg);
        }
    }

    private void validateEncryptedData(final byte[] data)
            throws GeneralCryptoLibException {

        if (data.length <= _initVectorByteLength) {
            throw new GeneralCryptoLibException(
                "Byte length "
                    + data.length
                    + " of bitwise concatenation of initialization vector and encrypted data is less than or equal to byte length "
                    + _initVectorByteLength + " of initialization vector.");
        }
    }

    /**
     * Checks whether an authentication tag is to be generated (when encrypting)
     * or retrieved from the encrypted data (when decrypting).
     * 
     * @return true if using a tag, false otherwise.
     */
    private boolean usingAuthenticationTag() {

        return _authTagBitLength > 0;
    }
}
