/*
    This file is part of Kismet

    Kismet is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 2 of the License, or
    (at your option) any later version.

    Kismet is distributed in the hope that it will be useful,
      but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with Kismet; if not, write to the Free Software
    Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
*/

#include "config.h"

#include <stdio.h>
#include <time.h>
#include <list>
#include <map>
#include <vector>
#include <algorithm>
#include <string>
#include <sstream>

#include <sys/types.h>
#include <unistd.h>
#include <errno.h>

#include "ipc_remote.h"

void IPC_MessageClient::ProcessMessage(string in_msg, int in_flags) {
	// Build it into a msgbus ipc block ... is there a smarter way to
	// build these things?
	ipc_packet *pack = 
		(ipc_packet *) malloc(sizeof(ipc_packet) + 8 + in_msg.length() + 1);
	ipc_msgbus_pass *msgb = (ipc_msgbus_pass *) pack->data;

	msgb->msg_flags = in_flags;
	msgb->msg_len = in_msg.length() + 1;
	snprintf(msgb->msg, msgb->msg_len, in_msg.c_str());

	pack->data_len = 8 + in_msg.length() + 1;

	pack->ipc_cmdnum = MSG_CMD_ID;
	pack->ipc_ack = 0;

	// If it's a fatal frame, push it as a shutdown
	if ((in_flags & MSGFLAG_FATAL)) {
		((IPCRemote *) auxptr)->ShutdownIPC(pack);
		return;
	}
	
	// Push it via the IPC
	((IPCRemote *) auxptr)->SendIPC(pack);

	// It gets freed once its sent so don't free it ourselves here
}

int ipc_msg_callback(IPC_CMD_PARMS) {
	(void) auxptr;

	// Child IPC does nothing with a MSG frame
	if (parent == 0)
		return 0;

	// Blow up on short IPC
	if (len < 5) {
		_MSG("IPC messagebus handler got a short message block",
			 MSGFLAG_ERROR);
		return -1;
	}

	// Strip it apart and inject it into the message bus
	ipc_msgbus_pass *pass = (ipc_msgbus_pass *) data;

	_MSG(pass->msg, pass->msg_flags);

	return 0;
}

int ipc_die_callback(IPC_CMD_PARMS) {
	(void) data;
	(void) len;

	// Child receiving DIE command shuts down
	if (parent == 0) {
		// Call the internal shutdown process
		((IPCRemote *) auxptr)->IPCDie();
		// Exit entirely, if we didn't already
		exit(1);
	}

	// Parent receiving DIE command knows child is dieing for some
	// reason, send a message note, we'll figure out later if this is
	// fatal
	ostringstream osstr;
	
	osstr << "IPC controller got notification that IPC child process " <<
		(int) ((IPCRemote *) auxptr)->FetchSpawnPid() << " is shutting down.";
	_MSG(osstr.str(), MSGFLAG_INFO);

	// Call the internal die sequence to make sure the child pid goes down
	((IPCRemote *) auxptr)->IPCDie();
	
	return 0;
}

int ipc_sync_callback(IPC_CMD_PARMS) {
	// Parent does nothing
	if (parent == 1)
		return 0;

	if (len < (int) sizeof(ipc_sync)) {
		_MSG("IPC sync handler got a short sync block", MSGFLAG_ERROR);
		return -1;
	}

	return ((IPCRemote *) auxptr)->SyncIPCCmd((ipc_sync *) data);
}

IPCRemote::IPCRemote() {
	fprintf(stderr, "FATAL OOPS:  IPCRemote called w/ no globalreg\n");
	exit(1);
}

IPCRemote::IPCRemote(GlobalRegistry *in_globalreg, string in_procname) {
	globalreg = in_globalreg;
	procname = in_procname;

	if (globalreg->messagebus == NULL) {
		fprintf(stderr, "FATAL OOPS:  IPCRemote called before messagebus\n");
		exit(1);
	}

	child_exec_mode = 0;

	next_cmdid = 0;
	ipc_pid = 0;
	ipc_spawned = 0;
	last_ack = 1; // Our "last" command was ack'd

	// Register builtin commands (in proper order!)
	RegisterIPCCmd(&ipc_die_callback, NULL, this, "DIE");
	RegisterIPCCmd(&ipc_msg_callback, NULL, this, "MSG");
	RegisterIPCCmd(&ipc_sync_callback, NULL, this, "SYNC");
}

int IPCRemote::SyncIPCCmd(ipc_sync *data) {
	// Search the map for something of this name
	for (map<unsigned int, ipc_cmd_rec *>::iterator x = ipc_cmd_map.begin();
		 x != ipc_cmd_map.end(); ++x) {
		ipc_cmd_rec *cr = x->second;
		string name = (char *) data->name;

		if (cr->name == name) {
			cr->id = data->ipc_cmdnum;
			ipc_cmd_map[data->ipc_cmdnum] = cr;
			ipc_cmd_map.erase(x);
			return 1;
		}
	}

	return 1;
}

int IPCRemote::SetChildExecMode(int argc, char *argv[]) {
	int tint;

	// Set us to child mode
	ipc_pid = 0;
	// Set our next cmd id to something big and negative, so that we can
	// stock cmds prior to a sync, but not interfere once the sync begins
	next_cmdid = -4098;

	child_exec_mode = 1;

	// Parse the FD out
	if (argc < 2)
		return -1;

	if (sscanf(argv[1], "%d", &tint) != 1) {
		return -1;
	}

	sockpair[0] = tint;

	return 1;
}

int IPCRemote::RegisterIPCCmd(IPCmdCallback in_callback, 
							  IPCmdCallback in_ackcallback,
							  void *in_aux, string in_name) {
	if (ipc_spawned) {
		_MSG("IPC_Remote - Tried to register a command after the IPC agent has "
			 "been spawned.  Commands must be registered before Spawn().",
			 MSGFLAG_FATAL);
		globalreg->fatal_condition = 1;
		return -1;
	}

	next_cmdid++;

	ipc_cmd_rec *rec = new ipc_cmd_rec;
	rec->auxptr = in_aux;
	rec->callback = in_callback;
	rec->ack_callback = in_ackcallback;
	rec->name = in_name;
	rec->id = next_cmdid;

	ipc_cmd_map[next_cmdid] = rec;

	return next_cmdid;
}

int IPCRemote::SpawnIPC() {
	// Don't build the socket pair if we're in exec child mode
	if (child_exec_mode == 0) {
		// Generate the socket pair before the split
		if (socketpair(PF_UNIX, SOCK_DGRAM, 0, sockpair) < 0) {
			_MSG("Unable to great socket pair for IPC communication: " +
				 string(strerror(errno)), MSGFLAG_FATAL);
			globalreg->fatal_condition = 1;
			return -1;
		}

		// Fork and launch the child control loop & set up the rest of our internal
		// info.
		if ((ipc_pid = fork()) < 0) {
			_MSG("Unable to fork() and create child process for IPC communication: " +
				 string(strerror(errno)), MSGFLAG_FATAL);
			globalreg->fatal_condition = 1;
			return -1;
		}

		// Split into the IPC child control loop if we've forked
		if (ipc_pid == 0) {
			// Run the client binary if we have one
			if (child_cmd != "") {
				char **cmdarg = new char *[3];
				cmdarg[0] = strdup(child_cmd.c_str());
				cmdarg[1] = new char[4];
				cmdarg[2] = NULL;

				snprintf(cmdarg[1], 4, "%d", sockpair[0]);

				execve(cmdarg[0], cmdarg, NULL);
			}

			IPC_Child_Loop();
			exit(0);
		}

		// Close the parent half of the socket pair
		close(sockpair[0]);

		// If we spawned something that needs to be synced, send all our protocols
		if (child_cmd != "") {
			for (map<unsigned int, ipc_cmd_rec *>::iterator x = ipc_cmd_map.begin();
				 x != ipc_cmd_map.end(); ++x) {

				if (x->first < LAST_BUILTIN_CMD_ID)
					continue;

				ipc_packet *pack = 
					(ipc_packet *) malloc(sizeof(ipc_packet) + sizeof(ipc_sync));

				ipc_sync *sync = (ipc_sync *) pack->data;

				sync->ipc_cmdnum = x->first;
				snprintf((char *) sync->name, 32, "%s", x->second->name.c_str());

				pack->data_len = sizeof(ipc_sync);
				pack->ipc_cmdnum = SYNC_CMD_ID;
				pack->ipc_ack = 0;
	
				// Push it via the IPC
				SendIPC(pack);
			}
		}
	}

	// We've spawned, can't set new commands anymore
	ipc_spawned = 1;

	return 1;
}

int IPCRemote::ShutdownIPC(ipc_packet *pack) {
	if (ipc_spawned == 0) 
		return 0;

	int sock;
	if (ipc_pid == 0)
		sock = sockpair[0];
	else
		sock = sockpair[1];

	// If we have a last frame, send it
	if (pack != NULL) {
		pack->sentinel = IPCRemoteSentinel;
		send(sock, pack, sizeof(ipc_packet) + pack->data_len, 0);
	}

	// Nothing special here, just a die frame.  Beauty is, we don't care
	// if we're the parent of the child, a clean shutdown will signal the other
	// side it's time to shuffle off
	ipc_packet *dpack = (ipc_packet *) malloc(sizeof(ipc_packet));
	dpack->data_len = 0;
	dpack->ipc_cmdnum = DIE_CMD_ID;
	dpack->ipc_ack = 0;
	dpack->sentinel = IPCRemoteSentinel;

	// Send it immediately
	send(sock, dpack, sizeof(ipc_packet) + dpack->data_len, 0);

	// Die fully
	IPCDie();

	return 1;
}

int IPCRemote::SendIPC(ipc_packet *pack) {
	// This is really just an enqueue system
	pack->sentinel = IPCRemoteSentinel;
	cmd_buf.push_back(pack);

	return 1;
}

int IPCRemote::FetchReadyState() {
	if (ipc_spawned == 0)
		return 0;

	if (last_ack == 0)
		return 0;

	if (cmd_buf.size() != 0)
		return 0;

	return 1;
}

void IPCRemote::IPC_Child_Loop() {
	fd_set rset, wset;

	// Obviously we're spawned
	ipc_spawned = 1;

	// Close the other half of the socket pair
	close(sockpair[1]);

	// Kluge in a new message bus to talk to our parent and give it only
	// the IPC client to replicate messages
	globalreg->messagebus = new MessageBus;
	IPC_MessageClient *ipcmc = new IPC_MessageClient(globalreg, this);
	globalreg->messagebus->RegisterClient(ipcmc, MSGFLAG_ALL);

	// ignore a bunch of signals
	signal(SIGINT, SIG_IGN);
	signal(SIGTERM, SIG_IGN);
	signal(SIGHUP, SIG_IGN);
	signal(SIGPIPE, SIG_IGN);

	// Set the process title
	init_proc_title(globalreg->argc, globalreg->argv, globalreg->envp);
	set_proc_title("%s", procname.c_str());

	while (1) {
		int max_fd = 0;

		FD_ZERO(&rset);
		FD_ZERO(&wset);

		FD_SET(sockpair[0], &rset);
		max_fd = sockpair[0];

		// Do we have data to send?
		if (cmd_buf.size() > 0) {
			FD_SET(sockpair[0], &wset);
		}

		struct timeval tm;
		tm.tv_sec = 1;
		tm.tv_usec = 0;

		// Timeout after 1 second if we stopped getting commands
		if (select(max_fd + 1, &rset, &wset, NULL, &tm) < 0) {
			// Die violently
			fprintf(stderr, "FATAL OOPS:  IPC command child %d got select() "
					"error and cannot continue cleanly: %s\n",
					getpid(), strerror(errno));
			exit(1);
		}

		// Handle in/out data
		if (Poll(rset, wset) < 0 || globalreg->fatal_condition) {
			exit(0);
		}

	}
}

void IPCRemote::IPCDie() {
	// If we're the child...
	if (ipc_pid == 0 && ipc_spawned) {
		// Shut down the child socket fd
		close(sockpair[0]);
		// and exit, we're done
		exit(0);
	} else if (ipc_spawned) {
		// otherwise if we're the parent...
		// Shut down the socket
		close(sockpair[1]);

		// Wait for the child process to be dead

		ostringstream osstr;

		osstr << "IPC controller waiting for IPC child process " <<
			(int) ipc_pid << " to end.";
		_MSG(osstr.str(), MSGFLAG_INFO);

		wait4(ipc_pid, NULL, 0, NULL);

		ipc_pid = 0;
		ipc_spawned = 0;

		// Flush all the queued packets
		while (cmd_buf.size() > 0) {
			ipc_packet *pack = cmd_buf.front();
			free(pack);
			cmd_buf.pop_front();
		}
		
	}
}

unsigned int IPCRemote::MergeSet(unsigned int in_max_fd, fd_set *out_rset,
								 fd_set *out_wset) {
	// Don't call this on the child, we have a micro-select loop in the child
	// process... also do nothing if we haven't spawned yet
	if (ipc_pid == 0 || ipc_spawned == 0)
		return in_max_fd;

	// Set the socket to be read
	FD_SET(sockpair[1], out_rset);
	
	// Set the write if we have data queued and our last command was
	// ack'd.  If it wasn't, rate limit ourselves down until it is.
	if (cmd_buf.size() > 0)
		FD_SET(sockpair[1], out_wset);

	if (in_max_fd < (unsigned int) sockpair[1])
		return sockpair[1];

	return in_max_fd;
}

int IPCRemote::Poll(fd_set& in_rset, fd_set& in_wset) {
	// This CAN be called by the parent or the child.  In the parent it's called
	// by the normal pollable architecture.  In the child we manually call it
	// from our micro-select loop.
	
	if (ipc_spawned == 0)
		return 0;

	ostringstream osstr;
	int sock;

	if (ipc_pid == 0)
		sock = sockpair[0];
	else
		sock = sockpair[1];

	// Process packets out
	if (FD_ISSET(sock, &in_wset)) {
		// Send as many frames as we have room for
		while (cmd_buf.size() > 0) {
			ipc_packet *pack = cmd_buf.front();

			// Send the frame
			if (send(sock, pack, sizeof(ipc_packet) + pack->data_len, 0) < 0) {
				if (errno == ENOBUFS)
					break;

				if (ipc_pid == 0) {
					// Blow up messily and spew into stderr
					fprintf(stderr, "IPC child %d got error writing packet to "
							"IPC socket: %s\n", getpid(), strerror(errno));
					globalreg->fatal_condition = 1;
					return -1;
				} else {
					osstr << "IPC controller got error writing data to IPC socket "
						"for IPC child pid " << ipc_pid << ": " << strerror(errno);
					_MSG(osstr.str(), MSGFLAG_FATAL);
					globalreg->fatal_condition = 1;
				}
			}

			// ACK frames themselves, msg frames, and death commands are not
			// expected to ack back that the command is done.  everything else
			// should.
			if (pack->ipc_cmdnum != DIE_CMD_ID &&
				pack->ipc_cmdnum != MSG_CMD_ID &&
				pack->ipc_ack == 0) {
				last_ack = 0;
			}

			// Finally delete it
			cmd_buf.pop_front();
			free(pack);
		}

	}

	// Process packets in
	if (FD_ISSET(sock, &in_rset)) {
		ipc_packet ipchdr;
		ipc_packet *fullpack = NULL;
		int ret;

		// Peek at the packet header
		if ((ret = recv(sock, &ipchdr, 
						sizeof(ipc_packet), MSG_PEEK)) < (int) sizeof(ipc_packet)) {
			if (ret < 0) {
				if (ipc_pid == 0) 
					osstr << "IPC child got error receiving packet header "
						"from controller: " << strerror(errno);
				else
					osstr << "IPC controller got error receiving packet header "
						"from IPC child pid " << ipc_pid << ": " << strerror(errno);
				_MSG(osstr.str(), MSGFLAG_FATAL);
				globalreg->fatal_condition = 1;
				return -1;
			} else {
				return 0;
			}
		}

		// validate the ipc header
		if (ipchdr.sentinel != IPCRemoteSentinel) {
			if (ipc_pid == 0) 
				osstr << "IPC child got error receiving packet header from "
					"controller: Invalid IPC sentinel value";
			else
				osstr << "IPC controller got error receiving packet header "
					"from IPC child pid " << ipc_pid << ": Invalid IPC "
					"sentinel value";
			_MSG(osstr.str(), MSGFLAG_FATAL);
			globalreg->fatal_condition = 1;
			return -1;
		}

		// See if its a command we understand
		map<unsigned int, ipc_cmd_rec *>::iterator cbitr = 
			ipc_cmd_map.find(ipchdr.ipc_cmdnum);
		if (cbitr == ipc_cmd_map.end()) {
			if (ipc_pid == 0) 
				osstr << "IPC child got error receiving packet header from "
					"controller: Unknown IPC command";
			else
				osstr << "IPC controller got error receiving packet header "
					"from IPC child pid " << ipc_pid << ": Unknown IPC command";
			_MSG(osstr.str(), MSGFLAG_FATAL);
			globalreg->fatal_condition = 1;
			return -1;
		}
		IPCmdCallback cback = cbitr->second->callback;
		IPCmdCallback ackcback = cbitr->second->ack_callback;
		void *cbackaux = cbitr->second->auxptr;

		// Get the full packet
		fullpack = (ipc_packet *) malloc(sizeof(ipc_packet) + ipchdr.data_len);

		if ((ret = recv(sock, fullpack, sizeof(ipc_packet) + ipchdr.data_len, 0)) < 
			(int) sizeof(ipc_packet) + (int) ipchdr.data_len) {
			if (ret < 0) {
				if (ipc_pid == 0) 
					osstr << "IPC child got error receiving packet "
						"from controller: " << strerror(errno);
				else
					osstr << "IPC controller got error receiving packet "
						"from IPC child pid " << ipc_pid << ": " << strerror(errno);
				_MSG(osstr.str(), MSGFLAG_FATAL);
				globalreg->fatal_condition = 1;
				return -1;
			} else {
				return 0;
			}
		}

		// If we've got an ack frame, and there is an ack callback for this 
		// command type, we send it on to them
		if (ipchdr.ipc_ack) {
			if (ackcback != NULL) {
				ret = (*ackcback)(globalreg, fullpack->data, fullpack->data_len,
								  cbackaux, ipc_pid);
			}
			last_ack = 1;
		} else {
			// We "know" the rest is valid, so call the handler w/ this function.
			// giving it the ipc pid lets us cheat and tell if its the parent or not,
			// since the child has a 0 pid
			ret = (*cback)(globalreg, fullpack->data, fullpack->data_len, 
						   cbackaux, ipc_pid);
			if (ret < 0) {
				if (ipc_pid == 0) 
					osstr << "IPC child got error executing command from controller.";
				else
					osstr << "IPC controller got error executing command "
						"from IPC child pid " << ipc_pid;
				_MSG(osstr.str(), MSGFLAG_FATAL);
				globalreg->fatal_condition = 1;
				return -1;
			}

			free(fullpack);

			// Queue a return ack frame that the command was received and processed
			// if it's not die, msg, or ack, and ret == 0, ie, send a generic ack
			if (ipchdr.ipc_cmdnum != DIE_CMD_ID &&
				ipchdr.ipc_cmdnum != MSG_CMD_ID &&
				ret == 0) {
				ipc_packet *ackpack = (ipc_packet *) malloc(sizeof(ipc_packet));
				ackpack->ipc_ack = 1;
				ackpack->ipc_cmdnum = ipchdr.ipc_cmdnum;
				ackpack->data_len = 0;
				SendIPC(ackpack);
			}
		}
	}

	
	return 1;
}