// ---------------------------------------------------------------------------
// - BlockCipher.cpp                                                         -
// - afnix:sec module - block 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 "Vector.hpp"
#include "Integer.hpp"
#include "QuarkZone.hpp"
#include "BlockCipher.hpp"

namespace afnix {

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

  // this procedure computes a block padding: bit padding
  static bool bc_pad_bitm (t_byte* bi, const long bisz, const long size) {
    // compute the padding count
    long pcnt = bisz - size;
    if (pcnt > 0) {
      // pad the first byte
      bi[size] = 0x80;
      // pad with the serie 80 00 00 00 ...
      for (long i = size+1; i < bisz; i++) bi[i] = 0x00;
      // mark as padded
      return true;
    }
    return false;
  }

  // this procedure compute the next block padding: bit padding
  static bool bc_nxp_bitm (t_byte* bi, const long bisz) {
    // pad the first byte
    bi[0] = 0x80;
    // pad the rest of the block
    for (long i = 1; i < bisz; i++) bi[i] = 0x00;
    // mark as padded
    return true;
  }

  // this procedure computes the number of valid bytes: bit padding
  static long bc_unp_bitm (const t_byte* bo, const long bosz) {
    long pidx = 0;
    for (long i = bosz-1; i >= 0; i--) {
      if (bo[i] == 0x80) {
	pidx = i;
	break;
      }
      if (bo[i] != 0x00) {
	throw Exception ("cipher-error", "invalid padded block");
      }
    }
    // check padded index
    if (bo[pidx] != 0x80) {
      throw Exception ("cipher-error", "invalid padded block");
    }
    // the size if the previous byte index + 1, so is pidx
    return pidx;
  }

  // this procedure computes a block padding: ANSI X.923
  static bool bc_pad_ansi (t_byte* bi, const long bisz, const long size) {
    // compute the padding count
    long pcnt = bisz - size;
    // check the padding
    if (pcnt > 256) {
      throw Exception ("cipher-error", "padding size is too long");
    }
    if (pcnt > 0) {
      // pad with the pad size
      for (long i = size; i < bisz-1; i++) bi[i] = 0x00;
      bi[bisz-1] = (t_byte) pcnt;
      // mark as padded
      return true;
    }
    return false;
  }

  // this procedure computes the next block padding: ANSI X.923
  static bool bc_nxp_ansi (t_byte* bi, const long bisz) {
    // check block size
    if (bisz > 256) {
      throw Exception ("cipher-error", "padding size is too long");
    }
    // pad with the block size
    for (long i = 0; i < bisz-1; i++) bi[i] = 0x00;
    bi[bisz-1] = (t_byte) bisz;
    // mark as padded
    return true;
  }

  // this procedure computes the number of valid bytes: ANSI X.923
  static long bc_unp_ansi (const t_byte* bo, const long bosz) {
    // initialize result
    long result = bosz;
    // extract valid bytes
    long pad = bo[bosz-1];
    if (pad > bosz) {
      throw Exception ("cipher-error", "block padding is too long");	
    }
    // update counter
    result -= pad;
    // number of valid bytes
    return result;
  }

  // this procedure computes a block padding: NIST 800-38A
  static bool bc_pad_nist (t_byte* bi, const long bisz, const long size) {
    // compute the padding count
    long pcnt = bisz - size;
    // check the padding
    if (pcnt > 256) {
      throw Exception ("cipher-error", "padding size is too long");
    }
    if (pcnt > 0) {
      // pad with the pad size
      for (long i = size; i < bisz; i++) bi[i] = (t_byte) pcnt;
      // mark as padded
      return true;
    }
    return false;
  }

  // this procedure computes the next block padding: NIST 800-38A
  static bool bc_nxp_nist (t_byte* bi, const long bisz) {
    // check block size
    if (bisz > 256) {
      throw Exception ("cipher-error", "padding size is too long");
    }
    // pad with the block size
    for (long i = 0; i < bisz; i++) bi[i] = (t_byte) bisz;
    // mark as padded
    return true;
  }

  // this procedure computes the number of valid bytes: NIST 800-38A
  static long bc_unp_nist (const t_byte* bo, const long bosz) {
    // initialize result
    long result = bosz;
    // extract valid bytes
    long pad = bo[bosz-1];
    if (pad > bosz) {
      throw Exception ("cipher-error", "block padding is too long");	
    }
    // update counter
    result -= pad;
    // number of valid bytes
    return result;
  }

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

  // create a cipher by name and key

  BlockCipher::BlockCipher (const String& name,
			    const Key& key) : Cipher (name) {
    // reset block size and padding
    d_cbsz = 0;
    d_pmod = PAD_N800;
    // set the key
    setkey (key);
  }

  // return the class name

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

  // set the block cipher key

  void BlockCipher::setkey (const Key& key) {
    wrlock ();
    try {
      // check the key for validity
      if (key.isbk () == false) {
	throw Exception ("key-error", "invalid key for block cipher");
      } 
      // set the key and return
      Cipher::setkey (key);
      unlock ();
    } catch (...) {
      unlock ();
      throw;
    }
  }

  // return the cipher block size

  long BlockCipher::getcbsz (void) const {
    rdlock ();
    try {
      long result = d_cbsz;
      unlock ();
      return result;
    } catch (...) {
      unlock ();
      throw;
    }
  }

  // set the block cipher padding mode

  void BlockCipher::setpmod (const t_pmod pmod) {
    wrlock ();
    try {
      d_pmod = pmod;
      unlock ();
    } catch (...) {
      unlock ();
      throw;
    }
  }

  // get the block cipher padding mode

  BlockCipher::t_pmod BlockCipher::getpmod (void) const {
    rdlock ();
    try {
      t_pmod result = d_pmod;
      unlock ();
      return result;
    } catch (...) {
      unlock ();
      throw;
    }
  }

  // get the normal waist from a file size

  t_long BlockCipher::waist (const t_long size) const {
    rdlock ();
    try {
      // compute waist size
      t_long result = (size / d_cbsz) * d_cbsz;
      if ((size == 0) || ((size % d_cbsz) != 0)) result+= d_cbsz;
      // eventually add a padding block
      switch (d_pmod) {
      case PAD_BITM:
      case PAD_X923:
      case PAD_N800:
	if ((size > 0) && ((size % d_cbsz) == 0)) result += d_cbsz;
	break;
      default:
	break;
      }
      unlock ();
      return result;
    } catch (...) {
      unlock ();
      throw;
    }
  }
  
  // process an input buffer into an output buffer

  t_long BlockCipher::stream (Buffer& ob, Buffer& ib) {
    // trivial check first
    if (d_cbsz == 0) return 0;
    // force the output buffer to be resizable
    ob.setrflg (true);
    // lock and process
    wrlock ();
    t_long  cc = 0;
    t_byte* bi = new t_byte[d_cbsz];
    t_byte* bo = new t_byte[d_cbsz];
    // loop as long as the buffer is not empty
    try {
      // initialize the block counter and input buffer
      long bcnt = 0;
      for (long i = 0; i < d_cbsz; i++) bi[i] = nilc;
      while (ib.empty () == false) {
	// load the buffer
	bi[bcnt++] = ib.read ();
	// check for full block
	if (bcnt < d_cbsz) continue;
	// process the block
	process (bo, bi);
	// check for unpadding in reverse mode
	if ((d_rflg == true) && (ib.empty () == true)) {
	  bcnt = getub (bo);
	}
	// update the buffer and counters
	if (bcnt > 0) ob.add ((char*) bo, bcnt);
	cc += bcnt;
	bcnt = 0;
	// clean input buffer
	for (long i = 0; i < d_cbsz; i++) bi[i] = nilc;
      }
      // check for last block
      if (bcnt != 0) {
	// check for normal mode
	if (d_rflg == true)  {
	  throw Exception ("cipher-error", "short block in reverse mode");
	}
	// check for padding
	mkpad (bi, bcnt);
	// process the block
	process (bo, bi);
	// update the buffer and counters
	ob.add ((char*) bo, d_cbsz);
	cc += bcnt;
      }
      // check for extra padded block
      if ((d_rflg == false) && (bcnt == 0)) {
	// fill next pad block
	bool pflg = mknxp (bi);
	// add it in the buffer
	if (pflg == true) {
	  // process the block
	  process (bo, bi);
	  // update the buffer and counters
	  ob.add ((char*) bo, d_cbsz);
	}
      }
      // clean and unlock
      delete [] bi;
      delete [] bo;
      unlock ();
    } catch (...) {
      delete [] bi;
      delete [] bo;
      unlock ();
      throw;
    }
    return cc;
  }

  // process an input buffer into an output buffer

  t_long BlockCipher::stream (Output& os, Input& is) {
    // trivial check first
    if (d_cbsz == 0) return 0;
    // lock and process
    wrlock ();
    t_long  cc = 0;
    t_byte* bi = new t_byte[d_cbsz];
    t_byte* bo = new t_byte[d_cbsz];
    // loop as long as the buffer is not empty
    try {
      // initialize the block counter and input buffer
      long bcnt = 0;
      for (long i = 0; i < d_cbsz; i++) bi[i] = nilc;
      while (is.valid () == true) {
	// load the buffer
	bi[bcnt++] = is.read ();
	// check for full block
	if (bcnt < d_cbsz) continue;
	// process the block
	process (bo, bi);
	// check for unpadding in reverse mode
	if ((d_rflg == true) && (is.iseos () == true)) {
	  bcnt = getub (bo);
	}
	// update the buffer and counters
	if (bcnt > 0) os.write ((char*) bo, bcnt);
	cc += bcnt;
	bcnt = 0;
	// clean input buffer
	for (long i = 0; i < d_cbsz; i++) bi[i] = nilc;
      }
      // check for last block
      if (bcnt != 0) {
	// check for normal mode
	if (d_rflg == true)  {
	  throw Exception ("cipher-error", "short block in reverse mode");
	}
	// check for padding
	mkpad (bi, bcnt);
	// process the block
	process (bo, bi);
	// update the buffer and counters
	os.write ((char*) bo, d_cbsz);
	cc += bcnt;
      }
      // check for extra padded block
      if ((d_rflg == false) && (bcnt == 0)) {
	// fill next pad block
	bool pflg = mknxp (bi);
	// add it in the buffer
	if (pflg == true) {
	  // process the block
	  process (bo, bi);
	  // update the buffer and counters
	  os.write ((char*) bo, d_cbsz);
	}
      }
      // clean and unlock
      delete [] bi;
      delete [] bo;
      unlock ();
    } catch (...) {
      delete [] bi;
      delete [] bo;
      unlock ();
      throw;
    }
    return cc;
  }

  // generate a block padding if any

  bool BlockCipher::mkpad (t_byte* bi, const long size) const {
    rdlock ();
    try {
      // check block size first
      if ((size < 0) || (size > d_cbsz)) {
	throw Exception ("cipher-error", "invalid block size to pad");
      }
      // check the cipher mode
      if (d_rflg == true) {
	throw Exception ("cipher-error", "invalid padding in reverse mode");
      }
      // by default do not pad
      bool result = false;
      // select from the padding mode
      switch (d_pmod) {
      case PAD_BITM:
	result = bc_pad_bitm (bi, d_cbsz, size);
	break;
      case PAD_X923:
	result = bc_pad_ansi (bi, d_cbsz, size);
	break;
      case PAD_N800:
	result = bc_pad_nist (bi, d_cbsz, size);
	break;
      default:
	result = false;
	break;
      }
      // unlock and return
      unlock ();
      return result;
    } catch (...) {
      unlock ();
      throw;
    }
  }

  // generate the next padded block if any

  bool BlockCipher::mknxp (t_byte* bi) const {
    rdlock ();
    try {
      // check the cipher mode
      if (d_rflg == true) {
	throw Exception ("cipher-error", "invalid padding in reverse mode");
      }
      // by default do not pad
      bool result = false;
      // select from the padding mode
      switch (d_pmod) {
      case PAD_BITM:
	result = bc_nxp_bitm (bi, d_cbsz);
	break;
      case PAD_X923:
	result = bc_nxp_ansi (bi, d_cbsz);
	break;
      case PAD_N800:
	result = bc_nxp_nist (bi, d_cbsz);
	break;
      default:
	result = false;
	break;
      }
      // unlock and return
      unlock ();
      return result;
    } catch (...) {
      unlock ();
      throw;
    }
  }

  // get the number of unpadded bytes
  
  long BlockCipher::getub (const t_byte* bo) const {
    rdlock ();
    try {
      // check input buffer
      if (bo == nilp) {
	throw Exception ("cipher-error", "invalid nil block to unpad");
      }
      // check the cipher mode
      if (d_rflg == false) {
	throw Exception ("cipher-error", "invalid unpadding in normal mode");
      }
      // initialize result
      long result = d_cbsz;
      // select from the padding mode
      switch (d_pmod) {
      case PAD_BITM:
	result = bc_unp_bitm (bo, d_cbsz);
	break;
      case PAD_X923:
	result = bc_unp_ansi (bo, d_cbsz);
	break;
      case PAD_N800:
	result = bc_unp_nist (bo, d_cbsz);
	break;
      default:
	result = d_cbsz;
	break;
      }
      // unlock and return
      unlock ();
      return result;
    } catch (...) {
      unlock ();
      throw;
    }
  }

  // process a block buffer into another one

  void BlockCipher::process (t_byte* bo, const t_byte* bi) {
    wrlock ();
    try {
      // encode/decode from reverse flag
      if (d_rflg == false)
	encode (bo, bi);
      else
	decode (bo, bi);
      // unlock and return
      unlock ();
    } catch (...) {
      unlock ();
      throw;
    }
  }

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

  // the object eval quarks
  static const long QUARK_NONE    = String::intern ("PAD-NONE");
  static const long QUARK_BITM    = String::intern ("PAD-BIT-MODE");
  static const long QUARK_X923    = String::intern ("PAD-ANSI-X923");
  static const long QUARK_N800    = String::intern ("PAD-NIST-800");
  static const long QUARK_BCIPHER = String::intern ("InputCipher");

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

  // the object supported quarks
  static const long QUARK_GETCBSZ = zone.intern ("get-block-size");
  static const long QUARK_SETPMOD = zone.intern ("set-padding-mode");
  static const long QUARK_GETPMOD = zone.intern ("get-padding-mode");

  // map an enumeration item to a block cipher padding mode
  static inline BlockCipher::t_pmod item_to_pmod (const Item& item) {
    // check for an input cipher item
    if (item.gettid () != QUARK_BCIPHER)
      throw Exception ("item-error", "item is not a block cipher item");
    // map the item to the enumeration
    long quark = item.getquark ();
    if (quark == QUARK_NONE) return BlockCipher::PAD_NONE;
    if (quark == QUARK_BITM) return BlockCipher::PAD_BITM;
    if (quark == QUARK_X923) return BlockCipher::PAD_X923;
    if (quark == QUARK_N800) return BlockCipher::PAD_N800;
    throw Exception ("item-error", 
		     "cannot map item to block cipher padding mode");
  }
  
  // map a block cipher padding mode to an item
  static inline Item* pmod_to_item (const BlockCipher::t_pmod pmod) {
    switch (pmod) {
    case BlockCipher::PAD_NONE:
      return new Item (QUARK_BCIPHER, QUARK_NONE);
      break;
    case BlockCipher::PAD_BITM:
      return new Item (QUARK_BCIPHER, QUARK_BITM);
      break;
    case BlockCipher::PAD_X923:
      return new Item (QUARK_BCIPHER, QUARK_X923);
      break;
    case BlockCipher::PAD_N800:
      return new Item (QUARK_BCIPHER, QUARK_N800);
      break;
    }
    return nilp;
  }
  
  // evaluate an object data member
  
  Object* BlockCipher::meval (Runnable* robj, Nameset* nset, 
			      const long quark) {
    if (quark == QUARK_NONE) 
      return new Item (QUARK_BCIPHER, QUARK_NONE);
    if (quark == QUARK_BITM) 
      return new Item (QUARK_BCIPHER, QUARK_BITM);
    if (quark == QUARK_X923) 
      return new Item (QUARK_BCIPHER, QUARK_X923);
    if (quark == QUARK_N800) 
      return new Item (QUARK_BCIPHER, QUARK_N800);
    throw Exception ("eval-error", "cannot evaluate member",
		     String::qmap (quark));
  }

  // return true if the given quark is defined

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

  // apply this object with a set of arguments and a quark
  
  Object* BlockCipher::apply (Runnable* robj, Nameset* nset, const long quark,
			      Vector* argv) {
    // get the number of arguments
    long argc = (argv == nilp) ? 0 : argv->length ();
    
    // check for 0 argument
    if (argc == 0) {
      if (quark == QUARK_GETCBSZ) return new Integer  (getcbsz ());
      if (quark == QUARK_GETPMOD) return pmod_to_item (getpmod ());
    }
    // check for 1 argument
    if (argc == 1) {
      if (quark == QUARK_SETPMOD) {
	Object* obj = argv->get (0);
	Item*  iobj = dynamic_cast <Item*> (obj);
	if (iobj == nilp) {
	  throw Exception ("type-error", "invalid object with set padding mode",
			   Object::repr (obj));
	}
	t_pmod pmod = item_to_pmod (*iobj);
	setpmod (pmod);
	return nilp;
      }
    }
    // call the cipher method
    return Cipher::apply (robj, nset, quark, argv);
  }
}
