#include <stdio.h>
#include <stdlib.h>
#include <strings.h>

extern const char *__progname;

#include "pp.h"
#include "bpp.h"
#include "rnd.h"
#include "str.h"
#include "algs.h"
#include "msgs.h"
#include "cmdline.h"
#include "pkt-util.h"
#include "transport.h"
#include "stdio-util.h"

typedef struct transport_state TRANSPORT_STATE;
typedef struct kex_state KEX_STATE;

struct transport_state {
  int first_kex;
  int stage;
#define TSS_REKEY        1
#define TSS_SENT_KEXINIT 2
#define TSS_RUNNING_KEX  3
#define TSS_NEWKEY_WAIT  4
#define TSS_DATA         5
  void *method_state;
  KEX_STATE *kex_state;
  PKTQ oq;
  int queuestate;
#define TQS_IDLE     1
#define TQS_QUEUEING 2
  } ;

struct kex_state {
  ENCALG *enc_c2s;
  ENCALG *enc_s2c;
  MACALG *mac_c2s;
  MACALG *mac_s2c;
  COMPALG *comp_c2s;
  COMPALG *comp_s2c;
  int ignorenext;
  } ;

static void *put_alg_list(void *opp, ALGLIST *algs, const char *(*namefn)(void *))
{
 char *buf;
 int n;
 FILE *f;

 static void foo(void *alg)
  { const char *name;
    name = (*namefn)(alg);
    fprintf(f,",%s",name);
  }

 f = fopen_alloc(&buf,&n);
 for (;algs;algs=algs->link) foo(algs->alg);
 fclose(f);
 if (n < 1) abort();
 opp = put_string(opp,buf+1,n-1);
 free(buf);
 return(opp);
}

static void print_offer(const char *tag, const void *name, int namelen, void *alg)
{
 static const char *prevtag = 0;

 if (tag != prevtag)
  { if (! prevtag) fprintf(stderr,"%s: offered algorithms (* = unrecognized):",__progname);
    fprintf(stderr,"\n");
    prevtag = tag;
    if (! tag) return;
    fprintf(stderr,"%12s:",tag);
  }
 if (tag) fprintf(stderr," %s%.*s",alg?"":"*",namelen,(const char *)name);
}

static void setalgs_list(ROSTR s, ALGLIST **listp, void *(*vfind)(ROSTR), const char *tag)
{
 static int foo(const void *name, int len)
  { ROSTR n;
    void *alg;
    ALGLIST *l;
    n.len = len;
    n.data = name;
    alg = (*vfind)(n);
    if (showoffer) print_offer(tag,name,len,alg);
    if (alg)
     { l = malloc(sizeof(ALGLIST));
       l->alg = alg;
       *listp = l;
       listp = &l->link;
     }
    return(0);
  }

 comma_list(s.data,s.len,foo);
 *listp = 0;
}

#define FOO(specific,short,str)\
	static void setalgs_##specific(ROSTR s, void *listpv)\
	{ setalgs_list(s,listpv,short##alg_vfind_rostr,str); }
FOO(kex,kex,"kex")
FOO(hk,hk,"hk")
FOO(enc_c2s,enc,"enc c->s")
FOO(enc_s2c,enc,"enc s->c")
FOO(mac_c2s,mac,"mac c->s")
FOO(mac_s2c,mac,"mac s->c")
FOO(comp_c2s,comp,"comp c->s")
FOO(comp_s2c,comp,"comp s->c")
#undef FOO

static void dump_alglists(ALGLIST *clist, ALGLIST *slist, const char *(*name)(void *))
{
 ALGLIST *l;

 fprintf(stderr,"Client list:");
 for (l=clist;l;l=l->link) fprintf(stderr," %s",(*name)(l->alg));
 fprintf(stderr,"\n");
 fprintf(stderr,"Server list:");
 for (l=slist;l;l=l->link) fprintf(stderr," %s",(*name)(l->alg));
 fprintf(stderr,"\n");
}

static void *derive_blob(BPP *b, char c, int len)
{
 unsigned char *t0;
 unsigned char *t;
 void *h;
 void (*last)(void);

 static void c_and_id(void)
  { (*b->kex_hash->process)(h,&c,1);
    (*b->kex_hash->process)(h,b->sessid,b->sessidlen);
  }

 static void blob_so_far(void)
  { (*b->kex_hash->process)(h,t0,t-t0);
  }

 last = &c_and_id;
 t0 = malloc(len);
 t = t0;
 while (1)
  { h = (*b->kex_hash->init)();
    (*b->kex_hash->process)(h,b->kex_k,b->kex_klen);
    (*b->kex_hash->process)(h,b->kex_h,b->kex_hlen);
    (*last)();
    if (len < b->kex_hash->hashlen)
     { (*b->kex_hash->done)(h,b->hashbuf);
       bcopy(b->hashbuf,t,len);
       return(t0);
     }
    else
     { (*b->kex_hash->done)(h,t);
       if (len == b->kex_hash->hashlen) return(t0);
       t += b->kex_hash->hashlen;
       last = &blob_so_far;
       len -= b->kex_hash->hashlen;
     }
  }
}

#if 0
static void dumpblob(const char *tag, const void *data, int len)
{
 int i;

 printf("%s",tag);
 for (i=0;i<len;i++) printf(" %02x",((const unsigned char *)data)[i]);
 printf("\n");
}
#endif

static void send_kexinit(LAYER *l, int firsttime)
{
 unsigned char *opp;
 unsigned char lclcookie[16];

 random_data(&lclcookie[0],16);
 opp = &l->b->opkt[0];
 *opp++ = SSH_MSG_KEXINIT;
 bcopy(&lclcookie[0],opp,16);
 opp += 16;
 opp = put_alg_list(opp,algs_kex,kexalg_name);
 opp = put_alg_list(opp,algs_hk,hkalg_name);
 opp = put_alg_list(opp,algs_enc_c2s,encalg_name);
 opp = put_alg_list(opp,algs_enc_s2c,encalg_name);
 opp = put_alg_list(opp,algs_mac_c2s,macalg_name);
 opp = put_alg_list(opp,algs_mac_s2c,macalg_name);
 opp = put_alg_list(opp,algs_comp_c2s,compalg_name);
 opp = put_alg_list(opp,algs_comp_s2c,compalg_name);
 opp = put_string(opp,lang_c2s?:"",-1);
 opp = put_string(opp,lang_s2c?:"",-1);
 *opp++ = 0;
 opp = put_uint32(opp,0);
 l->b->oplen = opp - &l->b->opkt[0];
 if (! firsttime) free(l->b->c_kexinit_payload);
 l->b->c_kexinit_paylen = l->b->oplen;
 l->b->c_kexinit_payload = malloc(l->b->c_kexinit_paylen);
 bcopy(&l->b->opkt[0],l->b->c_kexinit_payload,l->b->c_kexinit_paylen);
 below_opkt(l);
}

static KEX_STATE *recv_kexinit(LAYER *l, int firsttime)
{
 BPP *b;
 KEX_STATE *s;
 unsigned char remcookie[16];
 ALGLIST *remalgs_kex;
 ALGLIST *remalgs_hk;
 ALGLIST *remalgs_enc_c2s;
 ALGLIST *remalgs_enc_s2c;
 ALGLIST *remalgs_mac_c2s;
 ALGLIST *remalgs_mac_s2c;
 ALGLIST *remalgs_comp_c2s;
 ALGLIST *remalgs_comp_s2c;
 STR remlang_c2s;
 STR remlang_s2c;
 KEXALG *guessed_kex;
 int guessed_pkt;
 unsigned int extensions;

 b = l->b;
 s = malloc(sizeof(KEX_STATE));
 if (! firsttime) free(b->s_kexinit_payload);
 b->s_kexinit_paylen = b->iplen;
 b->s_kexinit_payload = malloc(b->s_kexinit_paylen);
 bcopy(&b->ipkt[0],b->s_kexinit_payload,b->s_kexinit_paylen);
 parse_packet(b,pp_fail,
	PP_IGNORE(1),
	PP_BLOB(16,&remcookie[0]),
	PP_STRING_CALL(setalgs_kex,&remalgs_kex),
	PP_STRING_CALL(setalgs_hk,&remalgs_hk),
	PP_STRING_CALL(setalgs_enc_c2s,&remalgs_enc_c2s),
	PP_STRING_CALL(setalgs_enc_s2c,&remalgs_enc_s2c),
	PP_STRING_CALL(setalgs_mac_c2s,&remalgs_mac_c2s),
	PP_STRING_CALL(setalgs_mac_s2c,&remalgs_mac_s2c),
	PP_STRING_CALL(setalgs_comp_c2s,&remalgs_comp_c2s),
	PP_STRING_CALL(setalgs_comp_s2c,&remalgs_comp_s2c),
	PP_STRING(&remlang_c2s),
	PP_STRING(&remlang_s2c),
	PP_BOOL(&guessed_pkt),
	PP_UINT32(&extensions),
	PP_ENDSHERE );
 if (showoffer) print_offer(0,0,0,0);
 /* Choose algorithms. */
 do <"kexfound">
  { ALGLIST *i1;
    ALGLIST *i2;
    ALGLIST *i3;
    /* Check for matching guesses */
    if (remalgs_kex)
     { guessed_kex = remalgs_kex->alg;
       if (algs_kex && (algs_kex->alg == guessed_kex))
	{ b->kex = guessed_kex;
	  break <"kexfound">;
	}
     }
    /* Kex algorithm is the first client kex algorithm for which */
    for <"nextkex"> (i1=algs_kex;i1;i1=i1->link)
     { /* + the server also supports the algorithm */
       do <"serversupport">
	{ for (i2=remalgs_kex;i2;i2=i2->link)
	   { if (i1->alg == i2->alg) break <"serversupport">;
	   }
	  continue <"nextkex">;
	} while (0);
       /* + if the algorithm requires an encryption-capable host key, */
       if (((KEXALG *)i1->alg)->need_enc)
	{ /* there is an encryption-capable algorithm on the server's
	     server_host_key_algorithms that is also supported by the
	     client */
	  do <"shkfound">
	   { for <"nextshk"> (i2=remalgs_hk;i2;i2=i2->link)
	      { do <"chkfound">
		 { for (i3=algs_hk;i3;i3=i3->link)
		    { if (i2->alg == i3->alg)
		       { if (! ((HKALG *)i3->alg)->can_enc) continue;
			 break <"chkfound">;
		       }
		    }
		   continue <"nextshk">;
		 } while (0);
		break <"shkfound">;
	      }
	     continue <"nextkex">;
	   } while (0);
	}
       /* + if the algorithm requires a signature-capable host key, */
       if (((KEXALG *)i1->alg)->need_sig)
	{ /* there is a signature-capable algorithm on the server's
	     server_host_key_algorithms that is also supported by the
	     client */
	  do <"shkfound">
	   { for <"nextshk"> (i2=remalgs_hk;i2;i2=i2->link)
	      { do <"chkfound">
		 { for (i3=algs_hk;i3;i3=i3->link)
		    { if (i2->alg == i3->alg)
		       { if (! ((HKALG *)i3->alg)->can_sig) continue;
			 break <"chkfound">;
		       }
		    }
		   continue <"nextshk">;
		 } while (0);
		break <"shkfound">;
	      }
	     continue <"nextkex">;
	   } while (0);
	}
       b->kex = i1->alg;
       break <"kexfound">;
     }
    send_disconnect(SSH_DISCONNECT_KEY_EXCHANGE_FAILED,"",0,"",0);
    fprintf(stderr,"%s: no usable key-exchange algorithm found\n",__progname);
    dump_alglists(algs_kex,remalgs_kex,kexalg_name);
    exit(1);
  } while (0);
printf("Using kex %s\n",b->kex->name);
 do <"hkfound">
  { ALGLIST *i1;
    ALGLIST *i2;
    for <"nexthk"> (i1=algs_hk;i1;i1=i1->link)
     { b->hk = i1->alg;
       if ((b->kex->need_enc && !b->hk->can_enc) ||
	   (b->kex->need_sig && !b->hk->can_sig) ) continue;
       do <"serversupport">
	{ for (i2=remalgs_hk;i2;i2=i2->link)
	   { if (i2->alg == i1->alg) break <"serversupport">;
	   }
	  continue <"nexthk">;
	} while (0);
       b->hk = i1->alg;
       break <"hkfound">;
     }
    send_disconnect(SSH_DISCONNECT_KEY_EXCHANGE_FAILED,"",0,"",0);
    fprintf(stderr,"%s: no usable server host key algorithm found\n",__progname);
    dump_alglists(algs_hk,remalgs_hk,hkalg_name);
    exit(1);
  } while (0);
printf("Using hk %s\n",b->hk->name);
#define FOO(clist,algvar,slist,namefn,msgtxt) do {\
 do <"found">								\
  { ALGLIST *i1;							\
    ALGLIST *i2;							\
    for <"nextc"> (i1=clist;i1;i1=i1->link)				\
     { do <"serversupport">						\
	{ for (i2=slist;i2;i2=i2->link)					\
	   { if (i1->alg == i2->alg) break <"serversupport">;		\
	   }								\
	  continue <"nextc">;						\
	} while (0);							\
       s->algvar = i1->alg;						\
       break <"found">;							\
     }									\
    send_disconnect(SSH_DISCONNECT_KEY_EXCHANGE_FAILED,"",0,"",0);	\
    fprintf(stderr,"%s: no usable server %s algorithm found\n",__progname,msgtxt);\
    dump_alglists(clist,slist,namefn);					\
    exit(1);								\
  } while (0);								\
printf("Using %s %s\n",msgtxt,s->algvar->name);				\
						} while (0)
 FOO(algs_enc_c2s,enc_c2s,remalgs_enc_c2s,encalg_name,"client->server encryption");
 FOO(algs_enc_s2c,enc_s2c,remalgs_enc_s2c,encalg_name,"server->client encryption");
 FOO(algs_mac_c2s,mac_c2s,remalgs_mac_c2s,macalg_name,"client->server MAC");
 FOO(algs_mac_s2c,mac_s2c,remalgs_mac_s2c,macalg_name,"server->client MAC");
 FOO(algs_comp_c2s,comp_c2s,remalgs_comp_c2s,compalg_name,"client->server compression");
 FOO(algs_comp_s2c,comp_s2c,remalgs_comp_s2c,compalg_name,"server->client compression");
#undef FOO
printf("Guessed kex packet sent by server: %s\n",guessed_pkt?"yes":"no");
 s->ignorenext = (guessed_pkt && (guessed_kex != b->kex));
 if (s->ignorenext) printf("...ignoring it\n");
 alglist_free(remalgs_kex);
 alglist_free(remalgs_hk);
 alglist_free(remalgs_enc_c2s);
 alglist_free(remalgs_enc_s2c);
 alglist_free(remalgs_mac_c2s);
 alglist_free(remalgs_mac_s2c);
 alglist_free(remalgs_comp_c2s);
 alglist_free(remalgs_comp_s2c);
 free_str(remlang_c2s);
 free_str(remlang_s2c);
 return(s);
}

static void kex_done(LAYER *l, int firsttime)
{
 if (firsttime)
  { l->b->sessidlen = l->b->kex_hlen;
    l->b->sessid = malloc(l->b->sessidlen);
    bcopy(l->b->kex_h,l->b->sessid,l->b->sessidlen);
  }
 if (! firsttime) free(l->b->hashbuf);
 l->b->hashbuf = malloc(l->b->kex_hash->hashlen);
 l->b->opkt[0] = SSH_MSG_NEWKEYS;
 l->b->oplen = 1;
 below_opkt(l);
}

static void install_new_keys_w(BPP *b, KEX_STATE *s)
{
 (*b->w_enc->done)(b->w_encstate);
 (*b->w_comp->done)(b->w_compstate);
 (*b->w_mac->done)(b->w_macstate);
  { void *iv;
    void *key;
    iv = derive_blob(b,'A',s->enc_c2s->ivsize);
    key = derive_blob(b,'C',s->enc_c2s->keysize);
    b->w_enc = s->enc_c2s;
    b->w_encstate = encalg_init(b->w_enc,key,iv,'w');
    free(iv);
    free(key);
  }
  { void *key;
    key = derive_blob(b,'E',s->mac_c2s->keylen);
    b->w_mac = s->mac_c2s;
    b->w_macstate = macalg_init(b->w_mac,key);
    free(key);
  }
}

static void install_new_keys_r(BPP *b, KEX_STATE *s)
{
 (*b->r_enc->done)(b->r_encstate);
 (*b->r_comp->done)(b->r_compstate);
 (*b->r_mac->done)(b->r_macstate);
  { void *iv;
    void *key;
    iv = derive_blob(b,'B',s->enc_s2c->ivsize);
    key = derive_blob(b,'D',s->enc_s2c->keysize);
    b->r_enc = s->enc_s2c;
    b->r_encstate = encalg_init(b->r_enc,key,iv,'r');
    free(iv);
    free(key);
  }
  { void *key;
    key = derive_blob(b,'F',s->mac_s2c->keylen);
    b->r_mac = s->mac_s2c;
    b->r_macstate = macalg_init(b->r_mac,key);
    free(key);
  }
}

static void *transport_init(LAYER *l __attribute__((__unused__)))
{
 TRANSPORT_STATE *s;

 s = malloc(sizeof(TRANSPORT_STATE));
 s->stage = TSS_REKEY;
 s->first_kex = 1;
 pktq_init(&s->oq);
 s->queuestate = TQS_IDLE;
 return(s);
}

static void transport_i(LAYER *l, void *arg)
{
 TRANSPORT_STATE *s;
 BPP *b;

 s = arg;
 b = l->b;
 switch (s->stage)
  { case TSS_DATA:
       if (! b->want_rekey)
	{ switch (b->ipkt[0])
	   { case SSH_MSG_KEXINIT:
		s->kex_state = recv_kexinit(l,0);
		send_kexinit(l,0);
		s->stage = TSS_RUNNING_KEX;
		s->method_state = (*b->kex->init)(l);
		break;
	     case SSH_MSG_NEWKEYS:
		badmsg("DATA","MSG_NEWKEYS");
		break;
	     default:
		above_ipkt(l);
		break;
	   }
	  break;
	}
       /* s->stage = TSS_REKEY; overwritten just below */
       /* fall through */
    case TSS_REKEY:
       send_kexinit(l,0);
       s->stage = TSS_SENT_KEXINIT;
       /* fall through */
    case TSS_SENT_KEXINIT:
       switch (b->ipkt[0])
	{ case SSH_MSG_KEXINIT:
	     s->kex_state = recv_kexinit(l,s->first_kex);
	     s->stage = TSS_RUNNING_KEX;
	     s->method_state = (*b->kex->init)(l);
	     break;
	  case SSH_MSG_NEWKEYS:
	     fprintf(stderr,"%s: misplaced NEWKEYS (at SENT_KEXINIT)\n",__progname);
	     exit(1);
	     break;
	  default:
	     if (s->first_kex)
	      { fprintf(stderr,"%s: KEXINIT packet isn't\n",__progname);
		exit(1);
	      }
	     above_ipkt(l);
	     break;
	}
       break;
    case TSS_RUNNING_KEX:
       if (s->kex_state->ignorenext)
	{ /* XXX: should we ignore the next packet, or the next packet
	     with a type field in the key-exchange range?  We ignore
	     the next packet, since that's how the draft is worded.
	     Since only key-exchange and debug/ignore/disconnect packets
	     are allowed during key exchange anyway, receiving anything
	     else is a protocol error (we arguably should check). */
	  s->kex_state->ignorenext = 0;
	  return;
	}
       switch (b->ipkt[0])
	{ case SSH_MSG_KEXINIT:
	     badmsg("RUNNING_KEX","MSG_KEXINIT");
	     break;
	  case SSH_MSG_NEWKEYS:
	     badmsg("RUNNING_KEX","MSG_NEWKEYS");
	     break;
	  default:
	     if ((b->ipkt[0] >= 30) && (b->ipkt[0] <= 49))
	      { if ((*b->kex->run)(l,s->method_state))
		 { kex_done(l,s->first_kex);
		   install_new_keys_w(b,s->kex_state);
		   s->stage = TSS_NEWKEY_WAIT;
		 }
	      }
	     else
	      { badmsg("RUNNING_KEX","type %d",b->ipkt[0]);
	      }
	     break;
	}
       break;
    case TSS_NEWKEY_WAIT:
       switch (b->ipkt[0])
	{ case SSH_MSG_KEXINIT:
	     badmsg("NEWKEY_WAIT","MSG_NEWKEYS");
	     break;
	  case SSH_MSG_NEWKEYS:
	      { static int rekey_pkts(int bb)
		 { return(1<<((bb<8)?16:(bb<16)?bb*2:30));
		 }
		install_new_keys_r(b,s->kex_state);
		free(s->kex_state);
		s->stage = TSS_DATA;
		s->first_kex = 0;
		b->ibytes_to_rekey = 1<<30;
		b->obytes_to_rekey = 1<<30;
		b->ipkts_to_rekey = rekey_pkts(b->r_enc->blksize);
		b->opkts_to_rekey = rekey_pkts(b->w_enc->blksize);
		b->want_rekey = 0;
		if (s->queuestate == TQS_QUEUEING)
		 { while (pktq_get(&s->oq,&b->opkt[0],&b->oplen,sizeof(b->opkt)))
		    { below_opkt(l);
		    }
		   s->queuestate = TQS_IDLE;
		 }
	      }
	     break;
	  default:
	     badmsg("NEWKEY_WAIT","type %d",b->ipkt[0]);
	     break;
	}
       break;
  }
}

static void transport_o(LAYER *l, void *arg)
{
 TRANSPORT_STATE *s;

 s = arg;
 if (s->queuestate != TQS_IDLE)
  { pktq_put(&s->oq,&l->b->opkt[0],l->b->oplen);
    return;
  }
 switch (s->stage)
  { case TSS_REKEY:
       s->queuestate = TQS_QUEUEING;
       pktq_put(&s->oq,&l->b->opkt[0],l->b->oplen);
       send_kexinit(l,0);
       s->stage = TSS_SENT_KEXINIT;
       break;
    default:
       s->queuestate = TQS_QUEUEING;
       pktq_put(&s->oq,&l->b->opkt[0],l->b->oplen);
       return;
    case TSS_DATA:
       below_opkt(l);
       break;
  }
}

LAYERDESC layer_transport = { &transport_init, &transport_i, &transport_o };