/*
 * Copyright (c) 2008 Hypertriton, Inc. <http://hypertriton.com/>
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 * 
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE FOR
 * ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE
 * USE OF THIS SOFTWARE EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

#include <sys/types.h>
#include <sys/wait.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <sys/un.h>
#include <sys/time.h>

#include <stdio.h>
#include <unistd.h>
#include <errno.h>
#include <string.h>
#include <syslog.h>
#include <stdarg.h>
#include <netdb.h>
#include <signal.h>
#include <math.h>

#include <netinet/in.h>
#include <arpa/inet.h>

#include "../config/version.h"
#include "../config/sysconfdir.h"

#include "mailblockd.h"
#include "ports.h"
#include "config.h"

#define MBD_TOTSCORE_MAX 3.0e+38
#define MBD_COUNT_MAX    (0xffffffff-1)
#define MBD_LOGFILE

char *mbdHome, *mbdHost, *mbdPort, *pidFile, *logFile, *mbdPassword;
char *pathHosts, *pathScores;
FILE *mbdLog = NULL;
float minScore, maxScore;
u_int maxEntries, nExpireEntries;

TBL *tblHosts;				/* Allowed hosts */
TBL *tblScores;				/* Spam score table */
TBL_Entry **expireEntries;		/* For expiration process */
u_int      *expireHashes;		/* For expiration process */

static char mbdError[1024];
static int mbdSocks[4];
static int mbdSockCount;

static volatile int die_flag = 0;
static volatile int flush_flag = 0;
static volatile int chld_flag = 0;
static fd_set servfds;

void
MBD_SetError(const char *fmt, ...)
{
	va_list ap;
	va_start(ap, fmt);
	vsnprintf(mbdError, sizeof(mbdError), fmt, ap);
	va_end(ap);
}
char *
MBD_GetError(void)
{
	return (mbdError);
}

void
MBD_Fatal(const char *fmt, ...)
{
	char err[128];
	va_list ap;

	va_start(ap, fmt);
	vsnprintf(err, sizeof(err), fmt, ap);
	va_end(ap);
	fputs(err, stderr);
	fputc('\n', stderr);
	exit(1);
}

void *
MBD_Malloc(size_t size)
{
	void *p;

	if ((p = malloc(size)) == NULL) {
		MBD_Fatal("Out of memory");
	}
	return (p);
}

#ifdef DEBUG
void
MBD_Debug(const char *fmt, ...)
{
#ifdef MBD_LOGFILE
	static char msg[160];
	va_list ap;

	va_start(ap, fmt);
	vsnprintf(msg, sizeof(msg), fmt, ap);
	va_end(ap);
	fprintf(mbdLog, "[%d] %s\n", getpid(), msg);
	fflush(mbdLog);
#endif
}
#endif /* DEBUG */

static void
ReportExit(const char *what, pid_t rpid, int code)
{
	if (WIFSIGNALED(code)) {
		syslog(LOG_ERR, "[%d] %s exited (signal %d)",
		    (int)rpid, what, WTERMSIG(code));
	} else if (WIFEXITED(code) && WEXITSTATUS(code) != 0) {
		syslog(LOG_ERR, "[%d] %s exited due to failure (%d)",
		    (int)rpid, what, WEXITSTATUS(code));
	}
}

static void
Reap(void)
{
	int save_errno = errno;
	int rpid, code;

	do {
		rpid = waitpid(-1, &code, WNOHANG);
		if (rpid > 0) {
			ReportExit("mailblockd", rpid, code);
		}
	} while (rpid > 0 || (rpid == -1 && errno == EINTR));

	errno = save_errno;
}

static void master_sig_die(int sigraised) { die_flag = 1; }
static void master_sig_flush(int sigraised) { flush_flag = 1; }
static void master_sig_chld(int sigraised) { chld_flag = 1; }

static void
CloseFiles(void)
{
	int i;

	for (i = 0; i < mbdSockCount; i++)
		close(mbdSocks[i]);
}

#if 0
void
MBD_EnterServerProc(const char *name)
{
	struct sigaction sa;

	CloseFiles();

	sigemptyset(&sa.sa_mask);
	sa.sa_flags = SA_RESTART;
	sa.sa_handler = SIG_DFL;
	sigaction(SIGCHLD, &sa, NULL);
	sa.sa_handler = SIG_IGN;
	sigaction(SIGURG, &sa, NULL);
	sigaction(SIGHUP, &sa, NULL);
	sigaction(SIGPIPE, &sa, NULL);

	sigfillset(&sa.sa_mask);
	sa.sa_flags = SA_RESTART;
	sa.sa_handler = child_sig_die;
	sigaction(SIGINT, &sa, NULL);
	sigaction(SIGQUIT, &sa, NULL);
	sigaction(SIGTERM, &sa, NULL);
	sigaction(SIGUSR1, &sa, NULL);

	closelog();
	openlog(name, LOG_PID|LOG_NDELAY, LOG_LOCAL0);
	Setproctitle("%s", name);
}

void
MBD_ExitServerProc(int code)
{
	closelog();
	exit(code);
}
#endif

#ifdef MBD_LOGFILE
static void
OpenLog(void)
{
	int fd;

	if (mbdLog != NULL) {
		fclose(mbdLog);
	}
	if ((mbdLog = fopen(logFile, "a")) == NULL) {
		syslog(LOG_ERR, "%s: %m", logFile);
		fprintf(stderr, "%s: %s\n", logFile, strerror(errno));
		exit(1);
	}
	fd = fileno(mbdLog);
	fchmod(fd, 0640);
	fchown(fd, 0, 20);
}
#endif /* MBD_LOGFILE */

static __inline__ int
ProcessSignals(void)
{
	if (die_flag) {
		die_flag = 0;
#if 0
		kill(0, SIGUSR1);
#endif
		return (1);
	}
	if (chld_flag) {
		chld_flag = 0;
		Reap();
	}
	if (flush_flag) {
		flush_flag = 0;
#ifdef MBD_LOGFILE
		OpenLog();
#endif
		if (TBL_Save(tblScores, pathScores, 0600) == -1)
			MBD_Fatal("Saving scores: %s", MBD_GetError());
	}
	return (0);
}

static int
ValidIPv4(const char *s)
{
	const char *p;
	int nDigits = 0;
	int nBlocks = 0;
	size_t len = strlen(s);

	if (len < 7 || len > 15) {
		return (0);
	}
	for (p = &s[0]; *p != '\0'; p++) {
		if (isdigit(*p)) {
			if (++nDigits > 3) {
				return (0);
			}
		} else if (*p == '.') {
			if (++nBlocks > 4) {
				return (0);
			}
			nDigits = 0;
		} else {
			return (0);
		}
	}
	return (1);
}

static void
ExpireOldScores(void)
{
	TBL_Entry *ent;
	u_int nOldest = 0;
	u_int i, j;

	for (i = 0; i < nExpireEntries; i++) {
		expireEntries[i] = NULL;
	}
	for (i = 0; i < tblScores->nbuckets; i++) {
		SLIST_FOREACH(ent, &tblScores->buckets[i].ents, bents) {
			for (j = 0; j < nExpireEntries; j++) {
				if (expireEntries[j] == NULL ||
				    ent->stamp < expireEntries[j]->stamp) {
				    	expireEntries[j] = ent;
				    	expireHashes[j] = i;
					break;
				}
			}
		}
	}
	MBD_Debug("Expiring %d entries", (int)nExpireEntries);
	for (i = 0; i < nExpireEntries; i++) {
		if (expireEntries[i] == NULL) {
			continue;
		}
		MBD_Debug("Expiring: [%d] %s (hash=%u)", i,
		    expireEntries[i]->key, expireHashes[i]);
		TBL_DeleteEntry(tblScores, expireHashes[i], expireEntries[i]);
	}
	if (TBL_Save(tblScores, pathScores, 0600) == -1)
		MBD_Debug("Flush failed: %s", MBD_GetError());
}

static int
ProcessScore(const char *addr, char *data)
{
	TBL_Entry *ent;
	struct timeval tv;
	char *auth, *rcvd, *sScore;
	char *s;
	double score;
	u_int h;

	if (TBL_Lookup(tblHosts, addr, NULL) == -1) {
		MBD_SetError("Unknown host: %s", addr);
		return (-1);
	}
	s = data;
	auth = strsep(&s, " ");
	rcvd = strsep(&s, " ");
	sScore = strsep(&s, " ");
	if (auth == NULL || rcvd == NULL || sScore == NULL ||
	    auth[0] == '\0' || rcvd[0] == '\0' || sScore[0] == '\0') {
		MBD_SetError("Malformed packet");
		return (-1);
	}
	if (strcmp(auth, mbdPassword) != 0) {
		MBD_SetError("Bad password: %s", auth);
		return (-1);
	}
	if (!ValidIPv4(rcvd)) {
		MBD_SetError("Invalid IP: %s", rcvd);
		return (-1);
	}
	score = strtod(sScore, NULL);
	if (score == -HUGE_VAL || score == HUGE_VAL ||
	    score < minScore || score >= maxScore) {
		MBD_SetError("Illegal score: %f", score);
		return (-1);
	}

	h = TBL_Hash(tblScores,rcvd);
	if ((ent = TBL_LookupHash(tblScores, rcvd, h)) == NULL) {
		MBD_Debug("N %s (%f)", rcvd, score);
		if (tblScores->count >= maxEntries) {
			MBD_Debug("Maximum entries (%d), expiring", maxEntries);
			ExpireOldScores();
		}
		ent = TBL_Insert(tblScores, h, rcvd);
		ent->count = 1;
		ent->scoreTot = score;
		ent->scoreAvg = score;
	} else {
		MBD_Debug("E %s (%d/%f)", ent->key, ent->count, ent->scoreAvg);
		if (ent->count < MBD_COUNT_MAX) {
			ent->count++;
		}
		if ((ent->scoreTot + score) > MBD_TOTSCORE_MAX) {
			ent->scoreTot = MBD_TOTSCORE_MAX;
			ent->scoreAvg = MBD_TOTSCORE_MAX/ent->count;
		} else {
			ent->scoreTot += score;
			ent->scoreAvg = ent->scoreTot/ent->count;
		}
	}
	if (gettimeofday(&tv, NULL) != -1) {
		ent->stamp = tv.tv_sec;
	} else {
		ent->stamp = 0;
	}
	return (0);
}

int
main(int argc, char **argv, char **env)
{
	struct addrinfo hints, *res, *res0;
	const char *cause = NULL;
	int rv, i, maxfd = 0;
	struct sigaction sa;
	struct stat sb;
	FILE *f;
	CFG_File *cf;
	
	if ((cf = CFG_OpenFile(SYSCONFDIR, "mailblockd.conf")) == NULL) {
		fprintf(stderr, "Load configuration file: %s\n",
		    MBD_GetError());
		exit(1);
	}
	CFG_GetStr(cf, "home", &mbdHome, "/var/db/mailblockd/");
	CFG_GetStr(cf, "pid-file", &pidFile, "/var/run/mailblockd.pid");
	CFG_GetStr(cf, "log-file", &logFile, "/var/log/mailblockd");
	CFG_GetStr(cf, "hosts-file", &pathHosts, "/var/db/mailblockd/hosts");
	CFG_GetStr(cf, "scores-file", &pathScores, "/var/db/mailblockd/scores");
	CFG_GetStr(cf, "host", &mbdHost, "localhost");
	CFG_GetStr(cf, "port", &mbdPort, "925");
	CFG_GetStr(cf, "password", &mbdPassword, "secret");
	CFG_GetFlt(cf, "min-score", &minScore, -300.0);
	CFG_GetFlt(cf, "max-score", &maxScore, +300.0);
	CFG_GetUint(cf, "max-entries", &maxEntries, 50000);
	CFG_GetUint(cf, "expire-entries", &nExpireEntries, 100);
	CFG_CloseFile(cf);

#ifdef MBD_LOGFILE
	OpenLog();
#endif
	MBD_Debug("mailblockd %s", VERSION);
	syslog(LOG_INFO, "Server startup (%s)", VERSION);

	/* Load the tables */
	if (stat(mbdHome, &sb) != 0 &&
	    mkdir(mbdHome, 0755) == -1) {
		MBD_SetError("mkdir %s: %s", mbdHome, strerror(errno));
		goto fatal;
	}
	if ((tblHosts = TBL_Load(pathHosts, -1)) == NULL ||
	    (tblScores = TBL_Load(pathScores, maxEntries)) == NULL) {
		goto fatal;
	}

	expireEntries = Malloc(nExpireEntries*sizeof(TBL_Entry *));
	expireHashes = Malloc(nExpireEntries*sizeof(u_int));

	FD_ZERO(&servfds);

	/*
	 * Set up the server listening socket(s).
	 */
	memset(&hints, 0, sizeof(hints));
	hints.ai_family = PF_UNSPEC;
	hints.ai_socktype = SOCK_DGRAM;
	hints.ai_flags = AI_PASSIVE;
	if ((rv = getaddrinfo(mbdHost, mbdPort, &hints, &res0)) != 0) {
		syslog(LOG_ERR, "%s:%s: %s", mbdHost, mbdPort,
		    gai_strerror(rv));
		goto out;
	}
	for (mbdSockCount = 0, res = res0;
	     res != NULL && mbdSockCount < 4;
	     res = res->ai_next) {
		rv = socket(res->ai_family, res->ai_socktype, res->ai_protocol);
		if (rv == -1) {
			cause = "socket";
			continue;
		}
		i = 1;
		setsockopt(rv, SOL_SOCKET, SO_REUSEADDR, &i,
		    (my_socklen_t)sizeof(i));
		if (bind(rv, res->ai_addr, res->ai_addrlen) == -1) {
			cause = "bind";
			close(rv);
			continue;
		}
		mbdSocks[mbdSockCount++] = rv;
		FD_SET(rv, &servfds);
		if (rv > maxfd) { maxfd = rv; }
	}
	if (mbdSockCount == 0) {
		syslog(LOG_ERR, "%s: %m", cause);
		goto out;
	}
	freeaddrinfo(res0);

	sigemptyset(&sa.sa_mask);

	sa.sa_handler = master_sig_chld;
	sa.sa_flags = SA_RESTART;
	sigaction(SIGCHLD, &sa, NULL);
	
	sa.sa_handler = master_sig_flush;
	sa.sa_flags = SA_RESTART;
	sigaction(SIGHUP, &sa, NULL);
	
	sa.sa_handler = SIG_IGN;
	sigaction(SIGUSR1, &sa, NULL);

	sa.sa_handler = master_sig_die;
	sa.sa_flags = 0;
	sigaction(SIGINT, &sa, NULL);
	sigaction(SIGQUIT, &sa, NULL);
	sigaction(SIGTERM, &sa, NULL);

	/* Write the PID file */
	if (stat(pidFile, &sb) == 0) {
		MBD_Debug("WARNING: Overwriting existing PID file!");
		syslog(LOG_WARNING, "Overwriting %s!", pidFile);
	}
	if ((f = fopen(pidFile, "w")) == NULL) {
		MBD_Debug("%s: %s", pidFile, strerror(errno));
		syslog(LOG_ERR, "%s: %m", pidFile);
		goto out_close;
	}
	fprintf(f, "%lu\n", (unsigned long)getpid());
	fclose(f);

	for (;;) {
		fd_set rfds = servfds;

		rv = select(maxfd+1, &rfds, NULL, NULL, NULL);
		if (rv == -1) {
			if (errno == EINTR) {
				if (ProcessSignals() == 1) {
					goto out_close;
				}
				continue;
			} else {
				syslog(LOG_ERR, "select: %m; shutdown");
				exit(1);
			}
		}
		for (i = 0; i < mbdSockCount; i++) {
			char buf[40+12];
			struct sockaddr_in siPeer;
			int siPeerLen = sizeof(siPeer);
			ssize_t rv;

			if (!FD_ISSET(mbdSocks[i], &rfds)) {
				continue;
			}
			if ((rv = recvfrom(mbdSocks[i], buf, sizeof(buf), 0,
			    (struct sockaddr *)&siPeer, &siPeerLen)) == -1) {
				syslog(LOG_ERR, "recvfrom: %m; shutdown");
				exit(1);
			}
			if (rv > 0 && buf[rv-1] == '\n') {
				buf[rv-1] = '\0';
			} else {
				buf[rv] = '\0';
			}
			if (ProcessScore(inet_ntoa(siPeer.sin_addr), buf)
			    == -1) {
				MBD_Debug("ProcessScore: %s", MBD_GetError());
/*				syslog(LOG_ERR, "%s", MBD_GetError()); */
			}
		}
		if (ProcessSignals() == 1)
			break;
	}
out_close:
	syslog(LOG_NOTICE, "Server shutdown (signal); saving %d entries",
	    tblScores->count);
	if (TBL_Save(tblScores, pathScores, 0600) == -1)
		syslog(LOG_ERR, "Flush failed: %s", MBD_GetError());
out:
	TBL_Destroy(tblHosts);
	TBL_Destroy(tblScores);
	CloseFiles();
	unlink(pidFile);
	
	return (0);
fatal:
	fprintf(stderr, "FATAL: %s\n", MBD_GetError());
	syslog(LOG_CRIT, "FATAL: %s", MBD_GetError());
	return (1);
}