/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
package org.apache.shiro.crypto;

import org.apache.shiro.util.ByteSource;
import org.apache.shiro.util.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.crypto.CipherInputStream;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.security.Key;
import java.security.SecureRandom;
import java.security.spec.AlgorithmParameterSpec;

/**
 * Abstract {@code CipherService} implementation utilizing Java's JCA APIs.
 * <h2>Auto-generated Initialization Vectors</h2>
 * Shiro does something by default for all of its {@code CipherService} implementations that the JCA
 * {@link javax.crypto.Cipher Cipher} does not do:  by default,
 * <a href="http://en.wikipedia.org/wiki/Initialization_vector">initialization vector</a>s are automatically randomly
 * generated and prepended to encrypted data before returning from the {@code encrypt} methods.  That is, the returned
 * byte array or {@code OutputStream} is actually a concatenation of an initialization vector byte array plus the actual
 * encrypted data byte array.  The {@code decrypt} methods in turn know to read this prepended initialization vector
 * before decrypting the real data that follows.
 * <p/>
 * This is highly desirable because initialization vectors guarantee that, for a key and any plaintext, the encrypted
 * output will always be different <em>even if you call {@code encrypt} multiple times with the exact same arguments</em>.
 * This is essential in cryptography to ensure that data patterns cannot be identified across multiple input sources
 * that are the same or similar.
 * <p/>
 * You can turn off this behavior by setting the
 * {@link #setGenerateInitializationVectors(boolean) generateInitializationVectors} property to {@code false}, but it
 * is highly recommended that you do not do this unless you have a very good reason to do so, since you would be losing
 * a critical security feature.
 * <h3>Initialization Vector Size</h3>
 * This implementation defaults the {@link #setInitializationVectorSize(int) initializationVectorSize} attribute to
 * {@code 128} bits, a fairly common size.  Initialization vector sizes are very algorithm specific however, so subclass
 * implementations will often override this value in their constructor if necessary.
 * <p/>
 * Also note that {@code initializationVectorSize} values are specified in the number of
 * bits (not bytes!) to match common references in most cryptography documentation.  In practice though, initialization
 * vectors are always specified as a byte array, so ensure that if you set this property, that the value is a multiple
 * of {@code 8} to ensure that the IV can be correctly represented as a byte array (the
 * {@link #setInitializationVectorSize(int) setInitializationVectorSize} mutator method enforces this).
 *
 * @since 1.0
 */
public abstract class JcaCipherService implements CipherService {

    /**
     * Internal private log instance.
     */
    private static final Logger log = LoggerFactory.getLogger(JcaCipherService.class);

    /**
     * Default key size (in bits) for generated keys.
     */
    private static final int DEFAULT_KEY_SIZE = 128;

    /**
     * Default size of the internal buffer (in bytes) used to transfer data between streams during stream operations
     */
    private static final int DEFAULT_STREAMING_BUFFER_SIZE = 512;

    private static final int BITS_PER_BYTE = 8;

    /**
     * Default SecureRandom algorithm name to use when acquiring the SecureRandom instance.
     */
    private static final String RANDOM_NUM_GENERATOR_ALGORITHM_NAME = "SHA1PRNG";

    /**
     * The name of the cipher algorithm to use for all encryption, decryption, and key operations
     */
    private String algorithmName;

    /**
     * The size in bits (not bytes) of generated cipher keys
     */
    private int keySize;

    /**
     * The size of the internal buffer (in bytes) used to transfer data from one stream to another during stream operations
     */
    private int streamingBufferSize;

    private boolean generateInitializationVectors;
    private int initializationVectorSize;


    private SecureRandom secureRandom;

    /**
     * Creates a new {@code JcaCipherService} instance which will use the specified cipher {@code algorithmName}
     * for all encryption, decryption, and key operations.  Also, the following defaults are set:
     * <ul>
     * <li>{@link #setKeySize keySize} = 128 bits</li>
     * <li>{@link #setInitializationVectorSize(int) initializationVectorSize} = 128 bits</li>
     * <li>{@link #setStreamingBufferSize(int) streamingBufferSize} = 512 bytes</li>
     * </ul>
     *
     * @param algorithmName the name of the cipher algorithm to use for all encryption, decryption, and key operations
     */
    protected JcaCipherService(String algorithmName) {
        if (!StringUtils.hasText(algorithmName)) {
            throw new IllegalArgumentException("algorithmName argument cannot be null or empty.");
        }
        this.algorithmName = algorithmName;
        this.keySize = DEFAULT_KEY_SIZE;
        this.initializationVectorSize = DEFAULT_KEY_SIZE; //default to same size as the key size (a common algorithm practice)
        this.streamingBufferSize = DEFAULT_STREAMING_BUFFER_SIZE;
        this.generateInitializationVectors = true;
    }

    /**
     * Returns the cipher algorithm name that will be used for all encryption, decryption, and key operations (for
     * example, 'AES', 'Blowfish', 'RSA', 'DSA', 'TripleDES', etc).
     *
     * @return the cipher algorithm name that will be used for all encryption, decryption, and key operations
     */
    public String getAlgorithmName() {
        return algorithmName;
    }

    /**
     * Returns the size in bits (not bytes) of generated cipher keys.
     *
     * @return the size in bits (not bytes) of generated cipher keys.
     */
    public int getKeySize() {
        return keySize;
    }

    /**
     * Sets the size in bits (not bytes) of generated cipher keys.
     *
     * @param keySize the size in bits (not bytes) of generated cipher keys.
     */
    public void setKeySize(int keySize) {
        this.keySize = keySize;
    }

    public boolean isGenerateInitializationVectors() {
        return generateInitializationVectors;
    }

    public void setGenerateInitializationVectors(boolean generateInitializationVectors) {
        this.generateInitializationVectors = generateInitializationVectors;
    }

    /**
     * Returns the algorithm-specific size in bits of generated initialization vectors.
     *
     * @return the algorithm-specific size in bits of generated initialization vectors.
     */
    public int getInitializationVectorSize() {
        return initializationVectorSize;
    }

    /**
     * Sets the algorithm-specific initialization vector size in bits (not bytes!) to be used when generating
     * initialization vectors.  The  value must be a multiple of {@code 8} to ensure that the IV can be represented
     * as a byte array.
     *
     * @param initializationVectorSize the size in bits (not bytes) of generated initialization vectors.
     * @throws IllegalArgumentException if the size is not a multiple of {@code 8}.
     */
    public void setInitializationVectorSize(int initializationVectorSize) throws IllegalArgumentException {
        if (initializationVectorSize % BITS_PER_BYTE != 0) {
            String msg = "Initialization vector sizes are specified in bits, but must be a multiple of 8 so they " +
                    "can be easily represented as a byte array.";
            throw new IllegalArgumentException(msg);
        }
        this.initializationVectorSize = initializationVectorSize;
    }

    protected boolean isGenerateInitializationVectors(boolean streaming) {
        return isGenerateInitializationVectors();
    }

    /**
     * Returns the size in bytes of the internal buffer used to transfer data from one stream to another during stream
     * operations ({@link #encrypt(java.io.InputStream, java.io.OutputStream, byte[])} and
     * {@link #decrypt(java.io.InputStream, java.io.OutputStream, byte[])}).
     * <p/>
     * Default size is {@code 512} bytes.
     *
     * @return the size of the internal buffer used to transfer data from one stream to another during stream
     *         operations
     */
    public int getStreamingBufferSize() {
        return streamingBufferSize;
    }

    /**
     * Sets the size in bytes of the internal buffer used to transfer data from one stream to another during stream
     * operations ({@link #encrypt(java.io.InputStream, java.io.OutputStream, byte[])} and
     * {@link #decrypt(java.io.InputStream, java.io.OutputStream, byte[])}).
     * <p/>
     * Default size is {@code 512} bytes.
     *
     * @param streamingBufferSize the size of the internal buffer used to transfer data from one stream to another
     *                            during stream operations
     */
    public void setStreamingBufferSize(int streamingBufferSize) {
        this.streamingBufferSize = streamingBufferSize;
    }

    /**
     * Returns a source of randomness for encryption operations.  If one is not configured, and the underlying
     * algorithm needs one, the JDK {@code SHA1PRNG} instance will be used by default.
     *
     * @return a source of randomness for encryption operations.  If one is not configured, and the underlying
     *         algorithm needs one, the JDK {@code SHA1PRNG} instance will be used by default.
     */
    public SecureRandom getSecureRandom() {
        return secureRandom;
    }

    /**
     * Sets a source of randomness for encryption operations.  If one is not configured, and the underlying
     * algorithm needs one, the JDK {@code SHA1PRNG} instance will be used by default.
     *
     * @param secureRandom a source of randomness for encryption operations.  If one is not configured, and the
     *                     underlying algorithm needs one, the JDK {@code SHA1PRNG} instance will be used by default.
     */
    public void setSecureRandom(SecureRandom secureRandom) {
        this.secureRandom = secureRandom;
    }

    protected static SecureRandom getDefaultSecureRandom() {
        try {
            return java.security.SecureRandom.getInstance(RANDOM_NUM_GENERATOR_ALGORITHM_NAME);
        } catch (java.security.NoSuchAlgorithmException e) {
            log.debug("The SecureRandom SHA1PRNG algorithm is not available on the current platform.  Using the " +
                    "platform's default SecureRandom algorithm.", e);
            return new java.security.SecureRandom();
        }
    }

    protected SecureRandom ensureSecureRandom() {
        SecureRandom random = getSecureRandom();
        if (random == null) {
            random = getDefaultSecureRandom();
        }
        return random;
    }

    /**
     * Returns the transformation string to use with the {@link javax.crypto.Cipher#getInstance} invocation when
     * creating a new {@code Cipher} instance.  This default implementation always returns
     * {@link #getAlgorithmName() getAlgorithmName()}.  Block cipher implementations will want to override this method
     * to support appending cipher operation modes and padding schemes.
     *
     * @param streaming if the transformation string is going to be used for a Cipher for stream-based encryption or not.
     * @return the transformation string to use with the {@link javax.crypto.Cipher#getInstance} invocation when
     *         creating a new {@code Cipher} instance.
     */
    protected String getTransformationString(boolean streaming) {
        return getAlgorithmName();
    }

    protected byte[] generateInitializationVector(boolean streaming) {
        int size = getInitializationVectorSize();
        if (size <= 0) {
            String msg = "initializationVectorSize property must be greater than zero.  This number is " +
                    "typically set in the " + CipherService.class.getSimpleName() + " subclass constructor.  " +
                    "Also check your configuration to ensure that if you are setting a value, it is positive.";
            throw new IllegalStateException(msg);
        }
        if (size % BITS_PER_BYTE != 0) {
            String msg = "initializationVectorSize property must be a multiple of 8 to represent as a byte array.";
            throw new IllegalStateException(msg);
        }
        int sizeInBytes = size / BITS_PER_BYTE;
        byte[] ivBytes = new byte[sizeInBytes];
        SecureRandom random = ensureSecureRandom();
        random.nextBytes(ivBytes);
        return ivBytes;
    }

    public ByteSource encrypt(byte[] plaintext, byte[] key) {
        byte[] ivBytes = null;
        boolean generate = isGenerateInitializationVectors(false);
        if (generate) {
            ivBytes = generateInitializationVector(false);
            if (ivBytes == null || ivBytes.length == 0) {
                throw new IllegalStateException("Initialization vector generation is enabled - generated vector" +
                        "cannot be null or empty.");
            }
        }
        return encrypt(plaintext, key, ivBytes, generate);
    }

    private ByteSource encrypt(byte[] plaintext, byte[] key, byte[] iv, boolean prependIv) throws CryptoException {

        final int MODE = javax.crypto.Cipher.ENCRYPT_MODE;

        byte[] output;

        if (prependIv && iv != null && iv.length > 0) {

            byte[] encrypted = crypt(plaintext, key, iv, MODE);

            output = new byte[iv.length + encrypted.length];

            //now copy the iv bytes + encrypted bytes into one output array:

            // iv bytes:
            System.arraycopy(iv, 0, output, 0, iv.length);

            // + encrypted bytes:
            System.arraycopy(encrypted, 0, output, iv.length, encrypted.length);
        } else {
            output = crypt(plaintext, key, iv, MODE);
        }

        if (log.isTraceEnabled()) {
            log.trace("Incoming plaintext of size " + (plaintext != null ? plaintext.length : 0) + ".  Ciphertext " +
                    "byte array is size " + (output != null ? output.length : 0));
        }

        return ByteSource.Util.bytes(output);
    }

    public ByteSource decrypt(byte[] ciphertext, byte[] key) throws CryptoException {

        byte[] encrypted = ciphertext;

        //No IV, check if we need to read the IV from the stream:
        byte[] iv = null;

        if (isGenerateInitializationVectors(false)) {
            try {
                //We are generating IVs, so the ciphertext argument array is not actually 100% cipher text.  Instead, it
                //is:
                // - the first N bytes is the initialization vector, where N equals the value of the
                // 'initializationVectorSize' attribute.
                // - the remaining bytes in the method argument (arg.length - N) is the real cipher text.

                //So we need to chunk the method argument into its constituent parts to find the IV and then use
                //the IV to decrypt the real ciphertext:

                int ivSize = getInitializationVectorSize();
                int ivByteSize = ivSize / BITS_PER_BYTE;

                //now we know how large the iv is, so extract the iv bytes:
                iv = new byte[ivByteSize];
                System.arraycopy(ciphertext, 0, iv, 0, ivByteSize);

                //remaining data is the actual encrypted ciphertext.  Isolate it:
                int encryptedSize = ciphertext.length - ivByteSize;
                encrypted = new byte[encryptedSize];
                System.arraycopy(ciphertext, ivByteSize, encrypted, 0, encryptedSize);
            } catch (Exception e) {
                String msg = "Unable to correctly extract the Initialization Vector or ciphertext.";
                throw new CryptoException(msg, e);
            }
        }

        return decrypt(encrypted, key, iv);
    }

    private ByteSource decrypt(byte[] ciphertext, byte[] key, byte[] iv) throws CryptoException {
        if (log.isTraceEnabled()) {
            log.trace("Attempting to decrypt incoming byte array of length " +
                    (ciphertext != null ? ciphertext.length : 0));
        }
        byte[] decrypted = crypt(ciphertext, key, iv, javax.crypto.Cipher.DECRYPT_MODE);
        return decrypted == null ? null : ByteSource.Util.bytes(decrypted);
    }

    /**
     * Returns a new {@link javax.crypto.Cipher Cipher} instance to use for encryption/decryption operations.  The
     * Cipher's {@code transformationString} for the {@code Cipher}.{@link javax.crypto.Cipher#getInstance getInstance}
     * call is obtaind via the {@link #getTransformationString(boolean) getTransformationString} method.
     *
     * @param streaming {@code true} if the cipher instance will be used as a stream cipher, {@code false} if it will be
     *                  used as a block cipher.
     * @return a new JDK {@code Cipher} instance.
     * @throws CryptoException if a new Cipher instance cannot be constructed based on the
     *                         {@link #getTransformationString(boolean) getTransformationString} value.
     */
    private javax.crypto.Cipher newCipherInstance(boolean streaming) throws CryptoException {
        String transformationString = getTransformationString(streaming);
        try {
            return javax.crypto.Cipher.getInstance(transformationString);
        } catch (Exception e) {
            String msg = "Unable to acquire a Java JCA Cipher instance using " +
                    javax.crypto.Cipher.class.getName() + ".getInstance( \"" + transformationString + "\" ). " +
                    getAlgorithmName() + " under this configuration is required for the " +
                    getClass().getName() + " instance to function.";
            throw new CryptoException(msg, e);
        }
    }

    /**
     * Functions as follows:
     * <ol>
     * <li>Creates a {@link #newCipherInstance(boolean) new JDK cipher instance}</li>
     * <li>Converts the specified key bytes into an {@link #getAlgorithmName() algorithm}-compatible JDK
     * {@link Key key} instance</li>
     * <li>{@link #init(javax.crypto.Cipher, int, java.security.Key, AlgorithmParameterSpec, SecureRandom) Initializes}
     * the JDK cipher instance with the JDK key</li>
     * <li>Calls the {@link #crypt(javax.crypto.Cipher, byte[]) crypt(cipher,bytes)} method to either encrypt or
     * decrypt the data based on the specified Cipher behavior mode
     * ({@link javax.crypto.Cipher#ENCRYPT_MODE Cipher.ENCRYPT_MODE} or
     * {@link javax.crypto.Cipher#DECRYPT_MODE Cipher.DECRYPT_MODE})</li>
     * </ol>
     *
     * @param bytes the bytes to crypt
     * @param key   the key to use to perform the encryption or decryption.
     * @param iv    the initialization vector to use for the crypt operation (optional, may be {@code null}).
     * @param mode  the JDK Cipher behavior mode (Cipher.ENCRYPT_MODE or Cipher.DECRYPT_MODE).
     * @return the resulting crypted byte array
     * @throws IllegalArgumentException if {@code bytes} are null or empty.
     * @throws CryptoException          if Cipher initialization or the crypt operation fails
     */
    private byte[] crypt(byte[] bytes, byte[] key, byte[] iv, int mode) throws IllegalArgumentException, CryptoException {
        if (key == null || key.length == 0) {
            throw new IllegalArgumentException("key argument cannot be null or empty.");
        }
        javax.crypto.Cipher cipher = initNewCipher(mode, key, iv, false);
        return crypt(cipher, bytes);
    }

    /**
     * Calls the {@link javax.crypto.Cipher#doFinal(byte[]) doFinal(bytes)} method, propagating any exception that
     * might arise in an {@link CryptoException}
     *
     * @param cipher the JDK Cipher to finalize (perform the actual cryption)
     * @param bytes  the bytes to crypt
     * @return the resulting crypted byte array.
     * @throws CryptoException if there is an illegal block size or bad padding
     */
    private byte[] crypt(javax.crypto.Cipher cipher, byte[] bytes) throws CryptoException {
        try {
            return cipher.doFinal(bytes);
        } catch (Exception e) {
            String msg = "Unable to execute 'doFinal' with cipher instance [" + cipher + "].";
            throw new CryptoException(msg, e);
        }
    }

    /**
     * Initializes the JDK Cipher with the specified mode and key.  This is primarily a utility method to catch any
     * potential {@link java.security.InvalidKeyException InvalidKeyException} that might arise.
     *
     * @param cipher the JDK Cipher to {@link javax.crypto.Cipher#init(int, java.security.Key) init}.
     * @param mode   the Cipher mode
     * @param key    the Cipher's Key
     * @param spec   the JDK AlgorithmParameterSpec for cipher initialization (optional, may be null).
     * @param random the SecureRandom to use for cipher initialization (optional, may be null).
     * @throws CryptoException if the key is invalid
     */
    private void init(javax.crypto.Cipher cipher, int mode, java.security.Key key,
                      AlgorithmParameterSpec spec, SecureRandom random) throws CryptoException {
        try {
            if (random != null) {
                if (spec != null) {
                    cipher.init(mode, key, spec, random);
                } else {
                    cipher.init(mode, key, random);
                }
            } else {
                if (spec != null) {
                    cipher.init(mode, key, spec);
                } else {
                    cipher.init(mode, key);
                }
            }
        } catch (Exception e) {
            String msg = "Unable to init cipher instance.";
            throw new CryptoException(msg, e);
        }
    }


    public void encrypt(InputStream in, OutputStream out, byte[] key) throws CryptoException {
        byte[] iv = null;
        boolean generate = isGenerateInitializationVectors(true);
        if (generate) {
            iv = generateInitializationVector(true);
            if (iv == null || iv.length == 0) {
                throw new IllegalStateException("Initialization vector generation is enabled - generated vector" +
                        "cannot be null or empty.");
            }
        }
        encrypt(in, out, key, iv, generate);
    }

    private void encrypt(InputStream in, OutputStream out, byte[] key, byte[] iv, boolean prependIv) throws CryptoException {
        if (prependIv && iv != null && iv.length > 0) {
            try {
                //first write the IV:
                out.write(iv);
            } catch (IOException e) {
                throw new CryptoException(e);
            }
        }

        crypt(in, out, key, iv, javax.crypto.Cipher.ENCRYPT_MODE);
    }

    public void decrypt(InputStream in, OutputStream out, byte[] key) throws CryptoException {
        decrypt(in, out, key, isGenerateInitializationVectors(true));
    }

    private void decrypt(InputStream in, OutputStream out, byte[] key, boolean ivPrepended) throws CryptoException {

        byte[] iv = null;
        //No Initialization Vector provided as a method argument - check if we need to read the IV from the stream:
        if (ivPrepended) {
            //we are generating IVs, so we need to read the previously-generated IV from the stream before
            //we decrypt the rest of the stream (we need the IV to decrypt):
            int ivSize = getInitializationVectorSize();
            int ivByteSize = ivSize / BITS_PER_BYTE;
            iv = new byte[ivByteSize];
            int read;

            try {
                read = in.read(iv);
            } catch (IOException e) {
                String msg = "Unable to correctly read the Initialization Vector from the input stream.";
                throw new CryptoException(msg, e);
            }

            if (read != ivByteSize) {
                throw new CryptoException("Unable to read initialization vector bytes from the InputStream.  " +
                        "This is required when initialization vectors are autogenerated during an encryption " +
                        "operation.");
            }
        }

        decrypt(in, out, key, iv);
    }

    private void decrypt(InputStream in, OutputStream out, byte[] decryptionKey, byte[] iv) throws CryptoException {
        crypt(in, out, decryptionKey, iv, javax.crypto.Cipher.DECRYPT_MODE);
    }

    private void crypt(InputStream in, OutputStream out, byte[] keyBytes, byte[] iv, int cryptMode) throws CryptoException {
        if (in == null) {
            throw new NullPointerException("InputStream argument cannot be null.");
        }
        if (out == null) {
            throw new NullPointerException("OutputStream argument cannot be null.");
        }

        javax.crypto.Cipher cipher = initNewCipher(cryptMode, keyBytes, iv, true);

        CipherInputStream cis = new CipherInputStream(in, cipher);

        int bufSize = getStreamingBufferSize();
        byte[] buffer = new byte[bufSize];

        int bytesRead;
        try {
            while ((bytesRead = cis.read(buffer)) != -1) {
                out.write(buffer, 0, bytesRead);
            }
        } catch (IOException e) {
            throw new CryptoException(e);
        }
    }

    private javax.crypto.Cipher initNewCipher(int jcaCipherMode, byte[] key, byte[] iv, boolean streaming)
            throws CryptoException {

        javax.crypto.Cipher cipher = newCipherInstance(streaming);
        java.security.Key jdkKey = new SecretKeySpec(key, getAlgorithmName());
        AlgorithmParameterSpec ivSpec = null;

        if (iv != null && iv.length > 0) {
            ivSpec = createParameterSpec(iv, streaming);
        }

        init(cipher, jcaCipherMode, jdkKey, ivSpec, getSecureRandom());

        return cipher;
    }

    protected AlgorithmParameterSpec createParameterSpec(byte[] iv, boolean streaming) {
        return new IvParameterSpec(iv);
    }
}