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

#define _XOPEN_SOURCE 600

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

#include <signal.h>
#include <fcntl.h>

struct {
	struct sockaddr_in6 a;
	struct winsize ws;
	u32 flags, notconnected;
	s32 ptm, pid;
	u32 lptime; /* time of last packet */
	/* stream state */
	u32 rcvd;   /* recieved data bytes from client */
	u32 rcvd_c; /* client version */
	u32 rb_r;   /* readed from ring buffer (recieved data by client) */
	u32 rb_r_c; /* client version */
	/* ring buffer */
	u32 rb_w;  /* written to ring buffer */
	u32 rb_wm; /* written masked (mod RB_SZ) */
	u32 rb_rm; /* readed masked */
	u32 rb_free, rb_free_for_read;
	u8 rb[RB_SZ];
} peer[MAX_PEERS];

u32 ccnt;

static void disconnect(s32 skt, s32 n){
	u8 pkt[4 + NONCE_SZ + HMAC_SZ];

	*(u32*)pkt = 0;
	send_pkt(skt, &pkt, 4, (struct sockaddr_storage*) &peer[n].a);
	kill(peer[n].pid, SIGKILL);
	close(peer[n].ptm);
	if (n != ccnt-1) memcpy(&peer[n], &peer[ccnt-1], sizeof(peer[0]));
	ccnt--;
	memset(&peer[ccnt], 0, sizeof(peer[0]));
	return;
}

int main(int argc, char *argv[]){
	u8  term[CONNECT];
	u8  *home;
	u16 *pkt;
	u16 pos;
	u16 port;
	s32 skt, res, fdmax, pts, i;
	s32 pkt_int = PKT_INT;
	u32 arcv;
	u32 rb_mask = RB_SZ - 1;
	struct sockaddr_in6 c;
	struct sigaction sa;
	struct timeval   slp;
	struct timespec  t1, t2;
	struct termios   attr;
	fd_set rset;

	/* check shell */
	res = access(SHELL, F_OK); ERRDIE(res, SHELL);
	res = access(SHELL, X_OK); ERRDIE(res, SHELL);

	/* arguments */
	if (argc != 2 && argc != 3){
		fprintf(stderr, "Using:\n\t%s [address] <port>\n\n"
		        "Example:\n\t%s 2222\n\t%s 127.0.0.1 2222\n\n",
		        argv[0], argv[0], argv[0]);
		exit(1);
	}

	memset(&c, 0, sizeof(c));
	port = (u16)atoi(argv[argc-1]);
	if (argc == 3){
		res = inet_pton(AF_INET6, argv[1], c.sin6_addr.s6_addr);
		if (res != 1){
			res = inet_pton(AF_INET, argv[1], &arcv);
			if (res != 1){
				fprintf(stderr, "wrong addr: %s\n",argv[1]);
				exit(1);
			}
			c.sin6_addr.s6_addr[10] = 0xFF;
			c.sin6_addr.s6_addr[11] = 0xFF;
			memcpy(&c.sin6_addr.s6_addr[12], &arcv, 4);
		}
	}

	/* get term attrs */
	res = tcgetattr(STDOUT_FILENO, &attr); ERRDIE(res, "tcgetattr");

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

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

	/* echo on */
	attr.c_lflag |= ECHO;
	res = tcsetattr(STDOUT_FILENO, TCSANOW, &attr);
	ERRDIE(res, "tcsetattr");

	/* 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);

	/* daemonize: redirect std streams */
	close(STDIN_FILENO);
	close(STDOUT_FILENO);
	close(STDERR_FILENO);
	res = open("/dev/null", O_RDONLY); ERRDIE(res, "open");
	res = open(LOGFILE, O_CREAT | O_WRONLY, S_IRUSR | S_IWUSR);
	ERRDIE(res, "open");
	res = open(LOGFILE, O_CREAT | O_WRONLY, S_IRUSR | S_IWUSR);
	ERRDIE(res, "open");

	/* daemonize: fork 1 */
	res = fork(); ERRDIE(res, "fork 1");
	if (res != 0) exit(0);

	/* daemonize: become session leader */
	setsid();

	/* daemonize: fork 2 */
	res = fork(); ERRDIE(res, "fork 2");
	if (res != 0) exit(0); /* kill session leader */

	/* daemonized! */

	/* chdir to HOME */
	home = getenv("HOME");
	if (home != NULL) chdir(home);

	/* kill zombie */
	memset(&sa, 0, sizeof(sa));
	sa.sa_handler = SIG_IGN;
	sa.sa_flags = SA_NOCLDWAIT;
	sigaction(SIGCHLD, &sa, NULL);

	/* init some vars */
	memset(peer, 0, sizeof(peer));
	ccnt = 0;

	/* cur pkt buf */
	pkt = malloc(MAX_PKT_SZ);
	if (pkt == NULL) fputs("pkt = malloc() failed\n",stderr), exit(1);

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

	/* bind */
	c.sin6_family = AF_INET6;
	c.sin6_port = htons(port);
	res = bind(skt, (struct sockaddr*) &c, sizeof(c));
	ERRDIE(res, "bind");

inf_loop:
	/* select */
	fdmax = skt + 1;
	FD_ZERO(&rset);
	FD_SET(skt, &rset);
	for (i=0; i < ccnt; i++){
		if (peer[i].rb_free == 0) continue;
		FD_SET(peer[i].ptm, &rset);
		if (fdmax <= peer[i].ptm) fdmax = peer[i].ptm + 1;
	}

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

	/* recv packet from client */
	if (FD_ISSET(skt, &rset)){
		res = recv_pkt(skt, pkt, MAX_PKT_SZ,
		                             (struct sockaddr_storage*) &c);

		pkt[0] = ntohs(pkt[0]);
		pkt[1] = ntohs(pkt[1]);
		if (res < 4 || res != 4 + pkt[0] + pkt[1]) goto ignore_p;
		for (i=0; i < ccnt; i++){
			res = memcmp(peer[i].a.sin6_addr.s6_addr,
			                           c.sin6_addr.s6_addr, 16);
			if (res == 0 && c.sin6_port == peer[i].a.sin6_port)
				break;
		}

		if (i == ccnt){
			/* client not connected */
			if (ccnt == MAX_PEERS || (pkt[0] & CONNECT) == 0)
				goto ignore_p;

			/* add peer to array */
			memcpy(&peer[i].a, &c, sizeof(c));

			/* peer is not fully connected yet */
			peer[i].notconnected = NOT_CONNECTED;

			/* init ring buffer */
			peer[i].rb_free = peer[i].rb_free_for_read = RB_SZ;

			/* pseudoterminal */
			peer[i].ptm = posix_openpt(O_RDWR);
			ERRDIE(peer[i].ptm, "posix_openpt");
			res = grantpt (peer[i].ptm); ERRDIE(res,"grantpt");
			res = unlockpt(peer[i].ptm); ERRDIE(res,"unlockpt");
        		pts = open(ptsname(peer[i].ptm), O_RDWR);
			ERRDIE(pts, "pts = open()");

			/* read TERM from packet */
			memset(term, '\0', CONNECT);
			memcpy(term, &pkt[2], CONNECT-1);

			/* fork */
			peer[i].pid = fork(); ERRDIE(peer[i].pid, "fork");
			if (peer[i].pid == 0){ /* child */
				close(peer[i].ptm);
				close(0), close(1), close(2);
				dup(pts), dup(pts), dup(pts);
				close(pts);
				setsid();
				res = ioctl(0, TIOCSCTTY, 1);
				ERRDIE(res, "TIOCSCTTY");
				setenv("TERM", term, 1);

				res = execlp(SHELL, SHELL, NULL);
				ERRDIE(res, "execlp");
			}
			close(pts);
			ccnt++;
		}

		/* client connected */
		clock_gettime(CLOCK_MONOTONIC, &t2);
		duplex257_prng_feed(prng_state, t2.tv_nsec);
		peer[i].lptime = t2.tv_sec;
		pos = 2;
		if (pkt[0] & CONNECT){
			if (!peer[i].notconnected){
				/* fully connected client sent packet with
				 * CONNECT flag, disconnect him */
				disconnect(skt, i);
				goto ignore_p;
			}
			peer[i].notconnected = NOT_CONNECTED;
			pos += CONNECT >> 1;
		} else if (peer[i].notconnected) peer[i].notconnected--;

		/* copy flags for responce */
		peer[i].flags |= pkt[0];

		/* ping, term sizes */
		if (pkt[0] & PING){
			peer[i].ws.ws_row = ntohs(pkt[pos++]);
			peer[i].ws.ws_col = ntohs(pkt[pos++]);
			ioctl(peer[i].ptm, TIOCSWINSZ, &peer[i].ws);
		}

		/* stream state */
		if (pkt[0] & SSTATE){
			peer[i].rb_r_c = ntohl(*(u32*)&pkt[pos]), pos += 2;
			peer[i].rcvd_c = ntohl(*(u32*)&pkt[pos]), pos += 2;

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

			/* rb_r_c is ok, update state of ring buffer */
			peer[i].rb_r = peer[i].rb_r_c;
			peer[i].rb_free = RB_SZ -
			                      (peer[i].rb_w - peer[i].rb_r);
			peer[i].rb_rm = peer[i].rb_r & rb_mask;
			peer[i].rb_free_for_read = peer[i].rb_free;
			if (peer[i].rb_free > peer[i].rb_rm)
				peer[i].rb_free_for_read -= peer[i].rb_rm;

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

	/* read ptms to ring buffers */
	for (i=0; i < ccnt; i++){
		if (0 == FD_ISSET(peer[i].ptm, &rset)) continue;

		res = read(peer[i].ptm, peer[i].rb + peer[i].rb_wm,
			   peer[i].rb_free_for_read);

		if (res < 1){/* disconnect */
			disconnect(skt, i);
			if (i != ccnt-1) i--;
			continue;
		}

		peer[i].rb_w += res;
		peer[i].rb_wm = peer[i].rb_w & rb_mask;
		peer[i].rb_free -= res;
		peer[i].rb_free_for_read = peer[i].rb_free;
		if (peer[i].rb_free > peer[i].rb_rm)
			peer[i].rb_free_for_read -= peer[i].rb_rm;
	}

	/* send packets to clients every PKT_INT (+- MIN_SLEEP) microsecs */
	clock_gettime(CLOCK_MONOTONIC, &t2);
	pkt_int -= abs((s32)t2.tv_nsec - (s32)t1.tv_nsec) / 1000;
	if (pkt_int < MIN_SLEEP)
	for (pkt_int=PKT_INT, i=0; i < ccnt; i++){
		/* make packet for peer i */
		pos = 2;
		pkt[0] = 0;

		/* disconnect on timeout */
		if (t2.tv_sec - peer[i].lptime >= TIMEOUT_D){
			disconnect(skt, i);
			continue;
		}

		/* set PING bit after getting packet with PING from peer */
		if (peer[i].flags & PING){
			pkt[0] |= PING;
			pkt[pos++] = htons(peer[i].ws.ws_row);
			pkt[pos++] = htons(peer[i].ws.ws_col);
		}

		/* add streams state and data if data exist */
		pkt[1] = peer[i].rb_w - peer[i].rb_r;
		if (pkt[1]) {
			/* add state of streams */
			pkt[0] |= SSTATE;
			*(u32*)&pkt[pos] = htonl(peer[i].rb_r), pos += 2;
			*(u32*)&pkt[pos] = htonl(peer[i].rcvd), 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, peer[i].rb + peer[i].rb_rm,
			                                            pkt[1]);
			pos += pkt[1];
		} else { /* no data to send */
			/* send streams state if it differ from peer state*/
			if (peer[i].rb_r != peer[i].rb_r_c ||
			    peer[i].rcvd != peer[i].rcvd_c){
				pkt[0] |= SSTATE;
				*(u32*)&pkt[pos] = htonl(peer[i].rb_r);
				pos += 2;
				*(u32*)&pkt[pos] = htonl(peer[i].rcvd);
				pos += 2;

				/* send state only once */
				peer[i].rb_r_c = peer[i].rb_r;
				peer[i].rcvd_c = peer[i].rcvd;
			}
			pos <<= 1;
		}

		peer[i].flags = 0;
		if (pkt[0]){
			/* add CONNECT bit if peer is not fully connected */
			if (peer[i].notconnected) pkt[0] |= CONNECT;
			pkt[0] = htons(pkt[0]);
			pkt[1] = htons(pkt[1]);
			send_pkt(skt, pkt, pos,
			             (struct sockaddr_storage*) &peer[i].a);
		}
	}

	goto inf_loop;

	return 0;
}