#include <sha1.h>
#include <stdio.h>
#include <errno.h>
#include <stdlib.h>
#include <stdarg.h>
#include <unistd.h>
#include <sys/un.h>
#include <strings.h>
#include <sys/uio.h>
#include <sys/socket.h>

extern const char *__progname;

#include "pp.h"
#include "b64.h"
#include "msgs.h"
#include "cmdline.h"
#include "keyfiles.h"
#include "alg-util.h"
#include "pkt-util.h"
#include "stdio-util.h"
#include "priv-crypto.h"

#include "agent-client.h"

static int tried = 0;
static int conn;

static inline unsigned int scr_to_uint32(unsigned int v) { return(v); }
static inline const void *scr_to_const_void_p(const void *v) { return(v); }
static inline int scr_to_int(int v) { return(v); }

static void agent_exchange(void **rdp, int *rlp, int pkttype, ...)
#define T_END T__END
#define T__END 1
#define T_UINT32(v) T__UINT32,scr_to_uint32(v)
#define T__UINT32 2
#define T_STRING(d,l) T__STRING,scr_to_const_void_p(d),scr_to_int(l)
#define T__STRING 3
{
 va_list ap;
 char *obuf;
 int olen;
 FILE *f;
 int t;
 unsigned char ilenbuf[4];
 char *ibuf;
 unsigned int ilen;

 f = fopen_alloc(&obuf,&olen);
 fwrite("\0\0\0\0",1,4,f);
 putc(pkttype,f);
 va_start(ap,pkttype);
 while <"args"> (1)
  { t = va_arg(ap,int);
    switch (t)
     { case T__END:
	  break <"args">;
       case T__UINT32:
	  fput_uint32(f,va_arg(ap,unsigned int));
	  break;
       case T__STRING:
	   { const void *data;
	     int len;
	     data = va_arg(ap,const void *);
	     len = va_arg(ap,int);
	     fput_string(f,data,len);
	   }
	  break;
       default:
	  abort();
	  break;
     }
  }
 va_end(ap);
 fclose(f);
 if (olen < 5) abort();
 put_uint32(obuf,olen-4);
 write(conn,obuf,olen);
 t = recv(conn,&ilenbuf[0],4,MSG_WAITALL);
 if (t < 0)
  { fprintf(stderr,"%s: agent connection recv: %s\n",__progname,strerror(errno));
    exit(1);
  }
 if (t < 4)
  { fprintf(stderr,"%s: short agent connection read: %d < 4\n",__progname,t);
    exit(1);
  }
 ilen = get_uint32(&ilenbuf[0]);
 ibuf = malloc(ilen);
 if (! ibuf)
  { fprintf(stderr,"%s: agent error: reply malloc(%u) failed\n",__progname,ilen);
    exit(1);
  }
 t = recv(conn,ibuf,ilen,MSG_WAITALL);
 if (t < 0)
  { fprintf(stderr,"%s: agent connection recv: %s\n",__progname,strerror(errno));
    exit(1);
  }
 if (t < ilen)
  { fprintf(stderr,"%s: short agent connection read: %d < %u\n",__progname,t,ilen);
    exit(1);
  }
 *rdp = ibuf;
 *rlp = ilen;
}

static int connect_to_agent(int mustwork)
{
 __label__ failret;

 char *txt;
 FILE *f;
 void *blob;
 int blen;
 const unsigned char *cookie;
 int clen;
 unsigned char challenge[16];
 int plen;
 struct sockaddr_un *s_un;
 int slen;
 int s;
 int pipe[2];
 int n;
 void *h;
 unsigned char hash[20];
 void *reply;
 int rlen;
 unsigned int v;
 struct msghdr mh;
 struct cmsghdr cmh;
 struct iovec iov;
 unsigned char ctlbuf[sizeof(struct cmsghdr)+sizeof(int)];

 static void fail(const char *, ...)
		__attribute__((__format__(__printf__,1,2),__noreturn__));
 static void fail(const char *fmt, ...)
  { va_list ap;
    if (s >= 0) close(s);
    if (pipe[0] >= 0) close(pipe[0]);
    if (pipe[1] >= 0) close(pipe[1]);
    free(blob);
    free(s_un);
    free(reply);
    if (mustwork)
     { fprintf(stderr,"%s: can't connect to agent: ",__progname);
       va_start(ap,fmt);
       vfprintf(stderr,fmt,ap);
       va_end(ap);
       fprintf(stderr,"\n");
       exit(1);
     }
    goto failret;
  }

 if (0)
  {
failret:;
    conn = -1;
    return(0);
  }
 if (tried) return(conn>=0);
 tried = 1;
 s = -1;
 pipe[0] = -1;
 pipe[1] = -1;
 blob = 0;
 s_un = 0;
 reply = 0;
 txt = getenv("SSH_AGENT");
 if (txt == 0) fail("no SSH_AGENT in the environment");
 f = fopen_rstr(txt,-1);
 if (! b64_read_blob(f,&blob,&blen))
  { fclose(f);
    fail("can't parse $SSH_AGENT");
  }
 fclose(f);
 if (blen < 1) fail("$SSH_AGENT too short (blen %d)",blen);
 if (((unsigned char *)blob)[0] != 1) fail("wrong $SSH_AGENT version");
 if (blen < 2) fail("$SSH_AGENT too short (blen %d)",blen);
 clen = ((unsigned char *)blob)[1];
 cookie = ((unsigned char *)blob) + 2;
 plen = blen - (1+1+clen);
 if (plen < 1) fail("$SSH_AGENT too short (plen %d)",plen);
 slen = sizeof(*s_un) - sizeof(s_un->sun_path) + plen + 1;
 s_un = malloc(slen);
 s_un->sun_family = AF_LOCAL;
 s_un->sun_len = slen;
 bcopy(((unsigned char *)blob)+1+1+clen,&s_un->sun_path[0],plen);
 s_un->sun_path[plen] = '\0';
 s = socket(AF_LOCAL,SOCK_DGRAM,0);
 if (s < 0)
  { fprintf(stderr,"%s: socket (AF_LOCAL SOCK_DGRAM): %s\n",__progname,strerror(errno));
    exit(1);
  }
 if (socketpair(AF_LOCAL,SOCK_STREAM,0,&pipe[0]) < 0)
  { fprintf(stderr,"%s: socketpair (AF_LOCAL SOCK_STREAM): %s\n",__progname,strerror(errno));
    exit(1);
  }
 n = 0;
 iov.iov_base = &n;
 iov.iov_len = 1;
 cmh.cmsg_len = sizeof(struct cmsghdr) + sizeof(int);
 cmh.cmsg_level = SOL_SOCKET;
 cmh.cmsg_type = SCM_RIGHTS;
 bcopy(&cmh,&ctlbuf[0],sizeof(struct cmsghdr));
 bcopy(&pipe[1],&ctlbuf[sizeof(struct cmsghdr)],sizeof(int));
 mh.msg_name = (void *) s_un;
 mh.msg_namelen = slen;
 mh.msg_iov = &iov;
 mh.msg_iovlen = 1;
 mh.msg_control = (void *) &ctlbuf[0];
 mh.msg_controllen = sizeof(struct cmsghdr) + sizeof(int);
 mh.msg_flags = 0;
 if (sendmsg(s,&mh,0) < 0) fail("sendmsg: %s",strerror(errno));
 free(s_un);
 close(pipe[1]);
 close(s);
 s = pipe[0];
 pipe[0] = -1;
 pipe[1] = -1;
 n = recv(s,&challenge[0],16,MSG_WAITALL);
 if (n < 0) fail("challenge recv: %s",strerror(errno));
 if (n < 16) fail("challenge too short (%d < 16)",n);
 h = sha1_init();
 for (n=0;n<16;n++)
  { sha1_process_bytes(h,cookie,clen);
    sha1_process_bytes(h,&challenge[0],16);
  }
 sha1_result(h,&hash[0]);
 xor_block(&hash[0],&hash[10],&hash[0],10);
 write(s,&hash[0],10);
 conn = s;
 agent_exchange(&reply,&rlen,SSH_AGENT_REQUEST_VERSION,T_END);
 if (rlen != 5) fail("version response length wrong (%d != 5)",rlen);
 if (((unsigned char *)reply)[0] != SSH_AGENT_VERSION_RESPONSE)fail("version response message type wrong (%d != %d)",((unsigned char *)reply)[0],SSH_AGENT_VERSION_RESPONSE);
 v = get_uint32(((unsigned char *)reply)+1);
 if (v != 2) fail("version number wrong (%u != 2)",v);
 free(reply);
 return(1);
}

static int success_failure_response(void *reply, int rlen, const char *tag)
{
 if (rlen < 1)
  { fprintf(stderr,"%s: %s response too short (len %d)\n",__progname,tag,rlen);
    return(0);
  }
 switch (((unsigned char *)reply)[0])
  { case SSH_AGENT_SUCCESS:
       if (rlen != 1)
	{ fprintf(stderr,"%s: %s SUCCESS response length wrong (len %d)\n",__progname,tag,rlen);
	  return(0);
	}
       break;
    case SSH_AGENT_FAILURE:
       if (rlen != 5)
	{ fprintf(stderr,"%s: ADD_KEY FAILURE response length wrong (len %d)\n",__progname,rlen);
	  return(0);
	}
       else
	{ unsigned int v;
	  v = get_uint32(((unsigned char *)reply)+1);
	  fprintf(stderr,"%s: %s failed: ",__progname,tag);
	  switch (v)
	   { case SSH_AGENT_ERROR_TIMEOUT:
		fprintf(stderr,"timeout");
		break;
	     case SSH_AGENT_ERROR_KEY_NOT_FOUND:
		fprintf(stderr,"key not found");
		break;
	     case SSH_AGENT_ERROR_DECRYPT_FAILED:
		fprintf(stderr,"decrypt failed");
		break;
	     case SSH_AGENT_ERROR_SIZE_ERROR:
		fprintf(stderr,"size error");
		break;
	     case SSH_AGENT_ERROR_KEY_NOT_SUITABLE:
		fprintf(stderr,"key not suitable");
		break;
	     case SSH_AGENT_ERROR_DENIED:
		fprintf(stderr,"denied");
		break;
	     case SSH_AGENT_ERROR_FAILURE:
		fprintf(stderr,"failure");
		break;
	     case SSH_AGENT_ERROR_UNSUPPORTED_OP:
		fprintf(stderr,"unsupported op");
		break;
	     default:
		fprintf(stderr,"unknown error %u",v);
		break;
	   }
	  fprintf(stderr,"\n");
	  return(0);
	}
       break;
    default:
       fprintf(stderr,"%s: %s provoked unknown response %d\n",__progname,tag,((unsigned char *)reply)[0]);
       return(0);
       break;
  }
 return(1);
}

void agent_add(void)
{
 int i;
 void *pub;
 int publen;
 HKALG *alg;
 void *priv;
 int privlen;
 char *comment;
 char *file;
 const char *sshdir;

 if (! argsbegin)
  { /* XXX should iterate with hkalg_enum and use ->deffile */
    fprintf(stderr,"%s: must specify at least one argument keyfile\n",__progname);
    exit(1);
  }
 sshdir = 0;
 for (i=0;i<argscount;i++)
  { char *t;
    int atnl;
    FILE *errf;
    static int w(void *cookie __attribute__((__unused__)), const char *buf, int len)
     { int i;
       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 (! atnl) putc('\n',stderr);
       return(0);
     }
    file = argsbegin[i];
    if (index(argsbegin[i],'/'))
     { t = 0;
     }
    else
     { if (! sshdir) sshdir = get_sshdir("key-file directory");
       asprintf(&t,"%s/%s",sshdir,file);
       file = t;
     }
    atnl = 1;
    errf = funopen(0,0,&w,0,&c);
    if (load_key_pair(file,&pub,&publen,&alg,&priv,&privlen,&comment,0,errf))
     { void *privclear;
       int privclen;
       if (priv_unwrap(priv,privlen,comment,&privclear,&privclen,passphrase))
	{ void *reply;
	  int rlen;
	  connect_to_agent(1);
	  agent_exchange(&reply,&rlen,SSH_AGENT_ADD_KEY,
		T_STRING(privclear,privclen),
		T_STRING(pub,publen),
		T_STRING(comment?:"",comment?strlen(comment):0),
		/* XXX add constraints */
		T_END );
	  success_failure_response(reply,rlen,"ADD_KEY");
	  free(reply);
	  bzero(privclear,privlen);
	  free(privclear);
	}
       free(pub);
       bzero(priv,privlen);
       free(priv);
       free(comment);
     }
    fclose(errf);
    free(t);
  }
}

void agent_list(void)
{
 int n;

 static void printkey(HKALG *alg, void *pubdata, int publen, char *comment)
  { if (alg)
     { write_pub_to_fd("standard","output",1,alg,pubdata,publen,comment);
     }
    else
     { int i;
       printf("Unsupported: ");
       for (i=0;i<publen;i++) printf("%02x",((const unsigned char *)pubdata)[i]);
       printf("\n");
     }
    free(pubdata);
    free(comment);
    n ++;
  }

 if (argsbegin)
  { /* XXX should do some kind of matching */
    fprintf(stderr,"%s: must not specify any arguments\n",__progname);
    exit(1);
  }
 connect_to_agent(1);
 n = 0;
 agent_list_keys(printkey);
 if (n == 0) printf("The agent holds no keys.\n");
}

void agent_delete(void)
{
 int i;
 void *pub;
 int publen;
 HKALG *alg;
 char *comment;
 char *file;
 const char *sshdir;

 if (! argsbegin)
  { /* XXX should iterate with hkalg_enum and use ->deffile */
    fprintf(stderr,"%s: must specify at least one argument keyfile\n",__progname);
    exit(1);
  }
 sshdir = 0;
 for (i=0;i<argscount;i++)
  { char *t;
    int atnl;
    FILE *errf;
    static int w(void *cookie __attribute__((__unused__)), const char *buf, int len)
     { int i;
       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 (! atnl) putc('\n',stderr);
       return(0);
     }
    file = argsbegin[i];
    if (index(argsbegin[i],'/'))
     { t = 0;
     }
    else
     { if (! sshdir) sshdir = get_sshdir("key-file directory");
       asprintf(&t,"%s/%s",sshdir,file);
       file = t;
     }
    atnl = 1;
    errf = funopen(0,0,&w,0,&c);
    if (load_key_pair(file,&pub,&publen,&alg,0,0,&comment,0,errf))
     { void *reply;
       int rlen;
       connect_to_agent(1);
       agent_exchange(&reply,&rlen,SSH_AGENT_DELETE_KEY,
	     T_STRING(pub,publen),
	     T_STRING(comment?:"",comment?strlen(comment):0),
	     T_END );
       success_failure_response(reply,rlen,"DELETE_KEY");
       free(reply);
       free(pub);
       free(comment);
     }
    fclose(errf);
    free(t);
  }
}

void agent_delete_all(void)
{
 void *reply;
 int rlen;

 if (argsbegin)
  { fprintf(stderr,"%s: must not specify any arguments\n",__progname);
    exit(1);
  }
 connect_to_agent(1);
 agent_exchange(&reply,&rlen,SSH_AGENT_DELETE_ALL_KEYS,T_END);
 success_failure_response(reply,rlen,"DELETE_ALL_KEYS");
 free(reply);
}

void agent_list_keys(void (*fn)(HKALG *, void *, int, char *))
{
 int i;
 HKALG *alg;
 void *reply;
 int rlen;
 const void *rest;
 int restlen;
 unsigned int nkeys;

 static void fail(const void *dp __attribute__((__unused__)), const char *fmt, ...)
  { va_list ap;
    fprintf(stderr,"%s: LIST_KEYS response invalid: ",__progname);
    va_start(ap,fmt);
    vfprintf(stderr,fmt,ap);
    va_end(ap);
    fprintf(stderr,"\n");
    exit(1);
  }

 if (! connect_to_agent(0)) return;
 agent_exchange(&reply,&rlen,SSH_AGENT_LIST_KEYS,T_END);
 if (rlen < 1)
  { fprintf(stderr,"%s: LIST_KEYS response too short (len %d)\n",__progname,rlen);
    exit(1);
  }
 if (((unsigned char *)reply)[0] != SSH_AGENT_KEY_LIST)
  { fprintf(stderr,"%s: LIST_KEYS response type wrong (%d)\n",__progname,((unsigned char *)reply)[0]);
    exit(1);
  }
 parse_data(reply,rlen,fail,
	PP_IGNORE(1),
	PP_UINT32(&nkeys),
	PP_REST(&rest,&restlen) );
 for (;nkeys>0;nkeys--)
  { STR pkblob;
    STR desc;
    parse_data(rest,restlen,fail,
	PP_STRING(&pkblob),
	PP_STRING(&desc),
	PP_REST(&rest,&restlen) );
    do <"handled">
     { char *t;
       t = malloc(desc.len+1);
       bcopy(desc.data,t,desc.len);
       t[desc.len] = '\0';
       for (i=0;(alg=hkalg_enum(i));i++)
	{ if ((*alg->checkpub)(pkblob.data,pkblob.len))
	   { (*fn)(alg,pkblob.data,pkblob.len,t);
	     break <"handled">;
	   }
	}
       (*fn)(0,pkblob.data,pkblob.len,t);
     } while (0);
    free_str(desc);
  }
}

int agent_sign(const void *keyd, int keyl, const void *hashd, int hashl, void **sigdp, int *siglp)
{
 void *reply;
 int rlen;
 STR sig;

 static void fail(const void *dp __attribute__((__unused__)), const char *fmt, ...)
  { va_list ap;
    fprintf(stderr,"%s: PRIVATE_KEY_OP response invalid: ",__progname);
    va_start(ap,fmt);
    vfprintf(stderr,fmt,ap);
    va_end(ap);
    fprintf(stderr,"\n");
    exit(1);
  }

 agent_exchange(&reply,&rlen,SSH_AGENT_PRIVATE_KEY_OP,
	T_STRING("sign",-1),
	T_STRING(keyd,keyl),
	T_STRING(hashd,hashl),
	T_END );
 if (rlen < 1)
  { fprintf(stderr,"%s: PRIVATE_KEY_OP response too short (len %d)\n",__progname,rlen);
    exit(1);
  }
 switch (((unsigned char *)reply)[0])
  { case SSH_AGENT_OPERATION_COMPLETE:
       parse_data(reply,rlen,fail,
		PP_BACKOUT(reply),
		PP_IGNORE(1),
		PP_STRING(&sig),
		PP_ENDSHERE );
       *sigdp = sig.data;
       *siglp = sig.len;
       return(1);
       break;
    case SSH_AGENT_FAILURE:
       bzero(reply,rlen);
       free(reply);
       return(0);
       break;
  }
 fail(0,"wrong message type %d",((unsigned char *)reply)[0]);
 return(0);
}