/*

 * Copyright (c) 2003, 2011, 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 sun.security.pkcs11;

import java.nio.ByteBuffer;

import java.util.Arrays;

import java.util.Locale;

import java.security.*;

import java.security.spec.*;

import javax.crypto.*;

import javax.crypto.spec.*;

import sun.nio.ch.DirectBuffer;

import sun.security.pkcs11.wrapper.*;

import static sun.security.pkcs11.wrapper.PKCS11Constants.*;

/**

 * Cipher implementation class. This class currently supports

 * DES, DESede, AES, ARCFOUR, and Blowfish.

 *

 * This class is designed to support ECB, CBC, CTR with NoPadding

 * and ECB, CBC with PKCS5Padding. It will use its own padding impl

 * if the native mechanism does not support padding.

 *

 * Note that PKCS#11 currently only supports ECB, CBC, and CTR.

 * There are no provisions for other modes such as CFB, OFB, and PCBC.

 *

 * @author  Andreas Sterbenz

 * @since   1.5

 */

final class P11Cipher extends CipherSpi {

    // mode constant for ECB mode

    private final static int MODE_ECB = 3;

    // mode constant for CBC mode

    private final static int MODE_CBC = 4;

    // mode constant for CTR mode

    private final static int MODE_CTR = 5;

    // padding constant for NoPadding

    private final static int PAD_NONE = 5;

    // padding constant for PKCS5Padding

    private final static int PAD_PKCS5 = 6;

    private static interface Padding {

        // ENC: format the specified buffer with padding bytes and return the

        // actual padding length

        int setPaddingBytes(byte[] paddingBuffer, int padLen);

        // DEC: return the length of trailing padding bytes given the specified

        // padded data

        int unpad(byte[] paddedData, int len)

                throws BadPaddingException, IllegalBlockSizeException;

    }

    private static class PKCS5Padding implements Padding {

        private final int blockSize;

        PKCS5Padding(int blockSize)

                throws NoSuchPaddingException {

            if (blockSize == 0) {

                throw new NoSuchPaddingException

                        ("PKCS#5 padding not supported with stream ciphers");

            }

            this.blockSize = blockSize;

        }

        public int setPaddingBytes(byte[] paddingBuffer, int padLen) {

            Arrays.fill(paddingBuffer, 0, padLen, (byte) (padLen & 0x007f));

            return padLen;

        }

        public int unpad(byte[] paddedData, int len)

                throws BadPaddingException, IllegalBlockSizeException {

            if ((len < 1) || (len % blockSize != 0)) {

                throw new IllegalBlockSizeException

                    ("Input length must be multiples of " + blockSize);

            }

            byte padValue = paddedData[len - 1];

            if (padValue < 1 || padValue > blockSize) {

                throw new BadPaddingException("Invalid pad value!");

            }

            // sanity check padding bytes

            int padStartIndex = len - padValue;

            for (int i = padStartIndex; i < len; i++) {

                if (paddedData[i] != padValue) {

                    throw new BadPaddingException("Invalid pad bytes!");

                }

            }

            return padValue;

        }

    }

    // token instance

    private final Token token;

    // algorithm name

    private final String algorithm;

    // name of the key algorithm, e.g. DES instead of algorithm DES/CBC/...

    private final String keyAlgorithm;

    // mechanism id

    private final long mechanism;

    // associated session, if any

    private Session session;

    // key, if init() was called

    private P11Key p11Key;

    // flag indicating whether an operation is initialized

    private boolean initialized;

    // falg indicating encrypt or decrypt mode

    private boolean encrypt;

    // mode, one of MODE_* above (MODE_ECB for stream ciphers)

    private int blockMode;

    // block size, 0 for stream ciphers

    private final int blockSize;

    // padding type, on of PAD_* above (PAD_NONE for stream ciphers)

    private int paddingType;

    // when the padding is requested but unsupported by the native mechanism,

    // we use the following to do padding and necessary data buffering.

    // padding object which generate padding and unpad the decrypted data

    private Padding paddingObj;

    // buffer for holding back the block which contains padding bytes

    private byte[] padBuffer;

    private int padBufferLen;

    // original IV, if in MODE_CBC or MODE_CTR

    private byte[] iv;

    // number of bytes buffered internally by the native mechanism and padBuffer

    // if we do the padding

    private int bytesBuffered;

    P11Cipher(Token token, String algorithm, long mechanism)

            throws PKCS11Exception, NoSuchAlgorithmException {

        super();

        this.token = token;

        this.algorithm = algorithm;

        this.mechanism = mechanism;

        String algoParts[] = algorithm.split("/");

        keyAlgorithm = algoParts[0];

        if (keyAlgorithm.equals("AES")) {

            blockSize = 16;

        } else if (keyAlgorithm.equals("RC4") ||

                keyAlgorithm.equals("ARCFOUR")) {

            blockSize = 0;

        } else { // DES, DESede, Blowfish

            blockSize = 8;

        }

        this.blockMode =

                (algoParts.length > 1 ? parseMode(algoParts[1]) : MODE_ECB);

        String defPadding = (blockSize == 0 ? "NoPadding" : "PKCS5Padding");

        String paddingStr =

                (algoParts.length > 2 ? algoParts[2] : defPadding);

        try {

            engineSetPadding(paddingStr);

        } catch (NoSuchPaddingException nspe) {

            // should not happen

            throw new ProviderException(nspe);

        }

    }

    protected void engineSetMode(String mode) throws NoSuchAlgorithmException {

        // Disallow change of mode for now since currently it's explicitly

        // defined in transformation strings

        throw new NoSuchAlgorithmException("Unsupported mode " + mode);

    }

    private int parseMode(String mode) throws NoSuchAlgorithmException {

        mode = mode.toUpperCase(Locale.ENGLISH);

        int result;

        if (mode.equals("ECB")) {

            result = MODE_ECB;

        } else if (mode.equals("CBC")) {

            if (blockSize == 0) {

                throw new NoSuchAlgorithmException

                        ("CBC mode not supported with stream ciphers");

            }

            result = MODE_CBC;

        } else if (mode.equals("CTR")) {

            result = MODE_CTR;

        } else {

            throw new NoSuchAlgorithmException("Unsupported mode " + mode);

        }

        return result;

    }

    // see JCE spec

    protected void engineSetPadding(String padding)

            throws NoSuchPaddingException {

        paddingObj = null;

        padBuffer = null;

        padding = padding.toUpperCase(Locale.ENGLISH);

        if (padding.equals("NOPADDING")) {

            paddingType = PAD_NONE;

        } else if (padding.equals("PKCS5PADDING")) {

            if (this.blockMode == MODE_CTR) {

                throw new NoSuchPaddingException

                    ("PKCS#5 padding not supported with CTR mode");

            }

            paddingType = PAD_PKCS5;

            if (mechanism != CKM_DES_CBC_PAD && mechanism != CKM_DES3_CBC_PAD &&

                    mechanism != CKM_AES_CBC_PAD) {

                // no native padding support; use our own padding impl

                paddingObj = new PKCS5Padding(blockSize);

                padBuffer = new byte[blockSize];

            }

        } else {

            throw new NoSuchPaddingException("Unsupported padding " + padding);

        }

    }

    // see JCE spec

    protected int engineGetBlockSize() {

        return blockSize;

    }

    // see JCE spec

    protected int engineGetOutputSize(int inputLen) {

        return doFinalLength(inputLen);

    }

    // see JCE spec

    protected byte[] engineGetIV() {

    }

    // see JCE spec

    protected AlgorithmParameters engineGetParameters() {

        if (iv == null) {

            return null;

        }

        IvParameterSpec ivSpec = new IvParameterSpec(iv);

        try {

            AlgorithmParameters params =

                    AlgorithmParameters.getInstance(keyAlgorithm,

                    P11Util.getSunJceProvider());

            params.init(ivSpec);

            return params;

        } catch (GeneralSecurityException e) {

            // NoSuchAlgorithmException, NoSuchProviderException

            // InvalidParameterSpecException

            throw new ProviderException("Could not encode parameters", e);

        }

    }

    // see JCE spec

    protected void engineInit(int opmode, Key key, SecureRandom random)

            throws InvalidKeyException {

        try {

            implInit(opmode, key, null, random);

        } catch (InvalidAlgorithmParameterException e) {

            throw new InvalidKeyException("init() failed", e);

        }

    }

    // see JCE spec

    protected void engineInit(int opmode, Key key,

            AlgorithmParameterSpec params, SecureRandom random)

            throws InvalidKeyException, InvalidAlgorithmParameterException {

        byte[] ivValue;

        if (params != null) {

            if (params instanceof IvParameterSpec == false) {

                throw new InvalidAlgorithmParameterException

                        ("Only IvParameterSpec supported");

            }

            IvParameterSpec ivSpec = (IvParameterSpec) params;

            ivValue = ivSpec.getIV();

        } else {

            ivValue = null;

        }

        implInit(opmode, key, ivValue, random);

    }

    // see JCE spec

    protected void engineInit(int opmode, Key key, AlgorithmParameters params,

            SecureRandom random)

            throws InvalidKeyException, InvalidAlgorithmParameterException {

        byte[] ivValue;

        if (params != null) {

            try {

                        params.getParameterSpec(IvParameterSpec.class);

                ivValue = ivSpec.getIV();

            } catch (InvalidParameterSpecException e) {

                throw new InvalidAlgorithmParameterException

                        ("Could not decode IV", e);

            }

        } else {

            ivValue = null;

        }

        implInit(opmode, key, ivValue, random);

    }

    // actual init() implementation

    private void implInit(int opmode, Key key, byte[] iv,

            SecureRandom random)

            throws InvalidKeyException, InvalidAlgorithmParameterException {

        cancelOperation();

        switch (opmode) {

            case Cipher.ENCRYPT_MODE:

                encrypt = true;

                break;

            case Cipher.DECRYPT_MODE:

                encrypt = false;

                break;

            default:

                throw new InvalidAlgorithmParameterException

                        ("Unsupported mode: " + opmode);

        }

        if (blockMode == MODE_ECB) { // ECB or stream cipher

            if (iv != null) {

                if (blockSize == 0) {

                    throw new InvalidAlgorithmParameterException

                            ("IV not used with stream ciphers");

                } else {

                    throw new InvalidAlgorithmParameterException

                            ("IV not used in ECB mode");

                }

            }

        } else { // MODE_CBC or MODE_CTR

            if (iv == null) {

                if (encrypt == false) {

                    String exMsg =

                        (blockMode == MODE_CBC ?

                         "IV must be specified for decryption in CBC mode" :

                         "IV must be specified for decryption in CTR mode");

                    throw new InvalidAlgorithmParameterException(exMsg);

                }

                // generate random IV

                if (random == null) {

                    random = new SecureRandom();

                }

                iv = new byte[blockSize];

                random.nextBytes(iv);

            } else {

                if (iv.length != blockSize) {

                    throw new InvalidAlgorithmParameterException

                            ("IV length must match block size");

                }

            }

        }

        this.iv = iv;

        p11Key = P11SecretKeyFactory.convertKey(token, key, keyAlgorithm);

        try {

            initialize();

        } catch (PKCS11Exception e) {

            throw new InvalidKeyException("Could not initialize cipher", e);

        }

    }

    private void cancelOperation() {

        if (initialized == false) {

            return;

        }

        initialized = false;

        if ((session == null) || (token.explicitCancel == false)) {

            return;

        }

        // cancel operation by finishing it

        int bufLen = doFinalLength(0);

        byte[] buffer = new byte[bufLen];

        try {

            if (encrypt) {

                token.p11.C_EncryptFinal(session.id(), 0, buffer, 0, bufLen);

            } else {

                token.p11.C_DecryptFinal(session.id(), 0, buffer, 0, bufLen);

            }

        } catch (PKCS11Exception e) {

            throw new ProviderException("Cancel failed", e);

        } finally {

            reset();

        }

    }

    private void ensureInitialized() throws PKCS11Exception {

        if (initialized == false) {

            initialize();

        }

    }

    private void initialize() throws PKCS11Exception {

        if (session == null) {

            session = token.getOpSession();

        }

        CK_MECHANISM mechParams = (blockMode == MODE_CTR?

            new CK_MECHANISM(mechanism, new CK_AES_CTR_PARAMS(iv)) :

            new CK_MECHANISM(mechanism, iv));

        try {

            if (encrypt) {

                token.p11.C_EncryptInit(session.id(), mechParams, p11Key.keyID);

            } else {

                token.p11.C_DecryptInit(session.id(), mechParams, p11Key.keyID);

            }

        } catch (PKCS11Exception ex) {

            // release session when initialization failed

            session = token.releaseSession(session);

            throw ex;

        }

        bytesBuffered = 0;

        padBufferLen = 0;

        initialized = true;

    }

    // if update(inLen) is called, how big does the output buffer have to be?

    private int updateLength(int inLen) {

        if (inLen <= 0) {

            return 0;

        }

        int result = inLen + bytesBuffered;

        if (blockSize != 0) {

            // minus the number of bytes in the last incomplete block.

            result -= (result & (blockSize - 1));

        }

        return result;

    }

    // if doFinal(inLen) is called, how big does the output buffer have to be?

    private int doFinalLength(int inLen) {

        if (inLen < 0) {

            return 0;

        }

        int result = inLen + bytesBuffered;

        if (blockSize != 0 && encrypt && paddingType != PAD_NONE) {

            // add the number of bytes to make the last block complete.

            result += (blockSize - (result & (blockSize - 1)));

        }

        return result;

    }

    // reset the states to the pre-initialized values

    private void reset() {

        initialized = false;

        bytesBuffered = 0;

        padBufferLen = 0;

        if (session != null) {

            session = token.releaseSession(session);

        }

    }

    // see JCE spec

    protected byte[] engineUpdate(byte[] in, int inOfs, int inLen) {

        try {

            byte[] out = new byte[updateLength(inLen)];

            int n = engineUpdate(in, inOfs, inLen, out, 0);

            return P11Util.convert(out, 0, n);

        } catch (ShortBufferException e) {

            // convert since the output length is calculated by updateLength()

            throw new ProviderException(e);

        }

    }

    // see JCE spec

    protected int engineUpdate(byte[] in, int inOfs, int inLen, byte[] out,

            int outOfs) throws ShortBufferException {

        int outLen = out.length - outOfs;

        return implUpdate(in, inOfs, inLen, out, outOfs, outLen);

    }

    // see JCE spec

    @Override

    protected int engineUpdate(ByteBuffer inBuffer, ByteBuffer outBuffer)

            throws ShortBufferException {

        return implUpdate(inBuffer, outBuffer);

    }

    // see JCE spec

    protected byte[] engineDoFinal(byte[] in, int inOfs, int inLen)