/*
 * handshake.cpp
 *
 * Code used to initialize cutlass connections.
 *
 * Copyright (c) 2004,2005 Jack Lloyd <lloyd@randombit.net>
 *                         Todd MacDermid <tmacd@synacklabs.net>
 */

#include <cutlass.h>
#include <cutlass/kex.hpp>
#include <cutlass/filters.hpp>

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

/*************************************************
 * Temporary storage for kex data                 *
 *************************************************/
class Kex_State
   {
   public:
      ClientHello* client_hello;
      ServerHello* server_hello;
      ClientKeyExchange* client_dh;
      ServerKeyExchange* server_dh;

      RSA_PublicKey* peer_key;
      RSA_PrivateKey* local_key;

      Network_Layer network;

      SecureVector<byte> last_message;

      bool is_client;

      Kex_State(bool is_c)
         {
         client_hello = 0;
         server_hello = 0;
         client_dh = 0;
         server_dh = 0;
         peer_key = 0;
         local_key = 0;
         is_client = is_c;
         }
      ~Kex_State()
         {
         delete client_hello;
         delete server_hello;
         delete client_dh;
         delete server_dh;
         delete peer_key;
         delete local_key;
         }
   };

/* Have to use this because X509::copy_key is broken in Botan 1.4.0 */
static X509_PublicKey* copy_x509(const X509_PublicKey& key)
   {
   Pipe bits;
   bits.start_msg();
   X509::encode(key, bits);
   bits.end_msg();
   DataSource_Memory source(bits.read_all());
   return X509::load_key(source);
   }

static RSA_PublicKey* extract_pub(cutlass_public_key key_handle)
   {
   if(key_handle.key == 0)
      return 0;

   try {
      X509_PublicKey* x509 = static_cast<X509_PublicKey*>(key_handle.key);

      X509_PublicKey* copy = copy_x509(*x509);

      if(dynamic_cast<RSA_PublicKey*>(copy))
         return dynamic_cast<RSA_PublicKey*>(copy);
      else
         {
         delete copy;
         return 0;
         }
   }
   catch(std::exception& e)
      {
      fprintf(stderr, "extract_pub: %s\n", e.what());
      return 0;
      }
   }

/*
 * cutlass_set_keys turns our kex material goop into the actual keys
 * used during a conversation.
 */

int
cutlass_set_keys(cutlass_t *cut_handle, conn_t *conn,
                 const DH_PrivateKey& dh_me,
                 const DH_PublicKey& dh_otherguy,
                 bool is_client)
   {
   const std::string CIPHER = "AES-256";
   const std::string HMAC_HASH = "SHA-1";
  
   const u32bit CIPHER_KEYLEN = max_keylength_of(CIPHER);
   const u32bit MAC_KEYLEN = output_length_of(HMAC_HASH);
  
   // FIXME: should match handshake_hash() in kex.cpp
   const std::string HASH_FN = "SHA-256";
  
   try {
      Kex_State* state = static_cast<Kex_State*>(conn->kex_info);
    
      Pipe pipe(new Hash_Filter(HASH_FN));
      pipe.start_msg();
      pipe.write(dh_me.derive_key(dh_otherguy));
      pipe.write(state->client_hello->random());
      pipe.write(state->server_hello->random());
      pipe.end_msg();
    
      SecureVector<byte> master_secret = pipe.read_all();
    
      Keyed_Filter* cipher_c2s = get_cipher(CIPHER + "/CTR-BE", ENCRYPTION);
      Keyed_Filter* cipher_s2c = get_cipher(CIPHER + "/CTR-BE", ENCRYPTION);
    
      KDF2 kdf(HASH_FN);

      // Derive the session keys
      SymmetricKey key_c2s_cipher, key_s2c_cipher, key_c2s_mac, key_s2c_mac;
      key_c2s_cipher = kdf.derive_key(CIPHER_KEYLEN, master_secret,
                                      "client cipher");
      key_s2c_cipher = kdf.derive_key(CIPHER_KEYLEN, master_secret,
                                      "server cipher");
      key_c2s_mac = kdf.derive_key(MAC_KEYLEN, master_secret, "client MAC");
      key_s2c_mac = kdf.derive_key(MAC_KEYLEN, master_secret, "server MAC");

      cipher_c2s->set_key(key_c2s_cipher);
      cipher_s2c->set_key(key_s2c_cipher);
    
      Keyed_Filter* mac_c2s =
         new MAC_Filter("HMAC(" + HMAC_HASH + ")", key_c2s_mac);

      Keyed_Filter* mac_s2c =
         new MAC_Filter("HMAC(" + HMAC_HASH + ")", key_s2c_mac);

      cutlass_keyset keyset = empty_keyset();

      if(is_client)
         {
         keyset.opaque_cipher_filter_send = cipher_c2s;
         keyset.opaque_cipher_filter_recv = cipher_s2c;
        
         keyset.opaque_cipher_send = new Pipe(cipher_c2s);
         keyset.opaque_cipher_recv = new Pipe(cipher_s2c);
        
         keyset.opaque_mac_send = new Pipe(mac_c2s);
         keyset.opaque_mac_recv = new Pipe(mac_s2c);

         keyset.send_proc = new Pipe(
            new Cutlass_Encryptor(CIPHER, HMAC_HASH,
                                  key_c2s_cipher, key_c2s_mac, conn->mtu)
            );

         keyset.recv_proc = new Pipe(
            new Cutlass_Decryptor(CIPHER, HMAC_HASH,
                                  key_s2c_cipher, key_s2c_mac)
            );

         conn->peer_key.key = copy_x509(state->server_hello->server_key());
         }
      else
         {
         keyset.opaque_cipher_filter_send = cipher_s2c;
         keyset.opaque_cipher_filter_recv = cipher_c2s;
        
         keyset.opaque_cipher_send = new Pipe(cipher_s2c);
         keyset.opaque_cipher_recv = new Pipe(cipher_c2s);
        
         keyset.opaque_mac_send = new Pipe(mac_s2c);
         keyset.opaque_mac_recv = new Pipe(mac_c2s);

         keyset.send_proc = new Pipe(
            new Cutlass_Encryptor(CIPHER, HMAC_HASH,
                                  key_s2c_cipher, key_s2c_mac, conn->mtu)
            );
         keyset.recv_proc = new Pipe(
            new Cutlass_Decryptor(CIPHER, HMAC_HASH,
                                  key_c2s_cipher, key_c2s_mac)
            );

         conn->peer_key.key = copy_x509(state->client_hello->client_key());
         }

      fingerprint_key_bin(conn->peer_key, conn->fingerprint);

      free_keyset(conn->keyset);
      conn->keyset = keyset;
   }
   catch(std::exception& e)
      {
      cutlass_sysmsg(cut_handle, CUT_ERROR, "%s", e.what());
      return 1;
      }
   return 0;
   }

static RSA_PrivateKey* extract_priv(cutlass_private_key key_handle)
   {
   if(key_handle.key == 0)
      {
      cutlass_sysmsg(NULL, CUT_ERROR, "extract_priv: Arg was NULL");
      return 0;
      }

   try {
      PKCS8_PrivateKey* pkcs8 = static_cast<PKCS8_PrivateKey*>(key_handle.key);

      /* PKCS8::copy_key is broken in 1.4.0 */
      Pipe bits;
      bits.start_msg();
      PKCS8::encode(*pkcs8, bits);
      bits.end_msg();
      DataSource_Memory source(bits.read_all());
      PKCS8_PrivateKey* copy = PKCS8::load_key(source);

      if(dynamic_cast<RSA_PrivateKey*>(copy))
         return dynamic_cast<RSA_PrivateKey*>(copy);
      else
         {
         delete copy;
         return 0;
         }
   }
   catch(std::exception& e)
      {
      fprintf(stderr, "extract_priv: %s\n", e.what());
      return 0;
      }
   }

static int bail_out(cutlass_t* cut_handle, conn_t* conn,
                    const char* msg)
   {
   Kex_State* state = static_cast<Kex_State*>(conn->kex_info);
   delete state;
   conn->kex_info = 0;

   cutlass_sysmsg(cut_handle, CUT_ERROR, msg);
   return -1;
   }

/* Create temp state, and send the client hello */
static int cutlass_do_init_handshake(cutlass_t *cut_handle, conn_t *conn,
                                     RSA_PrivateKey* local_key,
                                     RSA_PublicKey* peer_key = 0)
   {
   // FIXME: this leaks if state transition goes south
   conn->kex_info = new Kex_State(true);
   Kex_State* state = static_cast<Kex_State*>(conn->kex_info);

   if(!local_key)
      {
      cutlass_sysmsg(cut_handle, CUT_ERROR,
                     "cutlass_do_init_handshake: No local key!");
      return -1;
      }

   if(peer_key)
      state->client_hello = new ClientHello(*local_key, *peer_key);
   else
      state->client_hello = new ClientHello(*local_key);

   state->peer_key = peer_key;
   state->local_key = local_key;

   conn->conn_state = CUTSTATE_C_SENT_HELLO;
   state->network.write(*(state->client_hello), cut_handle, conn);

   return 0;
   }

int cutlass_handshake_init(cutlass_t *cut_handle, conn_t *conn,
                           cutlass_public_key peer_key_handle,
                           cutlass_private_key local_key_handle)
   {
   RSA_PrivateKey* local_rsa = extract_priv(local_key_handle);
   RSA_PublicKey* peer_rsa = extract_pub(peer_key_handle);

   if(!local_rsa || !peer_rsa)
      {
      cutlass_sysmsg(cut_handle, CUT_ERROR,
                     "Someone foolishly handed cutlass_handshake_init "
                     "a non-RSA key.");
      return -1;
      }
   
   return cutlass_do_init_handshake(cut_handle, conn, local_rsa, peer_rsa);
   }

int cutlass_handshake_init_no_server(cutlass_t *cut_handle, conn_t *conn,
                                     cutlass_private_key local_key_handle)
   {
   RSA_PrivateKey* local_rsa = extract_priv(local_key_handle);

   if(!local_rsa)
      {
      cutlass_sysmsg(cut_handle, CUT_ERROR,
                     "Someone foolishly handed cutlass_handshake_init "
                     "a non-RSA key.");
      return -1;
      }

   return cutlass_do_init_handshake(cut_handle, conn, local_rsa, 0);
   }

/*
 * cutlass_process_handshake processes the received handshake packet,
 * updates the kex state, and if appropriate, emits a new packet for
 * the next stage of handshaking (or activates the connection).
 *
 * Returns 0 on normal success, 1 on success that activates the
 * connection, or -1 on error.
 */

int
cutlass_process_handshake(cutlass_t *cut_handle, conn_t *conn,
                          const uint8_t* in,
                          struct cutlass_packet_hdr *head,
                          struct timespec *now)
   {
   if(conn->kex_info == 0)
      conn->kex_info = new Kex_State(false);

   Kex_State* state = static_cast<Kex_State*>(conn->kex_info);
  
   if(NULL == in && NULL == head)
      {
      state->network.resend_last(cut_handle, conn);
      return 0;
      }
  
   if(!state->local_key)
      state->local_key = extract_priv(cut_handle->local_key);

   //  printf("Got some stuff in cutlass_handshake_update: %d/%d\n",
   //         head->cut_type, head->length);
   //  printf("&(state->network) = %p\n", &(state->network));

   Handshake_Message* new_message = 
      state->network.read(head->cut_type, in, head->length);
  
   if(!new_message) /* We got just a fragment, no full message yet */
      return 0;

   switch(conn->conn_state)
      {
      case CUTSTATE_UNINITIALZED:
         if(new_message->type() == CLIENT_HELLO)
            {
            state->client_hello = dynamic_cast<ClientHello*>(new_message);
            state->peer_key =
               new RSA_PublicKey(state->client_hello->client_key());
      
            state->server_hello =
               new ServerHello(*(state->local_key),
                               state->client_hello->random());
      
            state->network.write(*(state->server_hello), cut_handle, conn);
            conn->conn_state = CUTSTATE_S_SENT_HELLO;
            return(0);
            }
         else
            {
            cutlass_sysmsg(cut_handle, CUT_INFO,
                           "cutlass_process_handshake: "
                           "Non-C_HELLO received in uninitialized state: %d\n",
                           new_message->type());
            return(-1);
            }
         break;
      case CUTSTATE_C_SENT_HELLO:
         if(new_message->type() == SERVER_HELLO)
            {
            state->server_hello = dynamic_cast<ServerHello*>(new_message);
            /* check server's key against what we expected */
            if(state->peer_key)
               {
               RSA_PublicKey key_got = state->server_hello->server_key();
        
               if(state->peer_key->get_e() != key_got.get_e() ||
                  state->peer_key->get_n() != key_got.get_n())
                  {
                  // Keys didn't match. Your trust model asplode
                  // TODO: call an action handler or some shit like that
                  }
               }
            else
               state->peer_key =
                  new RSA_PublicKey(state->server_hello->server_key());

            if(state->client_hello->random() != state->server_hello->client_random())
               return bail_out(cut_handle, conn, 
                               "Mismatch of nonces in hello messages");
       
            state->client_dh = new ClientKeyExchange(*(state->local_key), DH_1024,
                                                     state->client_hello->random(),
                                                     state->server_hello->random());
      
            state->network.write(*(state->client_dh), cut_handle, conn);
            conn->conn_state = CUTSTATE_C_SENT_DH;
            }
         else
            {
            cutlass_sysmsg(cut_handle, CUT_INFO,
                           "cutlass_process_handshake: "
                           "Non-S_HELLO received in C_SENT state\n");
            return(-1);
            }
         break;
      case CUTSTATE_S_SENT_HELLO:
         if(new_message->type() == CLIENT_KEX)
            {
            state->client_dh = dynamic_cast<ClientKeyExchange*>(new_message);

            if(!state->client_dh->check_sig(*(state->peer_key),
                                            state->client_hello->random(),
                                            state->server_hello->random()))
               return bail_out(cut_handle, conn, 
                               "Client key exchange signature failed");

            state->server_dh = 
               new ServerKeyExchange(*(state->local_key),
                                     state->client_dh->client_dh_group(),
                                     state->client_dh->client_dh_val(),
                                     state->client_hello->random(),
                                     state->server_hello->random());

            state->network.write(*(state->server_dh), cut_handle, conn);
            conn->conn_state = CUTSTATE_S_SENT_DH;

            /* We're the server */
            cutlass_set_keys(cut_handle, conn,
                             state->server_dh->server_dh_priv(
                                state->client_dh->client_dh_group()
                                ),
                             state->client_dh->client_dh(),
                             false); 
            }
         else
            {
            cutlass_sysmsg(cut_handle, CUT_INFO,
                           "cutlass_process_handshake: "
                           "Non-C_KEX received in S_SENT state\n");
            return(-1);
            }
         break;
      case CUTSTATE_C_SENT_DH:
         if(new_message->type() == SERVER_KEX)
            {
            state->server_dh = dynamic_cast<ServerKeyExchange*>(new_message);

            if(!state->server_dh->check_sig(*(state->peer_key),
                                            state->client_hello->random(),
                                            state->server_hello->random(),
                                            state->client_dh->client_dh_group(),
                                            state->client_dh->client_dh_val()))
               return bail_out(cut_handle, conn, 
                               "Server key exchange signature failed");

            /* We're the client */
            cutlass_set_keys(cut_handle, conn,
                             state->client_dh->client_dh_priv(),
                             state->server_dh->server_dh(
                                state->client_dh->client_dh_group()
                                ),
                             true); 
            cutlass_conn_activate(cut_handle, conn, now);

            cutlass_sysmsg(cut_handle, CUT_DEBUG, 
                           "cutlass_process_handshake: "
                           "Kex complete, sending info\n");
            return(1);
            }
         else
            {
            cutlass_sysmsg(cut_handle, CUT_INFO,
                           "cutlass_process_handshake: "
                           "Non-S_KEX received in C_SENT_DH state\n");
            return(-1);
            }
         break;
      case CUTSTATE_S_SENT_DH:
         cutlass_sysmsg(cut_handle, CUT_ERROR,
                        "cutlass_process_handshake: "
                        "in CUTSTATE_S_SENT_DH state. "
                        "Should only be retransmitting\n");
         return(-1);
         break;
      case CUTSTATE_ACTIVE:
         cutlass_sysmsg(cut_handle, CUT_INFO,
                        "cutlass_process_handshake: "
                        "Got here while in ACTIVE state\n");
         return(-1);
         break;
      default:
         cutlass_sysmsg(cut_handle, CUT_INFO,
                        "cutlass_process_handshake: "
                        "I'm in a non-KEXy state, here...\n");
         return(-1);
         break;
      }
  
   return(-1); /* should never reach here, normal return at end of try block */
   }

/* 
 * cutlass_conn_activate is what gets called when key exchange is complete,
 * and we use it to derive and set up any crypto things out of the
 * state. 
 */
int
cutlass_conn_activate(cutlass_t *cut_handle, conn_t *conn, 
                      struct timespec *now)
{
  struct cut_action_obj action_obj;
  uint8_t tmp_fingerprint[CUT_ID_LEN];
  conn_t *existing_conn;
  
  if((NULL == cut_handle) || (NULL == conn)) {
    cutlass_sysmsg(cut_handle, CUT_ERROR,
                   "cutlass_conn_activate: Was passed a NULL pointer\n");
    return(-1);
  }
  
  Kex_State* state = static_cast<Kex_State*>(conn->kex_info);
  delete state;
  conn->kex_info = 0;

  cutlass_sysmsg(cut_handle, CUT_DEBUG,
                 "cutlass_conn_activate: New connection complete\n");

  conn->conn_state = CUTSTATE_ACTIVE;
  cutlass_info_send(cut_handle, conn, now);

  if(cut_handle->action_handlers[CUT_USER_CONN] != NULL) {
    memset(&action_obj, 0, sizeof(struct cut_action_obj));
    action_obj.action_type = CUT_USER_CONN;
    memcpy(&(action_obj.remote_addr), 
           &(conn->remote_addr), 
           sizeof(struct sockaddr_in));
    action_obj.remote_key = copy_public_key(conn->peer_key);
    cutlass_action(cut_handle, conn, &action_obj);
  }

  memcpy(tmp_fingerprint, conn->fingerprint, CUT_ID_LEN);
  pthread_mutex_unlock(&(conn->conn_mutex));
  
  existing_conn = hashtable_fingerprint_find(cut_handle, tmp_fingerprint);
  if(existing_conn != NULL) {
    cutlass_sysmsg(cut_handle, CUT_DEBUG,
                   "cutlass_conn_activate: Existing connection - resetting\n");
    cutlass_conn_del(cut_handle, existing_conn);
  }
  hashtable_fingerprint_add(cut_handle, conn, tmp_fingerprint);

  return(0);
}