package tech.loedige;

import javax.crypto.*;
import javax.crypto.spec.SecretKeySpec;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;

public class KeyStreamGenerator_AES_CTR extends KeyStreamGenerator {

    private Cipher cipher;
    private final byte[] state = new byte[16];
    private final byte[] ctr = new byte[16];
    private int currentByte=0;

    public void initState(byte[] key, byte[] iv)
            throws NoSuchPaddingException, NoSuchAlgorithmException, InvalidKeyException {
        if((iv.length!=16)||(key.length!=16)){
            throw new IllegalArgumentException("Invalid input length");
        }
        System.arraycopy(iv,0,ctr,0,16);

        SecretKey secretKey = new SecretKeySpec(key,"AES");
        cipher = Cipher.getInstance("AES/ECB/NOPADDING");
        cipher.init(Cipher.ENCRYPT_MODE,secretKey);
    }

    @Override
    byte getNextKeyStreamByte(byte mByte) {
        currentByte++;
        if(currentByte>=16){
            try {
                updateStateBlock();
            } catch (BadPaddingException | IllegalBlockSizeException e) {
                e.printStackTrace();
            }
            currentByte=0;
        }
        return state[currentByte];
    }

    private void updateStateBlock()
            throws BadPaddingException, IllegalBlockSizeException {
        incrementCtr();
        System.arraycopy(cipher.doFinal(ctr),0,state,0,16);
    }

    /**
     * inkrementiert den counter
     */
    private void incrementCtr(){
        for(int i = 0; i<16; i++){
            //falls das Byte die volle Länge erreicht hat findet ein Übertrag statt
            if(ctr[i]==0xff){
                ctr[i]=0x00;
            }
            else{
                //falls das Byte nicht die volle Länge hat endet die Inkrementierung
                ctr[i]++;
                return;
            }
        }
    }
}