/*
 *  LIVED - Live Daemon.
 *  v0.01
 *  
 *  By Carl Ritson - 2001
 *  Licensed under GPL2 - www.gnu.org
 *
 *  Works on a Master+Slave System, sends heartbeats, when the master stops
 *  Sending Heartbeats or sends an explicit ASSERT signal to the slave,
 *  the slave takes over.  When the master starts sending heartbeats again,
 *  or the slave receives an explicit DROP signal it drops and lets the 
 *  master do the work.
 *  The DROP and ASSERT behaviors are defined by user scripts,
 *  see lived.conf.
 *
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

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

#include <sys/wait.h>

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

#include <asm/errno.h>
#include <signal.h>
#include <unistd.h>

#include <syslog.h>

#define DEBUG 0
#define BUFSIZE 1024

struct in_addr peeraddr;
int maxtokens 	= 0;
int tokens	= 0;
int heartrate 	= 0;
u_int16_t port 	= 0;
int master	= 0;

char *password		= NULL;
char *assert_script 	= NULL;
char *drop_script 	= NULL;
char *listenaddr	= NULL;

int spin	= 0;
int main_soc	= 0;

enum { HEARTBEAT = 101, ASSERT = 202, DROP = 303 };

int get_line(FILE *file, char *buffer, int length) {
	int c = fgetc(file);
	int r = 0;
	while(c != EOF && r < (length - 1) && (char)c != '\n') {
		(*buffer) = (char)c;
		++buffer;
		++r;
		c = fgetc(file);
	}
	*buffer = '\0';
	if (DEBUG) {fprintf(stderr,"\tGet_line\n");}
	if (c == EOF) {
		return c;
	} else {
		return r + 1;
	}
}
void load_config(char *filename) {
	char buffer[BUFSIZE];
	char *item;
	char *var;
	FILE *config;
	config = fopen(filename,"r");
	if(config == NULL) {
		fprintf(stderr,"Failed to config open %s\n",filename);
		exit(1);
	} else {
		while(get_line(config,buffer,BUFSIZE) != EOF) {
			if(*buffer != '#' && *buffer != '\0') {
				item = strtok(buffer," ");
				var = strtok(NULL," ");
				if(strcasecmp(item,"heartrate") == 0) {
					heartrate = atoi(var);
				} else if(strcasecmp(item,"peeraddr") == 0) {
					inet_aton(var,&peeraddr);
				} else if(strcasecmp(item,"maxtokens") == 0) {
					maxtokens = atoi(var);
				} else if(strcasecmp(item,"port") == 0) {
					port = htons(atoi(var));
				} else if(strcasecmp(item,"password") == 0) {
					password = strdup(var);
				} else if(strcasecmp(item,"assert_script") == 0) {
					assert_script = strdup(var);
				} else if(strcasecmp(item,"drop_script") == 0) {
					drop_script = strdup(var);
				} else if(strcasecmp(item,"master") == 0) {
					master = atoi(var);
				} else if(strcasecmp(item,"listenaddr") == 0) {
					listenaddr = strdup(var);
				}
			}
		}
		fclose(config);
	}
}

int process_message(char *buffer) {
	int i;
	char *temp;
	temp = strtok(buffer," ");
	if(strcmp(temp,password) == 0) {
		temp = strtok(NULL," ");
		i = atoi(temp);
	}
	return i;
}

void sleep_loop() {
	char buffer[BUFSIZE];
	struct sockaddr_in *from;
	struct sockaddr_in *expect;
	int in_len = sizeof(struct sockaddr_in);
	int r;
	int act;
	int gotbeat = 0;
	int gotassert = 0;
	int temp;
	
	/* Allocate and Zero 2 Address Structures */
	from = malloc(in_len);
	expect = malloc(in_len);
	memset(from,'\0',in_len);
	memset(expect,'\0',in_len);
	
	/* Setup the expect structure to what we expect */
	expect->sin_family = AF_INET;
	expect->sin_port = port;
	memcpy(&expect->sin_addr,&peeraddr,sizeof(struct in_addr));

	/* Fill the token bucket */
	tokens = maxtokens;
	
	for(;;) {
		sleep(heartrate);
		temp = in_len;
		gotbeat = 0;
		do {
			r = recvfrom(main_soc,buffer,BUFSIZE,MSG_DONTWAIT,(struct sockaddr *)from,&temp);
			if(r > 0 && r != EAGAIN) {
				if(memcmp(from,expect,temp) == 0) {
					act = process_message(buffer);
					if(act == HEARTBEAT) {
						gotbeat = 1;
					} else if(act == ASSERT) {
						gotassert = 1;
					} else if(act == DROP) {
						gotassert = 0;
					}
				}
			}
		} while(r > 0 && r != EAGAIN);
		if(gotbeat == 0) {
			--tokens;
		} else if(tokens < maxtokens) {
			++tokens;
		}
		if(gotassert || tokens <= 0) {
			if(gotassert) {
				syslog(LOG_WARNING,"Got assert signal from peer %s",inet_ntoa(peeraddr));
				sleep((int)(heartrate/2));
			} else {
				syslog(LOG_CRIT,"Token Bucket for peer %s Empty",inet_ntoa(peeraddr));
			}
			break;
		}
	}
	free(expect);
	free(from);
}

void awake_loop() {
	char buffer[BUFSIZE];
	char beat[BUFSIZE];
	struct sockaddr_in *from;
	struct sockaddr_in *expect;
	int in_len = sizeof(struct sockaddr_in);
	int beat_len;
	int r;
	int act;
	int gotdrop = 0;
	int temp;
	
	/* Prepare Heatbeat Message */
	sprintf(beat,"%s %d",password,HEARTBEAT);
	beat_len = strlen(beat);
	
	/* Allocate and Zero 2 Address Structures */
	from = malloc(in_len);
	expect = malloc(in_len);
	memset(from,'\0',in_len);
	memset(expect,'\0',in_len);
	
	/* Setup the expect structure to what we expect */
	/* Expect Structure is also used for sending heartbeats */
	expect->sin_family = AF_INET;
	expect->sin_port = port;
	memcpy(&expect->sin_addr,&peeraddr,sizeof(struct in_addr));

	for(;;) {
		temp = in_len;
		sleep(heartrate);
		sendto(main_soc,beat,beat_len,0,(struct sockaddr *)expect,in_len);
		do {
			r = recvfrom(main_soc,buffer,BUFSIZE,MSG_DONTWAIT,(struct sockaddr *)from,&temp);
			if(r > 0 && r != EAGAIN) {
				if(memcmp(from,expect,temp) == 0) {
					act = process_message(buffer);
					if(act == DROP) {
						gotdrop = 1;
					} else if(act == ASSERT) {
						gotdrop = 0;
					} else if(act == HEARTBEAT && !master) {
						gotdrop = 1;
					}
				}
			}
		} while(r > 0 && r != EAGAIN);
		if(gotdrop) {
			syslog(LOG_WARNING,"Got drop signal from peer %s",inet_ntoa(peeraddr));
			break;
		}
	}
	free(expect);
	free(from);
}	

void assert() {
	pid_t p;
	syslog(LOG_CRIT,"Asserting (%s)",assert_script);
	p = fork();
	if(p == 0) {
		execve(assert_script,NULL,NULL);
		exit(0);
	} else {
		waitpid(p,NULL,0);
	}
}

void drop() {
	pid_t p;
	syslog(LOG_CRIT,"Dropping (%s)",drop_script);
	p = fork();
	if(p == 0) {
		execve(drop_script,NULL,NULL);
		exit(0);
	} else {
		waitpid(p,NULL,0);
	}
}

void send_assert() {
	char buffer[BUFSIZE];
	struct sockaddr_in *peer = malloc(sizeof(struct sockaddr_in));	
	/* Prepare Address Structure */
	peer->sin_family = AF_INET;
	peer->sin_port = port;
	memcpy(&peer->sin_addr,&peeraddr,sizeof(struct in_addr));
	/* Prepare ASSERT Message */
	sprintf(buffer,"%s %d",password,ASSERT);
	/* Send 2 Copies 2 Seconds Apart */
	sendto(main_soc,buffer,strlen(buffer),0,(struct sockaddr *)peer,sizeof(struct sockaddr_in));
	sleep(2);
	sendto(main_soc,buffer,strlen(buffer),0,(struct sockaddr *)peer,sizeof(struct sockaddr_in));
	free(peer);
	syslog(LOG_WARNING,"Sent Assert Signal to Peer %s",inet_ntoa(peeraddr));
}

void send_drop() {
	char buffer[BUFSIZE];
	struct sockaddr_in *peer = malloc(sizeof(struct sockaddr_in));	
	/* Prepare Address Structure */
	peer->sin_family = AF_INET;
	peer->sin_port = port;
	memcpy(&peer->sin_addr,&peeraddr,sizeof(struct in_addr));
	/* Prepare Drop Message */
	sprintf(buffer,"%s %d",password,DROP);
	/* Send 2 Copies 2 Seconds Apart */
	sendto(main_soc,buffer,strlen(buffer),0,(struct sockaddr *)peer,sizeof(struct sockaddr_in));
	sleep(2);
	sendto(main_soc,buffer,strlen(buffer),0,(struct sockaddr *)peer,sizeof(struct sockaddr_in));
	free(peer);
	syslog(LOG_WARNING,"Sent Drop Signal to Peer %s",inet_ntoa(peeraddr));
}

void main_loop() {
	if(master) {
		for(;;) {
			send_drop();
			sleep(heartrate);
			assert();
			awake_loop();
		}
	} else {
		for(;;) {
			drop();
			sleep_loop();
			assert();
			awake_loop();
		}
	}
}

int sigtermhandler(int s) {
	syslog(LOG_WARNING,"Received Term Signal");
	if(master) {
		send_assert();
		drop();
	}
	closelog();
	exit(0);
	return 0;
}

void setup_signals() {
	signal(SIGTERM,(void *)&(sigtermhandler));
	signal(SIGSTOP,(void *)&(sigtermhandler));
	signal(SIGQUIT,(void *)&(sigtermhandler));
}

void setup_socket() {
	struct sockaddr_in addr;
	int ret;
	
	main_soc = socket(PF_INET,SOCK_DGRAM,0);
	
	addr.sin_family = AF_INET;
	addr.sin_port = port;
	if(listenaddr != NULL) {
		inet_aton(listenaddr,&addr.sin_addr);
	} else {
		inet_aton("0.0.0.0",&addr.sin_addr);
	}
	
	ret = bind(main_soc,(struct sockaddr *)&addr,sizeof(addr));
	if(ret != 0) {
		fprintf(stderr,"Bind Port %d Failed\n",ntohs(port));
		exit(1);
	}
}	

void drop_to_background() {
	freopen("/dev/null","r",stdin);
	freopen("/dev/null","a",stdout);
	freopen("/dev/null","a",stderr);
}


int signalusr2(int s) {
	spin = 0;
	return s;
}

void setup_log() {
	openlog("lived",0,LOG_DAEMON);
	syslog(LOG_NOTICE,"started pid:%d listening on:%s port:%d as %s",getpid(),listenaddr,ntohs(port),master?"Master":"Slave");
}

int main(int argc, char *argv[]) {
	pid_t p = 0;
	pid_t parent = getpid();
	if(argc < 2) {
		fprintf(stderr,"You must supply a config\n");
		exit(1);
	}
	load_config(argv[1]);
	if(assert_script == NULL || drop_script == NULL || password == NULL) {
		fprintf(stderr,"Config is incomplete\n");
		exit(1);
	}
	/* Register SIGUSR1 to spin unlock function */
	signal(SIGUSR2,(void *)&(signalusr2));
	
	p = fork();
	if(p == 0) {
		drop_to_background();
		setup_socket();
		setup_signals();
		setup_log();
		/* Send parent SIGUSR1 to say we are ready */
		kill(parent,SIGUSR2);
		main_loop();
	} else {
		printf("Dropping into Background\n");
		fflush(stdout);
		/* Child sends SIGUSR1 when ready */
		spin = 1;
		while(spin) {
			parent = waitpid(p,NULL,WNOHANG);
			if(parent != 0) {
				break;
			}
			usleep(10000);	
		}
		if(spin) {
			fprintf(stderr,"Fail to Drop to Background\n");
			return 1;
		} else {
			printf("Dropped into Background as : %d\n",p);
			return 0;
		}
	}
	/* Droppped off the end so return 1 */
	return 1;
}