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

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

#include "pp.h"
#include "rnd.h"
#include "algs.h"
#include "panic.h"
#include "nested.h"
#include "hostkey.h"
#include "pkt-util.h"
#include "algs-list.h"
#include "small-primes.h"

#define ENCPREF_LEN 15
static const unsigned char encpref[ENCPREF_LEN]
 = {	0x30,	/* universal, constructed, sequence */
	0x21,	/* length (33) */
		0x30,	/* universal, constructed, sequence */
		0x09,	/* length (9) */
			0x06,	/* universal, object ID */
			0x05,	/* length (5) */
				0x2b, 0x0e, 0x03, 0x02, 0x1a,
				/* 1.3.14.3.2.26 = id-SHA1 */
			0x05,	/* universal, null */
			0x00,	/* length (0) */
		0x04,	/* universal, octet string */
		0x14,	/* length (20) */ };

static void *rsahk_pub_data;
static int rsahk_pub_len;
static void *rsahk_priv_data;
static int rsahk_priv_len;
static int rsa_have_hk = 0;

typedef struct rsagen RSAGEN;
typedef struct primegen PRIMEGEN;

struct primegen {
  int step;
#define PGS_NONE  0
#define PGS_INIT  1
#define PGS_WHEEL 2
#define PGS_TRIAL 3
  MP_INT *loc;
  int nbits;
  int flags;
  MP_INT t1;
  MP_INT t2;
  unsigned int inc;
  int reps;
  int *moduli;
  } ;

struct rsagen {
  int step;
#define RGS_NONE 0
#define RGS_INIT 1
#define RGS_P    2
#define RGS_Q    3
#define RGS_E    4
  MP_INT *loc_e;
  MP_INT *loc_n;
  MP_INT *loc_d;
  int bits;
  unsigned int flags;
  int nbits;
  unsigned int rv;
  int rb;
  int etries;
  MP_INT p;
  MP_INT q;
  MP_INT p1;
  MP_INT q1;
  MP_INT t;
  PRIMEGEN pg;
  } ;

static int hk_ssh_rsa_canuse(int forserver)
{
 if (! forserver) return(1);
 if (! rsa_have_hk)
  { rsa_have_hk = read_host_key( &hkalg_ssh_rsa,
				 &rsahk_pub_data, &rsahk_pub_len,
				 &rsahk_priv_data, &rsahk_priv_len )
		  ? 1 : -1;
  }
 return(rsa_have_hk>0);
}

static void hk_ssh_rsa_host_key(BPP *b, void **privdp, int *privlp)
{
 if (rsa_have_hk <= 0) panic("don't have rsa hk");
 b->K_S.data = rsahk_pub_data;
 b->K_S.len = rsahk_pub_len;
 *privdp = rsahk_priv_data;
 *privlp = rsahk_priv_len;
}

static int hk_ssh_rsa_match(const void *d1, int l1, const void *d2, int l2)
{
 return((l1==l2)&&!bcmp(d1,d2,l1));
}

static HASHALG *hk_ssh_rsa_prehash(void)
{
 extern HASHALG hashalg_sha1;

 return(&hashalg_sha1);
}

static int hk_ssh_rsa_checksig(const void *pkdata, int pklen, const void *hash, const void *sig, int siglen)
{
 __label__ parsefail;

 MP_INT n;
 MP_INT e;
 MP_INT s;
 MP_INT m;
 MP_INT h;
 MP_INT tmp1;
 STR sstr;
 int nlen;
 int rv;

 NESTED void fail(const void *dp __attribute__((__unused__)), const char *fmt __attribute__((__unused__)), ...)
  { goto parsefail;
  }

 NESTED void clears(void)
  { mpz_clear(&n);
    mpz_clear(&e);
    mpz_clear(&s);
    mpz_clear(&m);
    mpz_clear(&h);
    mpz_clear(&tmp1);
  }

 if (pklen < 11+4+4) return(0);
 if (bcmp(pkdata,"\0\0\0\7ssh-rsa",11)) return(0);
 if (siglen < 11+4) return(0);
 if (bcmp(sig,"\0\0\0\7ssh-rsa",11)) return(0);
 mpz_init(&n);
 mpz_init(&e);
 mpz_init(&s);
 mpz_init(&m);
 mpz_init(&h);
 mpz_init(&tmp1);
 parse_data(pkdata,pklen,fail,
	PP_IGNORE(11),
	PP_MPINT(&e),
	PP_MPINT(&n),
	PP_ENDSHERE );
 nlen = (mpz_sizeinbase(&n,2) + 7) >> 3;
 if (nlen < ENCPREF_LEN+20+10+1) return(0);
 parse_data(sig,siglen,fail,
	PP_IGNORE(11),
	PP_STRING(&sstr),
	PP_ENDSHERE );
 if (0)
  {
parsefail:;
    clears();
    return(0);
  }
 data_to_mpint(&s,sstr.data,sstr.len);
 if (mpz_cmp(&s,&n) >= 0) fail(0,0);
 mpz_powm(&m,&s,&e,&n);
 mpz_set_ui(&h,1);
 mpz_mul_2exp(&h,&h,((nlen-(ENCPREF_LEN+20+3))*8)+1);
 mpz_sub_ui(&h,&h,1);
 mpz_mul_2exp(&h,&h,(ENCPREF_LEN+20+1)*8);
 data_to_mpint(&tmp1,&encpref[0],ENCPREF_LEN);
 mpz_mul_2exp(&tmp1,&tmp1,20*8);
 mpz_add(&h,&h,&tmp1);
 data_to_mpint(&tmp1,hash,20);
 mpz_add(&h,&h,&tmp1);
 rv = (mpz_cmp(&m,&h) == 0);
 if (! rv)
  { printf("hk_ssh_rsa_checksig failed:\n");
    printf("H(m) =");
     { int i;
       for (i=0;i<20;i++) printf(" %02x",((const unsigned char *)hash)[i]);
     }
    printf("\n");
    printf("n = ");
    mpz_out_str(stdout,16,&n);
    printf("\n");
    printf("e = ");
    mpz_out_str(stdout,16,&e);
    printf("\n");
    printf("s = ");
    mpz_out_str(stdout,16,&s);
    printf("\n");
    printf("m = ");
    mpz_out_str(stdout,16,&m);
    printf("\n");
    printf("h = ");
    mpz_out_str(stdout,16,&h);
    printf("\n");
  }
 clears();
 return(rv);
}

static int hk_ssh_rsa_checkpub(const void *blob, int len)
{
 return((len >= 11) && !bcmp(blob,"\0\0\0\7ssh-rsa",11));
}

static int hk_ssh_rsa_sign(STR *into, const void *pubdata, int publen, const void *privdata, int privlen, const void *hash)
{
 __label__ parsefail;

 MP_INT n;
 MP_INT e;
 MP_INT d;
 MP_INT s;
 MP_INT m;
 MP_INT h;
 MP_INT tmp1;
 int nlen;
 unsigned char *res;
 unsigned char *rp;
 int rlen;

 NESTED void fail(const void *dp __attribute__((__unused__)), const char *fmt __attribute__((__unused__)), ...)
  { goto parsefail;
  }

 NESTED void clears(void)
  { mpz_clear(&n);
    mpz_clear(&e);
    mpz_clear(&d);
    mpz_clear(&s);
    mpz_clear(&m);
    mpz_clear(&h);
    mpz_clear(&tmp1);
  }

 if (publen < 11+4+4) return(0);
 if (bcmp(pubdata,"\0\0\0\7ssh-rsa",11)) return(0);
 mpz_init(&n);
 mpz_init(&e);
 mpz_init(&d);
 mpz_init(&s);
 mpz_init(&m);
 mpz_init(&h);
 mpz_init(&tmp1);
 parse_data(pubdata,publen,fail,
	PP_IGNORE(11),
	PP_MPINT(&e),
	PP_MPINT(&n),
	PP_ENDSHERE );
 if (0)
  {
parsefail:;
    clears();
    return(0);
  }
 nlen = (mpz_sizeinbase(&n,2) + 7) >> 3;
 if (nlen < ENCPREF_LEN+20+10+1) fail(0,0);
 data_to_mpint(&d,privdata,privlen);
 if (mpz_cmp(&d,&n) >= 0) fail(0,0);
 mpz_set_ui(&h,1);
 mpz_mul_2exp(&h,&h,((nlen-(ENCPREF_LEN+20+3))*8)+1);
 mpz_sub_ui(&h,&h,1);
 mpz_mul_2exp(&h,&h,(ENCPREF_LEN+20+1)*8);
 data_to_mpint(&tmp1,&encpref[0],ENCPREF_LEN);
 mpz_mul_2exp(&tmp1,&tmp1,20*8);
 mpz_add(&h,&h,&tmp1);
 data_to_mpint(&tmp1,hash,20);
 mpz_add(&h,&h,&tmp1);
 mpz_powm(&m,&h,&d,&n);
 rlen = 11 + 4 + ((mpz_sizeinbase(&m,2) + 7) >> 3);
 res = malloc(rlen);
 rp = put_string(res,"ssh-rsa",-1);
 rp = put_uint32(rp,rlen-11-4);
 mpint_to_data(&m,rp,rlen-11-4);
 if (mpz_cmp_ui(&m,0) != 0) panic("bits left over");
 into->data = res;
 into->len = rlen;
 clears();
 return(1);
}

static void random_prime_init(PRIMEGEN *g, MP_INT *p, int nbits, unsigned int flags)
{
 g->step = PGS_INIT;
 g->loc = p;
 g->nbits = nbits;
 g->flags = flags;
 mpz_init(&g->t1);
 mpz_init(&g->t2);
 g->moduli = malloc(n_small_primes*sizeof(int));
 g->inc = 0;
}

static int random_prime_step(PRIMEGEN *g)
{
 int i;
 unsigned char rbuf[(g->nbits+7)/8];

 switch (g->step)
  { default:
       panic("bad step");
       break;
    case PGS_INIT:
       random_data(&rbuf[0],sizeof(rbuf));
       if (g->nbits & 7) rbuf[0] &= 0xff >> (8 - (g->nbits & 7));
       rbuf[0] |= "\xc0\x01\x03\x06\x0c\x18\x30\x60"[g->nbits&7];
       if (sizeof(rbuf) > 1) rbuf[1] |= "\x00\x80\x00\x00\x00\x00\x00\x00"[g->nbits&7];
       rbuf[sizeof(rbuf)-1] |= 1;
       data_to_mpint(g->loc,&rbuf[0],sizeof(rbuf));
       for (i=n_small_primes-1;i>=0;i--) g->moduli[i] = mpz_mod_ui(&g->t1,g->loc,small_primes[i]);
       g->inc = -2U;
       g->step = PGS_WHEEL;
       return(0);
       break;
    case PGS_WHEEL:
       while <"inc"> (1)
	{ g->inc += 2;
	  if (g->inc >= 0x70000000)
	   { g->step = PGS_INIT;
	     return(0);
	   }
	  for (i=0;i<n_small_primes;i++)
	   { while (g->moduli[i]+g->inc >= small_primes[i]) g->moduli[i] -= small_primes[i];
	     if (g->moduli[i]+g->inc == 0) continue <"inc">;
	   }
	  mpz_add_ui(&g->t1,g->loc,g->inc);
	  if (g->flags & GEN_F_VERBOSE) write(1,".",1);
	  mpz_set_ui(&g->t2,2);
	  mpz_powm(&g->t2,&g->t2,&g->t1,&g->t1);
	  if (mpz_cmp_ui(&g->t2,2) != 0) return(0);
	  if (g->flags & GEN_F_VERBOSE) write(1,"+",1);
	  g->step = PGS_TRIAL;
	  g->reps = 32;
	  return(0);
	}
       break;
    case PGS_TRIAL:
       if (! mpz_probab_prime_p(&g->t1,1))
	{ g->inc += 2;
	  g->step = PGS_WHEEL;
	  return(0);
	}
       if (g->reps-- > 0) return(0);
       mpz_set(g->loc,&g->t1);
       free(g->moduli);
       mpz_clear(&g->t1);
       mpz_clear(&g->t2);
       g->step = PGS_NONE;
       return(1);
       break;
  }
}

static void random_prime_abort(PRIMEGEN *g)
{
 mpz_clear(&g->t1);
 mpz_clear(&g->t2);
 free(g->moduli);
 g->step = PGS_NONE;
}

/*
 * We want e to be large enough to make the low encryption exponent
 *  attack hopeless right out of the gate; but we also want it to have
 *  relatively few bits set, so as to speed up operation.  So we set e
 *  equal to 0x8000000000000001 plus up to two other bits in random
 *  positions, trying different random values until we get an e that's
 *  relatively prime to t=(p-1)(q-1).  (The code always sets two bits;
 *  I say "up to" two other bits because the code may hit the bits
 *  already set, or may pick the same bit for each of the two.)
 *
 * This way, it's unlikely that two different keys will have the same
 *  value of e, but even if they do, it is impossible to mount the
 *  attack, because it requires at least e different messages, and e is
 *  very large when considered as a count of things.
 */
#define ESIZE 8 /* size of e in bytes - ie, e is ESIZE*8 bits long */
#define LGESIZE 3 /* lg(ESIZE) */
#if ESIZE != (1 << LGESIZE)
#error "LGESIZE and ESIZE disagree"
#endif

static void rsa_gen_init(RSAGEN *r, MP_INT *ep, MP_INT *np, MP_INT *dp, int bits, unsigned int flags)
{
 if (bits < 256) bits = 1024;
 r->step = RGS_INIT;
 r->loc_e = ep;
 r->loc_n = np;
 r->loc_d = dp;
 r->bits = bits;
 r->flags = flags;
 mpz_init(&r->p);
 mpz_init(&r->q);
 mpz_init(&r->p1);
 mpz_init(&r->q1);
 mpz_init(&r->t);
}

static int rsa_gen_step(RSAGEN *r)
{
 switch (r->step)
  { default:
       panic("bad step");
       break;
    case RGS_INIT:
       if (r->flags & GEN_F_VERBOSE)
	{ fflush(stdout);
	  write(1,"p(",2);
	}
       random_prime_init(&r->pg,&r->p,r->bits>>1,r->flags);
       r->step = RGS_P;
       return(0);
       break;
    case RGS_P:
       if (! random_prime_step(&r->pg)) return(0);
       if (r->flags & GEN_F_VERBOSE) write(1,")q(",3);
       random_prime_init(&r->pg,&r->q,(r->bits+1)>>1,r->flags);
       r->step = RGS_Q;
       return(0);
       break;
    case RGS_Q:
       if (! random_prime_step(&r->pg)) return(0);
       if (r->flags & GEN_F_VERBOSE) write(1,")e(",3);
       mpz_sub_ui(&r->p1,&r->p,1);
       mpz_sub_ui(&r->q1,&r->q,1);
       mpz_mul(r->loc_n,&r->p,&r->q);
       mpz_mul(&r->t,&r->p1,&r->q1);
       r->nbits = mpz_sizeinbase(r->loc_n,2);
       r->rv = 0;
       r->rb = 0;
       r->step = RGS_E;
       r->etries = 0;
       return(0);
       break;
    case RGS_E:
	{ int bit;
	  unsigned char ebuf[ESIZE];
	  NESTED int ebitno(void)
	   { unsigned int b;
	     while (r->rb < LGESIZE+3)
	      { unsigned char v;
		random_data(&v,1);
		r->rv |= ((unsigned int)v) << r->rb;
		r->rb += 8;
	      }
	     b = r->rv & ((1 << (LGESIZE+3)) - 1);
	     r->rv >>= LGESIZE + 3;
	     r->rb -= LGESIZE + 3;
	     return(b);
	   }
	  if (r->etries > 32)
	   { if (r->flags & GEN_F_VERBOSE) write(1,"?)",2);
	     r->step = RGS_INIT;
	     return(0);
	   }
	  r->etries ++;
	  if (r->flags & GEN_F_VERBOSE) write(1,".",1);
	  bzero(&ebuf[0],sizeof(ebuf));
	  ebuf[0] = 0x80;
	  ebuf[ESIZE-1] = 0x01;
	  bit = ebitno();
	  ebuf[bit>>3] |= 0x01 << (bit & 7);
	  bit = ebitno();
	  ebuf[bit>>3] |= 0x01 << (bit & 7);
	  data_to_mpint(r->loc_e,&ebuf[0],sizeof(ebuf));
	  if (! mpz_invert(r->loc_d,r->loc_e,&r->t)) return(0);
	  mpz_mod(r->loc_d,r->loc_d,&r->t);
	  if (mpz_sizeinbase(r->loc_d,2)*2 < r->nbits) return(0);
	  if (r->flags & GEN_F_VERBOSE) write(1,")\n",2);
	  if (r->flags & GEN_F_VALUES)
	   { printf("n = ");
	     mpz_out_str(stdout,16,r->loc_n);
	     printf("\ne = ");
	     mpz_out_str(stdout,16,r->loc_e);
	     printf("\nd = ");
	     mpz_out_str(stdout,16,r->loc_d);
	     printf("\n");
	   }
	  mpz_clear(&r->p);
	  mpz_clear(&r->q);
	  mpz_clear(&r->p1);
	  mpz_clear(&r->q1);
	  mpz_clear(&r->t);
	  return(1);
	}
       break;
  }
}

static void rsa_gen_abort(RSAGEN *r)
{
 switch (r->step)
  { default:
       panic("bad step");
       break;
    case RGS_INIT:
    case RGS_E:
       break;
    case RGS_P:
    case RGS_Q:
       random_prime_abort(&r->pg);
       break;
  }
 mpz_clear(&r->p);
 mpz_clear(&r->q);
 mpz_clear(&r->p1);
 mpz_clear(&r->q1);
 mpz_clear(&r->t);
 r->step = RGS_NONE;
}

static void gen_rsa_pair(MP_INT *ep, MP_INT *np, MP_INT *dp, int bits, unsigned int flags)
{
 RSAGEN rg;

 rsa_gen_init(&rg,ep,np,dp,bits,flags);
 while (! rsa_gen_step(&rg)) ;
}

static void hk_ssh_rsa_genkey(FILE *pubf, FILE *privf, int bits, unsigned int flags)
{
 MP_INT n;
 MP_INT e;
 MP_INT d;
 int blobsize;
 unsigned char *blobbuf;
 unsigned char *bp;

 if (bits < 256) bits = 1024;
 mpz_init(&n);
 mpz_init(&e);
 mpz_init(&d);
 gen_rsa_pair(&e,&n,&d,bits,flags);
 blobsize = 11 + mpint_bytes(&e) + mpint_bytes(&n);
 blobbuf = malloc(blobsize+32);
 bp = &blobbuf[0];
 bp = put_string(bp,"ssh-rsa",-1);
 bp = put_mpint(bp,&e);
 bp = put_mpint(bp,&n);
 if (bp-blobbuf > blobsize) panic("blob overflow A");
 fwrite(blobbuf,1,bp-&blobbuf[0],pubf);
  { int i;
    i = (mpz_sizeinbase(&d,2) + 7) >> 3;
    if (i > blobsize) panic("blob overflow B"); /* can't happen: d < t < n */
    mpint_to_data(&d,blobbuf,i);
    fwrite(blobbuf,1,i,privf);
  }
 mpz_clear(&n);
 mpz_clear(&e);
 mpz_clear(&d);
}
#undef ESIZE

static void hk_ssh_rsa_private(const char *cmd, ...)
{
 va_list ap;

 va_start(ap,cmd);
 if (!strcmp(cmd,"gen-step"))
  { void *gp;
    gp = va_arg(ap,void *);
    if (! rsa_gen_step(gp)) return;
    *va_arg(ap,int *) = 1;
  }
 else if (!strcmp(cmd,"gen-init"))
  { void **gp;
    MP_INT *ep;
    MP_INT *np;
    MP_INT *dp;
    int bits;
    gp = va_arg(ap,void **);
    ep = va_arg(ap,MP_INT *);
    np = va_arg(ap,MP_INT *);
    dp = va_arg(ap,MP_INT *);
    bits = va_arg(ap,int);
    *gp = malloc(sizeof(RSAGEN));
    rsa_gen_init(*gp,ep,np,dp,bits,0);
  }
 else if (!strcmp(cmd,"gen-abort"))
  { void *gp;
    gp = va_arg(ap,void *);
    rsa_gen_abort(gp);
    free(gp);
  }
 else if (!strcmp(cmd,"gen"))
  { MP_INT *ep;
    MP_INT *np;
    MP_INT *dp;
    int bits;
    ep = va_arg(ap,MP_INT *);
    np = va_arg(ap,MP_INT *);
    dp = va_arg(ap,MP_INT *);
    bits = va_arg(ap,int);
    gen_rsa_pair(ep,np,dp,bits,0);
  }
 else if (!strcmp(cmd,"decode"))
  { ROSTR s;
    MP_INT *ep;
    MP_INT *np;
    void (*fail)(const void *, const char *, ...);
    s = va_arg(ap,ROSTR);
    ep = va_arg(ap,MP_INT *);
    np = va_arg(ap,MP_INT *);
    fail = va_arg(ap,__typeof__(fail));
    if ((s.len < 11) || bcmp(s.data,"\0\0\0\7ssh-rsa",11)) (*fail)(0,"public key data blob type wrong");
    parse_data(s.data,s.len,fail,
	PP_IGNORE(11),
	PP_MPINT(ep),
	PP_MPINT(np),
	PP_ENDSHERE );
  }
 else if (!strcmp(cmd,"encode"))
  { MP_INT *ep;
    MP_INT *np;
    STR *sp;
    unsigned char *bp;
    ep = va_arg(ap,MP_INT *);
    np = va_arg(ap,MP_INT *);
    sp = va_arg(ap,STR *);
    sp->len = 4+7 + mpint_bytes(ep) + mpint_bytes(np);
    sp->data = malloc(sp->len);
    bp = (void *)sp->data;
    bp = put_string(bp,"ssh-rsa",-1);
    bp = put_mpint(bp,ep);
    bp = put_mpint(bp,np);
    if (bp != sp->len+(unsigned char *)sp->data) panic("encoding overflow");
  }
 else
  { panic("bad cmd");
  }
 va_end(ap);
}

static const char *hk_ssh_rsa_altnames[]
 = { "rsa", 0 };

HKALG hkalg_ssh_rsa
 = { "ssh-rsa", 0, 0, 1,
     &hk_ssh_rsa_canuse,
     &hk_ssh_rsa_host_key,
     &hk_ssh_rsa_match,
     &hk_ssh_rsa_prehash,
     &hk_ssh_rsa_checksig,
     &hk_ssh_rsa_sign,
     &hk_ssh_rsa_genkey,
     &hk_ssh_rsa_checkpub,
     &hk_ssh_rsa_private,
     &hk_ssh_rsa_altnames[0],
     "identity.rsa",
     "hostkey.rsa" };