/*
 * 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.commons.crypto.stream;

import java.io.IOException;
import java.io.InputStream;
import java.nio.ByteBuffer;
import java.nio.channels.ReadableByteChannel;
import java.security.GeneralSecurityException;
import java.util.Properties;

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

import org.apache.commons.crypto.cipher.CryptoCipher;
import org.apache.commons.crypto.cipher.CryptoCipherFactory;
import org.apache.commons.crypto.stream.input.ChannelInput;
import org.apache.commons.crypto.stream.input.Input;
import org.apache.commons.crypto.stream.input.StreamInput;
import org.apache.commons.crypto.utils.AES;
import org.apache.commons.crypto.utils.Utils;

/**
 * <p>
 * CtrCryptoInputStream decrypts data. AES CTR mode is required in order to
 * ensure that the plain text and cipher text have a 1:1 mapping. CTR crypto
 * stream has stream characteristic which is useful for implement features like
 * random seek. The decryption is buffer based. The key points of the decryption
 * are (1) calculating the counter and (2) padding through stream position:
 * </p>
 * <p>
 * counter = base + pos/(algorithm blocksize); padding = pos%(algorithm
 * blocksize);
 * </p>
 * The underlying stream offset is maintained as state. It is not thread-safe.
 */
public class CtrCryptoInputStream extends CryptoInputStream {
    /**
     * <p>
     * This method is only for Counter (CTR) mode. Generally the CryptoCipher
     * calculates the IV and maintain encryption context internally.For example
     * a Cipher will maintain its encryption context internally when we do
     * encryption/decryption using the CryptoCipher#update interface.
     * </p>
     * <p>
     * Encryption/Decryption is not always on the entire file. For example, in
     * Hadoop, a node may only decrypt a portion of a file (i.e. a split). In
     * these situations, the counter is derived from the file position.
     * </p>
     * The IV can be calculated by combining the initial IV and the counter with
     * a lossless operation (concatenation, addition, or XOR).
     *
     * @see <a
     *      href="http://en.wikipedia.org/wiki/Block_cipher_mode_of_operation#Counter_.28CTR.29">
     *      http://en.wikipedia.org/wiki/Block_cipher_mode_of_operation#Counter_.28CTR.29</a>
     *
     * @param initIV initial IV
     * @param counter counter for input stream position
     * @param IV the IV for input stream position
     */
    static void calculateIV(final byte[] initIV, long counter, final byte[] IV) {
        int i = IV.length; // IV length

        Utils.checkArgument(initIV.length == CryptoCipherFactory.AES_BLOCK_SIZE);
        Utils.checkArgument(i == CryptoCipherFactory.AES_BLOCK_SIZE);

        int j = 0; // counter bytes index
        int sum = 0;
        while (i-- > 0) {
            // (sum >>> Byte.SIZE) is the carry for addition
            sum = (initIV[i] & 0xff) + (sum >>> Byte.SIZE); // NOPMD
            if (j++ < 8) { // Big-endian, and long is 8 bytes length
                sum += (byte) counter & 0xff;
                counter >>>= 8;
            }
            IV[i] = (byte) sum;
        }
    }

    /**
     * Underlying stream offset
     */
    private long streamOffset;

    /**
     * The initial IV.
     */
    private final byte[] initIV;

    /**
     * Initialization vector for the cipher.
     */
    private final byte[] iv;

    /**
     * Padding = pos%(algorithm blocksize); Padding is put into
     * {@link #inBuffer} before any other data goes in. The purpose of padding
     * is to put the input data at proper position.
     */
    private byte padding;

    /**
     * Flag to mark whether the cipher has been reset
     */
    private boolean cipherReset;

    /**
     * Constructs a {@link CtrCryptoInputStream}.
     *
     * @param input the input data.
     * @param cipher the CryptoCipher instance.
     * @param bufferSize the bufferSize.
     * @param key crypto key for the cipher.
     * @param iv Initialization vector for the cipher.
     * @throws IOException if an I/O error occurs.
     */
    protected CtrCryptoInputStream(final Input input, final CryptoCipher cipher,
            final int bufferSize, final byte[] key, final byte[] iv) throws IOException {
        this(input, cipher, bufferSize, key, iv, 0);
    }

    /**
     * Constructs a {@link CtrCryptoInputStream}.
     *
     * @param input the input data.
     * @param cipher the CryptoCipher instance.
     * @param bufferSize the bufferSize.
     * @param key crypto key for the cipher.
     * @param iv Initialization vector for the cipher.
     * @param streamOffset the start offset in the stream.
     * @throws IOException if an I/O error occurs.
     */
    protected CtrCryptoInputStream(final Input input, final CryptoCipher cipher,
            final int bufferSize, final byte[] key, final byte[] iv, final long streamOffset)
            throws IOException {
        super(input, cipher, bufferSize, AES.newSecretKeySpec(key),
                new IvParameterSpec(iv));

        this.initIV = iv.clone();
        this.iv = iv.clone();

        CryptoInputStream.checkStreamCipher(cipher);

        resetStreamOffset(streamOffset);
    }

    /**
     * Constructs a {@link CtrCryptoInputStream}.
     *
     * @param inputStream the input stream.
     * @param cipher the CryptoCipher instance.
     * @param bufferSize the bufferSize.
     * @param key crypto key for the cipher.
     * @param iv Initialization vector for the cipher.
     * @throws IOException if an I/O error occurs.
     */
    protected CtrCryptoInputStream(final InputStream inputStream, final CryptoCipher cipher,
            final int bufferSize, final byte[] key, final byte[] iv) throws IOException {
        this(inputStream, cipher, bufferSize, key, iv, 0);
    }

    /**
     * Constructs a {@link CtrCryptoInputStream}.
     *
     * @param inputStream the InputStream instance.
     * @param cipher the CryptoCipher instance.
     * @param bufferSize the bufferSize.
     * @param key crypto key for the cipher.
     * @param iv Initialization vector for the cipher.
     * @param streamOffset the start offset in the stream.
     * @throws IOException if an I/O error occurs.
     */
    @SuppressWarnings("resource") // Closing the instance closes the StreamInput
    protected CtrCryptoInputStream(final InputStream inputStream, final CryptoCipher cipher,
            final int bufferSize, final byte[] key, final byte[] iv, final long streamOffset)
            throws IOException {
        this(new StreamInput(inputStream, bufferSize), cipher, bufferSize, key, iv,
                streamOffset);
    }

    /**
     * Constructs a {@link CtrCryptoInputStream}.
     *
     * @param properties The {@code Properties} class represents a set of
     *        properties.
     * @param inputStream the input stream.
     * @param key crypto key for the cipher.
     * @param iv Initialization vector for the cipher.
     * @throws IOException if an I/O error occurs.
     */
    public CtrCryptoInputStream(final Properties properties, final InputStream inputStream, final byte[] key,
            final byte[] iv) throws IOException {
        this(properties, inputStream, key, iv, 0);
    }

    /**
     * Constructs a {@link CtrCryptoInputStream}.
     *
     * @param properties The {@code Properties} class represents a set of
     *        properties.
     * @param inputStream the InputStream instance.
     * @param key crypto key for the cipher.
     * @param iv Initialization vector for the cipher.
     * @param streamOffset the start offset in the stream.
     * @throws IOException if an I/O error occurs.
     */
    @SuppressWarnings("resource") // The CryptoCipher returned by getCipherInstance() is closed by CtrCryptoInputStream.
    public CtrCryptoInputStream(final Properties properties, final InputStream inputStream, final byte[] key,
            final byte[] iv, final long streamOffset) throws IOException {
        this(inputStream, Utils.getCipherInstance(
                AES.CTR_NO_PADDING, properties),
                CryptoInputStream.getBufferSize(properties), key, iv, streamOffset);
    }

    /**
     * Constructs a {@link CtrCryptoInputStream}.
     *
     * @param properties The {@code Properties} class represents a set of
     *        properties.
     * @param channel the ReadableByteChannel instance.
     * @param key crypto key for the cipher.
     * @param iv Initialization vector for the cipher.
     * @throws IOException if an I/O error occurs.
     */
    public CtrCryptoInputStream(final Properties properties, final ReadableByteChannel channel,
            final byte[] key, final byte[] iv) throws IOException {
        this(properties, channel, key, iv, 0);
    }

    /**
     * Constructs a {@link CtrCryptoInputStream}.
     *
     * @param properties The {@code Properties} class represents a set of
     *        properties.
     * @param in the ReadableByteChannel instance.
     * @param key crypto key for the cipher.
     * @param iv Initialization vector for the cipher.
     * @param streamOffset the start offset in the stream.
     * @throws IOException if an I/O error occurs.
     */
    @SuppressWarnings("resource") // The CryptoCipher returned by getCipherInstance() is closed by CtrCryptoInputStream.
    public CtrCryptoInputStream(final Properties properties, final ReadableByteChannel in,
            final byte[] key, final byte[] iv, final long streamOffset) throws IOException {
        this(in, Utils.getCipherInstance(
                AES.CTR_NO_PADDING, properties),
                CryptoInputStream.getBufferSize(properties), key, iv, streamOffset);
    }

    /**
     * Constructs a {@link CtrCryptoInputStream}.
     *
     * @param channel the ReadableByteChannel instance.
     * @param cipher the cipher instance.
     * @param bufferSize the bufferSize.
     * @param key crypto key for the cipher.
     * @param iv Initialization vector for the cipher.
     * @throws IOException if an I/O error occurs.
     */
    protected CtrCryptoInputStream(final ReadableByteChannel channel, final CryptoCipher cipher,
            final int bufferSize, final byte[] key, final byte[] iv) throws IOException {
        this(channel, cipher, bufferSize, key, iv, 0);
    }

    /**
     * Constructs a {@link CtrCryptoInputStream}.
     *
     * @param channel the ReadableByteChannel instance.
     * @param cipher the CryptoCipher instance.
     * @param bufferSize the bufferSize.
     * @param key crypto key for the cipher.
     * @param iv Initialization vector for the cipher.
     * @param streamOffset the start offset in the stream.
     * @throws IOException if an I/O error occurs.
     */
    @SuppressWarnings("resource") // Closing the instance closes the ChannelInput
    protected CtrCryptoInputStream(final ReadableByteChannel channel, final CryptoCipher cipher,
            final int bufferSize, final byte[] key, final byte[] iv, final long streamOffset)
            throws IOException {
        this(new ChannelInput(channel), cipher, bufferSize, key, iv, streamOffset);
    }

    /**
     * Does the decryption using inBuffer as input and outBuffer as output. Upon
     * return, inBuffer is cleared; the decrypted data starts at
     * outBuffer.position() and ends at outBuffer.limit().
     *
     * @throws IOException if an I/O error occurs.
     */
    @Override
    protected void decrypt() throws IOException {
        Utils.checkState(inBuffer.position() >= padding);
        if (inBuffer.position() == padding) {
            // There is no real data in inBuffer.
            return;
        }

        inBuffer.flip();
        outBuffer.clear();
        decryptBuffer(outBuffer);
        inBuffer.clear();
        outBuffer.flip();

        if (padding > 0) {
            /*
             * The plain text and cipher text have a 1:1 mapping, they start at
             * the same position.
             */
            outBuffer.position(padding);
        }
    }

    /**
     * Decrypts all data in buf: total n bytes from given start position. Output
     * is also buf and same start position. buf.position() and buf.limit()
     * should be unchanged after decryption.
     *
     * @param buf The buffer into which bytes are to be transferred.
     * @param offset the start offset in the data.
     * @param len the maximum number of decrypted data bytes to read.
     * @throws IOException if an I/O error occurs.
     */
    protected void decrypt(final ByteBuffer buf, final int offset, final int len)
            throws IOException {
        final int pos = buf.position();
        final int limit = buf.limit();
        int n = 0;
        while (n < len) {
            buf.position(offset + n);
            buf.limit(offset + n + Math.min(len - n, inBuffer.remaining()));
            inBuffer.put(buf);
            // Do decryption
            try {
                decrypt();
                buf.position(offset + n);
                buf.limit(limit);
                n += outBuffer.remaining();
                buf.put(outBuffer);
            } finally {
                padding = postDecryption(streamOffset - (len - n));
            }
        }
        buf.position(pos);
    }

    /**
     * Does the decryption using out as output.
     *
     * @param out the output ByteBuffer.
     * @throws IOException if an I/O error occurs.
     */
    protected void decryptBuffer(final ByteBuffer out) throws IOException {
        final int inputSize = inBuffer.remaining();
        try {
            final int n = cipher.update(inBuffer, out);
            if (n < inputSize) {
                /**
                 * Typically code will not get here. CryptoCipher#update will
                 * consume all input data and put result in outBuffer.
                 * CryptoCipher#doFinal will reset the cipher context.
                 */
                cipher.doFinal(inBuffer, out);
                cipherReset = true;
            }
        } catch (final GeneralSecurityException e) {
            throw new IOException(e);
        }
    }

    /**
     * Does the decryption using inBuffer as input and buf as output. Upon
     * return, inBuffer is cleared; the buf's position will be equal to
     * <i>p</i>&nbsp;{@code +}&nbsp;<i>n</i> where <i>p</i> is the position
     * before decryption, <i>n</i> is the number of bytes decrypted. The buf's
     * limit will not have changed.
     *
     * @param buf The buffer into which bytes are to be transferred.
     * @throws IOException if an I/O error occurs.
     */
    protected void decryptInPlace(final ByteBuffer buf) throws IOException {
        Utils.checkState(inBuffer.position() >= padding);
        Utils.checkState(buf.isDirect());
        Utils.checkState(buf.remaining() >= inBuffer.position());
        Utils.checkState(padding == 0);

        if (inBuffer.position() == padding) {
            // There is no real data in inBuffer.
            return;
        }
        inBuffer.flip();
        decryptBuffer(buf);
        inBuffer.clear();
    }

    /**
     * Decrypts more data by reading the under layer stream. The decrypted data
     * will be put in the output buffer.
     *
     * @return The number of decrypted data. -1 if end of the decrypted stream.
     * @throws IOException if an I/O error occurs.
     */
    @Override
    protected int decryptMore() throws IOException {
        final int n = input.read(inBuffer);
        if (n <= 0) {
            return n;
        }

        streamOffset += n; // Read n bytes
        decrypt();
        padding = postDecryption(streamOffset);
        return outBuffer.remaining();
    }

    /**
     * Gets the counter for input stream position.
     *
     * @param position the given position in the data.
     * @return the counter for input stream position.
     */
    protected long getCounter(final long position) {
        return position / cipher.getBlockSize();
    }

    /**
     * Gets the initialization vector.
     *
     * @return the initIV.
     */
    protected byte[] getInitIV() {
        return initIV;
    }

    /**
     * Gets the padding for input stream position.
     *
     * @param position the given position in the data.
     * @return the padding for input stream position.
     */
    protected byte getPadding(final long position) {
        return (byte) (position % cipher.getBlockSize());
    }

    /**
     * Gets the offset of the stream.
     *
     * @return the stream offset.
     */
    protected long getStreamOffset() {
        return streamOffset;
    }

    /**
     * Gets the position of the stream.
     *
     * @return the position of the stream.
     */
    protected long getStreamPosition() {
        return streamOffset - outBuffer.remaining();
    }

    /**
     * Overrides the {@link CtrCryptoInputStream#initCipher()}. Initializes the
     * cipher.
     */
    @Override
    protected void initCipher() {
        // Do nothing for initCipher
        // Will reset the cipher when reset the stream offset
    }

    /**
     * This method is executed immediately after decryption. Checks whether
     * cipher should be updated and recalculate padding if needed.
     *
     * @param position the given position in the data..
     * @return the byte.
     * @throws IOException if an I/O error occurs.
     */
    protected byte postDecryption(final long position) throws IOException {
        byte padding = 0;
        if (cipherReset) {
            /*
             * This code is generally not executed since the cipher usually
             * maintains cipher context (e.g. the counter) internally. However,
             * some implementations can't maintain context so a re-init is
             * necessary after each decryption call.
             */
            resetCipher(position);
            padding = getPadding(position);
            inBuffer.position(padding);
        }
        return padding;
    }

    /**
     * Overrides the {@link CtrCryptoInputStream#read(ByteBuffer)}. Reads a
     * sequence of bytes from this channel into the given buffer.
     *
     * @param buf The buffer into which bytes are to be transferred.
     * @return The number of bytes read, possibly zero, or {@code -1} if the
     *         channel has reached end-of-stream.
     * @throws IOException if an I/O error occurs.
     */
    @Override
    public int read(final ByteBuffer buf) throws IOException {
        checkStream();
        int unread = outBuffer.remaining();
        if (unread <= 0) { // Fill the unread decrypted data buffer firstly
            final int n = input.read(inBuffer);
            if (n <= 0) {
                return n;
            }

            streamOffset += n; // Read n bytes
            if (buf.isDirect() && buf.remaining() >= inBuffer.position()
                    && padding == 0) {
                // Use buf as the output buffer directly
                decryptInPlace(buf);
                padding = postDecryption(streamOffset);
                return n;
            }
            // Use outBuffer as the output buffer
            decrypt();
            padding = postDecryption(streamOffset);
        }

        // Copy decrypted data from outBuffer to buf
        unread = outBuffer.remaining();
        final int toRead = buf.remaining();
        if (toRead <= unread) {
            final int limit = outBuffer.limit();
            outBuffer.limit(outBuffer.position() + toRead);
            buf.put(outBuffer);
            outBuffer.limit(limit);
            return toRead;
        }
        buf.put(outBuffer);
        return unread;
    }

    /**
     * Calculates the counter and iv, resets the cipher.
     *
     * @param position the given position in the data.
     * @throws IOException if an I/O error occurs.
     */
    protected void resetCipher(final long position) throws IOException {
        final long counter = getCounter(position);
        CtrCryptoInputStream.calculateIV(initIV, counter, iv);
        try {
            cipher.init(Cipher.DECRYPT_MODE, key, new IvParameterSpec(iv));
        } catch (final GeneralSecurityException e) {
            throw new IOException(e);
        }
        cipherReset = false;
    }

    /**
     * Resets the underlying stream offset; clear {@link #inBuffer} and
     * {@link #outBuffer}. This Typically happens during {@link #skip(long)}.
     *
     * @param offset the offset of the stream.
     * @throws IOException if an I/O error occurs.
     */
    protected void resetStreamOffset(final long offset) throws IOException {
        streamOffset = offset;
        inBuffer.clear();
        outBuffer.clear();
        outBuffer.limit(0);
        resetCipher(offset);
        padding = getPadding(offset);
        inBuffer.position(padding); // Set proper position for input data.
    }

    /**
     * Seeks the stream to a specific position relative to start of the under
     * layer stream.
     *
     * @param position the given position in the data.
     * @throws IOException if an I/O error occurs.
     */
    public void seek(final long position) throws IOException {
        Utils.checkArgument(position >= 0, "Cannot seek to negative offset.");
        checkStream();
        /*
         * If data of target pos in the underlying stream has already been read
         * and decrypted in outBuffer, we just need to re-position outBuffer.
         */
        if (position >= getStreamPosition() && position <= getStreamOffset()) {
            final int forward = (int) (position - getStreamPosition());
            if (forward > 0) {
                outBuffer.position(outBuffer.position() + forward);
            }
        } else {
            input.seek(position);
            resetStreamOffset(position);
        }
    }

    /**
     * Sets the offset of stream.
     *
     * @param streamOffset the stream offset.
     */
    protected void setStreamOffset(final long streamOffset) {
        this.streamOffset = streamOffset;
    }

    /**
     * Overrides the {@link CryptoInputStream#skip(long)}. Skips over and
     * discards {@code n} bytes of data from this input stream.
     *
     * @param n the number of bytes to be skipped.
     * @return the actual number of bytes skipped.
     * @throws IOException if an I/O error occurs.
     */
    @Override
    public long skip(long n) throws IOException {
        Utils.checkArgument(n >= 0, "Negative skip length.");
        checkStream();

        if (n == 0) {
            return 0;
        }
        if (n <= outBuffer.remaining()) {
            final int pos = outBuffer.position() + (int) n;
            outBuffer.position(pos);
            return n;
        }
        /*
         * Subtract outBuffer.remaining() to see how many bytes we need to
         * skip in the underlying stream. Add outBuffer.remaining() to the
         * actual number of skipped bytes in the underlying stream to get
         * the number of skipped bytes from the user's point of view.
         */
        n -= outBuffer.remaining();
        long skipped = input.skip(n);
        if (skipped < 0) {
            skipped = 0;
        }
        final long pos = streamOffset + skipped;
        skipped += outBuffer.remaining();
        resetStreamOffset(pos);
        return skipped;
    }
}