/*
 * Transparent Imap Redirection Proxy
 *   version 20071008
 *
 * (c)2001, (c)2007 by Willem van Engen
 * BSD-style copyright
 *
 * To be used with inetd or something alike.
 *
 * Terms used:
 *   proxy   = this program, imap redirection proxy
 *   client  = imap client connecting to the proxy
 *   server  = imap server that the client wants to access
 *
 * This program uses POSIX threads. It can be compiled on FreeBSD like this
 *   gcc -o tirp main.c -pthread
 * On Linux, this should work
 *   gcc -lpthread -o tirp main.c
 * But first update the settings below (look for CHANGE ME).
 * Then install it somewhere and use something like inetd to start it. E.g.
 * add to /etc/inetd.conf:
 *   imap    stream  tcp     nowait  nobody  /usr/libexec/tirp       tirp
 */

/* ------ CHANGE ME ------- */
/* Set this to the imap server and port you would normally use in the client.
   Now the client connects to this proxy and this proxy connects to the
   server instead. */
#define SERVERNAME	"studentex3.student.tue.nl"
#define PORT		"imap"
/* --- END OF CHANGE ME --- */

/* #define DEBUG
#define DEBUGFILE "/tmp/tirpout" */

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

#include <netinet/in.h>
#include <netdb.h>

#include <stdio.h>
#include <stdlib.h>
#include <pthread.h>
#include <string.h>
#include <unistd.h>
#include <stdarg.h>
#include <fcntl.h>
#include <err.h>

#define IMAPBUFSIZE	8192		/* Grabbed from sylpheed's imap.h */
#define IMAPCMDSIZE	10
#define IMAPREFSIZE	100
#define SEPSTR		" \n\r"
#define HOSTNAMESIZE 100

struct serverdata {
	int socket;		/* Socket for communication with server */
	char *cmdlogin;	/* Command to log into server */
	char hold;		/* When hold==0, data is passed through */
					/* When hold==1, redirection is in progress */
					/* When hold==-1,the proxy can terminate */
	char hasrefs;	/* if hasrefs==1 the server can return referrals */
	char mailrefs;	/* if mailrefs==1 rlist and rsub commands should be used */
	char oldcmd[IMAPBUFSIZE];
	int oldcmdlen;
};
struct pthargs {
	int fd;
	struct serverdata *sd;
};

#define IMAPR_BYE		"bye"
#define IMAPR_BYE_SIZE		3
#define IMAPR_NOREF			"no [referral "
#define IMAPR_NOREF_SIZE	13
#define IMAPR_BYEREF		"bye [referral "
#define IMAPR_BYEREF_SIZE	14
#define IMAPCMD_LOGIN		"login "
#define IMAPCMD_LOGIN_SIZE	6
#define IMAPCMD_AUTH		"authenticate "
#define IMAPCMD_AUTH_SIZE	13
#define IMAPCMD_LIST		"list "
#define IMAPCMD_LIST_SIZE	5
#define IMAPCMD_SUB			"sub "
#define IMAPCMD_SUB_SIZE	4
#define IMAP_GETCAPS		"z8492 CAPABILITY\r\n"
#define IMAP_GETCAPS_SIZE	18
#define IMAP_LOGOUT			"z8493 LOGOUT\r\n"
#define IMAP_LOGOUT_SIZE	14

#define IMAP_CAP_LOGINREF	"LOGIN-REFERRALS"
#define IMAP_CAP_MAILREF	"MAILBOX-REFERRALS"

int socket_connect(const char*,const char*);
void pass_c2s(struct pthargs*);
void pass_s2c(struct pthargs*);
int dgets(char*,int,int);
void getcaps(struct serverdata*);
int newnstr(char**,char*,int);
void insertchar(char*,char,int);
int newconnection(struct serverdata*,char*);
void debug(char*, ...);

#ifdef DEBUG
FILE *debugfile;
#endif

#define ERRBUFSIZE	80
char errorbuf[ERRBUFSIZE];

int main(int argc,char *argv[]) {
	struct serverdata sd={0,NULL,0,0,0};
	pthread_t pthc2s,pths2c;
	struct pthargs ptas2c={STDOUT_FILENO,&sd},ptac2s={STDIN_FILENO,&sd};

#ifdef DEBUG
	if ( (debugfile=fopen("DEBUGFILE","w"))==NULL)
		err(1,"Couldn't open debugfile "DEBUGFILE);
#endif

	/* First setup connection with default server */
	if ( (sd.socket=socket_connect(SERVERNAME,PORT))==-1)
		err(1,errorbuf);

	/* Pass through greeting message from server */
	sd.oldcmdlen=dgets(sd.oldcmd,IMAPBUFSIZE,sd.socket);
	write(STDOUT_FILENO,sd.oldcmd,sd.oldcmdlen);
	/* Get capabilities */
	getcaps(&sd);

	/* Create threads for client->server and server->client */
	if (pthread_create(&pthc2s,NULL,(void*)pass_c2s,&ptac2s)!=0)
		err(1,"couldn't create thread pthc2s");
	if (pthread_create(&pths2c,NULL,(void*)pass_s2c,&ptas2c)!=0)
		err(1,"couldn't create thread pths2c");

	/* Wait for threads to finish */
	pthread_join(pthc2s,NULL);
	pthread_join(pths2c,NULL);

	/* Cleanup */
	close(sd.socket);
	if (sd.cmdlogin!=NULL) free(sd.cmdlogin);
#ifdef DEBUG
	fclose(debugfile);
#endif

	return 0;
}

/*
 * pass_c2s
 *
 * Passes data from client to server as long as hold==0. And catches
 * username and password.
 */
void pass_c2s(struct pthargs *pa) {
	char buf[IMAPBUFSIZE+1];
	int cmdstart;
	int buflen=0;

	while (pa->sd->hold!=-1 && buflen>=0) {
		if (pa->sd->hold==0) {
			buflen=dgets(buf,IMAPBUFSIZE,pa->fd);

			/* Command starts after tag */
			if ( (cmdstart=strcspn(buf,SEPSTR)+1)<buflen ) {
				/* save imap login */
				if (strncasecmp(buf+cmdstart,IMAPCMD_LOGIN,IMAPCMD_LOGIN_SIZE)==0) {
#ifdef DEBUG
					if (newnstr(&pa->sd->cmdlogin,buf,buflen)==-1)
						debug("--] get logincmd failed: %s\n",errorbuf);
					else
						debug("--] get logincmd = '%s'\n",pa->sd->cmdlogin);
#else
					newnstr(&pa->sd->cmdlogin,buf,buflen);
#endif
				} else
				/* list -> rlist */
				if (pa->sd->mailrefs == 1 && strncasecmp(buf+cmdstart,
										IMAPCMD_LIST,IMAPCMD_LIST_SIZE)==0) {
					insertchar(buf+cmdstart,'r',IMAPBUFSIZE+1-cmdstart);
					buflen++;
				} else
				/* sub -> rsub */
				if (pa->sd->mailrefs == 1 && strncasecmp(buf+cmdstart,
										IMAPCMD_SUB,IMAPCMD_SUB_SIZE)==0) {
					insertchar(buf+cmdstart,'r',IMAPBUFSIZE+1-cmdstart);
					buflen++;
				}
			}

			debug("--> %s\n",buf);

			memcpy(pa->sd->oldcmd,buf,buflen);
			pa->sd->oldcmdlen=buflen;
			write(pa->sd->socket,buf,buflen);
		} else usleep(100000);
	}
}

/*
 * pass_s2c
 *
 * Passes data from server to client. Catches redirections and handles
 * them so they become invisible for the client.
 * XXX at logout, proxy terminates at BYE and discards everything after
 * (like 'OK logout' message)
 */
void pass_s2c(struct pthargs *pa) {
	char buf[IMAPBUFSIZE];
	int cmdstart;
	int buflen=0;
	char *ref,*p;

	while (pa->sd->hold!=-1 && buflen>=0) {
		buflen=dgets(buf,IMAPBUFSIZE,pa->sd->socket);

		if ( (cmdstart=strcspn(buf,SEPSTR)+1)<buflen ) {
			if ( pa->sd->hasrefs==1 &&
					strncasecmp(buf+cmdstart,IMAPR_NOREF,IMAPR_NOREF_SIZE)==0){
				/* There is a referral !! */
				ref=buf+cmdstart+IMAPR_NOREF_SIZE;
				if ( (p=strchr(ref,']'))!=NULL) {
					pa->sd->hold=1;
					*p=0;
					debug("--> NO Referral: %s\n",ref);
					if (newconnection(pa->sd,ref)!=-1) {
						/*  And issue old command on new server */
						write(pa->sd->socket,pa->sd->oldcmd,pa->sd->oldcmdlen);
						pa->sd->hold=0;
						continue;
					} else {
						/* Connection failed: return error */
						snprintf(buf+cmdstart,IMAPBUFSIZE-cmdstart,
									"NO [REFERRAL %s] not available: %s\n",
									ref,errorbuf); 
						buflen=strlen(buf);
					}
					pa->sd->hold=0;
				}
			} else if (pa->sd->hasrefs==1 &&
				strncasecmp(buf+cmdstart,IMAPR_BYEREF,IMAPR_BYEREF_SIZE)==0){
				/* A BYE referral ! */
				ref=buf+cmdstart+IMAPR_BYEREF_SIZE;
				if ( (p=strchr(ref,']'))!=NULL) {
					pa->sd->hold=1;
					*p=0;
					debug("--> BYE Referral: %s\n",ref);
					/* Setup new connection */
					if (newconnection(pa->sd,ref)!=-1) {
						/*  And issue old command on new server */
						write(pa->sd->socket,pa->sd->oldcmd,pa->sd->oldcmdlen);
						pa->sd->hold=0;
						continue;
					} else
						/* Server disconn'd with BYE so proxy disconnects too */
						pa->sd->hold=-1;
				}
			} else if (strncasecmp(buf+cmdstart,IMAPR_BYE,IMAPR_BYE_SIZE)==0) {
				/* Server wants to disconnect */
				/* XXX only bye message is sent to client, not OK response */
				pa->sd->hold=-1;
			}
		}
			
		debug("--< %s\n",buf);
		write(pa->fd,buf,buflen);
	}
}

/*
 * Connect socket.
 *
 * After an example in "Internetworking with TCP/IP, Vol III", Douglas E. Comer
 * and David L. Stevens, Prentice Hall, ISBN 0-13-262148-7.
 * Returns the socket descriptor if succeeded. If failed, returns -1 and
 * failure string is in errorbuf.
 */
int socket_connect(const char *host, const char *service) {
	struct hostent *phe;
	struct servent *pse;
	struct sockaddr_in sin;
	int s;

	memset(&sin, 0, sizeof(sin));
	sin.sin_family = AF_INET;
#ifdef __LINUX__
	sin.sin_len=sizeof(struct sockaddr_in);
#endif

	/* Get portnumber from service name */
	if ( pse=getservbyname(service,"tcp") )
		sin.sin_port = pse->s_port;
	else if ( (sin.sin_port = htons((u_short)atoi(service))) == 0 ) {
		snprintf(errorbuf,ERRBUFSIZE,"don't know service \"%s\"",service);
		return -1;
	}

	/* Get ip address from hostname */
	if ( phe=gethostbyname(host) )
		memcpy(&sin.sin_addr, phe->h_addr, phe->h_length);
	else if ( (sin.sin_addr.s_addr=inet_addr(host)) == INADDR_NONE) {
		snprintf(errorbuf,ERRBUFSIZE,"don't know host \"%s\"",host);
		return -1;
	}

	/* Create socket */
	if ((s=socket(PF_INET, SOCK_STREAM, 0))==-1) {
		snprintf(errorbuf,ERRBUFSIZE,"can't create socket");
		return -1;
	}

	/* Connect socket */
	if (connect(s, (struct sockaddr*)&sin, sizeof(sin))==-1) {
		snprintf(errorbuf,ERRBUFSIZE,"can't connect to %s.%s",host,service);
		return -1;
	}

	return s;
}

/*
 * dgets
 *
 * fgets for a file descriptor/socket descriptor, but returns the number
 * of characters read.
 */
int dgets(char *str,int size,int fd) {
	int i;

	for (i=0;i<size-1;i++) {
		if (read(fd,str+i,1)<0) return -1;
		if ( *(str+i)=='\n' ) {
			i++;
			break;
		}
	}
	str[i]=0;
	return i;
}

/*
 * getcaps
 *
 * Gets capabilities from server and sets the struct serverdata's entries
 * according to them
 */
void getcaps(struct serverdata *sd) {
	char buf[IMAPBUFSIZE];
	int buflen;


	/* Ask for capabilities */
	write(sd->socket,IMAP_GETCAPS,IMAP_GETCAPS_SIZE);
	buflen=dgets(buf,IMAPBUFSIZE,sd->socket);

	/* Search if referrals are in caps */
	if (strstr(buf,IMAP_CAP_LOGINREF)!=NULL) {
		sd->hasrefs=1;
	}
	if (strstr(buf,IMAP_CAP_MAILREF)!=NULL) {
		sd->hasrefs=1;
		sd->mailrefs=1;
	}
	
	/* Read in "CAPABILITY COMPLETED" line */
	dgets(buf,IMAPBUFSIZE,sd->socket);
}

/*
 * newnstr
 *
 * allocates memory and copies string
 */
int newnstr(char **dest,char *src,int maxsize) {
	int size;

	if ( (size=strlen(src))>maxsize ) size=maxsize;

	if (*dest!=NULL) free(*dest);

	if ( (*dest=malloc(size))==NULL ) {
		snprintf(errorbuf,ERRBUFSIZE,
				"Couldn't allocate %d bytes of memory in newnstr",size);
		return -1;
	}

	strncpy(*dest,src,maxsize);
	dest[maxsize]=0;

	return 0;
}

/*
 * insertchar
 *
 * inserts a char into an array. Array length is including that one!
 */
void insertchar(char *buf,char ch,int maxlen) {
	int i;
	int size;

	/* Check length */
	if ( (size=(strlen(buf)+1))>maxlen ) size=maxlen;

	/* Shift everything */
	for (i=size-2;i>=0;i--) {
		buf[i+1]=buf[i];
	}
	/* And insert char */
	buf[0]=ch;
}

/*
 * newconnection
 *
 * Closes the old connection if needed and sets up a new connection with server
 * pointed to by supplied url (has to be imap://....). When new server cannot
 * be connected, connection with old server is retained and -1 is returned.
 * Also when the url is misformed (or not supported by the proxy), -1 is
 * returned.
 * BUGS (this function completely violates the rfc's):
 * - ignores everything in hostname part until '@'
 */
int newconnection(struct serverdata *sd,char *url) {
	int newsocket;
	char *ptr=url,*pt;
	char host[HOSTNAMESIZE];
	char buf[IMAPBUFSIZE];
	char *port=NULL;

	/* parse imap url; get servername, port if given. What else? */
	if (strncasecmp(url,"imap://",7)!=0) return -1; else ptr+=7;
	/* If there's an '@', skip everything before it. Now we're at hostname */
	if ( (pt=strchr(ptr,'@'))!=NULL) ptr=pt+1;
	/* Find end of hostname. If no slash at end, we only have a hostname */
	if ( (pt=strchr(ptr,'/'))==NULL) pt=strlen(url)+url-1;
	if ( (pt-ptr)>HOSTNAMESIZE ) pt=ptr+HOSTNAMESIZE;
	strncpy(host,ptr,pt-ptr);
	host[pt-ptr]=0;
	ptr=pt+1;
	if (ptr<buf+strlen(url)) {
		if ( (pt=strchr(host,':'))!=NULL) {
			*pt=0;
			port=pt+1;
		}
	}
	
	/* Connect */
	if ( (newsocket=socket_connect(host,port==NULL?"imap":port))==-1) {
		debug("--] connect failed: %s\n",errorbuf);
		return -1;
	}

	debug("--] connected to host \'%s\'\n",host);

	/* Get welcome message */
	dgets(buf,IMAPBUFSIZE,newsocket);
	debug("--< %s\n",buf);
	debug("--> %s\n",sd->cmdlogin);

	/* Login and such */
	if (sd->cmdlogin!=NULL) {
		write(newsocket,sd->cmdlogin,strlen(sd->cmdlogin));
		/* Get login answer */
		dgets(buf,IMAPBUFSIZE,newsocket);
		debug("--< %s\n",buf);
	}

	/* update sd with socket and servercaps */
	close(sd->socket);
	sd->socket=newsocket;
	getcaps(sd);
	
	return 0;
}

/*
 * DEBUG
 *
 * printf-like function for debug ouput
 */
void debug(char *str, ...) {
#ifdef DEBUG
	va_list ap;

	va_start(ap, str);
	vfprintf(debugfile,str,ap);
	va_end(ap);
	fflush(debugfile);
#endif
}