/*
    * 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.nio.ByteBuffer;
  import java.security.GeneralSecurityException;
  import java.util.Properties;
  import java.util.Queue;
  import java.util.concurrent.ConcurrentLinkedQueue;
  
  import javax.crypto.Cipher;
  import javax.crypto.spec.IvParameterSpec;
  
  import org.apache.commons.crypto.cipher.CryptoCipher;
  import org.apache.commons.crypto.stream.input.Input;
  import org.apache.commons.crypto.utils.AES;
  import org.apache.commons.crypto.utils.IoUtils;
  import org.apache.commons.crypto.utils.Utils;
  
  /**
   * PositionedCryptoInputStream provides the capability to decrypt the stream
   * starting at random position as well as provides the foundation for positioned
   * read for decrypting. This needs a stream cipher mode such as AES CTR mode.
   */
  public class PositionedCryptoInputStream extends CtrCryptoInputStream {
  
      private static class CipherState {
  
          private final CryptoCipher cryptoCipher;
          private boolean reset;
  
          /**
           * Constructs a new instance.
           *
           * @param cryptoCipher the CryptoCipher instance.
           */
          public CipherState(final CryptoCipher cryptoCipher) {
              this.cryptoCipher = cryptoCipher;
              this.reset = false;
          }
  
          /**
           * Gets the CryptoCipher instance.
           *
           * @return the cipher.
           */
          public CryptoCipher getCryptoCipher() {
              return cryptoCipher;
          }
  
          /**
           * Gets the reset.
           *
           * @return the value of reset.
           */
          public boolean isReset() {
              return reset;
          }
  
          /**
           * Sets the value of reset.
           *
           * @param reset the reset.
           */
          public void reset(final boolean reset) {
              this.reset = reset;
          }
      }
  
      /**
       * DirectBuffer pool
       */
      private final Queue<ByteBuffer> byteBufferPool = new ConcurrentLinkedQueue<>();
  
      /**
       * CryptoCipher pool
       */
      private final Queue<CipherState> cipherStatePool = new ConcurrentLinkedQueue<>();
  
      /**
       * properties for constructing a CryptoCipher
       */
      private final Properties properties;
 
     /**
      * Constructs a {@link PositionedCryptoInputStream}.
      *
      * @param properties The {@code Properties} class represents a set of
      *        properties.
      * @param in the input data.
      * @param key crypto key for the cipher.
      * @param iv Initialization vector for the cipher.
      * @param streamOffset the start offset in the data.
      * @throws IOException if an I/O error occurs.
      */
     @SuppressWarnings("resource") // The CryptoCipher returned by getCipherInstance() is closed by PositionedCryptoInputStream.
     public PositionedCryptoInputStream(final Properties properties, final Input in, final byte[] key,
             final byte[] iv, final long streamOffset) throws IOException {
         this(properties, in, Utils.getCipherInstance(AES.CTR_NO_PADDING, properties),
                 CryptoInputStream.getBufferSize(properties), key, iv, streamOffset);
     }
 
     /**
      * Constructs a {@link PositionedCryptoInputStream}.
      *
      * @param properties the properties of stream
      * @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 data.
      * @throws IOException if an I/O error occurs.
      */
     protected PositionedCryptoInputStream(final Properties properties, final Input input, final CryptoCipher cipher,
             final int bufferSize, final byte[] key, final byte[] iv, final long streamOffset)
             throws IOException {
         super(input, cipher, bufferSize, key, iv, streamOffset);
         this.properties = properties;
     }
 
     /** Cleans direct buffer pool */
     private void cleanByteBufferPool() {
         ByteBuffer buf;
         while ((buf = byteBufferPool.poll()) != null) {
             CryptoInputStream.freeDirectBuffer(buf);
         }
     }
 
     /** Cleans direct buffer pool */
     private void cleanCipherStatePool() {
         CipherState cs;
         while ((cs = cipherStatePool.poll()) != null) {
             try {
                 cs.getCryptoCipher().close();
             } catch (IOException ignored) {
                 // ignore
             }
         }
     }
 
     /**
      * Overrides the {@link CryptoInputStream#close()}. Closes this input stream
      * and releases any system resources associated with the stream.
      *
      * @throws IOException if an I/O error occurs.
      */
     @Override
     public void close() throws IOException {
         if (!isOpen()) {
             return;
         }
 
         cleanByteBufferPool();
         cleanCipherStatePool();
         super.close();
     }
 
     /**
      * 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().
      *
      * @param state the CipherState instance.
      * @param inByteBuffer the input buffer.
      * @param outByteBuffer the output buffer.
      * @param padding the padding.
      * @throws IOException if an I/O error occurs.
      */
     private void decrypt(final CipherState state, final ByteBuffer inByteBuffer,
             final ByteBuffer outByteBuffer, final byte padding) throws IOException {
         Utils.checkState(inByteBuffer.position() >= padding);
         if (inByteBuffer.position() == padding) {
             // There is no real data in inBuffer.
             return;
         }
         inByteBuffer.flip();
         outByteBuffer.clear();
         decryptBuffer(state, inByteBuffer, outByteBuffer);
         inByteBuffer.clear();
         outByteBuffer.flip();
         if (padding > 0) {
             /*
              * The plain text and cipher text have a 1:1 mapping, they start at
              * the same position.
              */
             outByteBuffer.position(padding);
         }
     }
 
     /**
      * Decrypts length bytes in buffer starting at offset. Output is also put
      * into buffer starting at offset. It is thread-safe.
      *
      * @param buffer the buffer into which the data is read.
      * @param offset the start offset in the data.
      * @param position the offset from the start of the stream.
      * @param length the maximum number of bytes to read.
      * @throws IOException if an I/O error occurs.
      */
     protected void decrypt(final long position, final byte[] buffer, final int offset, final int length)
             throws IOException {
         final ByteBuffer inByteBuffer = getBuffer();
         final ByteBuffer outByteBuffer = getBuffer();
         CipherState state = null;
         try {
             state = getCipherState();
             final byte[] iv = getInitIV().clone();
             resetCipher(state, position, iv);
             byte padding = getPadding(position);
             inByteBuffer.position(padding); // Set proper position for input data.
 
             int n = 0;
             while (n < length) {
                 final int toDecrypt = Math.min(length - n, inByteBuffer.remaining());
                 inByteBuffer.put(buffer, offset + n, toDecrypt);
 
                 // Do decryption
                 decrypt(state, inByteBuffer, outByteBuffer, padding);
 
                 outByteBuffer.get(buffer, offset + n, toDecrypt);
                 n += toDecrypt;
                 padding = postDecryption(state, inByteBuffer, position + n, iv);
             }
         } finally {
             returnToPool(inByteBuffer);
             returnToPool(outByteBuffer);
             returnToPool(state);
         }
     }
 
     /**
      * Does the decryption using inBuffer as input and outBuffer as output.
      *
      * @param state the CipherState instance.
      * @param inByteBuffer the input buffer.
      * @param outByteBuffer the output buffer.
      * @throws IOException if an I/O error occurs.
      */
     @SuppressWarnings("resource") // getCryptoCipher does not allocate
     private void decryptBuffer(final CipherState state, final ByteBuffer inByteBuffer,
             final ByteBuffer outByteBuffer) throws IOException {
         final int inputSize = inByteBuffer.remaining();
         try {
             final int n = state.getCryptoCipher().update(inByteBuffer, outByteBuffer);
             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.
                  */
                 state.getCryptoCipher().doFinal(inByteBuffer, outByteBuffer);
                 state.reset(true);
             }
         } catch (final GeneralSecurityException e) {
             throw new IOException(e);
         }
     }
 
     /**
      * Gets direct buffer from pool. Caller MUST also call {@link #returnToPool(ByteBuffer)}.
      *
      * @return the buffer.
      * @see #returnToPool(ByteBuffer)
      */
     private ByteBuffer getBuffer() {
         final ByteBuffer buffer = byteBufferPool.poll();
         return buffer != null ? buffer : ByteBuffer.allocateDirect(getBufferSize());
     }
 
     /**
      * Gets CryptoCipher from pool. Caller MUST also call {@link #returnToPool(CipherState)}.
      *
      * @return the CipherState instance.
      * @throws IOException if an I/O error occurs.
      */
     @SuppressWarnings("resource") // Caller calls #returnToPool(CipherState)
     private CipherState getCipherState() throws IOException {
         final CipherState state = cipherStatePool.poll();
         return state != null ? state : new CipherState(Utils.getCipherInstance(AES.CTR_NO_PADDING, properties));
     }
 
     /**
      * This method is executed immediately after decryption. Check whether
      * cipher should be updated and recalculate padding if needed.
      *
      * @param state the CipherState instance.
      * @param inByteBuffer the input buffer.
      * @param position the offset from the start of the stream.
      * @param iv the iv.
      * @return the padding.
      */
     private byte postDecryption(final CipherState state, final ByteBuffer inByteBuffer,
             final long position, final byte[] iv) {
         byte padding = 0;
         if (state.isReset()) {
             /*
              * 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(state, position, iv);
             padding = getPadding(position);
             inByteBuffer.position(padding);
         }
         return padding;
     }
 
     /**
      * Reads up to the specified number of bytes from a given position within a
      * stream and return the number of bytes read. This does not change the
      * current offset of the stream, and is thread-safe.
      *
      * @param buffer the buffer into which the data is read.
      * @param length the maximum number of bytes to read.
      * @param offset the start offset in the data.
      * @param position the offset from the start of the stream.
      * @throws IOException if an I/O error occurs.
      * @return int the total number of decrypted data bytes read into the
      *         buffer.
      */
     public int read(final long position, final byte[] buffer, final int offset, final int length)
             throws IOException {
         checkStream();
         final int n = input.read(position, buffer, offset, length);
         if (n > 0) {
             // This operation does not change the current offset of the file
             decrypt(position, buffer, offset, n);
         }
         return n;
     }
 
     /**
      * Reads the specified number of bytes from a given position within a
      * stream. This does not change the current offset of the stream and is
      * thread-safe.
      *
      * @param position the offset from the start of the stream.
      * @param buffer the buffer into which the data is read.
      * @throws IOException if an I/O error occurs.
      */
     public void readFully(final long position, final byte[] buffer) throws IOException {
         readFully(position, buffer, 0, buffer.length);
     }
 
     /**
      * Reads the specified number of bytes from a given position within a
      * stream. This does not change the current offset of the stream and is
      * thread-safe.
      *
      * @param buffer the buffer into which the data is read.
      * @param length the maximum number of bytes to read.
      * @param offset the start offset in the data.
      * @param position the offset from the start of the stream.
      * @throws IOException if an I/O error occurs.
      */
     public void readFully(final long position, final byte[] buffer, final int offset, final int length)
             throws IOException {
         checkStream();
         IoUtils.readFully(input, position, buffer, offset, length);
         if (length > 0) {
             // This operation does not change the current offset of the file
             decrypt(position, buffer, offset, length);
         }
     }
 
     /**
      * Calculates the counter and iv, reset the cipher.
      *
      * @param state the CipherState instance.
      * @param position the offset from the start of the stream.
      * @param iv the iv.
      */
     @SuppressWarnings("resource") // getCryptoCipher does not allocate
     private void resetCipher(final CipherState state, final long position, final byte[] iv) {
         final long counter = getCounter(position);
         CtrCryptoInputStream.calculateIV(getInitIV(), counter, iv);
         try {
             state.getCryptoCipher().init(Cipher.DECRYPT_MODE, key, new IvParameterSpec(iv));
         } catch (final GeneralSecurityException e) {
             // Ignore
         }
         state.reset(false);
     }
 
     /**
      * Returns direct buffer to pool.
      *
      * @param buf the buffer.
      */
     private void returnToPool(final ByteBuffer buf) {
         if (buf != null) {
             buf.clear();
             byteBufferPool.add(buf);
         }
     }
 
     /**
      * Returns CryptoCipher to pool.
      *
      * @param state the CipherState instance.
      */
     private void returnToPool(final CipherState state) {
         if (state != null) {
             cipherStatePool.add(state);
         }
     }
 }