/*
   Place this file to compiled Dovecot sources' root directory and run:

   gcc imaptest.c -o imaptest -Wall -W -I. -Isrc/lib -DHAVE_CONFIG_H src/lib/liblib.a
*/
#include "lib.h"
#include "base64.h"
#include "buffer.h"
#include "str.h"
#include "network.h"
#include "istream.h"
#include "ostream.h"

#include <stdio.h>
#include <stdlib.h>
#include <ctype.h>
#include <fcntl.h>
#include <unistd.h>

/* IP / port where to connect to */
#define IP "127.0.0.1"
#define PORT 14300
/* Username. You can give either a single user, or a number of users in which
   case the username gets randomized at each connection. */
//#define USERNAME_TEMPLATE "u%04d@d%04d.domain.org"
//#define USERNAME_TEMPLATE "test%ld"
#define USERNAME_TEMPLATE "tsiraine"
#define USER_RAND 99
#define DOMAIN_RAND 99
/* Password (for all users) */
#define PASSWORD "pass"
/* Number of simultaneous client connections */
#define CLIENTS_COUNT 20
/* Try to keep around this many messages in mailbox (in expunge + append) */
#define MESSAGE_COUNT_THRESHOLD 100
/* Append messages from this mbox file to mailboxes */
#define MBOX_PATH "/home/tsiraine/mail/dovecot.mbox"

/* Use AUTHENTICATE command instead of LOGIN command */
#define USE_AUTHENTICATE
/* Append mails */
#define APPENDS
/* Expunge mails */
#define EXPUNGES
/* Randomly disconnect clients */
#define DISCONNECTS 20

struct {
  char *ip, *username_template, *password, *mbox_path;
  int port;
  unsigned int clients_count;
  int use_authenticate;
  int message_count_threshold;
} conf;

enum client_state {
	STATE_BANNER,
	STATE_AUTHENTICATE,
	STATE_LOGIN,
	STATE_SELECT,
	STATE_FETCH,
	STATE_FETCH2,
	STATE_STORE,
	STATE_EXPUNGE,
	STATE_APPEND,
        STATE_LOGOUT,

        STATE_COUNT
};

char *states[STATE_COUNT] = { "BANNER", "AUTHENTICATE", "LOGIN", "SELECT",
			      "FETCH", "FETCH2", "STORE", "EXPUNGE",
			      "APPEND", "LOGOUT" };

struct client {
        unsigned int idx;
        unsigned int cur, messages, deleted;
	struct istream *input;
	struct ostream *output;
	struct io *io;
	enum client_state state;
        struct mbox *mbox;

	uoff_t literal_length;
        int literal_nextline;
        time_t last_io;

	char *username;
};

struct mbox {
	int fd;
	char *path;
	struct istream *input;
};

static int clients_count = 0;
static struct ioloop *ioloop;
static unsigned int counters[STATE_COUNT], disconnects;
static struct client **clients;

void client_free(struct client *client);

static struct mbox *mbox_open(const char *path)
{
	struct mbox *mbox;

	mbox = i_new(struct mbox, 1);
	mbox->path = i_strdup(path);
	mbox->fd = open(path, O_RDONLY);
	if (mbox->fd == -1)
		i_fatal("open(%s) failed: %m", path);
	mbox->input =
		i_stream_create_file(mbox->fd, default_pool, (size_t)-1, FALSE);
	return mbox;
}

static void mbox_close(struct mbox *mbox)
{
	i_stream_unref(&mbox->input);
	(void)close(mbox->fd);
	i_free(mbox->path);
	i_free(mbox);
}

static uoff_t mbox_get_next_size(struct mbox *mbox)
{
	const char *line;
	uoff_t offset, last_offset, size;

	line = i_stream_read_next_line(mbox->input);
	if (line == NULL) {
		if (mbox->input->v_offset == 0)
			i_fatal("Empty mbox file: %s", mbox->path);
		i_stream_seek(mbox->input, 0);
		return mbox_get_next_size(mbox);
	}

	/* should be From-line */
	if (strncmp(line, "From ", 5) != 0) {
		if (mbox->input->v_offset == 0)
			i_fatal("Not a valid mbox file: %s", mbox->path);
		i_panic("From-line not found at %"PRIuUOFF_T,
			mbox->input->v_offset);
	}

        offset = last_offset = mbox->input->v_offset;
        while ((line = i_stream_read_next_line(mbox->input)) != NULL) {
                if (strncmp(line, "From ", 5) == 0) {
                        if (offset != last_offset)
                                break;

                        /* empty body */
                        offset = last_offset;
                }
                last_offset = mbox->input->v_offset;
        }
        if (offset == last_offset)
                i_fatal("mbox file ends with From-line: %s", mbox->path);

        size = last_offset - offset;
        i_stream_seek(mbox->input, offset);
	return size;
}

static int client_parse_literal(struct client *client)
{
	const unsigned char *data;
	size_t size;

	data = i_stream_get_data(client->input, &size);
	if (size < client->literal_length) {
		i_stream_skip(client->input, size);
		client->literal_length -= size;
		return FALSE;
	}

	i_stream_skip(client->input, client->literal_length);
	client->literal_length = 0;
	client->literal_nextline = TRUE;
	return TRUE;
}

static int client_skip_line(struct client *client, const char *line)
{
	size_t len = strlen(line);

	if (line[len-1] != '}')
		return TRUE;
	len--;

	/* skip literal */
	while (len > 0 && i_isdigit(line[len-1]))
		len--;
	if (len == 0 || line[len-1] != '{') {
		i_error("Invalid literal in: %s", line);
		client_free(client);
		return FALSE;
	}
	client->literal_length = strtoul(t_strcut(line+len, '}'), NULL, 10);
	return client_parse_literal(client);
}

static void client_handle_untagged(struct client *client, const char *line)
{
	if (i_isdigit(line[3])) {
		const char *p = strchr(line+2, ' ');
		unsigned int num;

		if (p++ == NULL) {
			i_error("%s: Unexpected: %s", client->username, line);
			return;
		}

		num = strtoul(t_strcut(line+2, ' '), NULL, 10);
		if (strcasecmp(p, "EXISTS") == 0)
			client->messages = num;

                if (num > client->messages &&
		    client->state > STATE_SELECT) {
			i_fatal("%s: seq too high (%u > %u, state=%d): %s",
                                client->username, num, client->messages,
                                client->state, line);
		}

		if (strcasecmp(p, "EXPUNGE") == 0)
			client->messages--;
	} else if (strncasecmp(line+3, "BYE ", 4) == 0) {
		i_info("%s: %s", client->username, line + 3);
	}
}

static int client_append(struct client *client)
{
	const char *cmd;
	struct istream *input;
	uoff_t size, next_offset;
	off_t ret;

	size = mbox_get_next_size(client->mbox);
	cmd = t_strdup_printf("a APPEND INBOX {%"PRIuUOFF_T"+}\r\n", size);
        o_stream_send_str(client->output, cmd);

        next_offset = client->mbox->input->v_offset + size;
	input = i_stream_create_limit(default_pool, client->mbox->input,
				      client->mbox->input->v_offset, size);
	ret = o_stream_send_istream(client->output, input);
        i_stream_unref(&input);
        o_stream_send_str(client->output, "\r\n");

	if (ret != (off_t)size) {
		if (ret < 0)
			i_error("APPEND failed: %m");
		else {
			i_error("APPEND failed: Sent only %"
				PRIuUOFF_T" of %"PRIuUOFF_T,
				ret, size);
                }
                i_stream_seek(client->mbox->input, next_offset);
		return -1;
	}
	return 0;
}

static void client_input(void *context)
{
	struct client *client = context;
	const char *line, *str;
	string_t *cmd, *buf;
	unsigned int i;

        client->last_io = ioloop_time;

	switch (i_stream_read(client->input)) {
	case 0:
		return;
	case -1:
		/* disconnected */
		client_free(client);
		return;
	case -2:
		/* buffer full */
		i_error("line too long");
		client_free(client);
		return;
	}

	if (client->literal_length > 0) {
		if (!client_parse_literal(client))
			return;
	}

	while ((line = i_stream_next_line(client->input)) != NULL) {
		if (client->literal_nextline) {
			client->literal_nextline = FALSE;
			if (!client_skip_line(client, line))
				return;
			continue;
		}

		if (strncmp(line, "* ", 2) == 0 &&
		    client->state != STATE_BANNER) {
			client_handle_untagged(client, line);
			if (!client_skip_line(client, line))
				return;
			continue;
		}

		switch (client->state) {
		case STATE_BANNER:
		  if(conf.use_authenticate)
		    {
			str = "l AUTHENTICATE plain\r\n";
                        client->state = STATE_AUTHENTICATE;
		    }
		  else
		    {
                        str = t_strdup_printf("l LOGIN \"%s\" \"%s\"\r\n",
                                              client->username, conf.password);
                        client->state = STATE_LOGIN;
		    }
			o_stream_send_str(client->output, str);
			break;
		case STATE_AUTHENTICATE:
			buf = t_str_new(512);
			cmd = t_str_new(512);
			str_append_c(buf, '\0');
			str_append(buf, client->username);
			str_append_c(buf, '\0');
			str_append(buf, conf.password);
			base64_encode(buf->data, buf->used, cmd);
			str_append(cmd, "\r\n");

			o_stream_send_str(client->output, str_c(cmd));
			client->state = STATE_LOGIN;
			break;
		case STATE_LOGIN:
			if (strncasecmp(line, "l OK", 4) != 0) {
				i_error("Login failed: %s", line);
				client_free(client);
				return;
			}
			o_stream_send_str(client->output, "s SELECT INBOX\r\n");
			client->state = STATE_SELECT;
			break;
		case STATE_SELECT:
			if (strncasecmp(line, "s OK", 4) != 0) {
				i_error("SELECT failed: %s", line);
				client_free(client);
				return;
			}
			/*i_info("%s: %u messages", client->username,
			       client->messages);*/

			if (client->messages == 0)
				goto __append;

			o_stream_send_str(client->output,
				"f FETCH 1:* (UID FLAGS ENVELOPE BODY)\r\n");
			client->state = STATE_FETCH;
			break;
		case STATE_FETCH:
			if (strncasecmp(line, "f OK", 4) != 0) {
				i_error("FETCH failed: %s", line);
				client_free(client);
				return;
			}
			o_stream_send_str(client->output,
				t_strdup_printf("f2 FETCH %lu (BODY[])\r\n",
					       (random() % client->messages) + 1));
			client->state = STATE_FETCH2;
			break;
		case STATE_FETCH2:
			if (strncasecmp(line, "f2 OK", 5) != 0) {
				i_error("FETCH BODY[] failed: %s", line);
				client_free(client);
				return;
                        }

#ifndef EXPUNGES
                        goto __append;
#endif

			/* delete them randomly */
			cmd = t_str_new(512);
			str_append(cmd, "s STORE ");
			for (i = 1; i <= client->messages; i++) {
				if ((random() % 10) == 0) {
					str_printfa(cmd, "%u,", i);
					client->deleted++;
				}
			}
			if (client->deleted == 0)
				goto __append;

			str_truncate(cmd, str_len(cmd) - 1);
			str_append(cmd, " +FLAGS.SILENT \\Deleted\r\n");
			o_stream_send_str(client->output, str_c(cmd));
			client->state = STATE_STORE;
			break;
		case STATE_STORE:
			if (strncasecmp(line, "s OK", 4) != 0) {
				i_error("STORE failed: %s", line);
				client_free(client);
				return;
			}

			o_stream_send_str(client->output, "e EXPUNGE\r\n");
			client->state = STATE_EXPUNGE;
			break;
		case STATE_EXPUNGE:
			if (strncasecmp(line, "e OK", 4) != 0) {
				i_error("EXPUNGE failed: %s", line);
				client_free(client);
				return;
			}

                __append:
#ifndef APPENDS
                        goto __logout;
#endif
			if (client->messages >= conf.message_count_threshold)
				goto __logout;

			if (client_append(client) < 0) {
				client_free(client);
				return;
			}
			client->state = STATE_APPEND;
			break;
		case STATE_APPEND:
			if (strncasecmp(line, "a OK", 4) != 0) {
				i_error("APPEND failed: %s", line);
				client_free(client);
				return;
			}

			if (client->messages < conf.message_count_threshold &&
			    (rand() % 5) == 0) {
				/* append more */
				if (client_append(client) < 0) {
					client_free(client);
					return;
				}
				break;
			}

		__logout:
			o_stream_send_str(client->output, "x LOGOUT\r\n");
			client->state = STATE_LOGOUT;
			break;
		case STATE_LOGOUT:
                        break;
                case STATE_COUNT:
                        i_unreached();
                }
                counters[client->state]++;
#ifdef DISCONNECTS
		if ((rand() % DISCONNECTS) == 0) {
			/* random disconnection */
                        disconnects++;
                        client_free(client);
			break;
                }
#endif
	}
}

static int client_output(void *context)
{
        struct client *client = context;
	int ret;

	ret = o_stream_flush(client->output);
        client->last_io = ioloop_time;
        return ret;
}

struct client *client_new(unsigned int idx, struct mbox *mbox)
{
	struct client *client;
	struct ip_addr ip;
	int fd;

	net_addr2ip(conf.ip, &ip);
	fd = net_connect_ip(&ip, conf.port, NULL);
	if (fd < 0) {
		i_error("connect() failed: %m");
		return NULL;
	}

	client = i_new(struct client, 1);
        client->idx = idx;
        client->mbox = mbox;
	client->input = i_stream_create_file(fd, default_pool, 65536, TRUE);
	client->output = o_stream_create_file(fd, default_pool, (size_t)-1, FALSE);
	client->io = io_add(fd, IO_READ, client_input, client);
	client->username = i_strdup_printf(conf.username_template,
					   (random() % USER_RAND) + 1,
					   (random() % DOMAIN_RAND) + 1);
        client->last_io = ioloop_time;
        o_stream_set_flush_callback(client->output, client_output, client);
	clients_count++;

        clients[idx] = client;
        return client;
}

void client_free(struct client *client)
{
	struct mbox *mbox = client->mbox;
        unsigned int idx = client->idx;

	--clients_count;
        clients[idx] = NULL;

	/*if (--clients_count == 0)
		io_loop_stop(ioloop);*/
	io_remove(&client->io);
	o_stream_unref(&client->output);
	i_stream_unref(&client->input);
	i_free(client->username);
	i_free(client);

	client_new(idx, mbox);
}

static void print_header(void)
{
        printf("Auth Logi Sele Fetc Fet2 Stor Expu Appe Logo Disc\n");
}

static void print_timeout(void *context __attr_unused__)
{
        static int rowcount = 0;
        unsigned int i;

        if ((rowcount++ % 10) == 0)
                print_header();

        for (i = 1; i < STATE_COUNT; i++) {
                printf("%4d ", counters[i]);
                counters[i] = 0;
        }
        printf("%4d\n", disconnects);

        for (i = 0; i < conf.clients_count; i++) {
                if (clients[i] != NULL &&
                    clients[i]->last_io < ioloop_time - 15) {
                        printf(" - %d. stalled for %u secs in %s\n", i,
                               (unsigned)(ioloop_time - clients[i]->last_io),
                               states[clients[i]->state]);
                }
        }

        disconnects = 0;
}

int main(int argc, char **argv)
{
	struct mbox *mbox;
	char *p, *s;
	unsigned int i;

	while(argc-- > 1)
	  {
	    if( (!strcmp(argv[argc], "-h")) || (!strcmp(argv[argc], "--help")))
	      {
		printf("imaptest [USER@IP:PORTNO] [pass=PASSWORD] [mbox=MBOX] [clients=CC] [msgs=NMSG] [use_authenticate] [PORTNO]\n");
		printf("USER = template for username. \"u%%04d\" will generate\nusers \"u0001\" to \"u0100\"\n");
		printf("MBOX = path to mbox from which we read mails to append.\n");
		printf("CC   = number of concurrent clients. [20]\n");
		printf("NMSG = target number of messages in INBOX. [100]\n");
		printf("If use_authenticate is given AUTHENTICATE will be used instead of LOGIN.\n");
		exit(0);
	      }

	    if(strcmp(argv[argc], "use_authenticate")==0)
	      {
		conf.use_authenticate = 1;
		continue;
	      }

	    /* pass=password */
	    if(strncmp(argv[argc], "pass=", 5)==0)
	      {
		conf.password = argv[argc] + 5;
		continue;
	      }

	    /* mbox=path */
	    if(strncmp(argv[argc], "mbox=", 5)==0)
	      {
		conf.mbox_path = argv[argc] + 5;
		continue;
	      }

	    /* clients=# */
	    if(strncmp(argv[argc], "clients=", 8)==0)
	      {
		conf.clients_count = atoi(argv[argc] + 8);
		continue;
	      }

	    /* msgs=# */
	    if(strncmp(argv[argc], "msgs=", 5)==0)
	      {
		conf.message_count_threshold = atoi(argv[argc] + 5);
		continue;
	      }

	    /* user@ip:port */
	    if( (p = strchr(argv[argc], '@')))
	      {
		if(p != argv[argc])
		  {
		    conf.username_template = strdup(argv[argc]);
		    s = strchr(conf.username_template, '@');
		    *s = 0;
		  }
		if(*(p+1) && (*(p+1) != ':'))
		  {
		    conf.ip = strdup(p+1);
		    s = strchr(conf.ip, ':');
		    if(s)
		      {
			*s = 0;
			conf.port = atoi(s+1);
		      }
		  }
		if(*(p+1) && (*(p+1) == ':') && *(p+2))
		  conf.port = atoi(p+2);
		continue;
	      }

	    /* 143 */
	    if(!conf.port)
	      for(p=argv[argc];*p;p++)
		{
		  if(!isdigit(*p))
		    break;
		  if(!*(p+1))
		    conf.port = atoi(argv[argc]);
		}
	  }

	if(!conf.password) conf.password = "pass";
	if(!conf.username_template) conf.username_template = "test";
	if(!conf.port) conf.port = 143;
	if(!conf.ip) conf.ip = "127.0.0.1";
	if(!conf.mbox_path) conf.mbox_path = "/root/imaptest.mbox";
	if(!conf.clients_count) conf.clients_count = 20;
	if(!conf.message_count_threshold) conf.message_count_threshold = 100;
	
	clients = malloc(conf.clients_count * sizeof(struct client*));
	memset(clients, 0, conf.clients_count * sizeof(struct client*));
	
	printf("Using %s@%s:%d\n pass=\"%s\"\n mbox=\"%s\"\n client_count=%d\n authcmd=%s\n messages in mbox (target)=%d\n",
	       conf.username_template,
	       conf.ip,
	       conf.port,
	       conf.password,
	       conf.mbox_path,
	       conf.clients_count,
	       conf.use_authenticate?"AUTHENTICATE":"LOGIN",
	       conf.message_count_threshold);
	
	lib_init();
	ioloop = io_loop_create(system_pool);

	mbox = mbox_open(conf.mbox_path);

	timeout_add(1000, print_timeout, NULL);
	for (i = 0; i < conf.clients_count; i++)
		client_new(i, mbox);
        io_loop_run(ioloop);

	mbox_close(mbox);
	io_loop_destroy(&ioloop);
	lib_deinit();
	return 0;
}