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

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

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

static AUTHKEY *authkeylist;

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, 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 ua_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 = 0;
	  ak->privdata = priv;
	  ak->privlen = privlen;
	  ak->comment = comment;
	  authkeylist = ak;
	  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 (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);
 if (authkeys)
  { 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(!!authkeys);
}

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

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

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

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

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

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_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 ua_publickey_retry(UASTATE *s)
{
 return(try_next_key(s));
}

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