#include <stdio.h>
#include <stdlib.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <unistd.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <syslog.h>
#include <signal.h>

#define GLBASE "/var/log/graylist"
#define WHITE_DAYS 3
#define GRAY_MINS 20
#define QUICKLIST_SMTP 1

/*
	graylist.c - quick and dirty graylisting implementation for qmail
	hdm [at] metasploit.com
*/

void tick_tock(int signo __attribute__ ((unused)) ) {
	return;
}

int createpath (char *path) {
	char cmd[1024];
	snprintf(cmd, 1024, "mkdir -p %s", path);
	system(cmd);
}

char *addr_to_path (char *base, char *ip) {
	struct in_addr addr;
	unsigned char a[4];
	char *p;
	
	inet_aton(ip, &addr);
	memcpy(a, &addr.s_addr, 4);
	
	p = malloc(32+strlen(base)+1);
	if (p == NULL) return NULL;
	
	/* we only graylist /24 blocks to speed up delivery from smtp farms */
	snprintf(p, 32+strlen(base), "%s/%d/%d/%d", base, a[0], a[1], a[2]);
	return p;
}


int is_mailserver(char *ip) {
	int s;
	fd_set fds;
	struct timeval tv;
	struct sockaddr_in addr;
	struct sigaction sa;
	
	memset(&addr, 0, sizeof(addr));
	inet_aton(ip, &addr.sin_addr);
	addr.sin_port = htons(25);
	addr.sin_family = AF_INET;
	
	memset(&sa, 0, sizeof(sa));
	sa.sa_handler = tick_tock;

	sigaction(SIGALRM, &sa, NULL);	
	
	s = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
	if (s < 0) return(0);

	alarm(10);
	if(connect(s, (struct sockaddr *) &addr, sizeof(addr)) < 0) {
		close(s);
		return(0);
	}
	
	write(s, "EHLO graylisting.engine\r\n", strlen("EHLO graylisting.engine\r\n"));
	write(s, "QUIT\r\n", 6);
	
	alarm(0);
	
	close(s);
	return(1);
}

int is_permalist(char *ip) {
	char *base = GLBASE "/perma";
	char *path = addr_to_path(base, ip);
	struct stat fs;
	
	if (path == NULL) {
		return 0;
	}
	
	/* return zero if the path does not exist */
	if( stat(path, &fs)) {
		free(path);
		return(0);
	}
	
	free(path);
	return(1);
}

int is_whitelist(char *ip) {
	char *base = GLBASE "/white";
	char *path = addr_to_path(base, ip);
	struct stat fs;
	signed long diff;
	
	/* fail over to allow everyone */
	if (path == NULL) {
		return 1;
	}
	
	/* return zero if the path does not exist */
	if( stat(path, &fs)) {
		free(path);
		return(0);
	}
	
	free(path);
	
	diff = time(NULL) - fs.st_ctime;
	
	/* whitelist expiration disabled (handled by cron)
	if (diff > (3600 * WHITE_DAYS)) {
		// syslog(LOG_INFO, "system %s was expired from the whitelist (%d seconds old)", ip, diff);
		rmdir(path);
		return(0);
	}
	*/
	
	return(1);
}

int is_graylist(char *ip) {
	char *base = GLBASE "/gray";
	char *path = addr_to_path(base, ip);
	struct stat fs;
	signed int diff;
	
	/* fail over to allow everyone */
	if (path == NULL) {
		return 1;
	}
	
	/* return zero if the path does not exist */
	if( stat(path, &fs)) {
		
		/* check to see if the source is running smtpd */
		if (QUICKLIST_SMTP && is_mailserver(ip)) {
			// syslog(LOG_INFO, "system %s was quicklisted out of the graylist", ip);
			return(1);
		}
		
		// syslog(LOG_INFO, "system %s was added to the graylist queue", ip);		
		free(path);
		return(0);
	}
	
	free(path);
	
	diff = time(NULL) - fs.st_ctime;
	/* see if enough time has elapsed since the last attempt */
	if (diff < (60 * GRAY_MINS)) {
		// syslog(LOG_INFO, "system %s reconnected before the timeout expired (%d seconds old)", ip, diff);
		return(0);
	}
	
	return(1);
}


/* 
	Remove the graylist entry if it exists
	Add a new whitelist entry
*/

void set_whitelist(char *ip) {
	char *wbase = GLBASE "/white";
	char *gbase = GLBASE "/gray";
	char *wpath = addr_to_path(wbase, ip);
	char *gpath = addr_to_path(gbase, ip);
	if (wpath == NULL || gpath == NULL) return;
	
	if (! createpath(wpath)) return;

	rmdir(gpath);

	free(wpath);
	free(gpath);
}

void set_graylist(char *ip) {
	char *gbase = GLBASE "/gray";
	char *gpath = addr_to_path(gbase, ip);
	if (gpath == NULL) return;	
	if (! createpath(gpath)) return;
	free(gpath);
}

int main (int argc, char **argv, char **envp) {
	char name[1024];
	gethostname(name, sizeof(name)-1);

	openlog("graylist", LOG_NDELAY | LOG_PID, LOG_DAEMON);
	
	char *ip  = getenv("TCPREMOTEIP");
	do
	{
		/* passthrough when the env is not set */
		if (ip == NULL) {
			syslog(LOG_INFO, "invalid ip address specified");
			ip = "UNKNOWN";
			break;
		}

		/* allow all permalisted systems */
		if (is_permalist(ip)) {
			// syslog(LOG_INFO, "system %s was found in the permalist", ip);
			break;
		}
				
		/* allow all whitelisted systems */
		if (is_whitelist(ip)) {
			// syslog(LOG_INFO, "system %s was found in the whitelist", ip);
			break;
		}
		
		/* allow all reconnecting graylisted system to whitelist */
		if (is_graylist(ip)) {
			// syslog(LOG_INFO, "system %s was moved to whitelist", ip);
			set_whitelist(ip);
			break;
		}
		
		/* add a graylist entry and delay the message */
		set_graylist(ip);

		/* display the temporary error message */
		printf("220 %s.metasploit.com 4.7.1 This system enforces graylisting...\r\n", name);
		sleep(5);
		
		printf("451 4.7.1 Come back later...\r\n");
		sleep(5);
		
		return(0);

	} while (0);

	closelog();
	
	argv[0] = "/var/qmail/bin/qmail-smtpd";
	execve(argv[0], argv, envp);
	return(0);
}