#include <stdio.h>
#include <errno.h>
#include <netdb.h>
#include <fcntl.h>
#include <stdlib.h>
#include <unistd.h>
#include <strings.h>
#include <sys/uio.h>
#include <sys/socket.h>
#include <sys/sysctl.h>

#include "pp.h"
#include "rnd.h"
#include "bpp.h"
#include "algs.h"
#include "msgs.h"
#include "remote.h"
#include "keygen.h"
#include "dequal.h"
#include "cmdline.h"
#include "pollloop.h"
#include "channels.h"
#include "pkt-util.h"
#include "transport.h"
#include "userauth-conn.h"

extern const char *__progname;

static BPP bpp;
static LAYER *disconn_layer;

static void openconn(void)
{
 struct addrinfo hints;
 struct addrinfo *ai0;
 struct addrinfo *ai;
 int e;
 int se;
 int fd;

 hints.ai_flags = 0;
 hints.ai_family = PF_UNSPEC;
 hints.ai_socktype = SOCK_STREAM;
 hints.ai_protocol = 0;
 hints.ai_addrlen = 0;
 hints.ai_canonname = 0;
 hints.ai_addr = 0;
 hints.ai_next = 0;
 if (! host)
  { fprintf(stderr,"%s: need a host to connect to\n",__progname);
    exit(1);
  }
 if (! port) port = "22";
 e = getaddrinfo(host,port,&hints,&ai0);
 if (e)
  { fprintf(stderr,"%s: %s/%s: %s\n",__progname,host,port,gai_strerror(e));
    exit(1);
  }
 if (! ai0)
  { fprintf(stderr,"%s: %s/%s: successful lookup but no addresses?\n",__progname,host,port);
    exit(1);
  }
 se = 0;
 for (ai=ai0;ai;ai=ai->ai_next)
  { fd = socket(ai->ai_family,ai->ai_socktype,ai->ai_protocol);
    if (fd < 0)
     { if (se == 0) se = errno;
       continue;
     }
    se = -1;
    if (connect(fd,ai->ai_addr,ai->ai_addrlen) < 0)
     { char hnbuf[NI_MAXHOST];
       char pnbuf[NI_MAXSERV];
       e = errno;
       if (getnameinfo(ai->ai_addr,ai->ai_addrlen,&hnbuf[0],NI_MAXHOST,&pnbuf[0],NI_MAXSERV,NI_NUMERICHOST|NI_NUMERICSERV))
	{ fprintf(stderr,"%s: %s%s%s: connect failed [%s], can't get numeric hostname info [%s]\n",__progname,host,port?"/":"",port?:"",strerror(e),strerror(errno));
	}
       else
	{ if (ai0->ai_next)
	   { fprintf(stderr,"%s: connect %s [%s/%s]: %s\n",__progname,host,&hnbuf[0],&pnbuf[0],strerror(e));
	   }
	  else
	   { fprintf(stderr,"%s: connect %s/%s: %s\n",__progname,&hnbuf[0],&pnbuf[0],strerror(e));
	   }
	}
       close(fd);
       continue;
     }
    freeaddrinfo(ai0);
    fcntl(fd,F_SETFL,fcntl(fd,F_GETFL,0)|O_NONBLOCK);
    bpp.fd = fd;
    return;
  }
 if (se > 0)
  { fprintf(stderr,"%s: socket: %s\n",__progname,strerror(se));
  }
 else if (se == 0)
  { fprintf(stderr,"%s: can't connect, can't tell why\n",__progname);
  }
 exit(1);
}

static void queue_block_common(BPP *b, OQB *q, int len)
{
 q->left = len;
 q->link = 0;
 *b->outtail = q;
 b->outtail = &q->link;
}

static void queue_block_point(BPP *b, const void *data, int len)
{
 OQB *q;

 q = malloc(sizeof(OQB));
 q->data = data;
 q->datafree = 0;
 queue_block_common(b,q,len);
}

static void queue_block_copy(BPP *b, const void *data, int len)
{
 OQB *q;

 q = malloc(sizeof(OQB)+len);
 q->data = (void *) (q+1);
 q->datafree = 0;
 bcopy(data,q+1,len);
 queue_block_common(b,q,len);
}

#if 0
static void queue_block_free(BPP *b, void *data, int len)
{
 OQB *q;

 q = malloc(sizeof(OQB));
 q->data = data;
 q->datafree = data;
 queue_block_common(b,q,len);
}
#endif

static void sendversion(void)
{
 strcpy(&bpp.c_version[0],"SSH-2.0-Moussh0.1 (fnord!)");
 bpp.c_vlen = strlen(&bpp.c_version[0]);
 queue_block_point(&bpp,&bpp.c_version[0],bpp.c_vlen);
 queue_block_point(&bpp,"\r\n",2);
}

static void svs_save(BPP *b, unsigned char ch)
{
 if (b->s_vlen >= SSH_MAX_VERSION_LEN)
  { fprintf(stderr,"%s: peer version string too long\n",__progname);
    exit(1);
  }
 b->s_version[b->s_vlen++] = ch;
}

static void bottom_i(BPP *b, const void *data, int len)
{
 int n;
 int prevril;

 static void wraw(const void *d, int l)
  { if (b->pkt.ril+len > SSH_MAX_PACKET_LEN) abort();
    bcopy(d,&b->rawipkt[b->pkt.ril],l);
    b->pkt.ril += l;
  }

 static void wpay(const void *d, int l)
  { if (b->iplen+len > SSH_MAX_PAYLOAD_LEN)
     { fprintf(stderr,"%s: bad packet on wire (payload overflow %d+%d > %d)\n",__progname,b->iplen,len,SSH_MAX_PAYLOAD_LEN);
       exit(1);
     }
    bcopy(d,&b->ipkt[b->iplen],l);
    b->iplen += l;
  }

 while (len > 0)
  { if (b->pkt.ril < 5)
     { n = encalg_i(b->r_enc,b->r_encstate,wraw,data,len,5-b->pkt.ril);
       data = ((const char *)data) + n;
       len -= n;
       if (b->pkt.ril >= 5)
	{ int pktlen;
	  int padlen;
	  int unit;
	  pktlen = get_uint32(&b->rawipkt[0]);
	  padlen = b->rawipkt[4];
	  if ( (pktlen < 5) ||
	       (pktlen > SSH_MAX_PACKET_LEN-4-b->r_mac->maclen) )
	   { fprintf(stderr,"%s: bad packet on wire (length %#x)\n",__progname,pktlen);
	     exit(1);
	   }
	  if (padlen < 4)
	   { fprintf(stderr,"%s: bad packet on wire (too little padding %u)\n",__progname,padlen);
	     exit(1);
	   }
	  if (padlen+1 >= pktlen)
	   { fprintf(stderr,"%s: bad packet on wire (no payload: pktlen %u padlen %u)\n",__progname,pktlen,padlen);
	     exit(1);
	   }
	  unit = b->r_enc->blksize;
	  if (unit < 8) unit = 8;
	  if ((4+pktlen) % unit)
	   { fprintf(stderr,"%s: bad packet on wire (length %u not a multiple of %d)\n",__progname,4+pktlen,unit);
	     exit(1);
	   }
	  b->pkt.riw = 4 + pktlen;
	  b->pkt.rit = b->pkt.riw + b->r_mac->maclen;
	  b->pkt.pay = pktlen - padlen - 1;
	}
     }
    else
     { prevril = b->pkt.ril;
       if (b->pkt.ril < b->pkt.riw)
	{ n = encalg_i(b->r_enc,b->r_encstate,wraw,data,len,b->pkt.riw-b->pkt.ril);
	  data = ((const char *)data) + n;
	  len -= n;
	}
       else if (b->pkt.ril < b->pkt.rit)
	{ n = b->pkt.rit - b->pkt.ril;
	  if (n > len) n = len;
	  wraw(data,n);
	  data = ((const char *)data) + n;
	  len -= n;
	}
       if ((prevril < b->pkt.riw) && (b->pkt.ril >= b->pkt.riw))
	{ b->iplen = 0;
	  (*b->r_comp->process)(b->r_compstate,&b->rawipkt[5],b->pkt.pay,wpay);
	  (*b->r_comp->flush)(b->r_compstate,wpay);
	}
       if ((prevril < b->pkt.rit) && (b->pkt.ril >= b->pkt.rit))
	{ unsigned char seqbuf[4];
	  b->r_seq ++;
	  put_uint32(&seqbuf[0],b->r_seq);
	  if (! (*b->r_mac->check)(b->r_macstate,&b->rawipkt[b->pkt.rit-b->r_mac->maclen],2,
				 (const void *)&seqbuf[0],4,
				 (const void *)&b->rawipkt[0],b->pkt.riw))
	   { fprintf(stderr,"%s: bad packet on wire (MAC failure)\n",__progname);
	     exit(1);
	   }
	  if (b->ipkts_to_rekey-- <= 0) set_rekey(b);
	  if ((b->ibytes_to_rekey-=b->pkt.riw) <= 0) set_rekey(b);
	  (*b->botlayer->ipkt)(b->botlayer,b->botlayer->arg);
	  b->pkt.ril = 0;
	}
     }
  }
}

void set_rekey(BPP *b)
{
 b->want_rekey = 1;
}

static void getversion(BPP *b, const void *data, int len)
{
 const unsigned char *dp;
 int c;

 do <"gotvers">
  { dp = data;
    while (len > 0)
     { c = *dp++;
       len --;
       switch <"svs.state"> (b->svs.state)
	{ case SVS_ATNL:
	     switch (c)
	      { case 'S':  b->svs.state = SVS_GOT_S;    break;
		case '\n': b->svs.state = SVS_ATNL;     break;
		default:   b->svs.state = SVS_FLUSHING; break;
	      }
	     break;
	  case SVS_GOT_S:
	     switch (c)
	      { case 'S':  b->svs.state = SVS_GOT_SS;   break;
		case '\n': b->svs.state = SVS_ATNL;     break;
		default:   b->svs.state = SVS_FLUSHING; break;
	      }
	     break;
	  case SVS_GOT_SS:
	     switch (c)
	      { case 'H':  b->svs.state = SVS_GOT_SSH;  break;
		case '\n': b->svs.state = SVS_ATNL;     break;
		default:   b->svs.state = SVS_FLUSHING; break;
	      }
	     break;
	  case SVS_GOT_SSH:
	     switch (c)
	      { case '-':  b->svs.state = SVS_GOT_SSH_; break;
		case '\n': b->svs.state = SVS_ATNL;     break;
		default:   b->svs.state = SVS_FLUSHING; break;
	      }
	     break;
	  case SVS_GOT_CR:
	     switch (c)
	      { case '\n': break <"gotvers">;
		default:
		   svs_save(b,'\r');
	  case <"svs.state"> SVS_GOT_SSH_:
		   switch (c)
		    { case '\n': break <"gotvers">;
		      case '\r': b->svs.state = SVS_GOT_CR; break;
		      default: svs_save(b,c); break;
		    }
		   break;
	      }
	     break;
	  case SVS_FLUSHING:
	     switch (c)
	      { case '\n': b->svs.state = SVS_ATNL; break;
	      }
	     break;
	}
     }
    return;
  } while (0);
 printf("remote version string: %.*s\n",b->s_vlen,&b->s_version[0]);
 if (! ( ((b->s_vlen >= 8) && !bcmp(&b->s_version[0],"SSH-2.0-",8)) ||
	 ((b->s_vlen >= 9) && !bcmp(&b->s_version[0],"SSH-1.99-",9)) ) )
  { fprintf(stderr,"%s: unrecognized peer version\n",__progname);
    exit(1);
  }
 b->input = &bottom_i;
 b->pkt.ril = 0;
 if (len) bottom_i(b,dp,len);
}

#if 0
static void rd_fp(void *fvp, void *buf, int len)
{
 fread(buf,1,len,fvp);
}
#endif

#if 0
static void wr_fp(void *fvp, const void *buf, int len)
{
 fwrite(buf,1,len,fvp);
}
#endif

#if 0
static void toplayer_r(BPP *b)
{
 return((*b->toplayer->rpkt)(b->toplayer,b->toplayer->arg));
}
#endif

#if 0
static void toplayer_w(BPP *b)
{
 (*b->toplayer->wpkt)(b->toplayer,b->toplayer->arg);
}
#endif

void above_ipkt(LAYER *l)
{
 (*l->above->ipkt)(l->above,l->above->arg);
}

void below_opkt(LAYER *l)
{
 (*l->below->opkt)(l->below,l->below->arg);
}

void send_disconnect(int reason, const void *desc, int desclen, const void *lang, int langlen)
{
 unsigned char *opp;

 opp = &bpp.opkt[0];
 *opp++ = SSH_MSG_DISCONNECT;
 opp = put_uint32(opp,reason);
 opp = put_string(opp,desc,desclen);
 opp = put_string(opp,lang,langlen);
 bpp.oplen = opp - &bpp.opkt[0];
 (*disconn_layer->opkt)(disconn_layer,disconn_layer->arg);
}

static void bypass_i(LAYER *l, void *arg __attribute__((__unused__)))
{
 above_ipkt(l);
}

static void bypass_o(LAYER *l, void *arg __attribute__((__unused__)))
{
 below_opkt(l);
}

static void push_layer(BPP *b, const LAYERDESC *ld)
{
 LAYER *l;

 l = malloc(sizeof(LAYER));
 l->b = b;
 l->ipkt = ld->ifn ?: bypass_i;
 l->opkt = ld->ofn ?: bypass_o;
 l->below = b->toplayer;
 l->above = 0;
 if (l->below) l->below->above = l; else b->botlayer = l;
 b->toplayer = l;
 l->arg = ld->init ? (*ld->init)(l) : 0;
}

static void *base_init(LAYER *l)
{
 disconn_layer = l;
 return(0);
}

static void base_o(LAYER *l, void *arg __attribute__((__unused__)))
{
 BPP *b;
 unsigned int padlen;
 unsigned int padmod;
 unsigned int paylen;
 unsigned int pktlen;
 unsigned int n;
 unsigned char seqbuf[4];

 static void wdata(const void *data, int len)
  { queue_block_copy(b,data,len);
  }

 static void writepay(const void *data, int len)
  { if (5+paylen+len+b->w_mac->maclen > SSH_MAX_PACKET_LEN)
     { fprintf(stderr,"%s: compressed data too long (%d+%d+%d+%d > %d)\n",__progname,5,paylen,len,b->w_mac->maclen,SSH_MAX_PACKET_LEN);
       exit(1);
     }
    bcopy(data,&b->rawopkt[5+paylen],len);
    paylen += len;
  }

 b = l->b;
 if (b->oplen > SSH_MAX_PAYLOAD_LEN)
  { fprintf(stderr,"%s: writing too-large packet (payload length %u)\n",__progname,b->oplen);
    exit(1);
  }
 paylen = 0;
 (*b->w_comp->process)(b->w_compstate,&b->opkt[0],b->oplen,writepay);
 (*b->w_comp->flush)(b->w_compstate,writepay);
 padmod = b->w_enc->blksize;
 if (padmod < 8) padmod = 8;
 padlen = 4;
 pktlen = 5 + paylen + padlen;
 n = pktlen % padmod;
 if (n) padlen += padmod - n;
 if (pktlen < 16)
  { padlen = 16 - (5 + paylen);
    pktlen = 5 + paylen + padlen;
    n = pktlen % padmod;
    if (n) padlen += padmod - n;
  }
 pktlen = 5 + paylen + padlen;
 random_data(&b->rawopkt[5+paylen],padlen);
 put_uint32(&b->rawopkt[0],pktlen-4);
 b->rawopkt[4] = padlen;
 b->w_seq ++;
 put_uint32(&seqbuf[0],b->w_seq);
 (*b->w_mac->gen)(b->w_macstate,&b->rawopkt[pktlen],2,
			(const void *)&seqbuf[0],4,
			(const void *)&b->rawopkt[0],pktlen);
 encalg_o(b->w_enc,b->w_encstate,wdata,&b->rawopkt[0],pktlen);
 if (b->w_mac->maclen) wdata(&b->rawopkt[pktlen],b->w_mac->maclen);
 if (b->opkts_to_rekey-- <= 0) set_rekey(b);
 if ((b->obytes_to_rekey-=pktlen) <= 0) set_rekey(b);
}

static LAYERDESC layer_base = { &base_init, 0, &base_o };

static void catchall_i(LAYER *l, void *arg __attribute__((__unused__)))
{
 fprintf(stderr,"%s: unhandled input packet, type %d\n",__progname,l->b->ipkt[0]);
 exit(1);
}

static LAYERDESC layer_catchall = { 0, &catchall_i, 0 };

static void d_i_d_i(LAYER *l, void *arg __attribute__((__unused__)))
{
 BPP *b;

 b = l->b;
 if (b->iplen < 1)
  { fprintf(stderr,"%s: ignoring empty protocol packet\n",__progname);
    return;
  }
 switch (b->ipkt[0])
  { case SSH_MSG_DISCONNECT:
	{ unsigned int reason;
	  STR desc;
	  STR lang;
	  parse_packet(b,pp_fail,
		PP_IGNORE(1),
		PP_UINT32(&reason),
		PP_STRING(&desc),
		PP_STRING(&lang),
		PP_ENDSHERE );
	  printf("Disconnect received: ");
	  switch (reason)
	   { case SSH_DISCONNECT_HOST_NOT_ALLOWED_TO_CONNECT:
		printf("Host not allowed to connect");
		break;
	     case SSH_DISCONNECT_PROTOCOL_ERROR:
		printf("Protocol error");
		break;
	     case SSH_DISCONNECT_KEY_EXCHANGE_FAILED:
		printf("Key exchange failed");
		break;
	     case SSH_DISCONNECT_RESERVED:
		printf("Reserved");
		break;
	     case SSH_DISCONNECT_MAC_ERROR:
		printf("Mac error");
		break;
	     case SSH_DISCONNECT_COMPRESSION_ERROR:
		printf("Compression error");
		break;
	     case SSH_DISCONNECT_SERVICE_NOT_AVAILABLE:
		printf("Service not available");
		break;
	     case SSH_DISCONNECT_PROTOCOL_VERSION_NOT_SUPPORTED:
		printf("Protocol version not supported");
		break;
	     case SSH_DISCONNECT_HOST_KEY_NOT_VERIFIABLE:
		printf("Host key not verifiable");
		break;
	     case SSH_DISCONNECT_CONNECTION_LOST:
		printf("Connection lost");
		break;
	     case SSH_DISCONNECT_BY_APPLICATION:
		printf("By application");
		break;
	     case SSH_DISCONNECT_TOO_MANY_CONNECTIONS:
		printf("Too many connections");
		break;
	     case SSH_DISCONNECT_AUTH_CANCELLED_BY_USER:
		printf("Auth cancelled by user");
		break;
	     case SSH_DISCONNECT_NO_MORE_AUTH_METHODS_AVAILABLE:
		printf("No more auth methods available");
		break;
	     case SSH_DISCONNECT_ILLEGAL_USER_NAME:
		printf("Illegal user name");
		break;
	     default:
		printf("unknown reason %u",reason);
		break;
	   }
	  printf("\n");
	  exit(0);
	}
       break;
    case SSH_MSG_IGNORE:
       break;
    case SSH_MSG_DEBUG:
	{ int disp;
	  STR msg;
	  STR lang;
	  parse_packet(b,pp_fail,
		PP_IGNORE(1),
		PP_BOOL(&disp),
		PP_STRING(&msg),
		PP_STRING(&lang),
		PP_ENDSHERE );
	  printf("Debug message [%.*s]: %.*s\n",lang.len,lang.data,msg.len,msg.data);
	}
       break;
    default:
       above_ipkt(l);
       break;
  }
}

static LAYERDESC layer_d_i_d = { 0, &d_i_d_i, 0 };

static void net_read(int id __attribute__((__unused__)), void *arg)
{
 BPP *b;
 char buf[8192];
 int n;

 b = arg;
 n = read(b->fd,&buf[0],sizeof(buf));
 if (n < 0)
  { switch (errno)
     { case EINTR:
       case EWOULDBLOCK:
	  return;
	  break;
     }
    fprintf(stderr,"%s: net read: %s\n",__progname,strerror(errno));
    exit(1);
  }
 if (n == 0)
  { fprintf(stderr,"%s: net read EOF\n",__progname);
    exit(1);
  }
 (*b->input)(b,&buf[0],n);
}

static void net_write(int id __attribute__((__unused__)), void *arg)
{
 static int maxiov = -1;
 static struct iovec *iov;
 static int iovn;
 BPP *b;
 int iovl;
 OQB *q;
 int w;

 b = arg;
 if (maxiov < 0)
  { int mib[2];
    size_t oldlen;
    mib[0] = CTL_KERN;
    mib[1] = KERN_IOV_MAX;
    oldlen = sizeof(maxiov);
    if (sysctl(&mib[0],2,&maxiov,&oldlen,0,0) < 0)
     { fprintf(stderr,"%s: can't get kern.iov_max: %s\n",__progname,strerror(errno));
       exit(1);
     }
    if (maxiov > 64) maxiov = 64;
    if (maxiov < 1) maxiov = 1;
    iov = 0;
    iovn = 0;
  }
 iovl = 0;
 for (q=b->output;q;q=q->link)
  { if (iovl >= maxiov) break;
    if (iovl >= iovn) iov = realloc(iov,(iovn=iovl+4)*sizeof(*iov));
    iov[iovl++] = (struct iovec){ .iov_base=dequal(q->data), .iov_len=q->left };
  }
 w = writev(b->fd,iov,iovl);
 if (w < 0)
  { switch (errno)
     { case EINTR:
       case EWOULDBLOCK:
	  return;
	  break;
     }
    fprintf(stderr,"%s: net write: %s\n",__progname,strerror(errno));
    exit(1);
  }
 if (w == 0)
  { fprintf(stderr,"%s: zero net write\n",__progname);
    exit(1);
  }
 while ((q=b->output) && (w >= q->left))
  { b->output = q->link;
    w -= q->left;
    free(q->datafree);
    free(q);
  }
 if (q)
  { if (w)
     { q->data += w;
       q->left -= w;
     }
  }
 else
  { b->outtail = &b->output;
  }
}

static int wtest_bpp(int id __attribute__((__unused__)), void *arg)
{
 return(((BPP *)arg)->output != 0);
}

static void bpp_setup(BPP *b)
{
 b->w_enc = encalg_vfind_c("none");
 b->w_encstate = encalg_init(b->w_enc,"","",'w');
 b->r_enc = b->w_enc;
 b->r_encstate = encalg_init(b->r_enc,"","",'r');
 b->w_comp = compalg_vfind_c("none");
 b->w_compstate = compalg_init(b->w_comp,'w');
 b->r_comp = b->w_comp;
 b->r_compstate = compalg_init(b->r_comp,'r');
 b->w_mac = macalg_vfind_c("none");
 b->w_macstate = macalg_init(b->w_mac,"");
 b->r_mac = b->w_mac;
 b->r_macstate = macalg_init(b->r_mac,"");
 b->w_seq = 0U - 1;
 b->r_seq = 0U - 1;
 b->ipkts_to_rekey = 1<<30;
 b->opkts_to_rekey = 1<<30;
 b->ibytes_to_rekey = 1<<30;
 b->obytes_to_rekey = 1<<30;
 b->output = 0;
 b->outtail = &b->output;
 b->input = &getversion;
 b->svs.state = SVS_ATNL;
 bcopy("SSH-",&b->s_version[0],4);
 b->s_vlen = 4;
 b->toplayer = 0;
 b->botlayer = 0;
}

int main(int, char **);
int main(int ac, char **av)
{
 handleargs(ac,av);
 switch (opmode)
  { case OPMODE_RUN:
       openconn();
       bpp_setup(&bpp);
       sendversion();
       push_layer(&bpp,&layer_base);
       push_layer(&bpp,&layer_d_i_d);
       push_layer(&bpp,&layer_transport);
       push_layer(&bpp,&layer_userauth_conn);
       push_layer(&bpp,&layer_channels);
       push_layer(&bpp,&layer_catchall);
       init_polling();
       start_remote_session();
       add_poll_fd(bpp.fd,rwtest_always,wtest_bpp,net_read,net_write,&bpp);
       while (1)
	{ pre_poll();
	  if (do_poll() < 0)
	   { if (errno == EINTR) continue;
	     fprintf(stderr,"%s: poll: %s\n",__progname,strerror(errno));
	     exit(1);
	   }
	  post_poll();
	}
       break;
    case OPMODE_KEYGEN:
       gen_key();
       break;
    default:
       abort();
       break;
  }
 exit(0);
}