/*
** SellaNMS
** $Id: icmp.cxx,v 1.17 2006/01/24 07:08:09 sella Exp $
** Copyright (c) 2001-2005 Digital Genesis Software, LLC. All Rights Reserved.
** Released under the GPL Version 2 License.
** http://www.digitalgenesis.com
*/

#ifdef HAVE_CONFIG_H
	#include "config.h"
#endif

static const char rcsid[] = "$Id: icmp.cxx,v 1.17 2006/01/24 07:08:09 sella Exp $";

#include "icmp.h"
#include "daemon.h"
#include "common/ip_util.h"
#include <fcntl.h>
#include <errno.h>
#include <unistd.h>
#include <pthread.h>
#include <sys/time.h>
#include <linux/if.h>
#include <sys/poll.h>
#include <sys/types.h>
#include <sys/socket.h>

/* FIXME: This module can be written efficient. Completed IP addresses should be removed from the tx list. */

extern ConfigNode* config;
extern bool interrupt;
extern int errno;

ICMP::ICMP() {
	pthread_mutex_init(&rx_hash_mutex, NULL);
}

ICMP::~ICMP() {
	while (EBUSY == pthread_mutex_destroy(&rx_hash_mutex)) {
		usleep(USLEEP_NORMLOOP_TIME);
	}
}

RESULT ICMP::readConfig() {
	RESULT result = SUCCESS;
	vector<string> v;

	string str;
	struct in_addr ip;

	ConfigNode node;

	if (!config->assignNode(getConfigPath(), node)) {
		LMSG("No '%s' section found within the config file.", getConfigPath());

		return FAILURE;
	}

	if (node.assignValue("source", str)) {
		inet_pton(AF_INET, str.c_str(), &(this->source));
	} else if (node.assignValue("interface", str)) {
		source_ip(str.c_str(), &ip);
		this->source = ip.s_addr;
	} else {
		this->source = htonl(INADDR_ANY);
	}

	node.assignValue("transmit", tx_count, DISCOVERY_ICMP_TRANSMITCOUNT_DEF);
	node.assignValue("require", rx_count, DISCOVERY_ICMP_REQUIREDCOUNT_DEF);
	/* Repair stupid config settings. */
	if (rx_count > tx_count) {
		rx_count = tx_count;
	}
	node.assignValue("timeout", timeout, DISCOVERY_ICMP_TIMEOUT_DEF);

	if (!node.assignValues("prefix", netblocks)) {
		LMSG("WARNING: Missing required values for 'prefix'");
		return FAILURE;
	}

	/* Clear any existing policies. */
	inputPolicy.clear();
	outputPolicy.clear();

	/* Add and compile input policy. */
	if (node.assignValues("policy input", v)) {
		if (inputPolicy.addPolicies(config, v) == FAILURE) {
			LMSG("NOTICE: Failed to add/compile input policies.");
			result = FAILURE;
		}
	}

	/* Add and compile output policy. */
	if (node.assignValues("policy output", v)) {
		if (outputPolicy.addPolicies(config, v) == FAILURE) {
			LMSG("NOTICE: Failed to add/compile output policies.");
			result = FAILURE;
		}
	}

	/* Configure Storage */
	if (storage.setConfig(*config) != SUCCESS) {
		result = FAILURE;
	}

	printConfig();

	return result;
}

void ICMP::printConfig() {
	string str;
	vector<string> v;
	vector<string>::iterator v_it;

	in_addr s;
	char	inet_buffer[16];

	if (!TRACE_ISSET(TRACE_CONFIG)) {
		return;
	}

	s.s_addr = this->source;

	TRACE(TRACE_CONFIG, "source: %s", inet_ntop(AF_INET, &s, inet_buffer, 16));
	TRACE(TRACE_CONFIG, "transmit: %lu (sec)", this->tx_count);
	TRACE(TRACE_CONFIG, "require: %lu", this->rx_count);
	TRACE(TRACE_CONFIG, "timeout: %lu", this->timeout);

	if (config->assignValues("discovery icmp policy input", v)) {
		str.erase();
		for (v_it = v.begin(); v_it != v.end(); ++v_it) {
			str.append(*v_it + " ");
		}
		TRACE(TRACE_CONFIG, "policy input: %s", str.c_str());
	} else {
		TRACE(TRACE_CONFIG, "policy input: none");
	}

	if (config->assignValues("discovery icmp policy output", v)) {
		str.erase();
		for (v_it = v.begin(); v_it != v.end(); ++v_it) {
			str.append(*v_it + " ");
		}
		TRACE(TRACE_CONFIG, "policy output: %s", str.c_str());
	} else {
		TRACE(TRACE_CONFIG, "policy output: none");
	}

	return;
}

RESULT ICMP::execute() {
	RESULT result = SUCCESS;

	int err;
	pthread_t thread_id = 0;
	this->module_id = (uint16_t) getpid();

	slist<pthread_t> t_list;
	slist<pthread_t>::iterator it;

	TRACE(TRACE_THREAD, "ICMP::execute() starting");

	/* Create the storage thread */
	if ((err = pthread_create(&thread_id, NULL, (void* (*)(void*))ICMP::storage_thread, (void*) this)) == 0) {
		t_list.push_front(thread_id);
	}

	/* Create the receiver thread */
	if (result == SUCCESS) {
		if ((err = pthread_create(&thread_id, NULL, (void* (*)(void*))ICMP::receive_thread, (void*) this)) == 0) {
			t_list.push_front(thread_id);
		}
	}

	/* Create the transmitter thread */
	if (result == SUCCESS) {
		if ((err = pthread_create(&thread_id, NULL, (void* (*)(void*))ICMP::transmit_thread, (void*) this)) == 0) {
			t_list.push_front(thread_id);
		}
	}

	/* Join each thread to free its resources. */
	for (it = t_list.begin(); it != t_list.end(); ++it) {
		/* Wait for thread to complete. */
		if ((err = pthread_join(*it, NULL)) != 0) {
			LMSG("pthread_join returned %d", err);
		}
	}

	TRACE(TRACE_THREAD, "ICMP::execute() ending");

	return result;
}

RESULT ICMP::transmit_thread(void *args) {
	/* FIXME: This should be coded to randomize targets, rather than in order. */
	RESULT result = SUCCESS;

	ICMP *instance = (ICMP*) args;

	unsigned long sleep_count = 0;

	Netblock netblock;
	char *ip_address = NULL;
	sockaddr_in dst_address;
	sockaddr_in src_address;
		
	char buffer[PACKET_BUFFER_SIZE];
	char inet_buffer[16];

	struct icmp *request = (icmp*)buffer;
	char *payload = (char*)(request) + sizeof(struct icmp);

	TRACE(TRACE_THREAD,"ICMP::transmit_thread() starting main loop");

	grantPriv();

	int s = socket(PF_INET, SOCK_RAW, IPPROTO_ICMP);
	fcntl(s, F_SETFL, O_NONBLOCK | O_ASYNC);

	src_address.sin_family = AF_INET;
	src_address.sin_port = 0;
	src_address.sin_addr.s_addr = instance->source;

	if (bind(s, (sockaddr *)&src_address, sizeof(sockaddr_in)) == -1) {
		LMSG("WARNING: Unable to bind to local address '%s'.\n", inet_ntop(AF_INET, &(src_address.sin_addr), inet_buffer, 16));
		close(s);
		s = -1;
	}

	dropPriv();

	if (s < 0) {
		LMSG("WARNING: Failed to get a usable socket");
	
		result = FAILURE;
		interrupt = true;
	} else {
		memset(buffer, 0, PACKET_BUFFER_SIZE);
		dst_address.sin_family = AF_INET;
		dst_address.sin_port = 0;

		request->icmp_id = instance->module_id; 
		request->icmp_type = ICMP_ECHO;
		memset(payload, 0, PAYLOAD_SIZE);
		strncpy(payload, PAYLOAD, PAYLOAD_SIZE);

		/* Send tx_count packets to each address. */
		for (unsigned long i = 0; i < instance->tx_count && !interrupt; i++) {
			request->icmp_seq = i;
			//early_exit = true;

			/* Grab the next netblock. */
			for (unsigned long j = 0; j < instance->netblocks.size(); j++) {		
				netblock.setNetblock(instance->netblocks[j]);

				if (!netblock.isValid()) {
					LMSG("Unable to parse the netblock %s.", instance->netblocks[j].c_str());
	
					result = FAILURE;
				} else {
					/* Use all addresses within netblock. */
					for (unsigned long k = 0; (ip_address = netblock[k]) != NULL && !interrupt; k++) {
						inet_pton(AF_INET, ip_address, &dst_address.sin_addr);	
						delete[] ip_address;

						/* Check if we can skip this TX. */
						if (instance->rx_hash.find(dst_address.sin_addr.s_addr) != instance->rx_hash.end()) {
							if (instance->rx_hash[dst_address.sin_addr.s_addr] == 0) {
								continue;
							}
						}

//						early_exit = false;
						request->icmp_cksum = 0;
						request->icmp_cksum = in_cksum((unsigned short *)request, sizeof(struct icmp) + PAYLOAD_SIZE);
	
						if (sendto(s, request, sizeof(struct icmp) + PAYLOAD_SIZE, 0, (sockaddr*)&dst_address, sizeof(struct sockaddr_in)) == -1) {
							if (errno == EAGAIN) {
#ifdef _POSIX_PRIORITY_SCHEDULING
								sched_yield();
#else
								usleep(0);
#endif
							}
						}
						// Without this, we use too much CPU and rx_thread loses packets. This should be configurable.
						if (++sleep_count % DISCOVERY_ICMP_SLEEPCOUNTINT_DEF == 0) {
							usleep(USLEEP_FASTLOOP_TIME);
						}
					}
				}
			}
			/* Wait timeout between volleys. */
			usleep(instance->timeout * 1000);
		}
		close(s);
	}

	if (!interrupt) {
		/* Shutdown the other threads. */
		usleep(instance->timeout * 1000 + USLEEP_LONGLOOP_TIME);
		interrupt = true;
	}

	TRACE(TRACE_THREAD, "ICMP::transmit_thread() ending main loop");

	pthread_exit(NULL);
}

RESULT ICMP::receive_thread(void* args) {
	ICMP *instance = (ICMP*) args;

	int size = 0;
	int count = 0;

	char buffer[PACKET_BUFFER_SIZE];
	memset(buffer, 0, PACKET_BUFFER_SIZE);
	struct ip *ip_header = (ip*)buffer;
	struct icmp *reply = NULL;

	TRACE(TRACE_THREAD, "ICMP::receive_thread() starting main loop");

	grantPriv();

	struct pollfd sock_events;
	sock_events.fd = socket(PF_INET, SOCK_RAW, IPPROTO_ICMP);
	sock_events.events = POLLIN;

	long max_buffer_size = 64000;	

	size = sizeof(long);

	setsockopt(sock_events.fd, SOL_SOCKET, SO_RCVBUF, &max_buffer_size, sizeof(max_buffer_size));
	fcntl(sock_events.fd, F_SETFL, O_NONBLOCK);	

	dropPriv();

	while (!interrupt) {
		count = poll(&sock_events, 1, instance->timeout);

		if (count > 0) {
			switch (sock_events.revents) {
				case POLLIN:
					size = recvfrom(sock_events.fd, (void *)ip_header, PACKET_BUFFER_SIZE, 0, NULL, 0);
					if (size != -1) {	
						reply = (icmp *) (buffer + (ip_header->ip_hl << 2));

						if (reply->icmp_type == ICMP_ECHOREPLY && reply->icmp_id == instance->module_id) {
							if (instance->rx_hash.find(ip_header->ip_src.s_addr) == instance->rx_hash.end()) {
								pthread_mutex_lock(&(instance->rx_hash_mutex));
									instance->rx_hash[ip_header->ip_src.s_addr] = 1;
								pthread_mutex_unlock(&(instance->rx_hash_mutex));
							} else {
								if (instance->rx_hash[ip_header->ip_src.s_addr] != 0) {
									pthread_mutex_lock(&(instance->rx_hash_mutex));				
										instance->rx_hash[ip_header->ip_src.s_addr]++;
									pthread_mutex_unlock(&(instance->rx_hash_mutex));
								}
							}
						}
					}
					break;

				case POLLERR:
				case POLLHUP:
				case POLLNVAL:
					break;
			}
		} else if (count < 0) {
			if (errno != EINTR) {
				LMSG("Error: Call to poll() returned %d - %s. Exiting thread.", count, strerror(errno));
				interrupt = true;
			}
		}
	}

	close(sock_events.fd);

	TRACE(TRACE_THREAD, "ICMP::receive_thread() ending main loop");

	pthread_exit(NULL);
}

RESULT ICMP::storage_thread(void *args) {
	ICMP *instance = (ICMP*) args;

	string sql;
	char buf[132];
	char inet_buffer[16];

	struct in_addr adr;

	bool write = false;
	StorageData data;

	TRACE(TRACE_THREAD, "ICMP::storage_thread() starting main loop");

	/* Clear database of previous discovery information */
	if (instance->storage.query(data, "DELETE FROM discovery_icmp") != SUCCESS) {
		LMSG("ICMP::storage_thread() failed to delete previous discovery information.");
		interrupt = true;
	}
	data.clear();

	while (!interrupt) {
		for (received_hash_t::iterator it = instance->rx_hash.begin(); it != instance->rx_hash.end(); ++it) {
			if (it->second >= instance->rx_count) {
				adr.s_addr = it->first;	

				snprintf(buf, 132, "('%s'),", inet_ntop(AF_INET, &adr, inet_buffer, 16));
				sql.append(buf);

				pthread_mutex_lock(&(instance->rx_hash_mutex));
					instance->rx_hash[it->first] = 0;
				pthread_mutex_unlock(&(instance->rx_hash_mutex));
				write = true;
			}
		}

		if (write) {
			sql.erase(sql.size() - 1);
			sql.insert(0, "INSERT INTO discovery_icmp (ip) VALUES ");
			if (instance->storage.query(data, sql.c_str()) == SUCCESS) {
				sql.erase();
				data.clear();
				write = false;
			} else {
				LMSG("NOTICE: Failed to write data in ICMP.");
			}
		}
		usleep(USLEEP_NORMLOOP_TIME);
	}

	TRACE(TRACE_THREAD, "ICMP::storage_thread() ending main loop");

	pthread_exit(NULL);
}

/*
** Local Variables:
** c-basic-offset: 3
** tab-width: 3
** End:
** vim:
** vim: noet ts=3 sw=3
*/
