/*
 * Copyright (c) 2014, Oracle and/or its affiliates. All rights reserved.
 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
 *
 * This code is free software; you can redistribute it and/or modify it
 * under the terms of the GNU General Public License version 2 only, as
 * published by the Free Software Foundation.  Oracle designates this
 * particular file as subject to the "Classpath" exception as provided
 * by Oracle in the LICENSE file that accompanied this code.
 *
 * This code 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 General Public License
 * version 2 for more details (a copy is included in the LICENSE file that
 * accompanied this code).
 *
 * You should have received a copy of the GNU General Public License version
 * 2 along with this work; if not, write to the Free Software Foundation,
 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
 *
 * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
 * or visit www.oracle.com if you need additional information or have any
 * questions.
 */

package com.oracle.security.ucrypto;

import java.nio.ByteBuffer;
import java.util.Set;
import java.util.Arrays;
import java.util.concurrent.ConcurrentSkipListSet;
import java.lang.ref.*;

import java.security.AlgorithmParameters;
import java.security.GeneralSecurityException;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.Key;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;


import java.security.spec.AlgorithmParameterSpec;
import java.security.spec.InvalidParameterSpecException;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.CipherSpi;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.ShortBufferException;

import javax.crypto.spec.IvParameterSpec;

/**
 * Wrapper class which uses NativeCipher class and Java impls of padding scheme.
 * This class currently supports
 * - AES/ECB/PKCS5PADDING
 * - AES/CBC/PKCS5PADDING
 * - AES/CFB128/PKCS5PADDING
 *
 * @since 1.9
 */
public class NativeCipherWithJavaPadding extends CipherSpi {

    private static interface Padding {
        // ENC: generate and return the necessary padding bytes
        int getPadLen(int dataLen);

        // ENC: generate and return the necessary padding bytes
        byte[] getPaddingBytes(int dataLen);

        // DEC: process the decrypted data and buffer up the potential padding
        // bytes
        byte[] bufferBytes(byte[] intermediateData);

        // DEC: return the length of internally buffered pad bytes
        int getBufferedLength();

        // DEC: unpad and place the output in 'out', starting from outOfs
        // and return the number of bytes unpadded into 'out'.
        int unpad(byte[] paddedData, byte[] out, int outOfs)
                throws BadPaddingException, IllegalBlockSizeException,
                ShortBufferException;

        // DEC: Clears the padding object to the initial state
        void clear();
    }

    private static class PKCS5Padding implements Padding {
        private final int blockSize;
        // buffer for storing the potential padding bytes
        private ByteBuffer trailingBytes = null;

        PKCS5Padding(int blockSize)
            throws NoSuchPaddingException {
            if (blockSize == 0) {
                throw new NoSuchPaddingException
                        ("PKCS#5 padding not supported with stream ciphers");
            }
            this.blockSize = blockSize;
        }

        public int getPadLen(int dataLen) {
            return (blockSize - (dataLen & (blockSize - 1)));
        }

        public byte[] getPaddingBytes(int dataLen) {
            byte padValue = (byte) getPadLen(dataLen);
            byte[] paddingBytes = new byte[padValue];
            Arrays.fill(paddingBytes, padValue);
            return paddingBytes;
        }

        public byte[] bufferBytes(byte[] dataFromUpdate) {
            if (dataFromUpdate == null || dataFromUpdate.length == 0) {
                return null;
            }
            byte[] result = null;
            if (trailingBytes == null) {
                trailingBytes = ByteBuffer.wrap(new byte[blockSize]);
            }
            int tbSize = trailingBytes.position();
            if (dataFromUpdate.length > trailingBytes.remaining()) {
                int totalLen = dataFromUpdate.length + tbSize;
                int newTBSize = totalLen % blockSize;
                if (newTBSize == 0) {
                    newTBSize = blockSize;
                }
                if (tbSize == 0) {
                    result = Arrays.copyOf(dataFromUpdate, totalLen - newTBSize);
                } else {
                    // combine 'trailingBytes' and 'dataFromUpdate'
                    result = Arrays.copyOf(trailingBytes.array(),
                                           totalLen - newTBSize);
                    if (result.length != tbSize) {
                        System.arraycopy(dataFromUpdate, 0, result, tbSize,
                                         result.length - tbSize);
                    }
                }
                // update 'trailingBytes' w/ remaining bytes in 'dataFromUpdate'
                trailingBytes.clear();
                trailingBytes.put(dataFromUpdate,
                                  dataFromUpdate.length - newTBSize, newTBSize);
            } else {
                trailingBytes.put(dataFromUpdate);
            }
            return result;
        }

        public int getBufferedLength() {
            if (trailingBytes != null) {
                return trailingBytes.position();
            }
            return 0;
        }

        public int unpad(byte[] lastData, byte[] out, int outOfs)
                throws BadPaddingException, IllegalBlockSizeException,
                ShortBufferException {
            int tbSize = (trailingBytes == null? 0:trailingBytes.position());
            int dataLen = tbSize + lastData.length;
            // check total length
            if ((dataLen < 1) || (dataLen % blockSize != 0)) {
                UcryptoProvider.debug("PKCS5Padding: unpad, buffered " + tbSize +
                                 " bytes, last block " + lastData.length + " bytes");

                throw new IllegalBlockSizeException
                    ("Input length must be multiples of " + blockSize);
            }

            // check padding bytes
            if (lastData.length == 0) {
                if (tbSize != 0) {
                    // work on 'trailingBytes' directly
                    lastData = Arrays.copyOf(trailingBytes.array(), tbSize);
                    trailingBytes.clear();
                    tbSize = 0;
                } else {
                    throw new BadPaddingException("No pad bytes found!");
                }
            }
            byte padValue = lastData[lastData.length - 1];
            if (padValue < 1 || padValue > blockSize) {
                UcryptoProvider.debug("PKCS5Padding: unpad, lastData: " + Arrays.toString(lastData));
                UcryptoProvider.debug("PKCS5Padding: unpad, padValue=" + padValue);
                throw new BadPaddingException("Invalid pad value: " + padValue);
            }

            // sanity check padding bytes
            int padStartIndex = lastData.length - padValue;
            for (int i = padStartIndex; i < lastData.length; i++) {
                if (lastData[i] != padValue) {
                    UcryptoProvider.debug("PKCS5Padding: unpad, lastData: " + Arrays.toString(lastData));
                    UcryptoProvider.debug("PKCS5Padding: unpad, padValue=" + padValue);
                    throw new BadPaddingException("Invalid padding bytes!");
                }
            }

            int actualOutLen = dataLen - padValue;
            // check output buffer capacity
            if (out.length - outOfs < actualOutLen) {
                throw new ShortBufferException("Output buffer too small, need " + actualOutLen +
                    ", got " + (out.length - outOfs));
            }
            try {
                if (tbSize != 0) {
                    trailingBytes.rewind();
                    if (tbSize < actualOutLen) {
                        trailingBytes.get(out, outOfs, tbSize);
                        outOfs += tbSize;
                    } else {
                        // copy from trailingBytes and we are done
                        trailingBytes.get(out, outOfs, actualOutLen);
                        return actualOutLen;
                    }
                }
                if (lastData.length > padValue) {
                    System.arraycopy(lastData, 0, out, outOfs,
                                     lastData.length - padValue);
                }
                return actualOutLen;
            } finally {
                clear();
            }
        }

        public void clear() {
            if (trailingBytes != null) trailingBytes.clear();
        }
    }

    public static final class AesEcbPKCS5 extends NativeCipherWithJavaPadding {
        public AesEcbPKCS5() throws NoSuchAlgorithmException, NoSuchPaddingException {
            super(new NativeCipher.AesEcbNoPadding(), "PKCS5Padding");
        }
    }

    public static final class AesCbcPKCS5 extends NativeCipherWithJavaPadding {
        public AesCbcPKCS5() throws NoSuchAlgorithmException, NoSuchPaddingException {
            super(new NativeCipher.AesCbcNoPadding(), "PKCS5Padding");
        }
    }

    public static final class AesCfb128PKCS5 extends NativeCipherWithJavaPadding {
        public AesCfb128PKCS5() throws NoSuchAlgorithmException, NoSuchPaddingException {
            super(new NativeCipher.AesCfb128NoPadding(), "PKCS5Padding");
        }
    }

    // fields (re)set in every init()
    private final NativeCipher nc;
    private final Padding padding;
    private final int blockSize;
    private int lastBlockLen = 0;

    // Only ECB, CBC, CTR, and CFB128 modes w/ NOPADDING for now
    NativeCipherWithJavaPadding(NativeCipher nc, String paddingScheme)
        throws NoSuchAlgorithmException, NoSuchPaddingException {
        this.nc = nc;
        this.blockSize = nc.engineGetBlockSize();
        if (paddingScheme.toUpperCase().equals("PKCS5PADDING")) {
            padding = new PKCS5Padding(blockSize);
        } else {
            throw new NoSuchAlgorithmException("Unsupported padding scheme: " + paddingScheme);
        }
    }

    void reset() {
        padding.clear();
        lastBlockLen = 0;
    }

    @Override
    protected synchronized void engineSetMode(String mode) throws NoSuchAlgorithmException {
        nc.engineSetMode(mode);
    }

    // see JCE spec
    @Override
    protected void engineSetPadding(String padding)
            throws NoSuchPaddingException {
        // Disallow change of padding for now since currently it's explicitly
        // defined in transformation strings
        throw new NoSuchPaddingException("Unsupported padding " + padding);
    }

    // see JCE spec
    @Override
    protected int engineGetBlockSize() {
        return blockSize;
    }

    // see JCE spec
    @Override
    protected synchronized int engineGetOutputSize(int inputLen) {
        int result = nc.engineGetOutputSize(inputLen);
        if (nc.encrypt) {
            result += padding.getPadLen(result);
        } else {
            result += padding.getBufferedLength();
        }
        return result;
    }

    // see JCE spec
    @Override
    protected synchronized byte[] engineGetIV() {
        return nc.engineGetIV();
    }

    // see JCE spec
    @Override
    protected synchronized AlgorithmParameters engineGetParameters() {
        return nc.engineGetParameters();
    }

    @Override
    protected int engineGetKeySize(Key key) throws InvalidKeyException {
        return nc.engineGetKeySize(key);
    }

    // see JCE spec
    @Override
    protected synchronized void engineInit(int opmode, Key key, SecureRandom random)
            throws InvalidKeyException {
        reset();
        nc.engineInit(opmode, key, random);
    }

    // see JCE spec
    @Override
    protected synchronized void engineInit(int opmode, Key key,
            AlgorithmParameterSpec params, SecureRandom random)
            throws InvalidKeyException, InvalidAlgorithmParameterException {
        reset();
        nc.engineInit(opmode, key, params, random);
    }

    // see JCE spec
    @Override
    protected synchronized void engineInit(int opmode, Key key, AlgorithmParameters params,
            SecureRandom random)
            throws InvalidKeyException, InvalidAlgorithmParameterException {
        reset();
        nc.engineInit(opmode, key, params, random);
    }

    // see JCE spec
    @Override
    protected synchronized byte[] engineUpdate(byte[] in, int inOfs, int inLen) {
        if (nc.encrypt) {
            lastBlockLen += inLen;
            lastBlockLen &= (blockSize - 1);
            return nc.engineUpdate(in, inOfs, inLen);
        } else {
            return padding.bufferBytes(nc.engineUpdate(in, inOfs, inLen));
        }
    }

    // see JCE spec
    @Override
    protected synchronized int engineUpdate(byte[] in, int inOfs, int inLen, byte[] out,
            int outOfs) throws ShortBufferException {
        if (nc.encrypt) {
            lastBlockLen += inLen;
            lastBlockLen &= (blockSize - 1);
            return nc.engineUpdate(in, inOfs, inLen, out, outOfs);
        } else {
            byte[] result = padding.bufferBytes(nc.engineUpdate(in, inOfs, inLen));
            if (result != null) {
                System.arraycopy(result, 0, out, outOfs, result.length);
                return result.length;
            } else return 0;
        }
    }

    // see JCE spec
    @Override
    protected synchronized byte[] engineDoFinal(byte[] in, int inOfs, int inLen)
            throws IllegalBlockSizeException, BadPaddingException {
        int estimatedOutLen = engineGetOutputSize(inLen);
        byte[] out = new byte[estimatedOutLen];
        try {
            int actualOut = this.engineDoFinal(in, inOfs, inLen, out, 0);
            // truncate off extra bytes
            if (actualOut != out.length) {
                out = Arrays.copyOf(out, actualOut);
            }
        } catch (ShortBufferException sbe) {
            throw new UcryptoException("Internal Error", sbe);
        } finally {
            reset();
        }
        return out;
    }

    // see JCE spec
    @Override
    protected synchronized int engineDoFinal(byte[] in, int inOfs, int inLen, byte[] out,
                                             int outOfs)
        throws ShortBufferException, IllegalBlockSizeException,
               BadPaddingException {
        int estimatedOutLen = engineGetOutputSize(inLen);

        if (out.length - outOfs < estimatedOutLen) {
            throw new ShortBufferException("Actual: " + (out.length - outOfs) +
                ". Estimated Out Length: " + estimatedOutLen);
        }
        try {
            if (nc.encrypt) {
                int k = nc.engineUpdate(in, inOfs, inLen, out, outOfs);
                lastBlockLen += inLen;
                lastBlockLen &= (blockSize - 1);
                byte[] padBytes = padding.getPaddingBytes(lastBlockLen);
                k += nc.engineDoFinal(padBytes, 0, padBytes.length, out, (outOfs + k));
                return k;
            } else {
                byte[] tempOut = nc.engineDoFinal(in, inOfs, inLen);
                int len = padding.unpad(tempOut, out, outOfs);
                return len;
            }
        } finally {
            reset();
        }
    }

    // see JCE spec
    @Override
    protected synchronized byte[] engineWrap(Key key) throws IllegalBlockSizeException,
                                                InvalidKeyException {
        byte[] result = null;
        try {
            byte[] encodedKey = key.getEncoded();
            if ((encodedKey == null) || (encodedKey.length == 0)) {
                throw new InvalidKeyException("Cannot get an encoding of " +
                                              "the key to be wrapped");
            }
            result = engineDoFinal(encodedKey, 0, encodedKey.length);
        } catch (BadPaddingException e) {
            // Should never happen for key wrapping
            throw new UcryptoException("Internal Error", e);
        }
        return result;
    }

    // see JCE spec
    @Override
    protected synchronized Key engineUnwrap(byte[] wrappedKey, String wrappedKeyAlgorithm,
                               int wrappedKeyType)
        throws InvalidKeyException, NoSuchAlgorithmException {

        byte[] encodedKey;
        try {
            encodedKey = engineDoFinal(wrappedKey, 0,
                                       wrappedKey.length);
        } catch (Exception e) {
            throw (InvalidKeyException)
                (new InvalidKeyException()).initCause(e);
        }

        return NativeCipher.constructKey(wrappedKeyType, encodedKey,
                                         wrappedKeyAlgorithm);
    }
}