// ---------------------------------------------------------------------------
// - TlsCkeyxh.cpp                                                           -
// - afnix:tls service - tls client key exchange 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-2023 amaury darsch                                   -
// ---------------------------------------------------------------------------

#include "Vector.hpp"
#include "Crypto.hpp"
#include "TlsTypes.hxx"
#include "TlsSuite.hxx"
#include "TlsUtils.hpp"
#include "TlsCkeyxh.hpp"
#include "QuarkZone.hpp"
#include "Exception.hpp"
#include "transient.tcc"

namespace afnix {

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

  // decode a handshake buffer to a premaster secret buffer with rsa
  static Buffer tls_rsa_prebuf (Buffer& hbuf, TlsState* ssta) {
    // get the cipher private key
    Key* prvk = (ssta == nullptr) ? nullptr : ssta->getprvk ();
    if (prvk == nullptr) {
      throw Exception ("tls-error", "nil rsa key during key exchange");
    }
    // create a decoding cipher by key
    t_transient<Cipher> cifr = Crypto::mkcipher (*prvk, true);
    if (cifr.valid() == false) {
      throw Exception ("tls-error", "cannot create key exchange cipher");
    }
    // the first 2 bytes should be the block size
    t_byte hbyt = (t_byte) hbuf.read ();
    t_byte lbyt = (t_byte) hbuf.read ();
    long   hlen = (((t_word) hbyt) << 8) + ((t_word) lbyt);
    // check the current buffer size
    if (hbuf.length () != hlen) {
      throw Exception ("tls-error", "inconsistent encrypted exchange size");
    }
    // the premaster result buffer
    Buffer result;
    long blen = cifr->stream (result, hbuf);
    if (blen != result.length ()) {
      throw Exception ("tls-error",
		       "internal error while decoding client key exchange");
    }
    return result;
  }

  // decode a handshake buffer to a premaster secret buffer with dhe
  static Buffer tls_dhe_prebuf (Buffer& hbuf, TlsState* ssta) {
    // collect dhe parameters
    TlsDhe* dhep = (ssta == nullptr) ? nullptr : ssta->getdhep ();
    if (dhep == nullptr) {
      throw Exception ("tls-error", "nil dhe during key exchange");
    }
    // the first 2 bytes should be the block size
    t_byte hbyt = (t_byte) hbuf.read ();
    t_byte lbyt = (t_byte) hbuf.read ();
    long   hlen = (((t_word) hbyt) << 8) + ((t_word) lbyt);
    // check the current buffer size
    if (hbuf.length () != hlen) {
      throw Exception ("tls-error", "inconsistent encrypted exchange size");
    }
    // rebuild client public key
    Relatif ykey (hbuf.tobyte (), hbuf.tosize());
    // compute common key
    Relatif kxch = dhep->tokxch (ykey);
    // get the relatif byte size
    long rmsb = kxch.getmsb ();
    long size = ((rmsb % 8) == 0) ? rmsb / 8 : rmsb / 8 + 1;
    if (size == 0L) {
      throw Exception ("tls-error", "null relatif view for key exchange");
    }
    // collect relatif byte array
    t_byte rbuf[size];
    if (kxch.toubuf (rbuf, size) != size) {
      throw Exception ("tls-error", "cannot map key exchange as byte buffer");
    }
    Buffer result (size); result.add ((char*) rbuf, size);
    return result;
  }
  
  // -------------------------------------------------------------------------
  // - class section                                                         -
  // -------------------------------------------------------------------------

  // create a default key exchange

  TlsCkeyxh::TlsCkeyxh (void) {
    d_mlen = 0L;
    p_mbuf = nullptr;
  }

  // create a key exchange by state

  TlsCkeyxh::TlsCkeyxh (TlsState* ssta) {
    // create a premaster secret
    d_mlen = TLS_SIZE_MSX;
    p_mbuf = TlsUtils::random (d_mlen, false);
    // set the state version
    if (ssta != nullptr) {
      // collect version
      p_mbuf[0] = ssta->getrmaj ();
      p_mbuf[1] = ssta->getrmin ();
      // set the premaster key
      ssta->setmkey(getmbuf());
    }
  }

  // create a key exchange by message block

  TlsCkeyxh::TlsCkeyxh (TlsHblock* hblk) {
    d_mlen = 0L;
    p_mbuf = nullptr;
    decode (hblk, nullptr);
  }
  
  // create a key exchange by message block and key

  TlsCkeyxh::TlsCkeyxh (TlsHblock* hblk, TlsState* ssta) {
    d_mlen = 0L;
    p_mbuf = nullptr;
    decode (hblk, ssta);
  }

  // destroy this  block
  
  TlsCkeyxh::~TlsCkeyxh (void) {
    delete [] p_mbuf; p_mbuf = nullptr;
    reset ();
  }

  // return the class name
  
  String TlsCkeyxh::repr (void) const {
    return "TlsCkeyxh";
  }

  // reset the key exchange

  void TlsCkeyxh::reset (void) {
    wrlock ();
    try {
      d_mlen = 0L;
      delete [] p_mbuf; p_mbuf = nullptr;
      unlock ();
    } catch (...) {
      unlock ();
      throw;
    }
  }

  // get the premaster buffer

  Buffer TlsCkeyxh::getmbuf (void) const {
    rdlock ();
    try {
      // check for valid buffer
      if (d_mlen == 0L) {
	throw Exception ("tls-error", "invalid null premaster buffer");
      }
      Buffer result;
      if (result.add ((char*) p_mbuf, d_mlen) != d_mlen) {
	throw Exception ("tls-error", "invalid premaster buffer");
      }
      unlock ();
      return result;
    } catch (...) {
      unlock ();
      throw;
    }
  }
  
  // get the key exchange info as a plist

  Plist TlsCkeyxh::getinfo (void) const {
    rdlock ();
    try {
      // create a result plist
      Plist plst;
      // here it is
      unlock ();
      return plst;
    } catch (...) {
      unlock ();
      throw;
    }
  }
  
  // decode the message block

  void TlsCkeyxh::decode (TlsHblock* hblk, TlsState* ssta) {
    wrlock ();
    try {
      // reset eveything
      reset ();
      if (hblk == nullptr) {
	unlock ();
	return;
      }
      // collect the cipher
      t_word cifr = (ssta == nullptr) ? 0x0000U : ssta->getcifr ();
      if (cifr == TLS_NULL_WITH_NULL_NULL) {
	unlock ();
	return;
      }
      // map the handshake to a buffer
      Buffer hbuf = hblk->tobuffer();
      // build the premaster buffer
      Buffer mbuf;
      switch (cifr) {
      case TLS_RSA_WITH_NULL_MD5:
      case TLS_RSA_WITH_NULL_SHA:
      case TLS_RSA_WITH_RC4_128_MD5:
      case TLS_RSA_WITH_RC4_128_SHA:
      case TLS_RSA_WITH_AES_128_CBC_SHA:
      case TLS_RSA_WITH_AES_256_CBC_SHA:
      case TLS_RSA_WITH_AES_128_CBC_SHA256:
      case TLS_RSA_WITH_AES_256_CBC_SHA256:
      case TLS_RSA_WITH_AES_128_GCM_SHA256:
      case TLS_RSA_WITH_AES_256_GCM_SHA384:
	mbuf = tls_rsa_prebuf (hbuf, ssta);
	break;
      case TLS_DHE_RSA_WITH_AES_128_CBC_SHA:
      case TLS_DHE_RSA_WITH_AES_256_CBC_SHA:
      case TLS_DHE_RSA_WITH_AES_128_CBC_SHA256:
      case TLS_DHE_RSA_WITH_AES_256_CBC_SHA256:
      case TLS_DHE_RSA_WITH_AES_128_GCM_SHA256:
      case TLS_DHE_RSA_WITH_AES_256_GCM_SHA384:
	mbuf = tls_dhe_prebuf (hbuf, ssta);
	break;
	break;
      default:
	mbuf = hbuf;
	break;
      }
      // copy the premaster secret
      d_mlen = mbuf.length ();
      delete [] p_mbuf; p_mbuf = (d_mlen > 0L) ? new t_byte[d_mlen] : nullptr;
      if (mbuf.copy ((char*) p_mbuf, d_mlen) != d_mlen) {
	throw Exception ("tls-error", "cannot extract pre-master secret");
      }
      mbuf.reset ();
      // check against malicious version
      //if (tls_vers_valid (p_mbuf[0], p_mbuf[1]) == false) {
      //throw Exception ("tls-error", "malicious version detected");
      //}
      // update the state premaster key
      if (ssta != nullptr) ssta->setmkey(getmbuf());
      // done
      unlock ();
    } catch (...) {
      reset ();
      unlock ();
      throw;
    }
  }
  
  // map a client key exchange to a chunk block

  TlsChunk TlsCkeyxh::tochunk (TlsState* sta) const {
    rdlock ();
    try {
      // get the public encoding key
      TlsCerts* cert = (sta == nullptr) ? nullptr : sta->getcert ();
      Key* pubk = (cert == nullptr) ? nullptr : cert->getpubk ();
      // create the pre-master buffer
      Buffer mbuf(d_mlen);
      if (mbuf.add ((char*) p_mbuf, d_mlen) != d_mlen) {
	throw Exception ("tls-error",
			 "cannot create client key exchange buffer");
      }
      TlsChunk result;
      if (pubk != nullptr) {
	// create an encoding cipher by key
	t_transient<Cipher> cifr = Crypto::mkcipher (*pubk, false);
	if (cifr.valid () == false) {
	  throw Exception ("tls-error", "cannot create key exchange cipher");
	}
	// stream into a buffer
	Buffer cbuf;
	cifr->stream (cbuf, mbuf);
	// add buffer size in the chunk
	long clen = cbuf.length ();
	if (clen  > 65536L) {
	  throw Exception ("tls-error", "inconsistent key exchange cipher");
	}
	t_byte hl = ((t_word) clen) >> 8;
	t_byte bl = ((t_word) clen) & 0x00FF;
	result.add (hl); result.add (bl);
	// fill the chunk
	result.add (cbuf);
      } else {
	if (d_mlen > 65536L) {
	  throw Exception ("tls-error", "inconsistent key exchange cipher");
	}
	// get the buffer length
	t_byte hl = ((t_word) d_mlen) >> 8;
	t_byte bl = ((t_word) d_mlen) & 0x00FF;
	result.add (hl); result.add (bl);
	result.add (mbuf);
      }
      unlock ();
      return result;
    } catch (...) {
      unlock ();
      throw;
    }
  }

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

  // create a new object in a generic way

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

    // check for 0 argument
    if (argc == 0) return new TlsCkeyxh;
    // check for 1 argument
    if (argc == 1) {
      Object*     obj = argv->get (0);
      TlsHblock* hblk = dynamic_cast<TlsHblock*> (obj);
      if ((hblk == nullptr) && (obj != nullptr)) {
	throw Exception ("type-error", 
			 "invalid object as tls client key exchange block",
			 Object::repr (obj));
      }
      return new TlsCkeyxh (hblk);
    }
    // too many arguments
    throw Exception ("argument-error", 
                     "too many argument with tls client key exchange");
  }
}
