/*
 * JBoss, Home of Professional Open Source.
 * See the COPYRIGHT.txt file distributed with this work for information
 * regarding copyright ownership.  Some portions may be licensed
 * to Red Hat, Inc. under one or more contributor license agreements.
 * 
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 * 
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Lesser General Public License for more details.
 * 
 * You should have received a copy of the GNU Lesser General Public
 * License along with this library; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
 * 02110-1301 USA.
 * Modified by
 * @author  Data Virtuality GmbH
 * @date    $Date$
 *  
 */
// https://issues.jboss.org/browse/TEIID-2302
package org.teiid.core.crypto;

import java.io.ByteArrayInputStream;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.io.Serializable;
import java.security.*;
import java.util.Arrays;

import javax.crypto.Cipher;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.SealedObject;
import javax.crypto.spec.IvParameterSpec;

import org.teiid.core.plugin.CorePlugin;
import org.teiid.core.util.AccessibleByteArrayOutputStream;
import org.teiid.core.util.ObjectInputStreamWithClassloader;


/**
 * <p>Public methods in this class throw only <code>CryptoException</code>s. </p>
 */
public class BasicCryptor implements Cryptor {
    /** The key to be used for decryption. */
    protected Key decryptKey;
    /** The <code>Cipher</code> to use for decryption. */
    private Cipher decryptCipher;
    /** The key to be used for encryption. */
    private final Key encryptKey;
    /** The <code>Cipher</code> to use for encryption. */
    protected Cipher encryptCipher;
    protected String cipherAlgorithm;

    private SecureRandom random;
    
    private ClassLoader classLoader = BasicCryptor.class.getClassLoader();

    private boolean useSealedObject = true;
    private IvParameterSpec iv;
    private byte[] randBuffer;
    public BasicCryptor(Key encryptKey, Key decryptKey, String algorithm) throws CryptoException {
        this(encryptKey, decryptKey, algorithm, null);
    }

    public BasicCryptor(Key encryptKey, Key decryptKey, String algorithm, IvParameterSpec iv) throws CryptoException {
        if (encryptKey == null) {
            throw new CryptoException(CorePlugin.Util.getString("crypto_encrypt_key_null"));
        }
        if (decryptKey == null) {
            throw new CryptoException(CorePlugin.Util.getString("crypto_decrypt_key_null"));
        }
        if (algorithm == null || algorithm.trim().isEmpty()) {
            throw new CryptoException(CorePlugin.Util.getString("crypto_algorithm_null_or_empty"));
        }

        this.encryptKey = encryptKey;
        this.cipherAlgorithm = algorithm;
        this.decryptKey = decryptKey;
        this.iv = iv;
        if (iv != null) {
            randBuffer = new byte[iv.getIV().length];
        }

        initializeSecureRandom();
        initEncryptCipher();
        initDecryptCipher();
    }

    private void initializeSecureRandom() throws CryptoException {
        random = FipsUtil.getSecureRandom();
    }

    public synchronized void setUseSealedObject(boolean useSealedObject) {
        this.useSealedObject = useSealedObject;
    }

    public synchronized void setClassLoader(ClassLoader classLoader) {
        this.classLoader = classLoader;
    }

    /**
     * Decrypt the ciphertext to yield the original cleartext.
     * In FIPS mode, this method will throw an exception if decryption fails, rather than
     * attempting to reinitialize the cipher and retry.
     * 
     * @param ciphertext The text to be decrypted, in byte form
     * @return The decrypted cleartext, in byte form
     * @throws CryptoException if decryption fails
     */
    public synchronized byte[] decrypt(byte[] ciphertext) throws CryptoException {
        try {
            byte[] result = decryptCipher.doFinal(ciphertext);
            if (iv != null) {
                //throw away the first block
                return Arrays.copyOfRange(result, iv.getIV().length, result.length);
            }
            return result;
        } catch (Exception e) {
            if (FipsUtil.isFipsModeEnabled()) {
                // In FIPS mode, we don't try to recover - throw an exception immediately
                throw new CryptoException(CorePlugin.Event.TEIID10006, 
                    CorePlugin.Util.gs(CorePlugin.Event.TEIID10006, e.getClass().getName(), e.getMessage()) + 
                    " " + CorePlugin.Util.getString("fips_no_fallback_suffix"));
            }
            
            // Only in non-FIPS mode, try to reinitialize the cipher
            try {
                initDecryptCipher();
            } catch (CryptoException err) {
                //shouldn't happen
            }
            throw new CryptoException(CorePlugin.Event.TEIID10006, 
                CorePlugin.Util.gs(CorePlugin.Event.TEIID10006, e.getClass().getName(), e.getMessage()));
        }
    }

    /**
     * Initialize the ciphers used for encryption and decryption.  The ciphers
     * define the algorithms to be used.  They are initialized with the
     * appropriate key to be used in the encryption or decryption operation.
     */
    protected void initDecryptCipher() throws CryptoException {
        
        // Create and initialize decryption cipher
        try {
            // Use FipsUtil to check if FIPS mode is enabled
            if (FipsUtil.isFipsModeEnabled() && FipsUtil.getCurrentFipsProvider() != null) {
                String provider = FipsUtil.getCurrentFipsProvider();
                
                // Validate the provider in FIPS mode
                FipsUtil.validateProvider(provider);
                
                decryptCipher = Cipher.getInstance(cipherAlgorithm, provider);
            } else {
                decryptCipher = Cipher.getInstance(cipherAlgorithm);
            }
            
            decryptCipher.init(Cipher.DECRYPT_MODE, decryptKey, iv);
        } catch (NoSuchAlgorithmException e) {
            throw new CryptoException(CorePlugin.Event.TEIID10009, e, CorePlugin.Util.gs(CorePlugin.Event.TEIID10009, cipherAlgorithm));
        } catch (NoSuchPaddingException e) {
            throw new CryptoException(CorePlugin.Event.TEIID10010, CorePlugin.Util.gs(CorePlugin.Event.TEIID10010, cipherAlgorithm, e.getClass().getName(), e.getMessage()));
        } catch (InvalidKeyException e) {
            throw new CryptoException(CorePlugin.Event.TEIID10011, e, CorePlugin.Util.gs(CorePlugin.Event.TEIID10011, e.getClass().getName(), e.getMessage()));
        } catch (InvalidAlgorithmParameterException e) {
            throw new CryptoException(CorePlugin.Event.TEIID10009, e, CorePlugin.Util.gs(CorePlugin.Event.TEIID10009, cipherAlgorithm));
        } catch (NoSuchProviderException e) {
            throw new CryptoException(CorePlugin.Event.TEIID10009, e, CorePlugin.Util.gs(CorePlugin.Event.TEIID10009, "Provider " + e.getMessage()));
        }
    }
    
    public synchronized Object unsealObject(Object object) throws CryptoException {
        if (useSealedObject) {
            if (!(object instanceof SealedObject)) {
                return object;
            }

            SealedObject so = (SealedObject) object;

            ClassLoader cl = Thread.currentThread().getContextClassLoader();
            try {
                if (cl != classLoader) {
                    // we can't take only BasicCryptor class loader, otherwise we will get dependency errors,
                    // for example, ClassNotFoundException errors, because so.getObject can use another classes
                    // from outside of BasicCryptor class loader
                    Thread.currentThread().setContextClassLoader(classLoader);
                }
                return so.getObject(decryptCipher);
            } catch (Exception e) {
                if (FipsUtil.isFipsModeEnabled()) {
                    // In FIPS mode, we don't try to recover - throw an exception immediately
                    throw new CryptoException(CorePlugin.Event.TEIID10006, 
                        CorePlugin.Util.gs(CorePlugin.Event.TEIID10006, e.getClass().getName(), e.getMessage()) +
                        " " + CorePlugin.Util.getString("fips_no_fallback_suffix"));
                }
                
                // Only in non-FIPS mode, try to reinitialize the cipher
                try {
                    initDecryptCipher();
                } catch (CryptoException err) {
                    //shouldn't happen
                }
                throw new CryptoException(CorePlugin.Event.TEIID10006, 
                    CorePlugin.Util.gs(CorePlugin.Event.TEIID10006, e.getClass().getName(), e.getMessage()));
            } finally {
                Thread.currentThread().setContextClassLoader(cl);
            }
        }
        
        if (!(object instanceof byte[])) {
            return object;
        }
        
        byte[] bytes = (byte[])object;
        bytes = decrypt(bytes);
        try {
            ObjectInputStream ois = new ObjectInputStreamWithClassloader(new ByteArrayInputStream(bytes), classLoader);
            return ois.readObject();
        } catch (Exception e) {
            throw new CryptoException(CorePlugin.Event.TEIID10006, 
                CorePlugin.Util.gs(CorePlugin.Event.TEIID10006, e.getClass().getName(), e.getMessage()));
        }
    }
            
    /**
     * Encrypt the cleartext in byte array format.
     * @param cleartext The text to be encrypted, in byte form
     */
    public byte[] encrypt(byte[] cleartext) throws CryptoException {
        return encrypt(cleartext, 0, cleartext.length);
    }

    public synchronized byte[] encrypt(byte[] buffer, int offset, int length)
 	            throws CryptoException {
        try {
            byte[] initBlock = null;
            if (iv != null) {
                random.nextBytes(randBuffer);
                initBlock = encryptCipher.update(randBuffer);
            }
            byte[] result = encryptCipher.doFinal(buffer, offset, length);
            if (initBlock != null) {
                byte[] newResult = Arrays.copyOf(initBlock, initBlock.length + result.length);
                System.arraycopy(result, 0, newResult, initBlock.length, result.length);
                return newResult;
            }
            return result;
        } catch (Exception e) {
            if (FipsUtil.isFipsModeEnabled()) {
                // In FIPS mode, we don't try to recover - throw an exception immediately
                throw new CryptoException(CorePlugin.Event.TEIID10013, 
                    CorePlugin.Util.gs(CorePlugin.Event.TEIID10013, e.getMessage()) +
                    " " + CorePlugin.Util.getString("fips_no_fallback_suffix"));
            }
            
            // Only in non-FIPS mode, try to reinitialize the cipher
            try {
                initEncryptCipher();
            } catch (CryptoException err) {
                //shouldn't happen
            }
            throw new CryptoException(CorePlugin.Event.TEIID10013, 
                CorePlugin.Util.gs(CorePlugin.Event.TEIID10013, e.getMessage()));
        }
    }

    /**
     * Initialize the cipher used for encryption.  The cipher defines the
     * algorithm to be used.  It is initialized with the appropriate key to
     * be used in the encryption operation.
     */
    protected void initEncryptCipher() throws CryptoException {

        // Create and initialize encryption cipher
        try {
            // Use FipsUtil to check if FIPS mode is enabled
            if (FipsUtil.isFipsModeEnabled() && FipsUtil.getCurrentFipsProvider() != null) {
                String provider = FipsUtil.getCurrentFipsProvider();
                
                // Validate the provider in FIPS mode
                FipsUtil.validateProvider(provider);
                
                encryptCipher = Cipher.getInstance(cipherAlgorithm, provider);
            } else {
                encryptCipher = Cipher.getInstance(cipherAlgorithm);
            }
            
            encryptCipher.init(Cipher.ENCRYPT_MODE, encryptKey, iv);
        } catch (NoSuchAlgorithmException e) {
            throw new CryptoException(CorePlugin.Event.TEIID10016, e, CorePlugin.Util.gs(CorePlugin.Event.TEIID10016, cipherAlgorithm));
        } catch (NoSuchPaddingException e) {
            throw new CryptoException(CorePlugin.Event.TEIID10017, e, CorePlugin.Util.gs(CorePlugin.Event.TEIID10017, cipherAlgorithm, e.getMessage()));
        } catch (InvalidKeyException e) {
            throw new CryptoException(CorePlugin.Event.TEIID10018, e, CorePlugin.Util.gs(CorePlugin.Event.TEIID10018, e.getMessage()));
        } catch (InvalidAlgorithmParameterException e) {
            throw new CryptoException(CorePlugin.Event.TEIID10016, e, CorePlugin.Util.gs(CorePlugin.Event.TEIID10016, cipherAlgorithm));
        } catch (NoSuchProviderException e) {
            throw new CryptoException(CorePlugin.Event.TEIID10016, e, CorePlugin.Util.gs(CorePlugin.Event.TEIID10016, "Provider " + e.getMessage()));
        }
    }
    
    public synchronized Object sealObject(Object object) throws CryptoException {
        try {
            if (useSealedObject) {
                return new SealedObject((Serializable) object, encryptCipher);
            }
            AccessibleByteArrayOutputStream baos = new AccessibleByteArrayOutputStream(1 << 13);
            ObjectOutputStream oos = new ObjectOutputStream(baos);
            oos.writeObject(object);
            oos.flush();
            oos.close();
            return encrypt(baos.getBuffer(), 0, baos.getCount());
        } catch (Exception e) {
            if (FipsUtil.isFipsModeEnabled()) {
                // In FIPS mode, we don't try to recover - throw an exception immediately
                throw new CryptoException(CorePlugin.Event.TEIID10013, 
                    CorePlugin.Util.gs(CorePlugin.Event.TEIID10013, e.getMessage()) +
                    " " + CorePlugin.Util.gs("fips_no_fallback_suffix"));
            }
            
            // Only in non-FIPS mode, try to reinitialize the cipher
            try {
                initEncryptCipher();
            } catch (CryptoException err) {
                //shouldn't happen
            }
            throw new CryptoException(CorePlugin.Event.TEIID10013, 
                CorePlugin.Util.gs(CorePlugin.Event.TEIID10013, e.getMessage()));
        }
    }

}
