// ---------------------------------------------------------------------------
// - InputCipher.cpp                                                         -
// - afnix:sec module - input cipher class implementation                    -
// ---------------------------------------------------------------------------
// - 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-2008 amaury darsch                                   -
// ---------------------------------------------------------------------------

#include "Item.hpp"
#include "Ascii.hpp"
#include "Vector.hpp"
#include "Utility.hpp"
#include "Unicode.hpp"
#include "QuarkZone.hpp"
#include "InputCipher.hpp"

namespace afnix {

  // -------------------------------------------------------------------------
  // - private section                                                       -
  // -------------------------------------------------------------------------

  // this procedure set the initial vector with random bytes
  static t_byte* preldiv_blk (const long ivsz) {
    if (ivsz <= 0) return nilp;
    t_byte* iv = new t_byte[ivsz];
    for (long i = 0; i < ivsz; i++) iv[i] = Utility::byternd ();
    return iv;
  }

  // this procedure reset the blocks with the initial vector
  static void resetiv_blk (const InputCipher::t_cmod cmod,
			   const long cbsz, const t_byte* iv,
			   t_byte* bo, t_byte* bi, t_byte* bl) {
    // check for cbc mode
    if (cmod == InputCipher::CBC) {
      for (long i = 0; i < cbsz; i++) bl[i] = iv[i];
    }
    // check for cfb mode
    if (cmod == InputCipher::CFB) {
      for (long i = 0; i < cbsz; i++) bl[i] = iv[i];
    }
  }

  // this procedure fill a block upto a block size
  static long preload_blk (t_byte* bi, Input* is, 
			   const long cbsz, const bool rflg) {
    // trivial check first
    if (is == nilp) return 0;
    // load the block
    long bcnt = 0;
    while (is->valid () == true) {
      bi[bcnt++] = is->read ();
      if (bcnt == cbsz) break;
    }
    // if the block is not full and we are in reverse mode
    // then, we do have a problem !
    if ((bcnt != cbsz) && (rflg == true)) {
      throw Exception ("cipher-error", "short block in reverse mode");
    }
    return bcnt;
  }

  // this procedure pad a block if needed
  static bool padload_blk (BlockCipher* cifr, t_byte* bi, const long bcnt) {
    // trivial check first
    if (cifr == nilp) return false;
    // extract parameters
    bool rflg = cifr->getrflg ();
    // pad the block in  normal mode
    if (rflg == false) return cifr->mkpad (bi, bcnt);
    // do not pad by default
    return false;
  }

  // this procedure process a block in ecb mode
  static void process_ecb (BlockCipher* cifr,
			   t_byte* bo, t_byte* bi, t_byte* bl) {
    // trivial check first
    if (cifr == nilp) return;
    if ((bo == nilp) || (bi == nilp) || (bl == nilp)) return;
    // extract parameters
    long cbsz = cifr->getcbsz ();
    bool rflg = cifr->getrflg ();
    // process the block
    if (rflg == false) {
      cifr->process (bo, bi);
      for (long i = 0; i < cbsz; i++) bl[i] = bo[i];
    } else {
      cifr->process (bo, bi);
      for (long i = 0; i < cbsz; i++) bl[i] = bi[i];
    }
  }

  // this procedure process a block in cbc mode
  static void process_cbc (BlockCipher* cifr,
			   t_byte* bo, t_byte* bi, t_byte* bl) {
    // trivial check first
    if (cifr == nilp) return;
    if ((bo == nilp) || (bi == nilp) || (bl == nilp)) return;
    // extract parameters
    long cbsz = cifr->getcbsz ();
    bool rflg = cifr->getrflg ();
    // process the block
    if (rflg == false) {
      for (long i = 0; i < cbsz; i++) bi[i] ^= bl[i];
      cifr->process (bo, bi);
      for (long i = 0; i < cbsz; i++) bl[i] = bo[i];
    } else {
      cifr->process (bo, bi);
      for (long i = 0; i < cbsz; i++) {
	bo[i] ^= bl[i];
	bl[i]  = bi[i];
      }
    }
  }

  // this procedure process a block in cfb mode
  static void process_cfb (BlockCipher* cifr,
			   t_byte* bo, t_byte* bi, t_byte* bl) {
    // trivial check first
    if (cifr == nilp) return;
    if ((bo == nilp) || (bi == nilp) || (bl == nilp)) return;
    // extract parameters
    long cbsz = cifr->getcbsz ();
    bool rflg = cifr->getrflg ();
    // process the block
    if (rflg == false) {
      cifr->encode (bo, bl);
      for (long i = 0; i < cbsz; i++) bo[i] ^= bi[i];
      for (long i = 0; i < cbsz; i++) bl[i]  = bo[i];
    } else {
      cifr->encode (bo, bl);
      for (long i = 0; i < cbsz; i++) bo[i] ^= bi[i];
      for (long i = 0; i < cbsz; i++) bl[i]  = bi[i];
    }
  }

  // this procedure process a block in ofb mode
  static void process_ofb (BlockCipher* cifr,
			   t_byte* bo, t_byte* bi, t_byte* bl) {
    // trivial check first
    if (cifr == nilp) return;
    if ((bo == nilp) || (bi == nilp) || (bl == nilp)) return;
    // extract parameters
    long cbsz = cifr->getcbsz ();
    // process the block
    cifr->encode (bo, bl);
    for (long i = 0; i < cbsz; i++) bl[i]  = bo[i];
    for (long i = 0; i < cbsz; i++) bo[i] ^= bi[i];
  }

  // this procedure process block based on the input cipher mode
  static void process_blk (BlockCipher* cifr, InputCipher::t_cmod cmod,
			   t_byte* bo, t_byte* bi, t_byte* bl) {
    // trivial checks
    if (cifr == nilp) return;
    // process the block with its mode
    switch (cmod) {
    case InputCipher::ECB:
      process_ecb (cifr, bo, bi, bl);
      break;
    case InputCipher::CBC:
      process_cbc (cifr, bo, bi, bl);
      break;
    case InputCipher::CFB:
      process_cfb (cifr, bo, bi, bl);
      break;
    case InputCipher::OFB:
      process_ofb (cifr, bo, bi, bl);
      break;
    }
  }

  // this procedure pushes the block in the stream buffer
  static void pushblk_buf (Buffer& sbuf, BlockCipher* cifr,
			   const InputCipher::t_cmod cmod,
			   const bool eosf, const bool pflg,
			   t_byte* bo, t_byte* bi, t_byte* bl) {
    // get the cipher parameters
    long cbsz = cifr->getcbsz ();
    bool rflg = cifr->getrflg ();
    long bcnt = cbsz;

    // in reverse mode we push only upto the pad if we are at the eos
    // this also depends on the padding mode
    if ((rflg == true) && (eosf == true)) {
      bcnt = cifr->getub (bo);
    }

    // add the block in the buffer - only add if not null - this takes
    // care of not pushing the extra pad block in reverse mode
    if (bcnt > 0) {
      sbuf.add ((const char*) bo, bcnt);
    }

    // in normal mode, if the last block was not padded, it might
    // be necessary to add an extra block which is fully padded
    if ((rflg == false) && (eosf == true) && (pflg == false)) {
      // set the next padding block
      bool nxpf = cifr->mknxp (bi);
      // if the next block is set process it
      if (nxpf == true) {
	// process the block
	process_blk (cifr, cmod, bo, bi, bl);
	// finally add it to the buffer
	sbuf.add ((const char*) bo, cbsz);
      }
    }
  }

  // this procedure update the stream buffer
  static void update_buf (Buffer& sbuf, BlockCipher* cifr, Input* is,
			  const InputCipher::t_cmod cmod,
			  t_byte* bo, t_byte* bi, t_byte* bl) {
    // check the input buffer first
    if (sbuf.empty () == false) return;
    // check if we can read the input stream
    if ((is == nilp) || (is->iseos () == true)) return;
    // check if we have a block cipher
    if (cifr == nilp) return;

    // get the cipher parameters
    long cbsz = cifr->getcbsz ();
    bool rflg = cifr->getrflg ();

    // fill in a block buffer
    long bcnt = preload_blk (bi, is, cbsz, rflg);
    // eventually pad the block
    bool pflg = padload_blk (cifr, bi, bcnt);
    // process the block
    process_blk (cifr, cmod, bo, bi, bl);
    // push the block buffer
    pushblk_buf (sbuf, cifr, cmod, is->iseos (), pflg, bo, bi, bl);
  }

  // -------------------------------------------------------------------------
  // - class section                                                         -
  // -------------------------------------------------------------------------

  // create a new input cipher stream by cipher only

  InputCipher::InputCipher (BlockCipher* cifr) {
    // bind object
    Object::iref (p_cifr = cifr);
    p_is = nilp;
    // grab cipher information
    d_cmod = ECB;
    // create buffer blocks
    long cbsz = (p_cifr == nilp) ? 0 : p_cifr->getcbsz ();
    p_bo = (cbsz > 0) ? new t_byte[cbsz] : nilp;
    p_bi = (cbsz > 0) ? new t_byte[cbsz] : nilp;
    p_bl = (cbsz > 0) ? new t_byte[cbsz] : nilp;
    // preset initial vector
    p_iv = preldiv_blk (cbsz);
    // reset eveything
    reset ();
  }

  // create a new input cipher stream by cipher and mode

  InputCipher::InputCipher (BlockCipher* cifr, const t_cmod cmod) {
    // bind object
    Object::iref (p_cifr = cifr);
    p_is = nilp;
    // grab cipher information
    d_cmod = cmod;
    // create buffer block
    long cbsz = (p_cifr == nilp) ? 0 : p_cifr->getcbsz ();
    p_bo = (cbsz > 0) ? new t_byte[cbsz] : nilp;
    p_bi = (cbsz > 0) ? new t_byte[cbsz] : nilp;
    p_bl = (cbsz > 0) ? new t_byte[cbsz] : nilp;
    // preset initial vector
    p_iv = preldiv_blk (cbsz);
    // reset eveything
    reset ();
  }

  // create a new input cipher stream by cipher and input stream

  InputCipher::InputCipher (BlockCipher* cifr, Input* is) {
    // bind object
    Object::iref (p_cifr = cifr);
    Object::iref (p_is = is);
    // grab cipher information
    d_cmod = ECB;
    // create buffer blocks
    long cbsz = (p_cifr == nilp) ? 0 : p_cifr->getcbsz ();
    p_bo = (cbsz > 0) ? new t_byte[cbsz] : nilp;
    p_bi = (cbsz > 0) ? new t_byte[cbsz] : nilp;
    p_bl = (cbsz > 0) ? new t_byte[cbsz] : nilp;
    // preset initial vector
    p_iv = preldiv_blk (cbsz);
    // reset eveything
    reset ();
  }

  // create a new input cipher stream by cipher, input stream and mode

  InputCipher::InputCipher (BlockCipher* cifr, Input* is, const t_cmod cmod) {
    // bind object
    Object::iref (p_cifr = cifr);
    Object::iref (p_is = is);
    // grab cipher information
    d_cmod = cmod;
    // create buffer block
    long cbsz = (p_cifr == nilp) ? 0 : p_cifr->getcbsz ();
    p_bo = (cbsz > 0) ? new t_byte[cbsz] : nilp;
    p_bi = (cbsz > 0) ? new t_byte[cbsz] : nilp;
    p_bl = (cbsz > 0) ? new t_byte[cbsz] : nilp;
    // preset initial vector
    p_iv = preldiv_blk (cbsz);
    // reset eveything
    reset ();
  }

  // destroy this input cipher

  InputCipher::~InputCipher (void) {
    Object::dref (p_cifr);
    Object::dref (p_is);
    delete [] p_bo;
    delete [] p_bi;
    delete [] p_bl;
    delete [] p_iv;
  }

  // return the class name

  String InputCipher::repr (void) const {
    return "InputCipher";
  }

  // reset the input cipher

  void InputCipher::reset (void) {
    wrlock ();
    try {
      // get the block size
      long cbsz = (p_cifr == nilp) ? 0 : p_cifr->getcbsz ();
      // reset buffer block 
      for (long i = 0; i < cbsz; i++) {
	p_bo[i] = nilc;
	p_bi[i] = nilc;
	p_bl[i] = nilc;
      }
      // reset with the initial vector
      resetiv_blk (d_cmod, cbsz, p_iv, p_bo, p_bi, p_bl);
      // done
      unlock ();
    } catch (...) {
      unlock ();
      throw;
    }
  }

  // set the initial vector by octet string

  void InputCipher::setiv (const String& ivs) {
    // convert the octet string
    long ivsz = 0;
    t_byte* sbuf = Unicode::stob (ivsz, ivs);
    // lock and set
    wrlock ();
    try {
      // check the size
      long cbsz = (p_cifr == nilp) ? 0 : p_cifr->getcbsz ();
      if (ivsz != cbsz) {
	throw Exception ("cipher-error", "invalid initial vector size");
      }
      for (long i = 0; i < cbsz; i++) p_iv[i] = sbuf[i];
      delete [] sbuf;
      reset ();
      unlock ();
    } catch (...) {
      delete [] sbuf;
      unlock ();
      throw;
    }
  }

  // set the initial vector

  void InputCipher::setiv (const Buffer& ivb) {
    wrlock ();
    try {
      long cbsz = (p_cifr == nilp) ? 0 : p_cifr->getcbsz ();
      for (long i = 0; i < cbsz; i++) p_iv[i] = ivb.get (i);
      reset ();
      unlock ();
    } catch (...) {
      unlock ();
      throw;
    }
  }

  // get the initial vector buffer

  String InputCipher::getiv (void) const {
    rdlock ();
    try {
      String result;
      long cbsz = (p_cifr == nilp) ? 0 : p_cifr->getcbsz ();
      if (cbsz > 0) result = Ascii::btos (p_iv, cbsz);
      unlock ();
      return result;
    } catch (...) {
      unlock ();
      throw;
    }
  }

  // set the input cipher input stream
  
  void InputCipher::setis (Input* is) {
    wrlock ();
    try {
      Object::iref (is);
      Object::dref (p_is);
      p_is = is;
      unlock ();
    } catch (...) {
      unlock ();
      throw;
    }
  }

  // get the input cipher mode

  InputCipher::t_cmod InputCipher::getmode (void) const {
    rdlock ();
    try {
      t_cmod result = d_cmod;
      unlock ();
      return result;
    } catch (...) {
      unlock ();
      throw;
    }
  }

  // read one byte from the input stream.

  char InputCipher::read (void) {
    wrlock ();
    try {
      // eventually update the buffer
      update_buf (d_sbuf, p_cifr, p_is, d_cmod, p_bo, p_bi, p_bl);
      // read a character
      char result = d_sbuf.read ();
      // eventually update again the buffer
      update_buf (d_sbuf, p_cifr, p_is, d_cmod, p_bo, p_bi, p_bl);
      // unlock and return
      unlock ();
      return result;
    } catch (...) {
      unlock ();
      throw;
    }
  }

  // return true if we are at the eos

  bool InputCipher::iseos (void) const {
    rdlock ();
    try {
      // eventually update the buffer
      update_buf (d_sbuf, p_cifr, p_is, d_cmod, p_bo, p_bi, p_bl);
      // check the buffer
      bool result = d_sbuf.empty ();
      unlock ();
      return result;
    } catch (...) {
      unlock ();
      throw;
    }
  }

  // return true if we can read a character

  bool InputCipher::valid (void) const {
    rdlock ();
    try {
      // eventually update the buffer
      update_buf (d_sbuf, p_cifr, p_is, d_cmod, p_bo, p_bi, p_bl);
      // check the buffer
      bool result = !d_sbuf.empty ();
      unlock ();
      return result;
    } catch (...) {
      unlock ();
      throw;
    }
  }

  // -------------------------------------------------------------------------
  // - object section                                                        -
  // -------------------------------------------------------------------------

  // the object eval quarks
  static const long QUARK_ECB     = String::intern ("ECB");
  static const long QUARK_CBC     = String::intern ("CBC");
  static const long QUARK_CFB     = String::intern ("CFB");
  static const long QUARK_OFB     = String::intern ("OFB");
  static const long QUARK_ICIPHER = String::intern ("InputCipher");

  // the quark zone
  static const long QUARK_ZONE_LENGTH = 5;
  static QuarkZone  zone (QUARK_ZONE_LENGTH);

  // the object supported quarks
  static const long QUARK_RESET   = zone.intern ("reset");
  static const long QUARK_SETIV   = zone.intern ("set-iv");
  static const long QUARK_GETIV   = zone.intern ("get-iv");
  static const long QUARK_SETIS   = zone.intern ("set-is");
  static const long QUARK_GETMODE = zone.intern ("get-mode");

  // map an enumeration item to an input cipher mode
  static inline InputCipher::t_cmod item_to_cmod (const Item& item) {
    // check for an input cipher item
    if (item.gettid () != QUARK_ICIPHER)
      throw Exception ("item-error", "item is not an input cipher item");
    // map the item to the enumeration
    long quark = item.getquark ();
    if (quark == QUARK_ECB) return InputCipher::ECB;
    if (quark == QUARK_CBC) return InputCipher::CBC;
    if (quark == QUARK_CFB) return InputCipher::CFB;
    if (quark == QUARK_OFB) return InputCipher::OFB;
    throw Exception ("item-error", "cannot map item to input cipher mode");
  }

  // map an input cipher mode to an item
  static inline Item* cmod_to_item (const InputCipher::t_cmod cmod) {
    switch (cmod) {
    case InputCipher::ECB:
      return new Item (QUARK_ICIPHER, QUARK_ECB);
      break;
    case InputCipher::CBC:
      return new Item (QUARK_ICIPHER, QUARK_CBC);
      break;
    case InputCipher::CFB:
      return new Item (QUARK_ICIPHER, QUARK_CFB);
      break;
    case InputCipher::OFB:
      return new Item (QUARK_ICIPHER, QUARK_OFB);
      break;
    }
    return nilp;
  }

  // evaluate an object data member

  Object* InputCipher::meval (Runnable* robj, Nameset* nset, 
			      const long quark) {
    if (quark == QUARK_ECB) 
      return new Item (QUARK_ICIPHER, QUARK_ECB);
    if (quark == QUARK_CBC) 
      return new Item (QUARK_ICIPHER, QUARK_CBC);
    if (quark == QUARK_CFB) 
      return new Item (QUARK_ICIPHER, QUARK_CFB);
    if (quark == QUARK_OFB) 
      return new Item (QUARK_ICIPHER, QUARK_OFB);
    throw Exception ("eval-error", "cannot evaluate member",
		     String::qmap (quark));
  }

  // create a new object in a generic way

  Object* InputCipher::mknew (Vector* argv) {
    // get the number of arguments
    long argc = (argv == nilp) ? 0 : argv->length ();

    // check for 1 argument
    if (argc == 1) {
      // get the cipher object
      Object* cobj = argv->get (0);
      BlockCipher* cifr = dynamic_cast <BlockCipher*> (cobj);
      if (cifr == nilp) {
	throw Exception ("argument-error", 
			 "invalid arguments with input cipher");
      }
      return new InputCipher (cifr);
    }
    // check for 2 arguments
    if (argc == 2) {
      // get the cipher object
      Object* cobj = argv->get (0);
      BlockCipher* cifr = dynamic_cast <BlockCipher*> (cobj);
      if (cifr == nilp) {
	throw Exception ("argument-error", 
			 "invalid arguments with input cipher");
      }
      // get the input stream object
      Object* iobj = argv->get (1);
      Input*  is   = dynamic_cast <Input*> (iobj);
      if (is != nilp) {
	return new InputCipher (cifr, is);
      }
      // check for a mode
      Item* item = dynamic_cast <Item*> (iobj);
      if (item != nilp) {
	t_cmod cmod = item_to_cmod (*item);
	return new InputCipher (cifr, cmod);
      }
      throw Exception ("argument-error", 
		       "invalid arguments with input cipher");
    }
    // check for 3 arguments
    if (argc == 3) {
      // get the cipher object
      Object* cobj = argv->get (0);
      BlockCipher* cifr = dynamic_cast <BlockCipher*> (cobj);
      if (cifr == nilp) {
	throw Exception ("argument-error", 
			 "invalid arguments with input cipher");
      }
      // get the input stream object
      Object* iobj = argv->get (1);
      Input*  is   = dynamic_cast <Input*> (iobj);
      if (is == nilp) {
	throw Exception ("argument-error", 
			 "invalid arguments with input cipher");
      }
      // get the mode as an item
      Object* mobj = argv->get (2);
      Item*   item = dynamic_cast <Item*> (mobj);
      if (item == nilp) {
	throw Exception ("argument-error", 
			 "invalid arguments with input cipher");
      }
      t_cmod cmod = item_to_cmod (*item);
      return new InputCipher (cifr, is, cmod);
    }
    throw Exception ("argument-error", 
		     "invalid arguments with with input cipher");
  }

  // return true if the given quark is defined

  bool InputCipher::isquark (const long quark, const bool hflg) const {
    rdlock ();
    if (zone.exists (quark) == true) {
      unlock ();
      return true;
    }
    bool result = hflg ? Input::isquark (quark, hflg) : false;
    unlock ();
    return result;
  }

  // apply this object with a set of arguments and a quark
  
  Object* InputCipher::apply (Runnable* robj, Nameset* nset, const long quark,
			      Vector* argv) {
    // get the number of arguments
    long argc = (argv == nilp) ? 0 : argv->length ();

    // dispatch 0 argument
    if (argc == 0) {
      if (quark == QUARK_GETIV)   return new String (getiv ());
      if (quark == QUARK_GETMODE) return cmod_to_item (getmode ());
      if (quark == QUARK_RESET) {
	reset ();
	return nilp;
      }
    }
    // dispatch 1 argument
    if (argc == 1) {
      if (quark == QUARK_SETIV) {
	Object* obj = argv->get (0);
	// check for a string
	String* ivs = dynamic_cast <String*> (obj);
	if (ivs != nilp) {
	  setiv (*ivs);
	  return nilp;
	}
	// check for a buffer
	Buffer* ivb = dynamic_cast <Buffer*> (obj);
	if (ivb != nilp) {
	  setiv (*ivb);
	  return nilp;
	}
	throw Exception ("type-error", "invalid input object for set-iv",
			 Object::repr (obj));
      }
      if (quark == QUARK_SETIS) {
	Object* obj = argv->get (0);
	Input*  is  = dynamic_cast <Input*> (obj);
	if (is == nilp) {
	  throw Exception ("type-error", 
			   "invalid input object for input cipher set-is",
			   Object::repr (obj));
	}
	setis (is);
	return nilp;
      }
    }
    // call the input stream method
    return Input::apply (robj, nset, quark, argv);
  }
}
