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

extern const char *__progname;

#include "pp.h"
#include "bpp.h"
#include "b64.h"
#include "str.h"
#include "algs.h"
#include "msgs.h"
#include "keygen.h"
#include "cmdline.h"
#include "keyfiles.h"
#include "pkt-util.h"
#include "userauth.h"
#include "priv-crypto.h"
#include "agent-client.h"

typedef struct authkey AUTHKEY;
typedef struct pkpriv PKPRIV;
typedef struct aktconc AKTCONC;

struct aktconc {
  AUTHKEY *list;
  AUTHKEY **tail;
  } ;

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

struct authkey {
  AUTHKEY *link;
  HKALG *alg;
  void *pubdata;
  int publen;
  unsigned int flags;
#define AKF_AGENT 0x00000001
#define AKF_NO_PP 0x00000002
  void *privdata;
  int privlen;
  char *comment;
  } ;

static AUTHKEY *authkeylist;

static AUTHKEY *sort_keys_by_preference(AUTHKEY *list)
{
 AKTCONC sub[2];
 AKTCONC rev;
 int i;

 static int prefer(AUTHKEY *a, AUTHKEY *b)
  { /* return true if a is preferred to b
       prefer passphrase-free file keys, then agent keys, then others
       arbitrarily prefer a over b if no basis for caring */
    if (a->flags & AKF_NO_PP) return(1);
    if (b->flags & AKF_NO_PP) return(0);
    if (a->flags & AKF_AGENT) return(1);
    if (b->flags & AKF_AGENT) return(0);
    return(1);
  }

 if (!list || !list->link) return(list);
 sub[0].tail = &sub[0].list;
 sub[1].tail = &sub[1].list;
 i = 0;
 while (list)
  { *sub[i].tail = list;
    sub[i].tail = &list->link;
    list = list->link;
    i = ! i;
  }
 *sub[0].tail = 0;
 *sub[1].tail = 0;
 sub[0].list = sort_keys_by_preference(sub[0].list);
 sub[1].list = sort_keys_by_preference(sub[1].list);
 rev.tail = &rev.list;
 while (sub[0].list || sub[1].list)
  { i = !sub[0].list || (sub[1].list && prefer(sub[1].list,sub[0].list));
    *rev.tail = sub[i].list;
    rev.tail = &sub[i].list->link;
    sub[i].list = sub[i].list->link;
  }
 *rev.tail = 0;
 return(rev.list);
}

static unsigned char *build_key_query(CUASTATE *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);
 if (verbose) printf("Trying key `%s'\n",p->k->comment);
 return(opp);
}

static unsigned char *build_key_sig(CUASTATE *s, AUTHKEY *k)
{
 PKPRIV *p;
 BPP *b;
 unsigned char *opp;
 STR sig;
 HASHALG *ha;
 void *hh;
 char *hash;
 void *privclr;
 int privlen;

 p = s->algpriv;
 b = s->layer->b;
 ha = (*p->k->alg->prehash)();
 hh = (*ha->init)();
 opp = put_string(&b->opkt[0],b->sessid,b->sessidlen);
 (*ha->process)(hh,&b->opkt[0],opp-&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);
 sig.data = 0;
 privclr = 0;
 do <"done">
  { if (k->flags & AKF_AGENT)
     { void *sigdata;
       int siglen;
       if (! agent_sign(p->k->pubdata,p->k->publen,hash,ha->hashlen,&sigdata,&siglen))
	{ fprintf(stderr,"%s: agent signing failed\n",__progname);
	  opp = 0;
	  break <"done">;
	}
       sig.data = sigdata;
       sig.len = siglen;
     }
    else
     { if (! priv_unwrap(p->k->privdata,p->k->privlen,p->k->comment,&privclr,&privlen,passphrase))
	{ fprintf(stderr,"%s: can't unwrap private key\n",__progname);
	  opp = 0;
	  break <"done">;
	}
       if (! (*p->k->alg->sign)(&sig,p->k->pubdata,p->k->publen,privclr,privlen,hash))
	{ fprintf(stderr,"%s: can't generate authentication signature\n",__progname);
	  opp = 0;
	  break <"done">;
	}
     }
    opp = put_string(opp,sig.data,sig.len);
  } while (0);
 if (privclr) bzero(privclr,privlen);
 free(privclr);
 bzero(hash,ha->hashlen);
 free(hash);
 bzero(sig.data,sig.len);
 free(sig.data);
 return(opp);
}

static int uac_publickey_canuse(BPP *b __attribute__((__unused__)))
{
 static int loadkey(const char *file, int gripe)
  { void *pub;
    int publen;
    void *priv;
    int privlen;
    char *comment;
    FILE *errf;
    HKALG *alg;
    int atnl;
    int rv;
    static int w(void *cookie __attribute__((__unused__)), const char *buf, int len)
     { int i;
       if (gripe)
	{ for (i=0;i<len;i++)
	   { if (atnl) fprintf(stderr,"%s: ",__progname);
	     putc(buf[i],stderr);
	     atnl = (buf[i] == '\n');
	   }
	}
       return(len);
     }
    static int c(void *cookie __attribute__((__unused__)))
     { if (gripe && !atnl) putc('\n',stderr);
       return(0);
     }
    atnl = 1;
    errf = funopen(0,0,&w,0,&c);
    rv = 0;
    if (load_key_pair(file,&pub,&publen,&alg,&priv,&privlen,&comment,0,errf))
     { if (! alg->can_sig)
	{ fprintf(errf,"can't use key in %s (algorithm `%s' can't do signatures)",file,alg->name);
	  free(pub);
	  free(priv);
	  free(comment);
	}
       else
	{ AUTHKEY *ak;
	  ak = malloc(sizeof(AUTHKEY));
	  ak->link = authkeylist;
	  ak->alg = alg;
	  ak->pubdata = pub;
	  ak->publen = publen;
	  ak->flags = priv_no_pp(priv,privlen) ? AKF_NO_PP : 0;
	  ak->privdata = priv;
	  ak->privlen = privlen;
	  ak->comment = comment;
	  authkeylist = ak;
	  if (verbose > 1) printf("saving file key `%s' (from %s)\n",comment,file);
	  rv = 1;
	}
     }
    fclose(errf);
    return(rv);
  }

 static void load_agent_key(HKALG *alg, void *data, int len, char *desc)
  { AUTHKEY *ak;
    ak = malloc(sizeof(AUTHKEY));
    ak->link = authkeylist;
    ak->alg = alg;
    ak->pubdata = data;
    ak->publen = len;
    ak->flags = AKF_AGENT;
    ak->privdata = 0;
    ak->privlen = 0;
    ak->comment = desc;
    authkeylist = ak;
    if (verbose > 1) printf("saving agent key `%s'\n",desc);
  }

 if (nokeyauth) return(0);
 if (nauthkeys)
  { int i;
    for (i=0;i<nauthkeys;i++)
     { do <"done">
	{ if (index(authkeys[i],'/'))
	   { loadkey(authkeys[i],1);
	   }
	  else
	   { char *t;
	     if (! authkeydir)
	      { if (! sshdir)
		 { loadkey(authkeys[i],1);
		   break <"done">;
		 }
		authkeydir = sshdir;
	      }
	     asprintf(&t,"%s/%s",authkeydir,authkeys[i]);
	     if (! loadkey(t,0))
	      { if (! loadkey(authkeys[i],0))
		 { fprintf(stderr,"%s: can't load `%s' - try specifying a full pathname\n",__progname,authkeys[i]);
		   exit(1);
		 }
	      }
	   }
	} while (0);
     }
  }
 else
  { int i;
    HKALG *a;
    static void try(const char *suf)
     { char *t;
       asprintf(&t,"%s/%s",authkeydir,suf);
       loadkey(t,0);
       free(t);
     }
    if (! authkeydir)
     { if (! sshdir)
	{ fprintf(stderr,"%s: nowhere to find key files\n",__progname);
	}
       else
	{ authkeydir = sshdir;
	}
     }
    if (authkeydir) for (i=0;(a=hkalg_enum(i));i++) try(a->deffile);
  }
 if (! noagent) agent_list_keys(load_agent_key);
 authkeylist = sort_keys_by_preference(authkeylist);
 if (verbose)
  { if (authkeylist)
     { AUTHKEY *ak;
       printf("publickey authentication possible:");
       for (ak=authkeylist;ak;ak=ak->link) printf(" %s",ak->alg->name);
       printf("\n");
     }
    else
     { printf("publickey authentication not possible\n");
     }
  }
 return(!!authkeylist);
}

static unsigned char *uac_publickey_start(CUASTATE *s)
{
 PKPRIV *p;

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

static int uac_publickey_partial(CUASTATE *s __attribute__((__unused__)))
{
 abort();
}

static int try_next_key(CUASTATE *s)
{
 PKPRIV *p;

 p = s->algpriv;
 if (! p->k) return(UA_NEXT);
 s->layer->b->oplen = build_key_query(s) - &s->layer->b->opkt[0];
 return(UA_SEND);
}

static int uac_publickey_run(CUASTATE *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_IGNORE(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);
	}
	{ unsigned char *sig;
	  sig = build_key_sig(s,p->k);
	  if (! sig) return(try_next_key(s));
	  b->oplen = sig - &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 uac_publickey_retry(CUASTATE *s)
{
 PKPRIV *p;

 p = s->algpriv;
 p->k = p->k->link;
 return(try_next_key(s));
}

static void uac_publickey_done(CUASTATE *s)
{
 PKPRIV *p;

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

static int uas_publickey_canuse(SUASTATE *s __attribute__((__unused__)))
{
 return(1);
}

static int uas_publickey_request(SUASTATE *s, const void *rest, int restlen)
{
 __label__ parsefail;

 int siginc;
 STR pkalg;
 HKALG *alg;
 STR pkblob;
 STR sigblob;
 const void *sigrest;
 int sigrlen;
 char *akfn;
 int rv;
 LAYER *l;
 unsigned char *opp;

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

 static void clears(void)
  { free_str(pkalg);
    free_str(pkblob);
    free_str(sigblob);
  }

 if (0)
  {
parsefail:;
    fprintf(stderr,"publickey extra data parse failure\n");
    send_disconnect(SSH_DISCONNECT_PROTOCOL_ERROR,0,0,0,0);
    exit(0);
  }
 l = s->layer;
 parse_data(rest,restlen,&fail,
	PP_BOOL(&siginc),
	PP_STRING(&pkalg),
	PP_STRING(&pkblob),
	PP_REST(&sigrest,&sigrlen) );
 if (siginc)
  { parse_data(sigrest,sigrlen,&fail,PP_STRING(&sigblob),PP_ENDSHERE);
  }
 else
  { parse_data(sigrest,sigrlen,&fail,PP_ENDSHERE);
    sigblob = STRZERO;
  }
 alg = hkalg_vfind_rostr(str_to_rostr(pkalg));
 if (!alg || !alg->can_sig || s->nosuchuser)
  { send_failure(s,0);
    clears();
    return(0);
  }
 rv = 0;
 asprintf(&akfn,"%s/.ssh/authorized-keys",s->homedir);
  { static int k(HKALG *a, const void *blob, int bloblen, const char *comment __attribute__((__unused__)))
     { HASHALG *ha;
       char *hashtmp;
       void *hashv;
       int rv;
       unsigned char *op0;
       unsigned char *opp;
       if (a != alg) return(0);
       if (! (*alg->match)(blob,bloblen,pkblob.data,pkblob.len)) return(0);
       if (! siginc) return(-1);
       ha = (*alg->prehash)();
       hashtmp = malloc(ha->hashlen);
       hashv = (*ha->init)();
       op0 = &l->b->opkt[0];
       opp = put_string(op0,l->b->sessid,l->b->sessidlen);
       (*ha->process)(hashv,op0,opp-op0);
       *op0 = SSH_MSG_USERAUTH_REQUEST;
       (*ha->process)(hashv,op0,1);
       opp = put_string(op0,s->curuser.data,s->curuser.len);
       (*ha->process)(hashv,op0,opp-op0);
       opp = put_string(op0,s->service.data,s->service.len);
       (*ha->process)(hashv,op0,opp-op0);
       opp = put_string(op0,"publickey",-1);
       (*ha->process)(hashv,op0,opp-op0);
       *op0 = 1;
       (*ha->process)(hashv,op0,1);
       opp = put_string(op0,pkalg.data,pkalg.len);
       (*ha->process)(hashv,op0,opp-op0);
       opp = put_string(op0,pkblob.data,pkblob.len);
       (*ha->process)(hashv,op0,opp-op0);
       (*ha->done)(hashv,hashtmp);
       rv = (*alg->checksig)(pkblob.data,pkblob.len,hashtmp,sigblob.data,sigblob.len) ? -3 : -2;
       bzero(hashtmp,ha->hashlen);
       free(hashtmp);
       return(rv);
     }
    switch (read_pubkeys_file(akfn,&k))
     { case 0: /* no matching key found */
	  if (verbose) printf("pkauth: no matching key found\n");
	  send_failure(s,0);
	  break;
       case -1: /* matching key found, no signature in packet */
	  if (verbose) printf("pkauth: matching key found, no sig in packet\n");
	  l->b->opkt[0] = SSH_MSG_USERAUTH_PK_OK;
	  opp = put_string(&l->b->opkt[1],pkalg.data,pkalg.len);
	  opp = put_string(opp,pkblob.data,pkblob.len);
	  l->b->oplen = opp - &l->b->opkt[0];
	  below_opkt(l);
	  break;
       case -2: /* matching key found, signature wrong */
	  if (verbose) printf("pkauth: matching key found, sig wrong\n");
	  send_failure(s,0);
	  break;
       case -3: /* matching key found, signature right */
	  if (verbose) printf("pkauth: matching key found, sig right\n");
	  rv = 1;
	  break;
     }
  }
 free_str(pkalg);
 free_str(pkblob);
 free_str(sigblob);
 return(rv);
}

UAALG uaalg_publickey
 = { "publickey",
     { &uac_publickey_canuse,
       &uac_publickey_start,
       &uac_publickey_partial,
       &uac_publickey_run,
       &uac_publickey_retry,
       &uac_publickey_done },
     { &uas_publickey_canuse,
       &uas_publickey_request } };