/* This file is in the public domain. */

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

extern const char *__progname;

#include "pp.h"
#include "bpp.h"
#include "str.h"
#include "rnd.h"
#include "algs.h"
#include "errf.h"
#include "msgs.h"
#include "hkdb.h"
#include "panic.h"
#include "pollloop.h"
#include "pkt-util.h"
#include "algs-list.h"

static const char *p1_hex =
	"FFFFFFFFFFFFFFFFC90FDAA22168C234C4C6628B80DC1CD1"
	"29024E088A67CC74020BBEA63B139B22514A08798E3404DD"
	"EF9519B3CD3A431B302B0A6DF25F14374FE1356D6D51C245"
	"E485B576625E7EC6F44C42E9A637ED6B0BFF5CB6F406B7ED"
	"EE386BFB5A899FA5AE9F24117C4B1FE649286651ECE65381"
	"FFFFFFFFFFFFFFFF";

static const char *p14_hex =
	"FFFFFFFFFFFFFFFFC90FDAA22168C234C4C6628B80DC1CD1"
	"29024E088A67CC74020BBEA63B139B22514A08798E3404DD"
	"EF9519B3CD3A431B302B0A6DF25F14374FE1356D6D51C245"
	"E485B576625E7EC6F44C42E9A637ED6B0BFF5CB6F406B7ED"
	"EE386BFB5A899FA5AE9F24117C4B1FE649286651ECE45B3D"
	"C2007CB8A163BF0598DA48361C55D39A69163FA8FD24CF5F"
	"83655D23DCA3AD961C62F356208552BB9ED529077096966D"
	"670C354E4ABC9804F1746C08CA18217C32905E462E36CE3B"
	"E39E772C180E86039B2783A2EC07A28FB5C55DF06F4C52C9"
	"DE2BCBF6955817183995497CEA956AE515D2261898FA0510"
	"15728E5A8AACAA68FFFFFFFFFFFFFFFF";

typedef struct dhgX_sha1_state DHGX_SHA1_STATE;

struct dhgX_sha1_state {
  int server;
  const char *phex;
  MP_INT p;
  union {
    struct {
      MP_INT x;
      MP_INT e;
      } c;
    struct {
      MP_INT y;
      MP_INT f;
      } s;
    } ;
  } ;

static void kexdh1_init(MP_INT *p, MP_INT *exp, MP_INT *pow, const char *phex)
{
 MP_INT q;
 int ptlen;
 int pblen;
 unsigned char *xbin;

 mpz_init_set_str(p,phex,16);
 mpz_init(&q);
 mpz_init(exp);
 mpz_init_set_ui(pow,2);
 mpz_tdiv_q_ui(&q,p,2);
 ptlen = strlen(phex);
 pblen = (ptlen + 1) / 2;
 xbin = malloc(pblen);
 while (1)
  { random_data(xbin,pblen);
    xbin[0] &= (ptlen & 1) ? 0x07 : 0x7f;
    data_to_mpint(exp,&xbin[0],pblen);
    if ((mpz_cmp(exp,&q) >= 0) || (mpz_cmp_ui(exp,1) <= 0)) continue;
    break;
  }
 mpz_powm(pow,pow,exp,p);
 mpz_clear(&q);
 free(xbin);
}

static void kexdh1_compute(MP_INT *K, MP_INT *e, MP_INT *f, BPP *b)
{
 void *hash;
 unsigned char *opp;

 b->kex_hlen = 20;
 b->kex_h = malloc(20);
 hash = sha1_init();
 opp = put_string(&b->opkt[0],&b->c_version[0],b->c_vlen);
 sha1_process_bytes(hash,&b->opkt[0],opp-&b->opkt[0]);
 opp = put_string(&b->opkt[0],&b->s_version[0],b->s_vlen);
 sha1_process_bytes(hash,&b->opkt[0],opp-&b->opkt[0]);
 opp = put_string(&b->opkt[0],b->c_kexinit_payload,b->c_kexinit_paylen);
 sha1_process_bytes(hash,&b->opkt[0],opp-&b->opkt[0]);
 opp = put_string(&b->opkt[0],b->s_kexinit_payload,b->s_kexinit_paylen);
 sha1_process_bytes(hash,&b->opkt[0],opp-&b->opkt[0]);
 opp = put_string(&b->opkt[0],b->K_S.data,b->K_S.len);
 sha1_process_bytes(hash,&b->opkt[0],opp-&b->opkt[0]);
 opp = put_mpint(&b->opkt[0],e);
 sha1_process_bytes(hash,&b->opkt[0],opp-&b->opkt[0]);
 opp = put_mpint(&b->opkt[0],f);
 sha1_process_bytes(hash,&b->opkt[0],opp-&b->opkt[0]);
 opp = put_mpint(&b->opkt[0],K);
 sha1_process_bytes(hash,&b->opkt[0],opp-&b->opkt[0]);
 b->kex_klen = opp - &b->opkt[0];
 b->kex_k = malloc(b->kex_klen);
 bcopy(&b->opkt[0],b->kex_k,b->kex_klen);
 sha1_result(hash,b->kex_h);
}

static const char *kex_d_h_gX_sha1_disabled(void)
{
 return(0);
}

static void *kex_d_h_gX_sha1_init_c(LAYER *l, const char *phex)
{
 DHGX_SHA1_STATE *s;
 unsigned char *opp;

 if (! l->b->hk->can_sig)
  { logmsg(LM_ERR|LM_PEER,"kex %s: server host key algorithm (%s) can't do signatures!",l->b->kex->name,l->b->hk->name);
    exit(1);
  }
 s = malloc(sizeof(DHGX_SHA1_STATE));
 s->server = 0;
 s->phex = phex;
 kexdh1_init(&s->p,&s->c.x,&s->c.e,phex);
 l->b->opkt[0] = SSH_MSG_KEXDH_INIT;
 opp = put_mpint(&l->b->opkt[1],&s->c.e);
 l->b->oplen = opp - &l->b->opkt[0];
 below_opkt(l);
 return(s);
}

static void *kex_d_h_g1_sha1_init_c(LAYER *l)
{
 return(kex_d_h_gX_sha1_init_c(l,&p1_hex[0]));
}

static void *kex_d_h_g14_sha1_init_c(LAYER *l)
{
 return(kex_d_h_gX_sha1_init_c(l,&p14_hex[0]));
}

static void *kex_d_h_gX_sha1_init_s(LAYER *l, const char *phex)
{
 DHGX_SHA1_STATE *s;

 if (! l->b->hk->can_sig)
  { logmsg(LM_ERR|LM_PEER,"kex %s: server host key algorithm (%s) can't do signatures!",l->b->kex->name,l->b->hk->name);
    exit(1);
  }
 s = malloc(sizeof(DHGX_SHA1_STATE));
 s->server = 1;
 s->phex = phex;
 return(s);
}

static void *kex_d_h_g1_sha1_init_s(LAYER *l)
{
 return(kex_d_h_gX_sha1_init_s(l,&p1_hex[0]));
}

static void *kex_d_h_g14_sha1_init_s(LAYER *l)
{
 return(kex_d_h_gX_sha1_init_s(l,&p14_hex[0]));
}

static int kex_d_h_gX_sha1_run_c(LAYER *l, DHGX_SHA1_STATE *s)
{
 BPP *b;
 MP_INT f;
 MP_INT K;
 STR Hsig;
 HASHALG *ha;
 void *hash;
 unsigned char sigbuf[20];

 b = l->b;
 mpz_init(&f);
 mpz_init(&K);
 if (b->ipkt[0] != SSH_MSG_KEXDH_REPLY)
  { logmsg(LM_WARN|LM_PEER,"KEXDH_REPLY packet isn't");
    exit(1);
  }
 parse_packet(b,&pp_fail,
	PP_IGNORE(1),
	PP_STRING(&b->K_S),
	PP_MPINT(&f),
	PP_STRING(&Hsig),
	PP_ENDSHERE );
 if ((mpz_cmp(&f,&s->p) >= 0) || (mpz_cmp_ui(&f,1) < 0))
  { logmsg(LM_WARN|LM_PEER,"server's f value is out of range");
    exit(1);
  }
 check_host_key(b->hk,b->K_S.data,b->K_S.len);
 mpz_powm(&K,&f,&s->c.x,&s->p);
 kexdh1_compute(&K,&s->c.e,&f,b);
 ha = (*b->hk->prehash)();
 if (sizeof(sigbuf) != ha->hashlen) panic("hash length disagreement");
 hash = (*ha->init)();
 (*ha->process)(hash,b->kex_h,20);
 (*ha->done)(hash,sigbuf);
 if (! (*b->hk->checksig)(b->K_S.data,b->K_S.len,&sigbuf[0],Hsig.data,Hsig.len))
  { logmsg(LM_WARN|LM_PEER,"kex hash signature wrong");
    exit(1);
  }
 mpz_clear(&s->p);
 mpz_clear(&s->c.x);
 mpz_clear(&s->c.e);
 mpz_clear(&f);
 mpz_clear(&K);
 b->kex_hash = &hashalg_sha1;
 free(s);
 free_str(b->K_S);
 return(1);
}

static int kex_d_h_gX_sha1_run_s(LAYER *l, DHGX_SHA1_STATE *s)
{
 BPP *b;
 MP_INT e;
 MP_INT K;
 HASHALG *ha;
 void *hash;
 unsigned char sigbuf[20];
 unsigned char *opp;
 STR sig;
 void *privdata;
 int privlen;

 b = l->b;
 if (! b->hk->can_sig)
  { logmsg(LM_ERR|LM_PEER,"kex %s: server host key algorithm (%s) can't do signatures!",l->b->kex->name,l->b->hk->name);
    exit(1);
  }
 mpz_init(&e);
 mpz_init(&K);
 if (b->ipkt[0] != SSH_MSG_KEXDH_INIT)
  { logmsg(LM_WARN|LM_PEER,"KEXDH_INIT packet isn't");
    exit(1);
  }
 parse_packet(b,&pp_fail,
	PP_IGNORE(1),
	PP_MPINT(&e),
	PP_ENDSHERE );
 kexdh1_init(&s->p,&s->s.y,&s->s.f,s->phex);
 (*b->hk->host_key)(b,&privdata,&privlen);
 mpz_powm(&K,&e,&s->s.y,&s->p);
 kexdh1_compute(&K,&e,&s->s.f,b);
 ha = (*b->hk->prehash)();
 if (sizeof(sigbuf) != ha->hashlen) panic("hash length disagreement");
 hash = (*ha->init)();
 (*ha->process)(hash,b->kex_h,20);
 (*ha->done)(hash,sigbuf);
 l->b->opkt[0] = SSH_MSG_KEXDH_REPLY;
 opp = put_string(&l->b->opkt[1],b->K_S.data,b->K_S.len);
 opp = put_mpint(opp,&s->s.f);
 if (! (*b->hk->sign)(&sig,b->K_S.data,b->K_S.len,privdata,privlen,&sigbuf[0]))
  { logmsg(LM_ERR|LM_PEER,"can't sign kex hash");
    exit(1);
  }
 opp = put_string(opp,sig.data,sig.len);
 l->b->oplen = opp - &l->b->opkt[0];
 below_opkt(l);
 mpz_clear(&s->p);
 mpz_clear(&s->s.y);
 mpz_clear(&s->s.f);
 mpz_clear(&e);
 mpz_clear(&K);
 free_str(sig);
 b->kex_hash = &hashalg_sha1;
 free(s);
 return(1);
}

static int kex_d_h_gX_sha1_run(LAYER *l, void *sv)
{
 DHGX_SHA1_STATE *s;

 s = sv;
 return(s->server?kex_d_h_gX_sha1_run_s(l,s):kex_d_h_gX_sha1_run_c(l,s));
}

static void kex_d_h_gX_sha1_servinit(KEXALG *a __attribute__((__unused__)))
{
}

static int kex_d_h_gX_sha1_servidle(KEXALG *a __attribute__((__unused__)))
{
 return(BLOCK_NIL);
}

static void kex_d_h_gX_sha1_servproc(KEXALG *a __attribute__((__unused__)))
{
}

KEXALG kexalg_diffie_hellman_group1_sha1
 = { "diffie-hellman-group1-sha1", 0, 0, 1,
     &kex_d_h_gX_sha1_disabled,
     &kex_d_h_g1_sha1_init_c,
     &kex_d_h_g1_sha1_init_s,
     &kex_d_h_gX_sha1_run,
     &kex_d_h_gX_sha1_servinit,
     &kex_d_h_gX_sha1_servidle,
     &kex_d_h_gX_sha1_servproc };

KEXALG kexalg_diffie_hellman_group14_sha1
 = { "diffie-hellman-group14-sha1", 0, 0, 1,
     &kex_d_h_gX_sha1_disabled,
     &kex_d_h_g14_sha1_init_c,
     &kex_d_h_g14_sha1_init_s,
     &kex_d_h_gX_sha1_run,
     &kex_d_h_gX_sha1_servinit,
     &kex_d_h_gX_sha1_servidle,
     &kex_d_h_gX_sha1_servproc };