/* FAt Remote SHell client by eeeee
 * Public Domain */

#include "config-common.h"
#include "config-farsh.h"
#include "common.h"

#include <sys/types.h>
#include <netdb.h>

struct termios old_attr;
static void exyt(){ tcsetattr(STDOUT_FILENO, TCSANOW, &old_attr); return; }

int main(int argc, char *argv[]){
	s32 skt, res, fdmax;
	s32 pkt_int = PKT_INT;
	u32 temp = 0;
	u32 connected = 0;
	u32 lptime; /* last packet time */
	u16 *pkt, pos;
	u8  term[CONNECT];
	u8  esc = 1;
	struct sockaddr_storage sender;
	struct sockaddr_in  *srv4, *snd4;
	struct sockaddr_in6 *srv6, *snd6;
	struct addrinfo *srv;
	struct termios  attr;
	struct timeval  slp;
	struct timespec t1, t2;
	struct winsize  wsc, wss; /* client, server */
	fd_set rset;

	/* ring buffer */
	u8 *rb;
	u32 rb_mask = RB_SZ - 1;
	u32 rb_w  = 0;  /* written to ring buffer     */
	u32 rb_wm = 0;  /* written masked (mod RB_SZ) */
	u32 rb_rm = 0;  /* readed masked              */
	u32 rb_free = RB_SZ;

	/* stream state */
	u32 rcvd   = 0; /* recieved data bytes by client */
	u32 rcvd_s = 0; /* recieved data bytes by client (server version) */
	u32 rb_r   = 0; /* readed from ring buffer by server */
	u32 rb_r_s = 0; /* server version */

	if (argc != 3){
		fprintf(stderr, "Using:\n\t"
		            "%s <host> <port>\n\n"
		        "Example:\n\t"
		            "%s 127.0.0.1 2222\n\t"
		            "%s example.org 2222\n\n",
		        argv[0], argv[0], argv[0]);
		exit(1);
	}

	do { /* dns, ip/port str to num */
		res = getaddrinfo(argv[1], argv[2], NULL, &srv);
		temp++;
		usleep(10000); /* 10 ms */
	} while (res == EAI_AGAIN && temp < 5);
	if (res){
		fprintf(stderr, "getaddrinfo: %s\n", gai_strerror(res));
		exit(1);
	}

	/* tty */
	if (isatty(STDOUT_FILENO) && isatty(STDIN_FILENO)){
		/* save term attrs */
		res = tcgetattr(STDOUT_FILENO, &old_attr);
		ERRDIE(res, "tcgetattr");

		/* restore term attrs at exit */
		res = atexit(exyt); ERRDIE(res, "atexit");

		/* echo off */
		memcpy(&attr, &old_attr, sizeof(attr));
		attr.c_lflag &= ~ECHO;
		res = tcsetattr(STDOUT_FILENO, TCSANOW, &attr);
		ERRDIE(res, "tcsetattr");

		/* get password */
		fputs("Password: ", stdout);
		rb = fgets(password, PASSWORD_BUF_SZ, stdin);
		fputs("\nconnecting...\n", stdout);
		if (rb == NULL)
			fputs("can't read password\n", stderr), exit(1);

		/* switch to raw mode */
		cfmakeraw(&attr);
		res = tcsetattr(STDOUT_FILENO, TCSANOW, &attr);
		ERRDIE(res, "tcsetattr");
	} else {
		fputs("this program can't work without tty\n", stderr);
		exit(1);
	}

	/* password without '\r\n' */
	res = 0;
	while (password[res] != '\0'
	    && password[res] != '\r'
	    && password[res] != '\n') res++;

	/* hash password */
	sponge256_hash32(key, password, res);
	memset(password, 0, PASSWORD_BUF_SZ);

	/* feed prng with key and current time */
	duplex257_prng_feed(prng_state, ((u32*)key)[0]);
	duplex257_prng_feed(prng_state, ((u32*)key)[1]);
	duplex257_prng_feed(prng_state, ((u32*)key)[2]);
	duplex257_prng_feed(prng_state, ((u32*)key)[3]);
	duplex257_prng_feed(prng_state, ((u32*)key)[4]);
	duplex257_prng_feed(prng_state, ((u32*)key)[5]);
	duplex257_prng_feed(prng_state, ((u32*)key)[6]);
	duplex257_prng_feed(prng_state, ((u32*)key)[7]);
	clock_gettime(CLOCK_MONOTONIC, &t1);
	duplex257_prng_feed(prng_state, (u32)t1.tv_sec);
	duplex257_prng_feed(prng_state, (u32)t1.tv_nsec);

	lptime = t1.tv_sec; /* init last packet time */

	/* get term size */
	res = ioctl(STDOUT_FILENO, TIOCGWINSZ, &wsc);
	ERRDIE(res, "TIOCGWINSZ");
	memset(&wss, 0, sizeof(wss));

	/* get TERM var */
	memset(term, '\0', CONNECT);
	rb = getenv("TERM");
	if (rb != NULL){
		res = strlen(rb);
		if (res >= CONNECT)
			fprintf(stderr, "strlen(TERM) > %d\n", CONNECT-1),
			exit(1);

		strcpy(term, rb);
	}

	/* packet buffer (recv/send) */
	pkt = malloc(MAX_PKT_SZ);
	if (pkt == NULL) fputs("pkt = malloc() failed\r\n",stderr), exit(1);

	/* ring buffer */
	rb = malloc(RB_SZ);
	if (rb == NULL) fputs("ringbuf malloc failed\r\n", stderr), exit(1);

	/* socket, setsockopt */
	skt = socket(srv->ai_family, SOCK_DGRAM, 0); ERRDIE(skt, "socket");
	res = 1;
	res = setsockopt(skt, SOL_SOCKET, SO_REUSEADDR, &res, sizeof(res));
	ERRDIE(res, "SO_REUSEADDR");

	/* sockaddr */
	srv4 = (struct sockaddr_in  *) srv->ai_addr;
	srv6 = (struct sockaddr_in6 *) srv->ai_addr;
	snd4 = (struct sockaddr_in  *) &sender;
	snd6 = (struct sockaddr_in6 *) &sender;

	/* fdmax for select */
	fdmax = skt;
	if (STDIN_FILENO > skt) fdmax = STDIN_FILENO;
	fdmax++;


inf_loop:
	/* select */
	FD_ZERO(&rset);
	FD_SET(skt, &rset);
	if (rb_free) FD_SET(STDIN_FILENO, &rset);

	slp.tv_sec = 0;
	slp.tv_usec = pkt_int;

	/* the contents of the slp structure after calling select() is
	 * undefined (not portable), so I manually count the sleep time */
	clock_gettime(CLOCK_MONOTONIC, &t1);
	select(fdmax, &rset, NULL, NULL, &slp);

	/* recieve packet */
	if (FD_ISSET(skt, &rset)){
		res = recv_pkt(skt, pkt, MAX_PKT_SZ, &sender);
		if (res < 4) goto ignore_p; /* ignore too small packets */

		/* check ipv4 packet addr */
		if (srv->ai_addr->sa_family == AF_INET){
			if (srv4->sin_port != snd4->sin_port ||
			    srv4->sin_addr.s_addr != snd4->sin_addr.s_addr)
				goto ignore_p;
		} else { /* check ipv6 packet addr */
			res = memcmp(srv6->sin6_addr.s6_addr,
			             snd6->sin6_addr.s6_addr, 16);
			if (0 != res || srv6->sin6_port != snd6->sin6_port)
				goto ignore_p;
		}

		pkt[0] = ntohs(pkt[0]);
		pkt[1] = ntohs(pkt[1]);
		if (pkt[0] == 0)
			fputs("\r\nserver closed connection\r\n", stdout),
			exit(0);

		/* ignore packet with wrong size */
		if (res != 4 + (pkt[0] & ~CONNECT) + pkt[1]) goto ignore_p;

		/* packet recv time for entropy and last packet time */
		clock_gettime(CLOCK_MONOTONIC, &t2);
		duplex257_prng_feed(prng_state, t2.tv_nsec);
		lptime = t2.tv_sec;

		pos = 2;
		connected |= pkt[0] & CONNECT;

		if (pkt[0] & PING){ /* read server term size */
			wss.ws_row = ntohs(pkt[pos++]);
			wss.ws_col = ntohs(pkt[pos++]);
		}

		if (pkt[0] & SSTATE){
			rcvd_s = ntohl(*(u32*)&pkt[pos]), pos += 2;
			rb_r_s = ntohl(*(u32*)&pkt[pos]), pos += 2;

			/* rb_w can overflow, so I can't ignore a packet if
			 * rb_w < rb_r_s. That's why I do these checks */
			if (RB_SZ < rb_w - rb_r_s ||
			    RB_SZ < rb_r_s - rb_r) goto ignore_p;

			/* rb_r_s is ok, update state of ring buffer */
			rb_r = rb_r_s;
			rb_free = RB_SZ - (rb_w - rb_r);
			rb_rm = rb_r & rb_mask;

			if (pkt[1]){ /* have data in packet */
				/* temp: already recieved amount of data*/
				temp = rcvd - rcvd_s;
				if (temp < pkt[1]){
					/* packet also have data that is not
					 * recieved yet */
					pos = (pos<<1) + temp;
					pkt[1] -= temp;
					write(STDOUT_FILENO, (u8*)pkt+pos,
					                            pkt[1]);
					rcvd += pkt[1];
				}
			}
		}
	}
ignore_p:

	/* read stdin to ringbuffer */
	if (FD_ISSET(STDIN_FILENO, &rset)){
		res = read(STDIN_FILENO, rb + rb_wm, 1);
		
		if (res != 1)
			fputs("\r\nreading stdin return res<1\r\n", stderr),
			exit(1);

		if (esc == 0 && rb[rb_wm] == DISCONNECT_KEY){
			memset(pkt, 0, 4 + CONNECT);
			pkt[0] = htons(CONNECT);
			send_pkt(skt, pkt, 4 + CONNECT,
			           (struct sockaddr_storage*) srv->ai_addr);
			fputs("\r\ndisconnected\r\n", stderr);
			exit(0);
		}

		esc = rb[rb_wm] - 0x1B;

		rb_w++;
		rb_wm = rb_w & rb_mask;
		rb_free--;
	}

	/* send packet every PKT_INT microseconds */
	clock_gettime(CLOCK_MONOTONIC, &t2);
	pkt_int -= abs((s32)t2.tv_nsec - (s32)t1.tv_nsec) / 1000;
	if (pkt_int < MIN_SLEEP){
		t2.tv_sec -= lptime;
		if (!connected && t2.tv_sec >= TIMEOUT_C)
			fputs("can't connect\r\n", stderr), exit(1);

		if (t2.tv_sec >= TIMEOUT_P){
			/* send PING if no packets from server for
			 * TIMEOUT_P seconds */
			wss.ws_row = 65535;
			if (t2.tv_sec >= TIMEOUT_D)
				fputs("\r\nserver timeout\r\n", stdout),
				exit(0);
		}

		pkt_int = PKT_INT;
		pos = 2; /* pos in array of u16 */
		pkt[0] = 0;

		/* add something to packet */
		if (!connected){
			/* add TERM var */
			pkt[0] |= CONNECT;
			memcpy(&pkt[pos], term, CONNECT);
			pos += CONNECT >> 1;
		}

		ioctl(STDOUT_FILENO, TIOCGWINSZ, &wsc);
		if ((wsc.ws_row - wss.ws_row) | (wsc.ws_col - wss.ws_col)){
			/* add term sizes */
			pkt[0] |= PING;
			pkt[pos++] = htons(wsc.ws_row);
			pkt[pos++] = htons(wsc.ws_col);
		}

		pkt[1] = rb_w - rb_r;
		if (pkt[1]){ /* have data */
			/* add state of streams */
			pkt[0] |= SSTATE;
			*(u32*)&pkt[pos] = htonl(rcvd), pos += 2;
			*(u32*)&pkt[pos] = htonl(rb_r), pos += 2;

			/* add data */
			pos <<= 1; /* now it pos in array of u8 */
			res = MAX_PKT_SZ - HMAC_SZ - NONCE_SZ - pos;
			if (res < pkt[1]) pkt[1] = res;
			memcpy((u8*)pkt + pos, rb + rb_rm, pkt[1]);
			pos += pkt[1];
		} else {
			if (rcvd != rcvd_s || rb_r != rb_r_s){
				/* add state of streams */
				pkt[0] |= SSTATE;
				*(u32*)&pkt[pos] = htonl(rcvd), pos += 2;
				*(u32*)&pkt[pos] = htonl(rb_r), pos += 2;

				/* send state only once */
				rcvd_s = rcvd;
				rb_r_s = rb_r;
			}
			pos <<= 1; /* now it pos in array of u8 */
		}

		if (pkt[0]){
			/* send packet */
			pkt[0] = htons(pkt[0]);
			pkt[1] = htons(pkt[1]);
			send_pkt(skt, pkt, pos,
			           (struct sockaddr_storage*) srv->ai_addr);
		}
	}

	goto inf_loop;

	return 0;
}