// ---------------------------------------------------------------------------
// - t_icipher.cpp                                                           -
// - afnix cryptographic library - aes class tester module                   -
// ---------------------------------------------------------------------------
// - This program is free software;  you can redistribute it  and/or  modify -
// - it provided that this copyright notice is kept intact.                  -
// -                                                                         -
// - This program  is  distributed in  the hope  that it will be useful, but -
// - without  any  warranty;  without  even   the   implied    warranty   of -
// - merchantability or fitness for a particular purpose.  In no event shall -
// - the copyright holder be liable for any  direct, indirect, incidental or -
// - special damages arising in any way out of the use of this software.     -
// ---------------------------------------------------------------------------
// - copyright (c) 1999-2007 amaury darsch                                   -
// ---------------------------------------------------------------------------

#include "Aes.hpp"
#include "Ascii.hpp"
#include "InputCipher.hpp"
#include "InputOutput.hpp"
  
namespace afnix {
  const t_byte FIPS_KEY_128 [] = {
    0x2B, 0x7E, 0x15, 0x16, 0x28, 0xAE, 0xD2, 0xA6, 
    0xAB, 0xF7, 0x15, 0x88, 0x09, 0xCF, 0x4F, 0x3C
  };
  
  // FIPS-197 16 bytes input test
  const char FIPS_TEST_BI [] = {
    0x32, 0x43, 0xF6, 0xA8, 0x88, 0x5A, 0x30, 0x8D,
    0x31, 0x31, 0x98, 0xA2, 0xE0, 0x37, 0x07, 0x34
  };
  
  // FIPS-197 16 bytes output result
  const char FIPS_TEST_BO [] = {
    0x39, 0x25, 0x84, 0x1D, 0x02, 0xDC, 0x09, 0xFB,
    0xDC, 0x11, 0x85, 0x97, 0x19, 0x6A, 0x0B, 0x32
  };
  
  // check ecb for one block
  static int check_ecb_fips (void) {
    // create a 128 bits key
    Key key (Key::K128, FIPS_KEY_128);
    // create an AES cipher with a 128 bits key
    Aes* aes = new Aes (key);    
    // create an input-output stream with the input test
    InputOutput* io = new InputOutput (FIPS_TEST_BI, 16);
    // create an input cipher
    InputCipher* ic = new InputCipher (aes, io);

    // create 32 bytes output buffer and process
    char bo[32];
    long count = 0;
    while (ic->valid (-1) == true) {
      if (count > 31) {
	delete ic;
	return 1;
      }
      bo[count++] = ic->read ();
    }
    if (count != 32) return 1;
    
    // verify first 16 bytes
    for (long i = 0; i < 16; i++) {
      if (bo[i] !=  FIPS_TEST_BO[i]) return 1;
    }
    // delete the input cipher
    delete ic;
    
    // create an input-output stream with the output test
    io  = new InputOutput (bo, 32);
    // create a new reverse cipher
    aes = new Aes (key, true);
    // create an input cipher
    ic  = new InputCipher (aes, io);

    // create a 16 bit input buffer and process
    char bi[16];
    count = 0;
    while (ic->valid (-1) == true) {
      if (count > 15) {
	delete ic;
	return 1;
      }
      bi[count++] = ic->read ();
    }
    if (count != 16) return 1;
    // verify the original 16 bytes
    for (long i = 0; i < 16; i++) {
      if (bi[i] !=  FIPS_TEST_BI[i]) return 1;
    }

    // success
    delete ic;
    return 0;
  }

  // check ecb for a very small block
  static int check_ecb_mesg (void) {
    // string message to crypt
    const char* mesg = "hello";
    long  mlen = Ascii::strlen (mesg);
    // create a 128 bits key
    Key key (Key::K128, FIPS_KEY_128);
    // create an AES cipher with a 128 bits key
    Aes* aes = new Aes (key);
    // create an input-output stream with the input test
    InputOutput* io = new InputOutput (mesg);    
    // create an input cipher
    InputCipher* ic = new InputCipher (aes, io);

    // create 16 bytes output buffer and process
    char bo[16];
    long count = 0;
    while (ic->valid (-1) == true) {
      if (count > 15) {
	delete ic;
	return 1;
      }
      bo[count++] = ic->read ();
    }
    // check block count
    if (count != 16) return 1;
    // delete the input cipher
    delete ic;
    
    // create an input-output stream with the output test
    io  = new InputOutput (bo, 16);
    // create a new reverse cipher
    aes = new Aes (key, true);
    // create an input cipher
    ic  = new InputCipher (aes, io);

    // create a 16 bit input buffer and process
    char bi[16];
    count = 0;
    while (ic->valid (-1) == true) {
      if (count > 15) {
	delete ic;
	return 1;
      }
      bi[count++] = ic->read ();
    }
    if (count != mlen) return 1;
    // verify the original 16 bytes
    for (long i = 0; i < mlen; i++) {
      if (bi[i] !=  mesg[i]) return 1;
    }

    // success
    delete ic;
    return 0;
  }

  // check cbc for one block
  static int check_cbc_fips (void) {
    // create a 128 bits key
    Key key (Key::K128, FIPS_KEY_128);
    // create an AES cipher with a 128 bits key
    Aes* aes = new Aes (key);
    // create an input-output stream with the input test
    InputOutput* io = new InputOutput (FIPS_TEST_BI, 16);    
    // create an input cipher
    InputCipher* ic = new InputCipher (aes, io, InputCipher::CBC);

    // create 48 bytes output buffer and process
    char bo[48];
    long count = 0;
    while (ic->valid (-1) == true) {
      if (count > 47) {
	delete ic;
	return 1;
      }
      bo[count++] = ic->read ();
    }
    // check counter
    if (count != 48) return 1;
    // delete the input cipher
    delete ic;
    
    // create an input-output stream with the output test
    io  = new InputOutput (bo, 48);
    // create a new reverse cipher
    aes = new Aes (key, true);
    // create an input cipher
    ic  = new InputCipher (aes, io, InputCipher::CBC);

    // create a 16 bit input buffer and process
    char bi[16];
    count = 0;
    while (ic->valid (-1) == true) {
      if (count > 15) {
	delete ic;
	return 1;
      }
      bi[count++] = ic->read ();
    }
    if (count != 16) return 1;
    // verify the original 16 bytes
    for (long i = 0; i < 16; i++) {
      if (bi[i] !=  FIPS_TEST_BI[i]) return 1;
    }

    // success
    delete ic;
    return 0;
  }

  // check cbc for a very small block
  static int check_cbc_mesg (void) {
    // string message to crypt
    const char* mesg = "hello";
    long  mlen  = Ascii::strlen (mesg);
    // create a 128 bits key
    Key key (Key::K128, FIPS_KEY_128);
    // create an AES cipher with a 128 bits key
    Aes* aes = new Aes (key);
    // create an input-output stream with the input test
    InputOutput* io = new InputOutput (mesg);    
    // create an input cipher
    InputCipher* ic = new InputCipher (aes, io, InputCipher::CBC);

    // create 32 bytes output buffer and process
    char bo[32];
    long count = 0;
    while (ic->valid (-1) == true) {
      if (count > 31) {
	delete ic;
	return 1;
      }
      bo[count++] = ic->read ();
    }
    // check block count
    if (count != 32) return 1;
    // delete the input cipher
    delete ic;
    
    // create an input-output stream with the output test
    io  = new InputOutput (bo, 32);
    // create a new reverse cipher
    aes = new Aes (key, true);
    // create an input cipher
    ic  = new InputCipher (aes, io, InputCipher::CBC);
    // create a 16 bit input buffer and process
    char bi[16];
    count = 0;
    while (ic->valid (-1) == true) {
      if (count > 15) {
	delete ic;
	return 1;
      }
      bi[count++] = ic->read ();
    }
    if (count != mlen) return 1;
    // verify the original 16 bytes
    for (long i = 0; i < mlen; i++) {
      if (bi[i] !=  mesg[i]) return 1;
    }

    // success
    delete ic;
    return 0;
  }
}

// main stuff starts here

int main (int, char**) {
  using namespace afnix;

  // check ecb mode
  if (check_ecb_fips () != 0) return 1;
  if (check_ecb_mesg () != 0) return 1;
  // check cbc mode
  if (check_cbc_fips () != 0) return 1;
  if (check_cbc_mesg () != 0) return 1;
  // success
  return 0;
}