#include <gmp.h>
#include <sha1.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <strings.h>

#include "pp.h"
#include "rnd.h"
#include "algs.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;

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

static void hk_ssh_rsa_host_key(BPP *b, void **privdp, int *privlp)
{
 if (! rsa_have_hk) abort();
 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;

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

 static 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;

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

 static 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) abort();
 into->data = res;
 into->len = rlen;
 clears();
 return(1);
}

static void random_prime(MP_INT *p, int nbits)
{
 int i;
 MP_INT t1;
 MP_INT t2;
 unsigned int inc;
 int moduli[n_small_primes];
 unsigned char rbuf[(nbits+7)/8];

 mpz_init(&t1);
 mpz_init(&t2);
 while <"newstart"> (1)
  { random_data(&rbuf[0],sizeof(rbuf));
    if (nbits & 7) rbuf[0] &= 0xff >> (8 - (nbits & 7));
    rbuf[0] |= "\xc0\x01\x03\x06\x0c\x18\x30\x60"[nbits&7];
    if (sizeof(rbuf) > 1) rbuf[1] |= "\x00\x80\x00\x00\x00\x00\x00\x00"[nbits&7];
    rbuf[sizeof(rbuf)-1] |= 1;
    data_to_mpint(p,&rbuf[0],sizeof(rbuf));
    for (i=n_small_primes-1;i>=0;i--) moduli[i] = mpz_mod_ui(&t1,p,small_primes[i]);
    for <"inc"> (inc=0;;inc+=2)
     { if (inc >= 0x70000000) continue <"newstart">;
       for (i=0;i<n_small_primes;i++)
	{ while (moduli[i]+inc >= small_primes[i]) moduli[i] -= small_primes[i];
	  if (moduli[i]+inc == 0) continue <"inc">;
	}
       mpz_add_ui(&t1,p,inc);
       write(1,".",1);
       mpz_set_ui(&t2,2);
       mpz_powm(&t2,&t2,&t1,&t1);
       if (mpz_cmp_ui(&t2,2) != 0) continue <"inc">;
       write(1,"+",1);
       if (! mpz_probab_prime_p(&t1,32)) continue <"inc">;
       mpz_set(p,&t1);
       return;
     }
  }
}

static void hk_ssh_rsa_genkey(FILE *pubf, FILE *privf, int bits)
{
 MP_INT p;
 MP_INT q;
 MP_INT n;
 MP_INT p1;
 MP_INT q1;
 MP_INT e;
 MP_INT d;
 MP_INT t;
 unsigned char *ebuf;
 int ebits;
 int ebytes;
 int blobsize;
 unsigned char *blobbuf;
 unsigned char *bp;
 int nbits;

 if (bits < 32) bits = 1024;
 mpz_init(&p);
 mpz_init(&q);
 mpz_init(&n);
 mpz_init(&p1);
 mpz_init(&q1);
 mpz_init(&e);
 mpz_init(&d);
 mpz_init(&t);
 write(1,"p(",2);
 random_prime(&p,bits>>1);
 write(1,")q(",3);
 random_prime(&q,(bits+1)>>1);
 write(1,")",1);
 mpz_sub_ui(&p1,&p,1);
 mpz_sub_ui(&q1,&q,1);
 mpz_mul(&n,&p,&q);
 mpz_mul(&t,&p1,&q1);
 nbits = mpz_sizeinbase(&n,2);
 write(1,"e(",2);
 ebits = bits - 1;
 ebytes = (ebits + 7) >> 3;
 ebuf = malloc(ebytes);
 while (1)
  { write(1,".",1);
    random_data(ebuf,ebytes);
    if (ebits & 7) ebuf[0] &= 0xff >> (8 - (ebits & 7));
    data_to_mpint(&e,&ebuf[0],ebytes);
    if (mpz_sizeinbase(&e,2)*2 < nbits) continue;
    mpz_gcd(&d,&e,&t);
    if (mpz_cmp_ui(&d,1) != 0) continue;
    mpz_invert(&d,&e,&t);
    mpz_mod(&d,&d,&t);
    if (mpz_sizeinbase(&d,2)*2 < nbits) continue;
    break;
  }
 write(1,")",1);
 printf("\nn = ");
 mpz_out_str(stdout,16,&n);
 printf("\ne = ");
 mpz_out_str(stdout,16,&e);
 printf("\nd = ");
 mpz_out_str(stdout,16,&d);
 printf("\n");
 blobsize = 11 + 4+(mpz_sizeinbase(&e,2)>>3)+1 + 4+(mpz_sizeinbase(&n,2)>>3)+1;
 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) abort();
 fwrite(blobbuf,1,bp-&blobbuf[0],pubf);
  { int i;
    i = (mpz_sizeinbase(&d,2) + 7) >> 3;
    if (i > blobsize) abort(); /* can't happen: d < n */
    mpint_to_data(&d,blobbuf,i);
    fwrite(blobbuf,1,i,privf);
  }
 mpz_clear(&p);
 mpz_clear(&q);
 mpz_clear(&n);
 mpz_clear(&p1);
 mpz_clear(&q1);
 mpz_clear(&e);
 mpz_clear(&d);
 mpz_clear(&t);
}
#undef BLOBSIZE

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

HKALG hkalg_ssh_rsa
 = { "ssh-rsa", 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_altnames[0],
     "identity.rsa",
     "hostkey.rsa" };