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

extern const char *__progname;

#include "pp.h"
#include "bpp.h"
#include "str.h"
#include "algs.h"
#include "msgs.h"
#include "pkt-util.h"
#include "userauth.h"

typedef struct authkey AUTHKEY;
typedef struct pkpriv PKPRIV;

struct pkpriv {
  AUTHKEY *k;
  int queried;
  } ;

struct authkey {
  AUTHKEY *link;
  HKALG *alg;
  void *pubdata;
  int publen;
  void *privdata;
  int privlen;
  char *comment;
  char *privname;
  } ;

static AUTHKEY *authkeys;

static unsigned char *build_key_query(UASTATE *s)
{
 PKPRIV *p;
 unsigned char *opp;

 p = s->algpriv;
 opp = (*s->pkthead)(s);
 *opp++ = 0;
 opp = put_string(opp,hkalg_name(p->k->alg),-1);
 opp = put_string(opp,p->k->pubdata,p->k->publen);
 return(opp);
}

static unsigned char *build_key_sig(UASTATE *s)
{
 PKPRIV *p;
 BPP *b;
 unsigned char *opp;
 STR sig;
 HASHALG *ha;
 void *hh;
 char *hash;

 p = s->algpriv;
 b = s->layer->b;
 ha = (*p->k->alg->prehash)();
 hh = (*ha->init)();
 opp = put_string(&s->layer->b->opkt[0],b->sessid,b->sessidlen);
 (*ha->process)(hh,&s->layer->b->opkt[0],opp-&s->layer->b->opkt[0]);
 opp = (*s->pkthead)(s);
 *opp++ = 1;
 opp = put_string(opp,hkalg_name(p->k->alg),-1);
 opp = put_string(opp,p->k->pubdata,p->k->publen);
 (*ha->process)(hh,&b->opkt[0],opp-&b->opkt[0]);
 hash = malloc(ha->hashlen);
 (*ha->done)(hh,hash);
 if (! (*p->k->alg->sign)(&sig,p->k->pubdata,p->k->publen,p->k->privdata,p->k->privlen,hash))
  { fprintf(stderr,"%s: can't generate authentication signature\n",__progname);
    exit(1);
  }
 opp = put_string(opp,sig.data,sig.len);
 free(sig.data);
 return(opp);
}

static char *read_comment(FILE *f, const char *def)
{
 char *cbuf;
 int l;
 int n;
 int c;

 static void savec(char c)
  { if (l >= n) cbuf = realloc(cbuf,n=l+16);
    cbuf[l++] = c;
  }

 cbuf = 0;
 l = 0;
 n = 0;
 while (1)
  { c = getc(f);
    switch (c)
     { case '\n': case EOF:
	  if (l == 0) return(strdup(def));
	  savec('\0');
	  return(cbuf);
	  break;
       case ' ':
	  if (l == 0) break;
       default:
	  savec(c);
	  if (l >= 80)
	   { savec('.');
	     savec('.');
	     savec('.');
	     savec('\0');
	     return(cbuf);
	   }
	  break;
     }
  }
}

static int ua_publickey_canuse(BPP *b __attribute__((__unused__)))
{
 char *home;
 char *keybuf;
 char *kfp;
 void *pub;
 int publen;
 void *priv;
 int privlen;
 char *comment;

 static void loadkey(const char *suf)
  { FILE *f;
    char kname[65];
    HKALG *alg;
    AUTHKEY *ak;
    sprintf(kfp,"%s.pub",suf);
    f = fopen(keybuf,"r");
    if (f == 0) return;
    if (fscanf(f,"%64s ",&kname[0]) != 1)
     { fclose(f);
       return;
     }
    alg = hkalg_vfind_c(&kname[0]);
    if ((alg == 0) || !alg->can_sig)
     { fclose(f);
       return;
     }
    if (! (*alg->read_pub)(f,&pub,&publen))
     { fclose(f);
       return;
     }
    comment = read_comment(f,suf);
    fclose(f);
    sprintf(kfp,"%s.priv",suf);
    f = fopen(keybuf,"r");
    if (f == 0)
     { free(pub);
       return;
     }
    if (! (*alg->read_priv)(f,pub,publen,&priv,&privlen))
     { fclose(f);
       free(pub);
       return;
     }
    fclose(f);
    ak = malloc(sizeof(AUTHKEY));
    ak->link = authkeys;
    ak->alg = alg;
    ak->pubdata = pub;
    ak->publen = publen;
    ak->privdata = priv;
    ak->privlen = privlen;
    strcpy(kfp,suf);
    ak->comment = comment;
    ak->privname = strdup(keybuf);
    authkeys = ak;
  }

 home = getenv("HOME");
 if (home == 0)
  { fprintf(stderr,"%s: no $HOME, public key authentication disabled\n",__progname);
    return(0);
  }
 keybuf = malloc(strlen(home)+64);
 kfp = keybuf + sprintf(keybuf,"%s/.ssh/",home);
 loadkey("identity.dsa");
 loadkey("identity.rsa");
 return(!!authkeys);
}

static unsigned char *ua_publickey_start(UASTATE *s)
{
 PKPRIV *p;

 p = malloc(sizeof(PKPRIV));
 s->algpriv = p;
 p->k = authkeys;
 p->queried = 1;
 return(build_key_query(s));
}

static int ua_publickey_partial(UASTATE *s __attribute__((__unused__)))
{
 abort();
}

static int ua_publickey_run(UASTATE *s)
{
 BPP *b;
 PKPRIV *p;
 STR algname;
 STR pkblob;

 b = s->layer->b;
 p = s->algpriv;
 switch (b->ipkt[0])
  { case SSH_MSG_USERAUTH_PK_OK:
       parse_packet(b,pp_fail,
		PP_SKIP(1),
		PP_STRING(&algname),
		PP_STRING(&pkblob),
		PP_ENDSHERE );
       if (! str_equalsC(algname,p->k->alg->name))
	{ fprintf(stderr,"%s: publickey authorization protocol error: algorithm name changed (sent %s, got %.*s)\n",__progname,p->k->alg->name,(int)algname.len,algname.data);
	  exit(1);
	}
       if (! str_equalsb(pkblob,p->k->pubdata,p->k->publen))
	{ int i;
	  fprintf(stderr,"%s: publickey authorization protocol error: key blob changed\n",__progname);
	  fprintf(stderr,"  sent:");
	  for (i=0;i<p->k->publen;i++) fprintf(stderr," %02x",((const unsigned char *)p->k->pubdata)[i]);
	  fprintf(stderr,"\n   got:");
	  for (i=0;i<pkblob.len;i++) fprintf(stderr," %02x",((const unsigned char *)pkblob.data)[i]);
	  fprintf(stderr,"\n");
	  exit(1);
	}
       b->oplen = build_key_sig(s) - &b->opkt[0];
       return(UA_SEND);
       break;
    default:
       fprintf(stderr,"%s: publickey authorization error: unexpected message (type %d)\n",__progname,b->ipkt[0]);
       exit(1);
       break;
  }
}

static int ua_publickey_retry(UASTATE *s __attribute__((__unused__)))
{
 abort();
}

static void ua_publickey_done(UASTATE *s)
{
 PKPRIV *p;

 p = s->algpriv;
 free(p);
 s->algpriv = 0;
}

UAALG uaalg_publickey
 = { "publickey",
     &ua_publickey_canuse,
     &ua_publickey_start,
     &ua_publickey_partial,
     &ua_publickey_run,
     &ua_publickey_retry,
     &ua_publickey_done };