/* Notification that our channel open request failed */
void recv_msg_channel_open_failure() {

	unsigned int chan;
	struct Channel * channel;
	chan = buf_getbyte(ses.payload);

	channel = getchannel(chan);
	if (channel == NULL) {
		dropbear_exit("Unknown channel");
	}

	closechannel(channel);
}

/* fails fatally */
static void check_signkey_bits(enum signkey_type type, int bits)
{
	switch (type) {
#ifdef DROPBEAR_RSA
		case DROPBEAR_SIGNKEY_RSA:
			if (bits < 512 || bits > 4096 || (bits % 8 != 0)) {
				dropbear_exit("Bits must satisfy 512 <= bits <= 4096, and be a"
				              " multiple of 8\n");
			}
			break;
#endif
#ifdef DROPEAR_DSS
		case DROPBEAR_SIGNKEY_DSS:
			if (bits != 1024) {
				dropbear_exit("DSS keys have a fixed size of 1024 bits\n");
				exit(EXIT_FAILURE);
			}
#endif
		default:
			(void)0; /* quiet, compiler. ecdsa handles checks itself */
	}
}

void m_mp_init_multi(mp_int *mp, ...) 
{
    mp_int* cur_arg = mp;
    va_list args;

    va_start(args, mp);        /* init args to next argument from caller */
    while (cur_arg != NULL) {
        if (mp_init(cur_arg) != MP_OKAY) {
			dropbear_exit("mem alloc error");
        }
        cur_arg = va_arg(args, mp_int*);
    }
    va_end(args);
}

/* process a packet, and also check that auth has been done */
static void process_postauth_packet(unsigned int type) {

	/* messages following here require userauth before use */
	if (!ses.authstate.authdone) {
		dropbear_exit("received message %d before userauth", type);
	}

	switch (type) {

		case SSH_MSG_CHANNEL_DATA:
			recv_msg_channel_data();
			break;

		case SSH_MSG_CHANNEL_WINDOW_ADJUST:
			recv_msg_channel_window_adjust();
			break;

		case SSH_MSG_CHANNEL_REQUEST:
			recv_msg_channel_request();
			break;

		case SSH_MSG_CHANNEL_OPEN:
			recv_msg_channel_open();
			break;

		case SSH_MSG_CHANNEL_EOF:
			recv_msg_channel_eof();
			break;

		case SSH_MSG_CHANNEL_CLOSE:
			recv_msg_channel_close();
			break;

#ifdef USE_LISTENERS /* for x11, tcp fwd etc */
		case SSH_MSG_CHANNEL_OPEN_CONFIRMATION:
			recv_msg_channel_open_confirmation();
			break;
			
		case SSH_MSG_CHANNEL_OPEN_FAILURE:
			recv_msg_channel_open_failure();
			break;
#endif
			
		default:
			TRACE(("unknown packet()"));
			recv_unimplemented();
			break;
	}
}

void svr_session(int sock, int childpipe, 
		char* remotehost, char *addrstring) {

	struct timeval timeout;

    reseedrandom();

	crypto_init();
	common_session_init(sock, remotehost);

	/* Initialise server specific parts of the session */
	svr_ses.childpipe = childpipe;
	svr_ses.addrstring = addrstring;
	svr_authinitialise();
	chaninitialise(svr_chantypes);
	svr_chansessinitialise();

	if (gettimeofday(&timeout, 0) < 0) {
		dropbear_exit("Error getting time");
	}

	ses.connecttimeout = timeout.tv_sec + AUTH_TIMEOUT;

	/* set up messages etc */
	ses.remoteclosed = svr_remoteclosed;

	/* packet handlers */
	ses.packettypes = svr_packettypes;
	ses.buf_match_algo = svr_buf_match_algo;

	ses.isserver = 1;

	/* We're ready to go now */
	sessinitdone = 1;

	/* exchange identification, version etc */
	session_identification();

	/* start off with key exchange */
	send_msg_kexinit();

	/* Run the main for loop. NULL is for the dispatcher - only the client
	 * code makes use of it */
	session_loop(NULL);

	/* Not reached */

}

/* Set up the general chansession environment, in particular child-exit
 * handling */
void chansessinitialise() {

	struct sigaction sa_chld;

	/* single child process intially */
	ses.childpids = (struct ChildPid*)m_malloc(sizeof(struct ChildPid));
	ses.childpids[0].pid = -1; /* unused */
	ses.childpids[0].chansess = NULL;
	ses.childpidsize = 1;
	sa_chld.sa_handler = sesssigchild_handler;
	sa_chld.sa_flags = SA_NOCLDSTOP;
	if (sigaction(SIGCHLD, &sa_chld, NULL) < 0) {
		dropbear_exit("signal() error");
	}
	
}

/* add a new environment variable, allocating space for the entry */
void addnewvar(const char* param, const char* var) {

	char* newvar;
	int plen, vlen;

	plen = strlen(param);
	vlen = strlen(var);

	newvar = m_malloc(plen + vlen + 2); /* 2 is for '=' and '\0' */
	memcpy(newvar, param, plen);
	newvar[plen] = '=';
	memcpy(&newvar[plen+1], var, vlen);
	newvar[plen+vlen+1] = '\0';
	if (putenv(newvar) < 0) {
		dropbear_exit("environ error");
	}
}

/* when we receive channel data, put it in a buffer for writing to the program/
 * shell etc */
void recv_msg_channel_data() {

	unsigned int chan;
	struct Channel * channel;
	unsigned int datalen;
	unsigned int pos;
	unsigned int maxdata;

	TRACE(("enter recv_msg_channel_data"));
	
	chan = buf_getint(ses.payload);
	channel = getchannel(chan);
	if (channel == NULL) {
		/* disconnect ? */
		dropbear_exit("Unknown channel");
	}

	assert(channel->infd != -1);

	datalen = buf_getint(ses.payload);

	/* if the client is going to send us more data than we've allocated, then 
	 * it has ignored the windowsize, so we "MAY ignore all extra data" */
	maxdata = channel->writebuf->size - channel->writebuf->pos;
	if (datalen > maxdata) {
		TRACE(("Warning: recv_msg_channel_data: extra data past window"));
		datalen = maxdata;
	}

	/* write to the buffer - we always append to the end of the buffer */
	pos = channel->writebuf->pos;
	buf_setpos(channel->writebuf, channel->writebuf->len);
	memcpy(buf_getwriteptr(channel->writebuf, datalen), 
			buf_getptr(ses.payload, datalen), datalen);
	buf_incrwritepos(channel->writebuf, datalen);
	buf_setpos(channel->writebuf, pos); /* revert pos */

	channel->recvwindow -= datalen;
/*	if (channel->recvwindow < RECV_MINWINDOW) {
		send_msg_channel_window_adjust(channel, 
				RECV_MAXWINDOW - channel->recvwindow);
		channel->recvwindow = RECV_MAXWINDOW;
	}*/

	TRACE(("leave recv_msg_channel_data"));
}

/* mostly taken from libtomcrypt's rsa key generation routine */
dropbear_rsa_key * gen_rsa_priv_key(unsigned int size) {

	dropbear_rsa_key * key;
	DEF_MP_INT(pminus);
	DEF_MP_INT(qminus);
	DEF_MP_INT(lcm);

	if (size < 512 || size > 4096 || (size % 8 != 0)) {
		dropbear_exit("Bits must satisfy 512 <= bits <= 4096, and be a"
			" multiple of 8");
	}

	key = m_malloc(sizeof(*key));
	m_mp_alloc_init_multi(&key->e, &key->n, &key->d, &key->p, &key->q, NULL);
	m_mp_init_multi(&pminus, &lcm, &qminus, NULL);

	if (mp_set_int(key->e, RSA_E) != MP_OKAY) {
		fprintf(stderr, "RSA generation failed\n");
		exit(1);
	}

	getrsaprime(key->p, &pminus, key->e, size/16);
	getrsaprime(key->q, &qminus, key->e, size/16);

	if (mp_mul(key->p, key->q, key->n) != MP_OKAY) {
		fprintf(stderr, "RSA generation failed\n");
		exit(1);
	}

	/* lcm(p-1, q-1) */
	if (mp_lcm(&pminus, &qminus, &lcm) != MP_OKAY) {
		fprintf(stderr, "RSA generation failed\n");
		exit(1);
	}

	/* de = 1 mod lcm(p-1,q-1) */
	/* therefore d = (e^-1) mod lcm(p-1,q-1) */
	if (mp_invmod(key->e, &lcm, key->d) != MP_OKAY) {
		fprintf(stderr, "RSA generation failed\n");
		exit(1);
	}

	mp_clear_multi(&pminus, &qminus, &lcm, NULL);

	return key;
}	

/* hash the current random pool with some unique identifiers
 * for this process and point-in-time. this is used to separate
 * the random pools for fork()ed processes. */
void reseedrandom() {

	pid_t pid;
	hash_state hs;
	struct timeval tv;

	if (!donerandinit) {
		dropbear_exit("seedrandom not done");
	}

	pid = getpid();
	gettimeofday(&tv, NULL);

	sha1_init(&hs);
	sha1_process(&hs, (void*)hashpool, sizeof(hashpool));
	sha1_process(&hs, (void*)&pid, sizeof(pid));
	sha1_process(&hs, (void*)&tv, sizeof(tv));
	sha1_done(&hs, hashpool);
}

/* Return a null-terminated string, it is malloced, so must be free()ed
 * Note that the string isn't checked for null bytes, hence the retlen
 * may be longer than what is returned by strlen */
char* buf_getstring(buffer* buf, unsigned int *retlen) {

	unsigned int len;
	char* ret;
	len = buf_getint(buf);
	if (len > MAX_STRING_LEN) {
		dropbear_exit("String too long");
	}

	if (retlen != NULL) {
		*retlen = len;
	}
	ret = m_malloc(len+1);
	memcpy(ret, buf_getptr(buf, len), len);
	buf_incrpos(buf, len);
	ret[len] = '\0';

	return ret;
}

static void cli_proxy_cmd(int *sock_in, int *sock_out, pid_t *pid_out) {
	char * ex_cmd = NULL;
	size_t ex_cmdlen;
	int ret;

	fill_passwd(cli_opts.own_user);

	ex_cmdlen = strlen(cli_opts.proxycmd) + 6; /* "exec " + command + '\0' */
	ex_cmd = m_malloc(ex_cmdlen);
	snprintf(ex_cmd, ex_cmdlen, "exec %s", cli_opts.proxycmd);

	ret = spawn_command(exec_proxy_cmd, ex_cmd,
			sock_out, sock_in, NULL, pid_out);
	m_free(ex_cmd);
	if (ret == DROPBEAR_FAILURE) {
		dropbear_exit("Failed running proxy command");
		*sock_in = *sock_out = -1;
	}
}

/* Confirmation that our channel open request (for forwardings) was 
 * successful*/
void recv_msg_channel_open_confirmation() {

	unsigned int chan;
	struct Channel * channel;

	TRACE(("enter recv_msg_channel_open_confirmation"));
	chan = buf_getint(ses.payload);

	channel = getchannel(chan);
	if (channel == NULL) {
		dropbear_exit("Unknown channel");
	}

	channel->remotechan =  buf_getint(ses.payload);
	channel->transwindow = buf_getint(ses.payload);
	channel->transmaxpacket = buf_getint(ses.payload);

	TRACE(("leave recv_msg_channel_open_confirmation"));
}

int cli_main(int argc, char ** argv) {
#else
int main(int argc, char ** argv) {
#endif

	int sock_in, sock_out;
	struct dropbear_progress_connection *progress = NULL;

	_dropbear_exit = cli_dropbear_exit;
	_dropbear_log = cli_dropbear_log;

	disallow_core();

	seedrandom();
	crypto_init();

	cli_getopts(argc, argv);

	TRACE(("user='%s' host='%s' port='%s'", cli_opts.username,
				cli_opts.remotehost, cli_opts.remoteport))

	if (signal(SIGPIPE, SIG_IGN) == SIG_ERR) {
		dropbear_exit("signal() error");
	}

#ifdef ENABLE_CLI_PROXYCMD
	if (cli_opts.proxycmd) {
		cli_proxy_cmd(&sock_in, &sock_out);
		m_free(cli_opts.proxycmd);
	} else
#endif
	{
		progress = connect_remote(cli_opts.ipfamily, cli_opts.remotehost,
				cli_opts.remoteport, cli_connected, &ses);
		sock_in = sock_out = -1;
	}

	cli_session(sock_in, sock_out, progress);

	/* not reached */
	return -1;
}

/* Handle a diffie-hellman key exchange initialisation. This involves
 * calculating a session key reply value, and corresponding hash. These
 * are carried out by send_msg_kexdh_reply(). recv_msg_kexdh_init() calls
 * that function, then brings the new keys into use */
void recv_msg_kexdh_init() {

	mp_int dh_e;

	TRACE(("enter recv_msg_kexdh_init"));
	if (!ses.kexstate.recvkexinit) {
		dropbear_exit("Premature kexdh_init message received");
	}

	m_mp_init(&dh_e);
	buf_getmpint(ses.payload, &dh_e);

	send_msg_kexdh_reply(&dh_e);

	mp_clear(&dh_e);

	send_msg_newkeys();
	ses.expecting = SSH_MSG_NEWKEYS;
	TRACE(("leave recv_msg_kexdh_init"));
}

/* compresses len bytes from src, outputting to dest (starting from the
 * respective current positions. */
static void buf_compress(buffer * dest, buffer * src, unsigned int len) {

	unsigned int endpos = src->pos + len;
	int result;

	TRACE(("enter buf_compress"));

	while (1) {

		ses.keys->trans_zstream->avail_in = endpos - src->pos;
		ses.keys->trans_zstream->next_in = 
			buf_getptr(src, ses.keys->trans_zstream->avail_in);

		ses.keys->trans_zstream->avail_out = dest->size - dest->pos;
		ses.keys->trans_zstream->next_out =
			buf_getwriteptr(dest, ses.keys->trans_zstream->avail_out);

		result = deflate(ses.keys->trans_zstream, Z_SYNC_FLUSH);

		buf_setpos(src, endpos - ses.keys->trans_zstream->avail_in);
		buf_setlen(dest, dest->size - ses.keys->trans_zstream->avail_out);
		buf_setpos(dest, dest->len);

		if (result != Z_OK) {
			dropbear_exit("zlib error");
		}

		if (ses.keys->trans_zstream->avail_in == 0) {
			break;
		}

		assert(ses.keys->trans_zstream->avail_out == 0);

		/* the buffer has been filled, we must extend. This only happens in
		 * unusual circumstances where the data grows in size after deflate(),
		 * but it is possible */
		buf_resize(dest, dest->size + ZLIB_COMPRESS_INCR);

	}
	TRACE(("leave buf_compress"));
}

dropbear_dss_key * gen_dss_priv_key(unsigned int size) {

	dropbear_dss_key *key;

	if (size != 1024) {
		dropbear_exit("DSS keys have a fixed size of 1024 bits");
	}

	key = m_malloc(sizeof(*key));

	m_mp_alloc_init_multi(&key->p, &key->q, &key->g, &key->y, &key->x, NULL);
	
	getq(key);
	getp(key, size/8);
	getg(key);
	getx(key);
	gety(key);

	return key;
	
}

/* Increment the outgoing data window for a channel - the remote end limits
 * the amount of data which may be transmitted, this window is decremented
 * as data is sent, and incremented upon receiving window-adjust messages */
void recv_msg_channel_window_adjust() {

	unsigned int chan;
	struct Channel * channel;
	unsigned int incr;
	
	chan = buf_getint(ses.payload);
	channel = getchannel(chan);

	if (channel == NULL) {
		dropbear_exit("Unknown channel"); /* TODO - disconnect */
	}
	
	incr = buf_getint(ses.payload);
	TRACE(("received window increment %d", incr));
	incr = MIN(incr, MAX_TRANS_WIN_INCR);
	
	channel->transwindow += incr;
	channel->transwindow = MIN(channel->transwindow, MAX_TRANS_WINDOW);

}

/* Create (malloc) a new buffer of size */
buffer* buf_new(unsigned int size) {

	buffer* buf;
	
	if (size > BUF_MAX_SIZE) {
		dropbear_exit("buf->size too big");
	}

	buf = (buffer*)m_malloc(sizeof(buffer)+size);

	if (size > 0) {
		buf->data = (unsigned char*)buf + sizeof(buffer);
	} else {
		buf->data = NULL;
	}

	buf->size = size;

	return buf;

}

/* non-blocking function writing out a current encrypted packet */
void write_packet() {

	int len, written;
	buffer * writebuf;
	
	TRACE(("enter write_packet"));
	assert(!isempty(&ses.writequeue));

	/* Get the next buffer in the queue of encrypted packets to write*/
	writebuf = (buffer*)examine(&ses.writequeue);

	len = writebuf->len - writebuf->pos;
	assert(len > 0);
	/* Try to write as much as possible */
	written = write(ses.sock, buf_getptr(writebuf, len), len);

	if (written < 0) {
		if (errno == EINTR) {
			TRACE(("leave writepacket: EINTR"));
			return;
		} else {
			dropbear_exit("error writing");
		}
	} 

	if (written == 0) {
		dropbear_close("remote host closed connection");
	}

	if (written == len) {
		/* We've finished with the packet, free it */
		dequeue(&ses.writequeue);
		buf_free(writebuf);
	} else {
		/* More packet left to write, leave it in the queue for later */
		buf_incrpos(writebuf, written);
	}

	TRACE(("leave write_packet"));
}