package com.trilead.ssh2.crypto.cipher;

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

/**
 * CipherInputStream.
 *
 * @author Christian Plattner, plattner@trilead.com
 * @version $Id: CipherInputStream.java,v 1.1 2007/10/15 12:49:55 cplattne Exp $
 */
public class CipherInputStream {
    BlockCipher currentCipher;
    InputStream bi;
    byte[] buffer;
    byte[] enc;
    int blockSize;
    int pos;

    /*
     * We cannot use java.io.BufferedInputStream, since that is not available in
     * J2ME. Everything could be improved alot here.
     */

    final int BUFF_SIZE = 2048;
    byte[] input_buffer = new byte[BUFF_SIZE];
    int input_buffer_pos = 0;
    int input_buffer_size = 0;

    public CipherInputStream(BlockCipher tc, InputStream bi) {
        this.bi = bi;
        changeCipher(tc);
    }

    private int fill_buffer() throws IOException {
        input_buffer_pos = 0;
        input_buffer_size = bi.read(input_buffer, 0, BUFF_SIZE);
        return input_buffer_size;
    }

    private int internal_read(byte[] b, int off, int len) throws IOException {
        if (input_buffer_size < 0)
            return -1;

        if (input_buffer_pos >= input_buffer_size) {
            if (fill_buffer() <= 0)
                return -1;
        }

        int avail = input_buffer_size - input_buffer_pos;
        int thiscopy = (len > avail) ? avail : len;
        System.arraycopy(input_buffer, input_buffer_pos, b, off, thiscopy);
        input_buffer_pos += thiscopy;
        return thiscopy;
    }

    public void changeCipher(BlockCipher bc) {
        this.currentCipher = bc;
        blockSize = bc.getBlockSize();
        buffer = new byte[blockSize];
        enc = new byte[blockSize];
        pos = blockSize;
    }

    private void getBlock() throws IOException {
        int n = 0;

        while (n < blockSize) {
            int len = internal_read(enc, n, blockSize - n);

            if (len < 0)
                throw new IOException("Cannot read full block, EOF reached.");

            n += len;
        }

        try {
            currentCipher.transformBlock(enc, 0, buffer, 0);
        }
        catch (Exception e) {
            throw new IOException("Error while decrypting block.");
        }

        pos = 0;
    }

    public int read(byte[] dst) throws IOException {
        return read(dst, 0, dst.length);
    }

    public int read(byte[] dst, int off, int len) throws IOException {
        int count = 0;

        while (len > 0) {
            if (pos >= blockSize)
                getBlock();

            int avail = blockSize - pos;
            int copy = Math.min(avail, len);
            System.arraycopy(buffer, pos, dst, off, copy);
            pos += copy;
            off += copy;
            len -= copy;
            count += copy;
        }

        return count;
    }

    public int read() throws IOException {
        if (pos >= blockSize) {
            getBlock();
        }

        return buffer[pos++] & 0xff;
    }

    public int readPlain(byte[] b, int off, int len) throws IOException {
        if (pos != blockSize)
            throw new IOException("Cannot read plain since crypto buffer is not aligned.");

        int n = 0;

        while (n < len) {
            int cnt = internal_read(b, off + n, len - n);

            if (cnt < 0)
                throw new IOException("Cannot fill buffer, EOF reached.");

            n += cnt;
        }

        return n;
    }
}