package org.apache.commons.jcs3.utils.serialization;

/*
 * 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.
 */

import java.io.IOException;
import java.nio.ByteBuffer;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.security.spec.InvalidKeySpecException;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.SecretKey;
import javax.crypto.SecretKeyFactory;
import javax.crypto.spec.GCMParameterSpec;
import javax.crypto.spec.PBEKeySpec;
import javax.crypto.spec.SecretKeySpec;

import org.apache.commons.jcs3.engine.behavior.IElementSerializer;

/**
 * Performs serialization and de-serialization. It encrypts and decrypts the
 * value.
 *
 * @since 3.1
 */
public class EncryptingSerializer extends StandardSerializer
{
    private static final String DEFAULT_SECRET_KEY_ALGORITHM = "PBKDF2WithHmacSHA256";
    private static final String DEFAULT_CIPHER = "AES/ECB/PKCS5Padding";
    private static final int KEYHASH_ITERATION_COUNT = 1000;
    private static final int KEY_LENGTH = 256;
    private static final int TAG_LENGTH = 128;
    private static final int IV_LENGTH = 12;
    private static final int SALT_LENGTH = 16;

    /** The pre-shared key */
    private String psk;

    /** The cipher transformation */
    private String cipherTransformation = DEFAULT_CIPHER;

    /** The random source */
    private final SecureRandom secureRandom;

    /** The secret-key factory */
    private final SecretKeyFactory secretKeyFactory;

    /** Wrapped serializer */
    private final IElementSerializer serializer;

    /**
     * Default constructor
     */
    public EncryptingSerializer()
    {
        this(new StandardSerializer());
    }

    /**
     * Wrapper constructor
     *
     * @param serializer
     *            the wrapped serializer
     */
    public EncryptingSerializer(IElementSerializer serializer)
    {
        this.serializer = serializer;

        try
        {
            this.secureRandom = new SecureRandom();
            this.secretKeyFactory = SecretKeyFactory.getInstance(DEFAULT_SECRET_KEY_ALGORITHM);
        }
        catch (NoSuchAlgorithmException e)
        {
            throw new RuntimeException("Could not set up encryption tools", e);
        }
    }

    /**
     * Set the pre-shared key for encryption and decryption
     *
     * @param psk the key
     */
    public void setPreSharedKey(String psk)
    {
        this.psk = psk;
    }

    /**
     * Set the cipher transformation for encryption and decryption
     * Default is AES/ECB/PKCS5Padding
     *
     * @param transformation the transformation
     */
    public void setAesCipherTransformation(String transformation)
    {
        this.cipherTransformation = transformation;
    }

    private byte[] getRandomBytes(int length)
    {
        byte[] bytes = new byte[length];
        secureRandom.nextBytes(bytes);

        return bytes;
    }

    private SecretKey createSecretKey(String password, byte[] salt) throws InvalidKeySpecException
    {
        /* Derive the key, given password and salt. */
        PBEKeySpec spec = new PBEKeySpec(password.toCharArray(), salt,
                KEYHASH_ITERATION_COUNT, KEY_LENGTH);
        SecretKey tmp = secretKeyFactory.generateSecret(spec);
        return new SecretKeySpec(tmp.getEncoded(), "AES");
    }

    private byte[] encrypt(byte[] source) throws IOException
    {
        try
        {
            byte[] salt = getRandomBytes(SALT_LENGTH);
            byte[] iv = getRandomBytes(IV_LENGTH);

            SecretKey secretKey = createSecretKey(psk, salt);
            Cipher cipher = Cipher.getInstance(cipherTransformation);
            if (cipher.getAlgorithm().startsWith("AES/GCM"))
            {
                cipher.init(Cipher.ENCRYPT_MODE, secretKey, new GCMParameterSpec(TAG_LENGTH, iv));
            }
            else
            {
                cipher.init(Cipher.ENCRYPT_MODE, secretKey);
            }

            byte[] encrypted = cipher.doFinal(source);

            // join initial vector, salt and encrypted data for later decryption
            return ByteBuffer.allocate(IV_LENGTH + SALT_LENGTH + encrypted.length)
                    .put(iv)
                    .put(salt)
                    .put(encrypted)
                    .array();
        }
        catch (NoSuchAlgorithmException | NoSuchPaddingException | BadPaddingException |
                IllegalBlockSizeException | InvalidKeyException | InvalidKeySpecException |
                InvalidAlgorithmParameterException e)
        {
            throw new IOException("Error while encrypting", e);
        }
    }

    private byte[] decrypt(byte[] source) throws IOException
    {
        try
        {
            // split data in initial vector, salt and encrypted data
            ByteBuffer wrapped = ByteBuffer.wrap(source);

            byte[] iv = new byte[IV_LENGTH];
            wrapped.get(iv);

            byte[] salt = new byte[SALT_LENGTH];
            wrapped.get(salt);

            byte[] encrypted = new byte[wrapped.remaining()];
            wrapped.get(encrypted);

            SecretKey secretKey = createSecretKey(psk, salt);
            Cipher cipher = Cipher.getInstance(cipherTransformation);

            if (cipher.getAlgorithm().startsWith("AES/GCM"))
            {
                cipher.init(Cipher.DECRYPT_MODE, secretKey, new GCMParameterSpec(TAG_LENGTH, iv));
            }
            else
            {
                cipher.init(Cipher.DECRYPT_MODE, secretKey);
            }

            return cipher.doFinal(encrypted);
        }
        catch (NoSuchAlgorithmException | NoSuchPaddingException | BadPaddingException |
                IllegalBlockSizeException | InvalidKeyException | InvalidKeySpecException |
                InvalidAlgorithmParameterException e)
        {
            throw new IOException("Error while decrypting", e);
        }
    }

    /**
     * Serializes an object using default serialization. Encrypts the byte array.
     * <p>
     * @param obj object
     * @return byte[]
     * @throws IOException on i/o problem
     */
    @Override
    public <T> byte[] serialize( final T obj )
        throws IOException
    {
        final byte[] unencrypted = serializer.serialize(obj);
        return encrypt(unencrypted);
    }

    /**
     * Uses default de-serialization to turn a byte array into an object. Decrypts the value
     * first. All exceptions are converted into IOExceptions.
     * <p>
     * @param data data bytes
     * @param loader class loader to use
     * @return Object
     * @throws IOException on i/o problem
     * @throws ClassNotFoundException if class is not found during deserialization
     */
    @Override
    public <T> T deSerialize( final byte[] data, final ClassLoader loader )
        throws IOException, ClassNotFoundException
    {
        if ( data == null )
        {
            return null;
        }

        final byte[] deccrypted = decrypt(data);
        return serializer.deSerialize(deccrypted, loader);
    }
}