package org.bouncycastle.crypto.io;

import java.io.FilterInputStream;
import java.io.IOException;
import java.io.InputStream;

import org.bouncycastle.crypto.BufferedBlockCipher;
import org.bouncycastle.crypto.InvalidCipherTextException;
import org.bouncycastle.crypto.SkippingCipher;
import org.bouncycastle.crypto.StreamCipher;
import org.bouncycastle.crypto.modes.AEADBlockCipher;
import org.bouncycastle.util.Arrays;

A CipherInputStream is composed of an InputStream and a cipher so that read() methods return data
that are read in from the underlying InputStream but have been additionally processed by the
Cipher. The cipher must be fully initialized before being used by a CipherInputStream.

For example, if the Cipher is initialized for decryption, the
CipherInputStream will attempt to read in data and decrypt them,
before returning the decrypted data.
/**
 * A CipherInputStream is composed of an InputStream and a cipher so that read() methods return data
 * that are read in from the underlying InputStream but have been additionally processed by the
 * Cipher. The cipher must be fully initialized before being used by a CipherInputStream.
 * <p>
 * For example, if the Cipher is initialized for decryption, the
 * CipherInputStream will attempt to read in data and decrypt them,
 * before returning the decrypted data.
 */
public class CipherInputStream
    extends FilterInputStream
{
    private static final int INPUT_BUF_SIZE = 2048;

    private SkippingCipher skippingCipher;
    private byte[] inBuf;

    private BufferedBlockCipher bufferedBlockCipher;
    private StreamCipher streamCipher;
    private AEADBlockCipher aeadBlockCipher;

    private byte[] buf;
    private byte[] markBuf;


    private int bufOff;
    private int maxBuf;
    private boolean finalized;
    private long markPosition;
    private int markBufOff;

    
Constructs a CipherInputStream from an InputStream and a
BufferedBlockCipher.
/**
     * Constructs a CipherInputStream from an InputStream and a
     * BufferedBlockCipher.
     */
    public CipherInputStream(
        InputStream is,
        BufferedBlockCipher cipher)
    {
        this(is, cipher, INPUT_BUF_SIZE);
    }

    
Constructs a CipherInputStream from an InputStream and a StreamCipher.
/**
     * Constructs a CipherInputStream from an InputStream and a StreamCipher.
     */
    public CipherInputStream(
        InputStream is,
        StreamCipher cipher)
    {
        this(is, cipher, INPUT_BUF_SIZE);
    }

    
Constructs a CipherInputStream from an InputStream and an AEADBlockCipher.
/**
     * Constructs a CipherInputStream from an InputStream and an AEADBlockCipher.
     */
    public CipherInputStream(
        InputStream is,
        AEADBlockCipher cipher)
    {
        this(is, cipher, INPUT_BUF_SIZE);
    }

    
Constructs a CipherInputStream from an InputStream, a
BufferedBlockCipher, and a specified internal buffer size.
/**
     * Constructs a CipherInputStream from an InputStream, a
     * BufferedBlockCipher, and a specified internal buffer size.
     */
    public CipherInputStream(
        InputStream is,
        BufferedBlockCipher cipher,
        int bufSize)
    {
        super(is);

        this.bufferedBlockCipher = cipher;
        this.inBuf = new byte[bufSize];
        this.skippingCipher = (cipher instanceof SkippingCipher) ? (SkippingCipher)cipher : null;
    }

    
Constructs a CipherInputStream from an InputStream, a StreamCipher, and a specified internal buffer size.
/**
     * Constructs a CipherInputStream from an InputStream, a StreamCipher, and a specified internal buffer size.
     */
    public CipherInputStream(
        InputStream is,
        StreamCipher cipher,
        int bufSize)
    {
        super(is);

        this.streamCipher = cipher;
        this.inBuf = new byte[bufSize];
        this.skippingCipher = (cipher instanceof SkippingCipher) ? (SkippingCipher)cipher : null;
    }

    
Constructs a CipherInputStream from an InputStream, an AEADBlockCipher, and a specified internal buffer size.
/**
     * Constructs a CipherInputStream from an InputStream, an AEADBlockCipher, and a specified internal buffer size.
     */
    public CipherInputStream(
        InputStream is,
        AEADBlockCipher cipher,
        int bufSize)
    {
        super(is);

        this.aeadBlockCipher = cipher;
        this.inBuf = new byte[bufSize];
        this.skippingCipher = (cipher instanceof SkippingCipher) ? (SkippingCipher)cipher : null;
    }

    
Read data from underlying stream and process with cipher until end of stream or some data is
available after cipher processing.
Returns:
-1 to indicate end of stream, or the number of bytes (> 0) available./**
     * Read data from underlying stream and process with cipher until end of stream or some data is
     * available after cipher processing.
     *
     * @return -1 to indicate end of stream, or the number of bytes (> 0) available.
     */
    private int nextChunk()
        throws IOException
    {
        if (finalized)
        {
            return -1;
        }

        bufOff = 0;
        maxBuf = 0;

        // Keep reading until EOF or cipher processing produces data
        while (maxBuf == 0)
        {
            int read = in.read(inBuf);
            if (read == -1)
            {
                finaliseCipher();
                if (maxBuf == 0)
                {
                    return -1;
                }
                return maxBuf;
            }

            try
            {
                ensureCapacity(read, false);
                if (bufferedBlockCipher != null)
                {
                    maxBuf = bufferedBlockCipher.processBytes(inBuf, 0, read, buf, 0);
                }
                else if (aeadBlockCipher != null)
                {
                    maxBuf = aeadBlockCipher.processBytes(inBuf, 0, read, buf, 0);
                }
                else
                {
                    streamCipher.processBytes(inBuf, 0, read, buf, 0);
                    maxBuf = read;
                }
            }
            catch (Exception e)
            {
                throw new CipherIOException("Error processing stream ", e);
            }
        }
        return maxBuf;
    }

    private void finaliseCipher()
        throws IOException
    {
        try
        {
            finalized = true;
            ensureCapacity(0, true);
            if (bufferedBlockCipher != null)
            {
                maxBuf = bufferedBlockCipher.doFinal(buf, 0);
            }
            else if (aeadBlockCipher != null)
            {
                maxBuf = aeadBlockCipher.doFinal(buf, 0);
            }
            else
            {
                maxBuf = 0; // a stream cipher
            }
        }
        catch (final InvalidCipherTextException e)
        {
            throw new InvalidCipherTextIOException("Error finalising cipher", e);
        }
        catch (Exception e)
        {
            throw new IOException("Error finalising cipher " + e);
        }
    }

    
Reads data from the underlying stream and processes it with the cipher until the cipher
outputs data, and returns the next available byte.

If the underlying stream is exhausted by this call, the cipher will be finalised.

Throws:
IOException – if there was an error closing the input stream.
InvalidCipherTextIOException – if the data read from the stream was invalid ciphertext
(e.g. the cipher is an AEAD cipher and the ciphertext tag check fails)./**
     * Reads data from the underlying stream and processes it with the cipher until the cipher
     * outputs data, and returns the next available byte.
     * <p>
     * If the underlying stream is exhausted by this call, the cipher will be finalised.
     * </p>
     * @throws IOException if there was an error closing the input stream.
     * @throws InvalidCipherTextIOException if the data read from the stream was invalid ciphertext
     * (e.g. the cipher is an AEAD cipher and the ciphertext tag check fails).
     */
    public int read()
        throws IOException
    {
        if (bufOff >= maxBuf)
        {
            if (nextChunk() < 0)
            {
                return -1;
            }
        }

        return buf[bufOff++] & 0xff;
    }

    
Reads data from the underlying stream and processes it with the cipher until the cipher
outputs data, and then returns up to b.length bytes in the provided array.

If the underlying stream is exhausted by this call, the cipher will be finalised.

Params:
b – the buffer into which the data is read.
Throws:
IOException – if there was an error closing the input stream.
InvalidCipherTextIOException – if the data read from the stream was invalid ciphertext
(e.g. the cipher is an AEAD cipher and the ciphertext tag check fails).
Returns:
the total number of bytes read into the buffer, or -1 if there is no
        more data because the end of the stream has been reached./**
     * Reads data from the underlying stream and processes it with the cipher until the cipher
     * outputs data, and then returns up to <code>b.length</code> bytes in the provided array.
     * <p>
     * If the underlying stream is exhausted by this call, the cipher will be finalised.
     * </p>
     * @param b the buffer into which the data is read.
     * @return the total number of bytes read into the buffer, or <code>-1</code> if there is no
     *         more data because the end of the stream has been reached.
     * @throws IOException if there was an error closing the input stream.
     * @throws InvalidCipherTextIOException if the data read from the stream was invalid ciphertext
     * (e.g. the cipher is an AEAD cipher and the ciphertext tag check fails).
     */
    public int read(
        byte[] b)
        throws IOException
    {
        return read(b, 0, b.length);
    }

    
Reads data from the underlying stream and processes it with the cipher until the cipher
outputs data, and then returns up to len bytes in the provided array.

If the underlying stream is exhausted by this call, the cipher will be finalised.

Params:
b –   the buffer into which the data is read.
off – the start offset in the destination array b
len – the maximum number of bytes read.
Throws:
IOException – if there was an error closing the input stream.
InvalidCipherTextIOException – if the data read from the stream was invalid ciphertext
(e.g. the cipher is an AEAD cipher and the ciphertext tag check fails).
Returns:
the total number of bytes read into the buffer, or -1 if there is no
        more data because the end of the stream has been reached./**
     * Reads data from the underlying stream and processes it with the cipher until the cipher
     * outputs data, and then returns up to <code>len</code> bytes in the provided array.
     * <p>
     * If the underlying stream is exhausted by this call, the cipher will be finalised.
     * </p>
     * @param b   the buffer into which the data is read.
     * @param off the start offset in the destination array <code>b</code>
     * @param len the maximum number of bytes read.
     * @return the total number of bytes read into the buffer, or <code>-1</code> if there is no
     *         more data because the end of the stream has been reached.
     * @throws IOException if there was an error closing the input stream.
     * @throws InvalidCipherTextIOException if the data read from the stream was invalid ciphertext
     * (e.g. the cipher is an AEAD cipher and the ciphertext tag check fails).
     */
    public int read(
        byte[] b,
        int off,
        int len)
        throws IOException
    {
        if (bufOff >= maxBuf)
        {
            if (nextChunk() < 0)
            {
                return -1;
            }
        }

        int toSupply = Math.min(len, available());
        System.arraycopy(buf, bufOff, b, off, toSupply);
        bufOff += toSupply;
        return toSupply;
    }

    public long skip(
        long n)
        throws IOException
    {
        if (n <= 0)
        {
            return 0;
        }

        if (skippingCipher != null)
        {
            int avail = available();
            if (n <= avail)
            {
                bufOff += n;

                return n;
            }

            bufOff = maxBuf;

            long skip = in.skip(n - avail);

            long cSkip = skippingCipher.skip(skip);

            if (skip != cSkip)
            {
                throw new IOException("Unable to skip cipher " + skip + " bytes.");
            }

            return skip + avail;
        }
        else
        {
            int skip = (int)Math.min(n, available());
            bufOff += skip;

            return skip;
        }
    }

    public int available()
        throws IOException
    {
        return maxBuf - bufOff;
    }

    
Ensure the cipher text buffer has space sufficient to accept an upcoming output.
Params:
updateSize – the size of the pending update.
finalOutput – true iff this the cipher is to be finalised./**
     * Ensure the cipher text buffer has space sufficient to accept an upcoming output.
     *
     * @param updateSize the size of the pending update.
     * @param finalOutput <code>true</code> iff this the cipher is to be finalised.
     */
    private void ensureCapacity(int updateSize, boolean finalOutput)
    {
        int bufLen = updateSize;
        if (finalOutput)
        {
            if (bufferedBlockCipher != null)
            {
                bufLen = bufferedBlockCipher.getOutputSize(updateSize);
            }
            else if (aeadBlockCipher != null)
            {
                bufLen = aeadBlockCipher.getOutputSize(updateSize);
            }
        }
        else
        {
            if (bufferedBlockCipher != null)
            {
                bufLen = bufferedBlockCipher.getUpdateOutputSize(updateSize);
            }
            else if (aeadBlockCipher != null)
            {
                bufLen = aeadBlockCipher.getUpdateOutputSize(updateSize);
            }
        }

        if ((buf == null) || (buf.length < bufLen))
        {
            buf = new byte[bufLen];
        }
    }

    
Closes the underlying input stream and finalises the processing of the data by the cipher.
Throws:
IOException – if there was an error closing the input stream.
InvalidCipherTextIOException – if the data read from the stream was invalid ciphertext
            (e.g. the cipher is an AEAD cipher and the ciphertext tag check fails)./**
     * Closes the underlying input stream and finalises the processing of the data by the cipher.
     *
     * @throws IOException if there was an error closing the input stream.
     * @throws InvalidCipherTextIOException if the data read from the stream was invalid ciphertext
     *             (e.g. the cipher is an AEAD cipher and the ciphertext tag check fails).
     */
    public void close()
        throws IOException
    {
        try
        {
            in.close();
        }
        finally
        {
            if (!finalized)
            {
                // Reset the cipher, discarding any data buffered in it
                // Errors in cipher finalisation trump I/O error closing input
                finaliseCipher();
            }
        }
        maxBuf = bufOff = 0;
        markBufOff = 0;
        markPosition = 0;
        if (markBuf != null)
        {
            Arrays.fill(markBuf, (byte)0);
            markBuf = null;
        }
        if (buf != null)
        {
            Arrays.fill(buf, (byte)0);
            buf = null;
        }
        Arrays.fill(inBuf, (byte)0);
    }

    
Mark the current position.

This method only works if markSupported() returns true - which means the underlying stream supports marking, and the cipher passed
in to this stream's constructor is a SkippingCipher (so capable of being reset to an arbitrary point easily).

Params:
readlimit – the maximum read ahead required before a reset() may be called./**
     * Mark the current position.
     * <p>
     * This method only works if markSupported() returns true - which means the underlying stream supports marking, and the cipher passed
     * in to this stream's constructor is a SkippingCipher (so capable of being reset to an arbitrary point easily).
     * </p>
     * @param readlimit the maximum read ahead required before a reset() may be called.
     */
    public void mark(int readlimit)
    {
        in.mark(readlimit);
        if (skippingCipher != null)
        {
            markPosition = skippingCipher.getPosition();
        }

        if (buf != null)
        {
            markBuf = new byte[buf.length];
            System.arraycopy(buf, 0, markBuf, 0, buf.length);
        }

        markBufOff = bufOff;
    }

    
Reset to the last marked position, if supported.
Throws:
IOException – if marking not supported by the cipher used, or the underlying stream./**
     * Reset to the last marked position, if supported.
     *
     * @throws IOException if marking not supported by the cipher used, or the underlying stream.
     */
    public void reset()
        throws IOException
    {
        if (skippingCipher == null)
        {
            throw new IOException("cipher must implement SkippingCipher to be used with reset()");
        }

        in.reset();

        skippingCipher.seekTo(markPosition);

        if (markBuf != null)
        {
            buf = markBuf;
        }

        bufOff = markBufOff;
     }

    
Return true if mark(readlimit) is supported. This will be true if the underlying stream supports marking and the
cipher used is a SkippingCipher,
Returns:
true if mark(readlimit) supported, false otherwise./**
     * Return true if mark(readlimit) is supported. This will be true if the underlying stream supports marking and the
     * cipher used is a SkippingCipher,
     *
     * @return true if mark(readlimit) supported, false otherwise.
     */
    public boolean markSupported()
    {
        if (skippingCipher != null)
        {
            return in.markSupported();
        }

        return false;
    }

}