/*
 * kex.cpp
 *
 * The key exchange stuff, called from handshake.c
 *
 * Copyright (c) 2004 Jack Lloyd <lloyd@randombit.net>
 *
 */

#include <cutlass/kex.hpp>

/* getting pulled in from somewhere, really fucks up C++ */
#undef min
#undef max

#include <botan/botan.h>
#include <botan/look_pk.h>
#include <botan/kdf.h>
using namespace Botan;

#include <iostream>
#include <iomanip>

#include <memory>

namespace {

SecureVector<byte> encode_key(const X509_PublicKey& key)
   {
   Pipe pipe;
   pipe.start_msg();
   X509::encode(key, pipe, RAW_BER);
   pipe.end_msg();
   return pipe.read_all();
   }

std::string handshake_hash()
   {
   return "SHA-256";
   }

u32bit handshake_hash_len()
   {
   return output_length_of(handshake_hash());
   }

u32bit nonce_len()
   {
   return 32;
   }

byte kex_version_code() { return 0; }

}

/*************************************************
* Serialize ClientHello                          *
*************************************************/
SecureVector<byte> ClientHello::serialize() const
   {
   SecureVector<byte> buffer;
   buffer.append(kex_version_code());
   buffer.append(c_random);
   buffer.append(hash);
   buffer.append(key);
   return buffer;
   }

/*************************************************
* Return the client's key                        *
*************************************************/
RSA_PublicKey ClientHello::client_key() const
   {
   std::auto_ptr<X509_PublicKey> key_ptr(X509::load_key(key));
   RSA_PublicKey* rsa_ptr = dynamic_cast<RSA_PublicKey*>(key_ptr.get());
   if(rsa_ptr)
      return RSA_PublicKey(*rsa_ptr);
   throw Decoding_Error("Key in ClientHello is not an RSA key");
   }

/*************************************************
* Check if the hash matches our key              *
*************************************************/
bool ClientHello::hash_matches(const RSA_PublicKey& server_key) const
   {
   Pipe pipe(new Hash_Filter(handshake_hash()));
   pipe.start_msg();
   pipe.write(c_random);
   pipe.write(key);
   X509::encode(server_key, pipe, RAW_BER);
   pipe.end_msg();

   if(hash == pipe.read_all())
      return true;
   return false;
   }

/*************************************************
* Create a new ClientHello                       *
*************************************************/
ClientHello::ClientHello(const RSA_PrivateKey& my_key,
                         const RSA_PublicKey& server_key)
   {
   c_random.create(nonce_len());
   Global_RNG::randomize(c_random, c_random.size(), Nonce);

   key = encode_key(my_key);

   Pipe pipe(new Hash_Filter(handshake_hash()));
   pipe.start_msg();
   pipe.write(c_random);
   pipe.write(key);
   X509::encode(server_key, pipe, RAW_BER);
   pipe.end_msg();

   hash = pipe.read_all();
   }

/*************************************************
* Create a new ClientHello (no server key known) *
*************************************************/
ClientHello::ClientHello(const RSA_PrivateKey& my_key)
   {
   c_random.create(nonce_len());
   Global_RNG::randomize(c_random, c_random.size(), Nonce);

   key = encode_key(my_key);

   /* A random nonce in replacement of the server key */
   SecureVector<byte> server_key(nonce_len());
   Global_RNG::randomize(server_key, server_key.size(), Nonce);

   Pipe pipe(new Hash_Filter(handshake_hash()));
   pipe.start_msg();
   pipe.write(c_random);
   pipe.write(key);
   pipe.write(server_key);
   pipe.end_msg();

   hash = pipe.read_all();
   }

/*************************************************
* Deserialize a ClientHello                      *
*************************************************/
ClientHello::ClientHello(const MemoryRegion<byte>& buf)
   {
   const u32bit HASH_LEN = handshake_hash_len();
   const u32bit NONCE_LEN = nonce_len();

   if(buf.size() <= 1 + NONCE_LEN + HASH_LEN)
      throw Decoding_Error("ClientHello deserialization failed");

   version = buf[0];
   if(version != kex_version_code())
      throw Decoding_Error("Bad version code in ClientHello");

   c_random.set(buf + 1, NONCE_LEN);
   hash.set(buf + 1 + NONCE_LEN, HASH_LEN);
   key.set(buf + 1 + NONCE_LEN + HASH_LEN, buf.size() - NONCE_LEN - HASH_LEN);
   }

/*************************************************
* Serialize a ServerHello                        *
*************************************************/
SecureVector<byte> ServerHello::serialize() const
   {
   SecureVector<byte> buffer;
   buffer.append(s_random);
   buffer.append(c_random);
   buffer.append(key);
   return buffer;
   }

/*************************************************
* Return the server's public key                 *
*************************************************/
RSA_PublicKey ServerHello::server_key() const
   {
   std::auto_ptr<X509_PublicKey> key_ptr(X509::load_key(key));
   RSA_PublicKey* rsa_ptr = dynamic_cast<RSA_PublicKey*>(key_ptr.get());
   if(rsa_ptr)
      return RSA_PublicKey(*rsa_ptr);
   throw Decoding_Error("Key in ServerHello is not an RSA key");
   }

/*************************************************
* Create a ServerHello                           *
*************************************************/
ServerHello::ServerHello(const RSA_PublicKey& s_key,
                         const MemoryRegion<byte>& c_rand)
   {
   key = encode_key(s_key);
   c_random = c_rand;
   s_random.create(nonce_len());
   cipher_algo = CIPHER_AES_256;
   hmac_hash_algo = HASH_SHA160;
   sig_hash_algo = HASH_SHA256;
   Global_RNG::randomize(s_random, s_random.size(), Nonce);
   }

/*************************************************
* Deserialize a ServerHello                      *
*************************************************/
ServerHello::ServerHello(const MemoryRegion<byte>& buf)
   {
   Pipe pipe;
   pipe.process_msg(buf);

   s_random.create(nonce_len());
   pipe.read(s_random, s_random.size());

   c_random.create(nonce_len());
   pipe.read(c_random, c_random.size());

   X509_PublicKey* rsa_key = X509::load_key(pipe);
   key = encode_key(*rsa_key);
   delete rsa_key;

   if(pipe.remaining())
      throw Decoding_Error("Extra bytes after ServerHello");

   cipher_algo = CIPHER_AES_128;
   hmac_hash_algo = HASH_SHA160;
   sig_hash_algo = HASH_SHA256;
   }

/*************************************************
* Serialize a ClientKeyExchange                  *
*************************************************/
SecureVector<byte> ClientKeyExchange::serialize() const
   {
   SecureVector<byte> buf;

   u16bit keylen = dh_c.size();
   u16bit siglen = sig.size();

   buf.append(dh_group);
   buf.append(get_byte(0, keylen));
   buf.append(get_byte(1, keylen));
   buf.append(dh_c);
   buf.append(get_byte(0, siglen));
   buf.append(get_byte(1, siglen));
   buf.append(sig);

   return buf;
   }

/*************************************************
* Check the signature on the stuff               *
*************************************************/
bool ClientKeyExchange::check_sig(const RSA_PublicKey& c_key,
                                  const MemoryRegion<byte>& c_random,
                                  const MemoryRegion<byte>& s_random) const
   {
   std::auto_ptr<PK_Verifier> verifier(
      get_pk_verifier(c_key, "EMSA4(" + handshake_hash() + ")")
      );

   verifier->update(s_random);
   verifier->update(c_random);
   verifier->update(&dh_group, 1); // FIXME: bad Botan interface
   verifier->update(dh_c);

   return verifier->check_signature(sig);
   }

/*************************************************
* Return the client DH key                       *
*************************************************/
DH_PublicKey ClientKeyExchange::client_dh() const
   {
   return DH_PublicKey(get_dh_group(dh_group), BigInt::decode(dh_c));
   }

/*************************************************
* Return the client private DH key               *
*************************************************/
DH_PrivateKey ClientKeyExchange::client_dh_priv() const
   {
   return DH_PrivateKey(get_dh_group(dh_group), BigInt::decode(dh_c_priv),
                        BigInt::decode(dh_c));
   }

/*************************************************
* Create a new ClientKeyExchange                 *
*************************************************/
ClientKeyExchange::ClientKeyExchange(const RSA_PrivateKey& rsa_key,
                                     DH_GROUP group_id,
                                     const MemoryRegion<byte>& c_random,
                                     const MemoryRegion<byte>& s_random)
   {
   dh_group = group_id;

   DL_Group group;
   if(dh_group == DH_1024)
      group = get_dl_group("IETF-1024");
   else if(dh_group == DH_1536)
      group = get_dl_group("IETF-1536");
   else
      throw Invalid_Argument("Unknown DH group id chosen in constructor");

   DH_PrivateKey emph_dh(group);
   dh_c = emph_dh.public_value();
   dh_c_priv = BigInt::encode(emph_dh.get_x());

   std::auto_ptr<PK_Signer> signer(
      get_pk_signer(rsa_key, "EMSA4(" + handshake_hash() + ")")
      );

   signer->update(s_random);
   signer->update(c_random);
   signer->update(&dh_group, 1); // FIXME: bad Botan interface
   signer->update(dh_c);
   sig = signer->signature();
   }

/*************************************************
* Deserialize a ClientKeyExchange                *
*************************************************/
ClientKeyExchange::ClientKeyExchange(const MemoryRegion<byte>& buf)
   {
   if(buf.size() < 3)
      throw Decoding_Error("ClientKeyExchange is way too short");

   dh_group = buf[0];
   u32bit keylen = make_u16bit(buf[1], buf[2]);

   if(buf.size() < 3 + keylen)
      throw Decoding_Error("ClientKeyExchange has bad DH length field");

   dh_c.set(buf + 3, keylen);

   u32bit siglen = make_u16bit(buf[3+keylen], buf[3+keylen+1]);
   if(buf.size() != siglen + keylen + 5)
      throw Decoding_Error("ClientKeyExchange has bad sig length");
   sig.set(buf + 5 + keylen, siglen);
   }

/*************************************************
* Serialize a ServerKeyExchange                  *
*************************************************/
SecureVector<byte> ServerKeyExchange::serialize() const
   {
   SecureVector<byte> buf;

   u16bit keylen = dh_s.size();
   u16bit siglen = sig.size();

   buf.append(get_byte(0, keylen));
   buf.append(get_byte(1, keylen));
   buf.append(dh_s);

   buf.append(get_byte(0, siglen));
   buf.append(get_byte(1, siglen));
   buf.append(sig);

   return buf;
   }

/*************************************************
* Return the DH key in the ServerKeyExchange     *
*************************************************/
DH_PublicKey ServerKeyExchange::server_dh(byte group_id) const
   {
   return DH_PublicKey(get_dh_group(group_id), BigInt::decode(dh_s));
   }

/*************************************************
* Return the DH key in the ServerKeyExchange     *
*************************************************/
DH_PrivateKey ServerKeyExchange::server_dh_priv(byte group_id) const
   {
   return DH_PrivateKey(get_dh_group(group_id), BigInt::decode(dh_s_priv),
                        BigInt::decode(dh_s));
   }

/*************************************************
* Check the signature on the stuff               *
*************************************************/
bool ServerKeyExchange::check_sig(const RSA_PublicKey& s_key,
                                  const MemoryRegion<byte>& c_random,
                                  const MemoryRegion<byte>& s_random,
                                  byte group_id,
                                  const MemoryRegion<byte>& dh_c) const
   {
   std::auto_ptr<PK_Verifier> verifier(
      get_pk_verifier(s_key, "EMSA4(" + handshake_hash() + ")")
      );

   verifier->update(c_random);
   verifier->update(s_random);
   verifier->update(&group_id, 1); // add this to Botan
   verifier->update(dh_c);
   verifier->update(dh_s);
   return verifier->check_signature(sig);
   }

/*************************************************
* Create a new ServerKeyExchange                 *
*************************************************/
ServerKeyExchange::ServerKeyExchange(const RSA_PrivateKey& server_rsa,
                                     byte dh_group,
                                     const MemoryRegion<byte>& dh_c,
                                     const MemoryRegion<byte>& c_random,
                                     const MemoryRegion<byte>& s_random)
   {
   DH_PrivateKey emph_dh(get_dh_group(dh_group));

   dh_s = emph_dh.public_value();
   dh_s_priv = BigInt::encode(emph_dh.get_x());

   std::auto_ptr<PK_Signer> signer(
      get_pk_signer(server_rsa, "EMSA4(" + handshake_hash() + ")")
      );

   signer->update(c_random);
   signer->update(s_random);
   signer->update(&dh_group, 1); // FIXME: bad Botan interface
   signer->update(dh_c);
   signer->update(dh_s);
   sig = signer->signature();
   }

/*************************************************
* Deserialize a ServerKeyExchange                *
*************************************************/
ServerKeyExchange::ServerKeyExchange(const MemoryRegion<byte>& buf)
   {
   if(buf.size() < 2)
      throw Decoding_Error("ServerKeyExchange is too small");
   u32bit dh_len = make_u16bit(buf[0], buf[1]);

   if(buf.size() < 2 + dh_len)
      throw Decoding_Error("ServerKeyExchange is too small");
   dh_s.set(buf + 2, dh_len);

   if(buf.size() < 4 + dh_len)
      throw Decoding_Error("ServerKeyExchange is too small");
   u32bit sig_len = make_u16bit(buf[dh_len+2], buf[dh_len+3]);

   if(buf.size() != 4 + dh_len + sig_len)
      throw Decoding_Error("ServerKeyExchange is too small");
   sig.set(buf + dh_len + 4, sig_len);
   }

/*************************************************
* Get the DL group matching this id              *
*************************************************/
DL_Group get_dh_group(byte group_id)
   {
   DL_Group group;

   if(group_id == DH_1024)
      group = get_dl_group("IETF-1024");
   else if(group_id == DH_1536)
      group = get_dl_group("IETF-1536");
   else
      throw Invalid_Argument("Unknown DH group id");

   return group;
   }