/*
** SellaNMS
** $Id: icmp.cxx,v 2.60 2005/03/31 18:57:12 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

#include "icmp.h"
#include "debug.h"
#include "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 int errno;

ICMP::ICMP() {
	this->name.append(" icmp");

	has_scheduler = false;
	has_execute = true;

	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;

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

	config->assignValue("packet-options transmit", tx_count, MODULE_DISCOVERY_ICMP_TRANSMITCOUNT_DEFAULT);
	config->assignValue("packet-options require", rx_count, MODULE_DISCOVERY_ICMP_REQUIREDCOUNT_DEFAULT);
	/* Repair stupid config settings. */
	if (rx_count > tx_count) {
		rx_count = tx_count;
	}
	config->assignValue("packet-options timeout", timeout, MODULE_DISCOVERY_ICMP_TIMEOUT_DEFAULT);
//	config->assignValue("packet-options source", source);

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

	/* Add and compile input policy. */
	if (config->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 (config->assignValues("policy output", v)) {
		if (outputPolicy.addPolicies(config, v) == FAILURE) {
			LMSG("NOTICE: Failed to add/compile output policies.");
			result = FAILURE;
		}
	}

#if 0
	/* Verify we have required policies. */
	if (!outputPolicy) {
		LMSG("NOTICE: A output policy is required for this module.");
		result = FAILURE;
	}
#endif

	/* Configure Storage */
   storage.setConfig(*config);

	return result;
}

RESULT ICMP::printConfig() {
	char	inet_buffer[16];

	in_addr s;
	s.s_addr = this->source;

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

	Discovery::printConfig();

	return SUCCESS;
}

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

	pthread_t storage_thread_id = 0;
	pthread_t rx_thread_id = 0;
	pthread_t tx_thread_id = 0;

	bool running = true;

	if (trace.thread) {
		LMSG("ICMP::execute() starting");
	}

	module_id = pthread_self();

	/* Create the storage thread */
	result = createThread(storage_thread_id, (void* (*)(void*))ICMP::storage_thread, (void*) this);

	/* Create the receiver thread */
	if (result == SUCCESS) {
		result = createThread(rx_thread_id, (void* (*)(void*))ICMP::receive_thread, (void*) this);
	}

	/* Create the transmitter thread */
	if (result == SUCCESS) {
		result = createThread(tx_thread_id, (void* (*)(void*))ICMP::transmit_thread, (void*) this);
	}

	/* Wait for tx thread to finish. */
	getThreadResult(tx_thread_id, t_result);
	/* Request rx thread kill itself. */
	this->killThread(rx_thread_id);
	/* Wait until rx thread is finished. */
	getThreadResult(rx_thread_id, t_result);
	/* Request storage thread kill itself. */
	this->killThread(storage_thread_id);
	/* Wait until storage thread is finished. */
	getThreadResult(storage_thread_id, t_result);

	if (trace.thread) {
		LMSG("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;
	pthread_t thread_id = pthread_self();

	ICMP *instance = (ICMP*) args;
	instance->setThreadStatus(thread_id, THREAD_RUNNING);

	bool early_exit = false;
	unsigned long sleep_count = 0;

	Netblock netblock;
	char *ip_address = NULL;
	sockaddr_in dst_address;
	sockaddr_in src_address;
		
	int recvbuffer = 0;

	char buffer[PACKET_BUFFER_SIZE];
	char inet_buffer[16];

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

	if (instance->trace.thread) {
		LMSG("ICMP::transmit_thread() starting main loop");
	}

	instance->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;
	}

	instance->dropPriv();

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

		request->icmp_id = (u_int16_t) 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 && !early_exit; 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]; 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 % MODULE_DISCOVERY_ICMP_SLEEPCOUNTINT_DEFAULT == 0) {
							usleep(USLEEP_FASTLOOP_TIME);
						}
					}
				}
			}
			/* Wait timeout between volleys. */
			usleep(instance->timeout * 1000);
		}
		close(s);
	}

	if (instance->trace.thread) {
		LMSG("ICMP::transmit_thread() ending main loop");
	}

	instance->setThreadStatus(thread_id, THREAD_COMPLETED);

	THREAD_RETURN(result);
}

RESULT ICMP::receive_thread(void* args) {
	RESULT result = SUCCESS;	
	STATUS status;
	pthread_t thread_id = pthread_self();

	ICMP *instance = (ICMP*) args;
	instance->setThreadStatus(thread_id, THREAD_RUNNING);

	int size = 0;
	int count = 0;
	bool running = true;

	struct sockaddr src_address;

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

	if (instance->trace.thread) {
		LMSG("ICMP::receive_thread() starting main loop");
	}

	instance->grantPriv();

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

	long max_buffer_size = 64000;	

	int so_rcvbuf = SO_RCVBUF;	
	size = sizeof(long);

	socklen_t addrsize = sizeof(sockaddr_in);

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

	instance->dropPriv();

	while (running) {
		instance->getThreadStatus(thread_id, status);
		if (status == THREAD_KILLED) {
			running = false;
		} else {
			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_id == (u_int16_t) instance->module_id && reply->icmp_type == ICMP_ECHOREPLY) {
								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) {
				LMSG("Error: Call to poll() returned %d. Exiting thread.", count);
				running = false;
			}
		}
	}

	close(sock_events.fd);

	if (instance->trace.thread) {
		LMSG("ICMP::receive_thread() ending main loop");
	}

	instance->setThreadStatus(thread_id, THREAD_COMPLETED);

	THREAD_RETURN(result);
}

RESULT ICMP::storage_thread(void *args) {
	RESULT result = SUCCESS;
	STATUS status;
	pthread_t thread_id = pthread_self();

	ICMP *instance = (ICMP*) args;
	instance->dropPriv();
	instance->setThreadStatus(thread_id, THREAD_RUNNING);
	string sql;
	char buf[132];
	char inet_buffer[16];

	struct in_addr adr;

	bool write = false;
	bool running = true;
	StorageData data;

	if (instance->trace.thread) {
		LMSG("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.");
	}
	data.clear();

	while (running) {
		instance->getThreadStatus(thread_id, status);
		if (status == THREAD_KILLED) {
			running = false;
		}
		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);
	}

	if (instance->trace.thread) {
		LMSG("ICMP::storage_thread() ending main loop");
	}

	instance->setThreadStatus(thread_id, THREAD_COMPLETED);

	THREAD_RETURN(result);
}

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