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.
/**
 * 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.OutputStream;
import java.nio.ByteBuffer;
import java.nio.channels.WritableByteChannel;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.Key;
import java.security.spec.AlgorithmParameterSpec;
import java.util.Properties;
import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.ShortBufferException;
import javax.crypto.spec.IvParameterSpec;

import org.apache.commons.crypto.cipher.CryptoCipher;
import org.apache.commons.crypto.stream.output.ChannelOutput;
import org.apache.commons.crypto.stream.output.Output;
import org.apache.commons.crypto.stream.output.StreamOutput;
import org.apache.commons.crypto.utils.Utils;

CryptoOutputStream encrypts data and writes to the under layer output. It supports any mode of operations such as AES CBC/CTR/GCM mode in concept. It is not thread-safe. /**
 * {@link CryptoOutputStream} encrypts data and writes to the under layer
 * output. It supports any mode of operations such as AES CBC/CTR/GCM mode in
 * concept. It is not thread-safe.
 */

public class CryptoOutputStream extends OutputStream implements
        WritableByteChannel {
    private final byte[] oneByteBuf = new byte[1];

    
The output. /** The output. */
    Output output;

    
the CryptoCipher instance /** the CryptoCipher instance */
    final CryptoCipher cipher;

    
The buffer size. /** The buffer size. */
    final int bufferSize;

    
Crypto key for the cipher. /** Crypto key for the cipher. */
    final Key key;

    
the algorithm parameters /** the algorithm parameters */
    final AlgorithmParameterSpec params;

    
Flag to mark whether the output stream is closed. /** Flag to mark whether the output stream is closed. */
    private boolean closed;

    
Input data buffer. The data starts at inBuffer.position() and ends at
inBuffer.limit().
/**
     * Input data buffer. The data starts at inBuffer.position() and ends at
     * inBuffer.limit().
     */
    ByteBuffer inBuffer;

    
Encrypted data buffer. The data starts at outBuffer.position() and ends
at outBuffer.limit().
/**
     * Encrypted data buffer. The data starts at outBuffer.position() and ends
     * at outBuffer.limit().
     */
    ByteBuffer outBuffer;

    
Constructs a CryptoOutputStream. 
Params:
transformation – the name of the transformation, e.g.,
AES/CBC/PKCS5Padding.
See the Java Cryptography Architecture Standard Algorithm Name Documentation
for information about standard transformation names.
props – The Properties class represents a set of
       properties.
out – the output stream.
key – crypto key for the cipher.
params – the algorithm parameters.
Throws:
IOException – if an I/O error occurs./**
     * Constructs a {@link CryptoOutputStream}.
     *
     * @param transformation the name of the transformation, e.g.,
     * <i>AES/CBC/PKCS5Padding</i>.
     * See the Java Cryptography Architecture Standard Algorithm Name Documentation
     * for information about standard transformation names.
     * @param props The <code>Properties</code> class represents a set of
     *        properties.
     * @param out the output stream.
     * @param key crypto key for the cipher.
     * @param params the algorithm parameters.
     * @throws IOException if an I/O error occurs.
     */

    public CryptoOutputStream(String transformation,
            Properties props, OutputStream out, Key key,
            AlgorithmParameterSpec params) throws IOException {
        this(out, Utils.getCipherInstance(transformation, props),
                CryptoInputStream.getBufferSize(props), key, params);

    }

    
Constructs a CryptoOutputStream. 
Params:
transformation – the name of the transformation, e.g.,
AES/CBC/PKCS5Padding.
See the Java Cryptography Architecture Standard Algorithm Name Documentation
for information about standard transformation names.
props – The Properties class represents a set of
       properties.
out – the WritableByteChannel instance.
key – crypto key for the cipher.
params – the algorithm parameters.
Throws:
IOException – if an I/O error occurs./**
     * Constructs a {@link CryptoOutputStream}.
     *
     * @param transformation the name of the transformation, e.g.,
     * <i>AES/CBC/PKCS5Padding</i>.
     * See the Java Cryptography Architecture Standard Algorithm Name Documentation
     * for information about standard transformation names.
     * @param props The <code>Properties</code> class represents a set of
     *        properties.
     * @param out the WritableByteChannel instance.
     * @param key crypto key for the cipher.
     * @param params the algorithm parameters.
     * @throws IOException if an I/O error occurs.
     */
    public CryptoOutputStream(String transformation,
            Properties props, WritableByteChannel out, Key key,
            AlgorithmParameterSpec params) throws IOException {
        this(out, Utils.getCipherInstance(transformation, props), CryptoInputStream
                .getBufferSize(props), key, params);

    }

    
Constructs a CryptoOutputStream. 
Params:
out – the output stream.
cipher – the CryptoCipher instance.
bufferSize – the bufferSize.
key – crypto key for the cipher.
params – the algorithm parameters.
Throws:
IOException – if an I/O error occurs./**
     * Constructs a {@link CryptoOutputStream}.
     *
     * @param out the output stream.
     * @param cipher the CryptoCipher instance.
     * @param bufferSize the bufferSize.
     * @param key crypto key for the cipher.
     * @param params the algorithm parameters.
     * @throws IOException if an I/O error occurs.
     */
    protected CryptoOutputStream(OutputStream out, CryptoCipher cipher,
            int bufferSize, Key key, AlgorithmParameterSpec params)
            throws IOException {
        this(new StreamOutput(out, bufferSize), cipher, bufferSize, key, params);
    }

    
Constructs a CryptoOutputStream. 
Params:
channel – the WritableByteChannel instance.
cipher – the cipher instance.
bufferSize – the bufferSize.
key – crypto key for the cipher.
params – the algorithm parameters.
Throws:
IOException – if an I/O error occurs./**
     * Constructs a {@link CryptoOutputStream}.
     *
     * @param channel the WritableByteChannel instance.
     * @param cipher the cipher instance.
     * @param bufferSize the bufferSize.
     * @param key crypto key for the cipher.
     * @param params the algorithm parameters.
     * @throws IOException if an I/O error occurs.
     */
    protected CryptoOutputStream(WritableByteChannel channel, CryptoCipher cipher,
            int bufferSize, Key key, AlgorithmParameterSpec params)
            throws IOException {
        this(new ChannelOutput(channel), cipher, bufferSize, key, params);
    }

    
Constructs a CryptoOutputStream. 
Params:
output – the output stream.
cipher – the CryptoCipher instance.
bufferSize – the bufferSize.
key – crypto key for the cipher.
params – the algorithm parameters.
Throws:
IOException – if an I/O error occurs./**
     * Constructs a {@link CryptoOutputStream}.
     *
     * @param output the output stream.
     * @param cipher the CryptoCipher instance.
     * @param bufferSize the bufferSize.
     * @param key crypto key for the cipher.
     * @param params the algorithm parameters.
     * @throws IOException if an I/O error occurs.
     */
    protected CryptoOutputStream(Output output, CryptoCipher cipher,
            int bufferSize, Key key, AlgorithmParameterSpec params)
            throws IOException {

        this.output = output;
        this.bufferSize = CryptoInputStream.checkBufferSize(cipher, bufferSize);
        this.cipher = cipher;

        this.key = key;
        this.params = params;

        if (!(params instanceof IvParameterSpec)) {
            // other AlgorithmParameterSpec such as GCMParameterSpec is not
            // supported now.
            throw new IOException("Illegal parameters");
        }

        inBuffer = ByteBuffer.allocateDirect(this.bufferSize);
        outBuffer = ByteBuffer.allocateDirect(this.bufferSize
                + cipher.getBlockSize());

        initCipher();
    }

    
Overrides the OutputStream.write(byte[]). Writes the specified byte to this output stream. 
Params:
b – the data.
Throws:
IOException – if an I/O error occurs./**
     * Overrides the {@link java.io.OutputStream#write(byte[])}. Writes the
     * specified byte to this output stream.
     *
     * @param b the data.
     * @throws IOException if an I/O error occurs.
     */
    @Override
    public void write(int b) throws IOException {
        oneByteBuf[0] = (byte) (b & 0xff);
        write(oneByteBuf, 0, oneByteBuf.length);
    }

    
Overrides the OutputStream.write(byte[], int, int). Encryption is buffer based. If there is enough room in inBuffer, then write to this buffer. If inBuffer is full, then do encryption and write data to the underlying stream. 
Params:
b – the data.
off – the start offset in the data.
len – the number of bytes to write.
Throws:
IOException – if an I/O error occurs./**
     * Overrides the {@link java.io.OutputStream#write(byte[], int, int)}.
     * Encryption is buffer based. If there is enough room in {@link #inBuffer},
     * then write to this buffer. If {@link #inBuffer} is full, then do
     * encryption and write data to the underlying stream.
     *
     * @param b the data.
     * @param off the start offset in the data.
     * @param len the number of bytes to write.
     * @throws IOException if an I/O error occurs.
     */
    @Override
    public void write(byte[] b, int off, int len) throws IOException {
        checkStream();
        if (b == null) {
            throw new NullPointerException();
        } else if (off < 0 || len < 0 || off > b.length || len > b.length - off) {
            throw new IndexOutOfBoundsException();
        }

        while (len > 0) {
            final int remaining = inBuffer.remaining();
            if (len < remaining) {
                inBuffer.put(b, off, len);
                len = 0;
            } else {
                inBuffer.put(b, off, remaining);
                off += remaining;
                len -= remaining;
                encrypt();
            }
        }
    }

    
Overrides the OutputStream.flush(). To flush, we need to encrypt the data in the buffer and write to the underlying stream, then do the flush. 
Throws:
IOException – if an I/O error occurs./**
     * Overrides the {@link OutputStream#flush()}. To flush, we need to encrypt
     * the data in the buffer and write to the underlying stream, then do the
     * flush.
     *
     * @throws IOException if an I/O error occurs.
     */
    @Override
    public void flush() throws IOException {
        checkStream();
        encrypt();
        output.flush();
        super.flush();
    }

    
Overrides the OutputStream.close(). Closes this output stream and releases any system resources associated with this stream. 
Throws:
IOException – if an I/O error occurs./**
     * Overrides the {@link OutputStream#close()}. Closes this output stream and
     * releases any system resources associated with this stream.
     *
     * @throws IOException if an I/O error occurs.
     */
    @Override
    public void close() throws IOException {
        if (closed) {
            return;
        }

        try {
            encryptFinal();
            output.close();
            freeBuffers();
            cipher.close();
            super.close();
        } finally {
            closed = true;
        }
    }

    
Overrides the Channel.isOpen(). Tells whether or not this channel is open. 
Returns:
true if, and only if, this channel is open/**
     * Overrides the {@link java.nio.channels.Channel#isOpen()}. Tells whether or not this channel
     * is open.
     *
     * @return <tt>true</tt> if, and only if, this channel is open
     */
    @Override
    public boolean isOpen() {
        return !closed;
    }

    
Overrides the WritableByteChannel.write(ByteBuffer). Writes a sequence of bytes to this channel from the given buffer. 
Params:
src – The buffer from which bytes are to be retrieved.
Throws:
IOException – if an I/O error occurs.
Returns:
The number of bytes written, possibly zero./**
     * Overrides the
     * {@link java.nio.channels.WritableByteChannel#write(ByteBuffer)}. Writes a
     * sequence of bytes to this channel from the given buffer.
     *
     * @param src The buffer from which bytes are to be retrieved.
     * @return The number of bytes written, possibly zero.
     * @throws IOException if an I/O error occurs.
     */
    @Override
    public int write(ByteBuffer src) throws IOException {
        checkStream();
        final int len = src.remaining();
        int remaining = len;
        while (remaining > 0) {
            final int space = inBuffer.remaining();
            if (remaining < space) {
                inBuffer.put(src);
                remaining = 0;
            } else {
                // to void copy twice, we set the limit to copy directly
                final int oldLimit = src.limit();
                final int newLimit = src.position() + space;
                src.limit(newLimit);

                inBuffer.put(src);

                // restore the old limit
                src.limit(oldLimit);

                remaining -= space;
                encrypt();
            }
        }

        return len;
    }

    
Initializes the cipher.
Throws:
IOException – if an I/O error occurs./**
     * Initializes the cipher.
     *
     * @throws IOException if an I/O error occurs.
     */
    protected void initCipher() throws IOException {
        try {
            cipher.init(Cipher.ENCRYPT_MODE, key, params);
        } catch (InvalidKeyException e) {
            throw new IOException(e);
        } catch (InvalidAlgorithmParameterException e) {
            throw new IOException(e);
        }
    }

    
Does the encryption, input is inBuffer and output is outBuffer. 
Throws:
IOException – if an I/O error occurs./**
     * Does the encryption, input is {@link #inBuffer} and output is
     * {@link #outBuffer}.
     *
     * @throws IOException if an I/O error occurs.
     */
    protected void encrypt() throws IOException {

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

        try {
            cipher.update(inBuffer, outBuffer);
        } catch (ShortBufferException e) {
            throw new IOException(e);
        }

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

        // write to output
        output.write(outBuffer);
    }

    
Does final encryption of the last data.
Throws:
IOException – if an I/O error occurs./**
     * Does final encryption of the last data.
     *
     * @throws IOException if an I/O error occurs.
     */
    protected void encryptFinal() throws IOException {
        inBuffer.flip();
        outBuffer.clear();

        try {
            cipher.doFinal(inBuffer, outBuffer);
        } catch (ShortBufferException e) {
            throw new IOException(e);
        } catch (IllegalBlockSizeException e) {
            throw new IOException(e);
        } catch (BadPaddingException e) {
            throw new IOException(e);
        }

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

        // write to output
        output.write(outBuffer);
    }

    
Checks whether the stream is closed.
Throws:
IOException – if an I/O error occurs./**
     * Checks whether the stream is closed.
     *
     * @throws IOException if an I/O error occurs.
     */
    protected void checkStream() throws IOException {
        if (closed) {
            throw new IOException("Stream closed");
        }
    }

    
Forcibly free the direct buffers. /** Forcibly free the direct buffers. */
    protected void freeBuffers() {
        CryptoInputStream.freeDirectBuffer(inBuffer);
        CryptoInputStream.freeDirectBuffer(outBuffer);
    }

    
Gets the outBuffer.
Returns:
the outBuffer./**
     * Gets the outBuffer.
     *
     * @return the outBuffer.
     */
    protected ByteBuffer getOutBuffer() {
        return outBuffer;
    }

    
Gets the internal Cipher.
Returns:
the cipher instance./**
     * Gets the internal Cipher.
     *
     * @return the cipher instance.
     */
    protected CryptoCipher getCipher() {
        return cipher;
    }

    
Gets the buffer size.
Returns:
the buffer size./**
     * Gets the buffer size.
     *
     * @return the buffer size.
     */
    protected int getBufferSize() {
        return bufferSize;
    }

    
Gets the inBuffer.
Returns:
the inBuffer./**
     * Gets the inBuffer.
     *
     * @return the inBuffer.
     */
    protected ByteBuffer getInBuffer() {
        return inBuffer;
    }
}