/*
** SellaNMS
** $Id: netflow.cxx,v 1.20 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 "netflow.h"
#include "debug.h"
#include "snmp_util.h"
#include <fcntl.h>
#include <unistd.h>
#include <pthread.h>
#include <sys/poll.h>
#include <sys/time.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <linux/if.h>

NetFlow::NetFlow(): proxy(NULL), accept(NULL), reject(NULL) {
	this->name.append(" netflow");

	has_scheduler = false;
	has_execute = true;
}

NetFlow::~NetFlow() {
	SAFE_DELETE(this->proxy);
	SAFE_DELETE(this->accept);
	SAFE_DELETE(this->reject);

}

RESULT NetFlow::readConfig() {
	RESULT result = SUCCESS;
	string str;
	vector<string> v;
	vector<string>::iterator v_it;

	if (config->assignValue("port", str)) {
		this->port = strtoul(str.c_str(), NULL, 10);
	} else {
		this->port = MODULE_MONITOR_NETFLOW_PORT_DEFAULT;
	}

	config->assignValues("proxy", this->proxy);

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

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

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

	return result;
}

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

	pthread_t proxy_thread_id = 0;
	pthread_t tree_thread_id = 0;
	pthread_t calc_thread_id = 0;
	pthread_t storage_thread_id = 0;
	pthread_t receive_thread_id = 0;

	netflow_queue = new netflow_queue_t();
	netflow_mutex = new pthread_mutex_t();
	tree_queue = new netflow_queue_t();
	tree_mutex = new pthread_mutex_t();
	proxy_queue = new netflow_queue_t();
	proxy_mutex = new pthread_mutex_t();

	bool running = true;
	struct snmp_session session;

	pthread_mutex_init(netflow_mutex, NULL);
	pthread_mutex_init(tree_mutex, NULL);
	pthread_mutex_init(proxy_mutex, NULL);

	result = FAILURE;
	
	/* Create the proxy thread (if needed) */
	if (proxy != NULL) {
		if (createThread(proxy_thread_id, (void* (*)(void*))NetFlow::proxy_thread, this) != SUCCESS) {
			result = FAILURE;
		}
	}

	/* Create the tree thread */
	if (createThread(tree_thread_id, (void* (*)(void*))NetFlow::tree_thread, this) != SUCCESS) {
		result = FAILURE;
	}

	/* Create the storage thread */
	if (createThread(storage_thread_id, (void* (*)(void*))NetFlow::storage_thread, this) != SUCCESS) {
		result = FAILURE;
	}

	/* Create the calc thread */
	if (createThread(calc_thread_id, (void* (*)(void*))NetFlow::calc_thread, this) != SUCCESS) {
		result = FAILURE;
	}

	/* Create the receiver thread */
	if (createThread(receive_thread_id, (void* (*)(void*))NetFlow::receive_thread, this) != SUCCESS) {
		result = FAILURE;
	}

	/* Wait for rx thread to complete. */
	this->getThreadResult(receive_thread_id, t_result);
	/* Request tree thread kill itself. */
	this->killThread(tree_thread_id);
	/* Wait for tree thread to complete. */
	this->getThreadResult(tree_thread_id, t_result);
	/* Request calc thread kill itself. */
	this->killThread(calc_thread_id);
	/* Wait for calc thread to complete. */
	this->getThreadResult(calc_thread_id, t_result);
	/* Request storage thread kill itself. */
	this->killThread(storage_thread_id);
	/* Wait for storage thread to complete. */
	this->getThreadResult(storage_thread_id, t_result);
	if (proxy_thread_id) {
		/* Request proxy thread kill itself. */
		this->killThread(proxy_thread_id);
		/* Wait for proxy thread to complete. */
		this->getThreadResult(proxy_thread_id, t_result);
	}

	pthread_mutex_destroy(netflow_mutex);
	pthread_mutex_destroy(tree_mutex);
	pthread_mutex_destroy(proxy_mutex);

	SAFE_DELETE(netflow_mutex);
	SAFE_DELETE(netflow_queue);
	SAFE_DELETE(proxy_mutex);
	SAFE_DELETE(proxy_queue);

	return result;
}

RESULT NetFlow::checkFilter(/* Something */) {
	RESULT result = SUCCESS;
	char buf[31];
	unsigned int len;
#if 0

	if (this->community) {
		result = FAILURE;
		len = (sizeof(buf) > pdu->community_len) ? pdu->community_len + 1 : sizeof(buf);
		snprintf(buf, len, "%s", (char*) pdu->community);
		for(vector<string>::iterator v_it = this->community->begin(); v_it != this->community->end(); ++v_it) {
			if (!v_it->compare(buf)) {
				result = SUCCESS;
				break;
			}
		}

		if (result == FAILURE) {
			return FAILURE;
		}
	}

	if (this->reject) {
		result = SUCCESS;
		for(vector<unsigned int>::iterator v_it = this->reject->begin(); v_it != this->reject->end(); ++v_it) {
			if (pdu->flow_type == *v_it) {
				result = FAILURE;
				break;
			}
		}
	} else if (this->accept) {
		result = FAILURE;
	}
	
	if (this->accept) {
		for(vector<unsigned int>::iterator v_it = this->accept->begin(); v_it != this->accept->end(); ++v_it) {
			if (pdu->flow_type == *v_it) {
				result = SUCCESS;
				break;
			}
		}
	}
#endif

	return result;
}


RESULT NetFlow::receive_thread(void *args) {
	RESULT result = SUCCESS;	
	STATUS status;

	pthread_t thread_id = pthread_self();

	int s, count;
	socklen_t len;
	struct sockaddr_in server;
	struct sockaddr_in client;
	fd_set fd;
	char buf[2024], *ptr;

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

	netflow_queue_t *netflow_queue = instance->netflow_queue;
	pthread_mutex_t *netflow_mutex = instance->netflow_mutex;
	netflow_queue_t *proxy_queue = instance->proxy_queue;
	pthread_mutex_t *proxy_mutex = instance->proxy_mutex;

	bool running = false;

	if ((s = socket(AF_INET, SOCK_DGRAM, 0)) == -1) {
		LMSG("NetFlow::receive_thread -Failed to create socket.");
	} else {
		bzero(&server, sizeof(struct sockaddr_in));
		server.sin_family = AF_INET;
		server.sin_addr.s_addr = htonl(INADDR_ANY);
		server.sin_port = htons(4444);
		if (bind(s, (struct sockaddr *) &server, sizeof(server)) == -1)  {
			LMSG("NetFlow::receive_thread - Failed to bind to port");
			
			close(s);
			s = -1;
		}

		running = true;
	}

	FD_ZERO(&fd);
	FD_SET(s, &fd);

	while (running) {
		instance->getThreadStatus(thread_id, status);
		if (status == THREAD_KILLED) {
			running = false;
		} else {
			select(s + 1, &fd, NULL, NULL, NULL);

			if	(FD_ISSET(s, &fd)) {
				len = sizeof(struct sockaddr);
				if ((count = recvfrom(s, buf, 2024 - 1, 0, (struct sockaddr *) &client, &len)) == -1) {
					LMSG("NetFlow::receive_thread - recvfrom() failed.");
				}

				if (instance->proxy) {
					ptr = (char *) malloc(count);
					if (ptr) {
						memcpy(ptr, buf, count);
						if (proxy_queue->size() > 1) {
							proxy_queue->push(ptr);
						} else {
							pthread_mutex_lock(proxy_mutex);
							proxy_queue->push(ptr);
							pthread_mutex_unlock(proxy_mutex);
						}
					}
				}
	
				ptr = (char *) malloc(count);
				if (ptr) {
					memcpy(ptr, buf, count);
					if (netflow_queue->size() > 1) {
						netflow_queue->push(ptr);
					} else {
						pthread_mutex_lock(netflow_mutex);
						netflow_queue->push(ptr);
						pthread_mutex_unlock(netflow_mutex);
					}
				}
			}
		}
	}

	close(s);

	instance->setThreadStatus(thread_id, THREAD_COMPLETED);

	THREAD_RETURN(result);
}

RESULT NetFlow::calc_thread(void *args) {
	RESULT result = SUCCESS;
	STATUS status;

	pthread_t thread_id = pthread_self();

	char *dgram, *ptr;
	netflow_v5_header header;
	int header_len = sizeof(netflow_v5_header);
	netflow_v5_record record;
	int record_len = sizeof(netflow_v5_record);

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

	netflow_queue_t *netflow_queue = instance->netflow_queue;
	pthread_mutex_t *netflow_mutex = instance->netflow_mutex;

	bool running = true;

	// Until we are asked to stop, store our successes.
	while (running) {
		instance->getThreadStatus(thread_id, status);
		if (status == THREAD_KILLED) {
			running = false;
		}
		while (!netflow_queue->empty()) {
			if (netflow_queue->size() <= 1) {
				dgram = netflow_queue->front();
				netflow_queue->pop();
			} else {
				pthread_mutex_lock(netflow_mutex);
				dgram = netflow_queue->front();
				netflow_queue->pop();
				pthread_mutex_unlock(netflow_mutex);
			}

			memcpy(&header, dgram, header_len);
			ptr = dgram + header_len;
			for(int i = 0; i < ntohs(header.count); i++) {
				assert(i <= 30); /* Shouldn't be more than 30 records per datagram */
				memcpy(&record, ptr, record_len);
				ptr += record_len;

//				LMSG("Souce ASN: %d, Destination ASN: %d", ntohs(record.src_as), ntohs(record.dst_as));
			}

			free(dgram);
#if 0
			if (instance->checkFilter(pdu) == FAILURE) {
				continue;
			}
#endif
		}

		usleep(USLEEP_NORMLOOP_TIME);
	}

	instance->setThreadStatus(thread_id, THREAD_COMPLETED);

	THREAD_RETURN(result);
}

RESULT NetFlow::tree_thread(void *args) {
	RESULT result = SUCCESS;
	STATUS status;

	pthread_t thread_id = pthread_self();

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

	netflow_queue_t *proxy_queue = instance->proxy_queue;
	pthread_mutex_t *proxy_mutex = instance->proxy_mutex;

	bool running = false;

	// Until we are asked to stop, store our successes.
	while (running) {
		instance->getThreadStatus(thread_id, status);
		if (status == THREAD_KILLED) {
			running = false;
		}
		while (!proxy_queue->empty()) {
		}

		usleep(USLEEP_NORMLOOP_TIME);
	}

	instance->setThreadStatus(thread_id, THREAD_COMPLETED);

	THREAD_RETURN(result);
}

RESULT NetFlow::proxy_thread(void *args) {
	/* NOTE:
	 *
	 * This needs to be rewritten as a RAW socket. NetFlow relies on the source of the packet,
	 * to know which router the data is coming from, so we need to spoof the source.
	 *
	 */
	
	RESULT result = SUCCESS;
	STATUS status;

	pthread_t thread_id = pthread_self();

	int *s, count, proxy_len;
	socklen_t len;
	struct sockaddr_in client;
	fd_set fd;
	char *dgram;
	netflow_v5_header header;
	int header_len = sizeof(netflow_v5_header);
	int record_len = sizeof(netflow_v5_record);

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

	netflow_queue_t *proxy_queue = instance->proxy_queue;
	pthread_mutex_t *proxy_mutex = instance->proxy_mutex;
	vector<string> *proxy = instance->proxy;

	bool running = true;

	proxy_len = proxy->size();
	s = new int[proxy_len];
	for(int i = 0; i < proxy_len; i++) {
		if ((s[i] = socket(AF_INET, SOCK_DGRAM, 0)) == -1) {
			LMSG("NetFlow::proxy_thread - Failed to create socket. (out of file descriptors?)");
			if (i < 1) { /* proxy_thread should run if it got at least 1 socket */
				LMSG("NetFlow::proxy_thread - proxy_thread will exit");
				running = false;
			} else {
				LMSG("NetFlow::proxy_thread - Will proxy to only %d of %d clients.", i + 1, proxy_len);
			}
			proxy_len = i;
			break;
		}
	}

	// Until we are asked to stop, store our successes.
	while (running) {
		instance->getThreadStatus(thread_id, status);
		if (status == THREAD_KILLED) {
			running = false;
		}
		while (!proxy_queue->empty()) {
			if (proxy_queue->size() <= 1) {
				dgram = proxy_queue->front();
				proxy_queue->pop();
			} else {
				pthread_mutex_lock(proxy_mutex);
				dgram = proxy_queue->front();
				proxy_queue->pop();
				pthread_mutex_unlock(proxy_mutex);
			}

			memcpy(&header, dgram, header_len);
			len = header_len + ntohs(header.count) * record_len; 

			header.version = ntohs(header.version);
			header.count = ntohs(header.count);
			header.uptime = ntohl(header.uptime);
			header.unix_secs = ntohl(header.unix_secs);
			
			for(int i = 0; i < proxy_len; i++) {
				bzero(&client, sizeof(struct sockaddr_in));
				client.sin_family = AF_INET;
				inet_pton(AF_INET, (*proxy)[i].c_str(), &(client.sin_addr));
				client.sin_port = htons(4444);

				if ((count = sendto(s[i], dgram, len, 0, (struct sockaddr *) &client, sizeof(struct sockaddr))) == -1) {
					LMSG("Failed to forward to %s", (*proxy)[i].c_str());
				}
			}

			free(dgram);
		}

		usleep(USLEEP_NORMLOOP_TIME);
	}

	for (int i = 0; i < proxy_len; i++) {
		close(s[i]);
	}
	delete[] s;

	instance->setThreadStatus(thread_id, THREAD_COMPLETED);

	THREAD_RETURN(result);
}

RESULT NetFlow::storage_thread(void *args) {
	RESULT result = SUCCESS;
	STATUS status;

	pthread_t thread_id = pthread_self();

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

	netflow_queue_t *netflow_queue = instance->netflow_queue;
	pthread_mutex_t *netflow_mutex = instance->netflow_mutex;

	struct in_addr adr;

	bool write = false;
	bool running = false;
	StorageData data;
	char buf[31];
	unsigned int len;
	string* last_id;
	snmp_pdu* pdu;
	struct sockaddr_in *source;
	variable_list *var;

	// Until we are asked to stop, store our successes.
	while (running) {
		instance->getThreadStatus(thread_id, status);
		if (status == THREAD_KILLED) {
			running = false;
		}
		while (!netflow_queue->empty()) {
			/* Fetch Data from queue */
#if 0
			pthread_mutex_lock(netflow_mutex);
			pdu = netflow_queue->front();
			netflow_queue->pop();
			pthread_mutex_unlock(netflow_mutex);
#endif

#if 0

			/* Conv PDU into a non-proxy V2 flow. */
			instance->conv_pdu(pdu, 2, false);
			
			if (instance->checkFilter(pdu) == FAILURE) {
				snmp_free_pdu(pdu);
				continue;
			}
			
			source = (struct sockaddr_in*) &(pdu->agent_addr);
			
			data->setFlag(STORAGEDATA_WRITE | STORAGEDATA_LASTID);
			data->setData("monitor.source", inet_ntoa(source->sin_addr));
			data->setData("monitor.protocol", "snmpflow");
			data->setData("monitor.layer", "3");

			data->setData("monitor.class", CLASS_INFORMATION);

			if (instance->writeStorage(data) != SUCCESS) {
				LMSG("Failed to write failures to monitor.");
			} else {
				// Grab the LAST_ID before erasing data.
				last_id = data->getLastId();

				data->clear();

				data->setData("monitor_data.monitor_id", *last_id);
				data->setData("monitor_data.variable", "community");
				len = (sizeof(buf) > pdu->community_len) ? pdu->community_len + 1 : sizeof(buf);
				snprintf(buf, len, "%s", (char*) pdu->community);
				data->setData("monitor_data.value", buf);

				/* Process variables */
				for(var = pdu->variables; var != NULL; var = var->next_variable) {
					data->setData("monitor_data.monitor_id", *last_id);
					convOidToString(buf, var->name, var->name_length);
					data->setData("monitor_data.variable", buf);

					switch (var->type) {
						case ASN_OBJECT_ID:
							convOidToString(buf, (oid*) var->val.str, (int) (var->val_len / sizeof(oid)));
							break;

						case ASN_IPADDRESS:
							convVarListToIP(buf, var);
							break;

						case ASN_INTEGER:
							snprintf(buf, MAX_OID_LEN, "%ld", (long *) *var->val.integer);
							break;

						case ASN_GAUGE:
						case ASN_COUNTER:
						case ASN_TIMETICKS:
							snprintf(buf, MAX_OID_LEN, "%lu", (long *) *var->val.integer);
							break;

						case ASN_NULL:
						case ASN_BOOLEAN:
						case ASN_COUNTER64:
						case ASN_OCTET_STR:
							sprint_value(buf, var->name, var->name_length, var);
							break;

						default:
							sprint_value(buf, var->name, var->name_length, var);
					}
					data->setData("monitor_data.value", buf);
				}

				if (instance->writeStorage(data) != SUCCESS) {
					LMSG("Failed to write to monitor_data.");
				}

				// Clear for next time thru.
				data->clear();

				delete last_id;
			}
#endif

			write = true;
		}

#if 0
		if (write) {
			if (instance->write(data) == SUCCESS) {
				data.clear();
				write = false;
			}
		}
#endif
		usleep(USLEEP_NORMLOOP_TIME);

	}

	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
*/
