/*
** SellaNMS
** $Id: snmp.cxx,v 1.89 2005/08/22 02:25:42 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 "snmp.h"
#include "debug.h"
#include "netblock.h"
#include "ip_util.h"
#include <errno.h>

#define BUF_SIZE 1024
#define USE_NEW_CODE

SNMP::SNMP(): maxThreadCount(0), timeout(0), retry(0), commit_count(0) {
	this->name.append(" snmp");

	has_scheduler = false;
	has_execute = true;

#ifdef USE_NEW_CODE
//	LMSG("YO!! IF THE DISCOVERY COMES OUT SCREWY, TRY TURNING OFF USE_NEW_CODE. I didn't verify the new function worked within this module.");
//	usleep(1000000);
#endif

	pthread_mutex_init(&to_do_queue_mutex, NULL);
	pthread_mutex_init(&result_queue_mutex, NULL);
}

SNMP::~SNMP() {
	while (EBUSY == pthread_mutex_destroy(&to_do_queue_mutex)) {
		usleep(USLEEP_NORMLOOP_TIME);
	}

	while (EBUSY == pthread_mutex_destroy(&result_queue_mutex)) {
		usleep(USLEEP_NORMLOOP_TIME);
	}

	/* Clean up real_oids */
	for (oid_vector_t::iterator it = real_oids.begin(); it != real_oids.end(); ++it) {
		SAFE_DELETE_ARRAY(it->first);
	}
}

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

	string str;
	struct in_addr ip;
	
	config->assignValue("threads", maxThreadCount, MODULE_DISCOVERY_SNMP_THREADCOUNT_DEFAULT);
	if (maxThreadCount > PTHREAD_THREADS_MAX) {
		maxThreadCount = PTHREAD_THREADS_MAX;
	}

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

	/* Used to limit the size of INSERTs, both for delay and size of SQL query. 300 gives ~16500 characters. */
	config->assignValue("commit-count", commit_count, MODULE_DISCOVERY_SNMP_COMMITCOUNT_DEFAULT);
	commit_count = (commit_count > 1) ? commit_count : 1;

	if (config->assignValue("timeout", timeout, MODULE_DISCOVERY_SNMP_TIMEOUT_DEFAULT)) {
		timeout *= 1000;
	}

	config->assignValue("retry", retry, MODULE_DISCOVERY_SNMP_RETRY_DEFAULT);

	if (!config->assignValues("community", communities)) {
		communities.push_back("public");
	}

	config->assignValues("variable", v);
	//FIXME: We need to shove some oids in here.
	
	/* Resolve all our str oids into real oids. */
	struct snmp_session session, *ss;
	void *opaque;
	pair<oid*, size_t> pair_ptr;
	oid root[MAX_OID_LEN];
	size_t root_length;
	char buf[255];
	bool done;

	for (vector<string>::iterator it = v.begin(); it != v.end(); ++it) {
		root_length = MAX_OID_LEN;
		if (read_objid(it->c_str(), root, &root_length) == 0) {
			if (trace.internal) {
				LMSG("read_objid() failed to resolve '%s'. Trying get_node().", it->c_str());
			}

			root_length = MAX_OID_LEN;
			if (get_node(it->c_str(), root, &root_length) == 0) {
				LMSG("WARNING: get_node() failed to resolve OID '%s'. Skipping this OID.", it->c_str());
				continue;
			}
		}

		pair_ptr.first = new oid[root_length];
		memmove(pair_ptr.first, root, sizeof(oid) * root_length);
		pair_ptr.second = root_length;
		real_oids.push_back(pair_ptr);
	}

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

	this->config->assignValues("prefix", netblocks);

	if (readKnownAddress() != SUCCESS) {
		LMSG("Failed to read known netblocks from database");
	}

	return result;
}

RESULT SNMP::printConfig() {
	LMSG("threads: %lu", this->maxThreadCount);
	LMSG("commit-count: %lu", this->commit_count);
	LMSG("timeout: %lu", this->timeout);
	LMSG("retry: %lu", this->retry);

	Discovery::printConfig();

	return SUCCESS;
}

RESULT SNMP::execute() {
	RESULT result = SUCCESS;
	RESULT thread_result;
	StorageData data;

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

	if (this->fillWorkQueue() == FAILURE) {
		LMSG("Warning: SNMP Discovery Module will not scan all wanted netblocks; check previous messages.");

		result = FAILURE;
	}

	if (this->spawnThreads() == FAILURE) {
		LMSG("Warning: SNMP Discovery Module was only able to spawn %i threads and was configured to spawn %i.", children.size(), maxThreadCount);

		result = FAILURE;
	}

	if (storage.query(data, "DELETE FROM discovery_snmp_oid") == FAILURE) {
		LMSG("Error: Failed to delete old info from discovery_snmp_oid.");
	}

	while (!children.empty() || !result_queue.empty()) {
		if (result_queue.empty()) {
			usleep(USLEEP_NORMLOOP_TIME);
		}

		if (this->queryThreads() == FAILURE) {
			DMSG(FLAG_ERROR, "Dependant function, SNMP::queryThreads(), failed.");

			result = FAILURE;
		}

		if (this->writeResults() == FAILURE) {
			DMSG(FLAG_ERROR, "Dependant function, SNMP::writeResults(), failed.");
	
			result = FAILURE;
		}
	}

	assert(result_queue.size() == 0);
	assert(children.empty());

	if (trace.thread) {
		LMSG("SNMP::execute() ending");
	}

	return result;
}

RESULT SNMP::fillWorkQueue() {
	RESULT result = SUCCESS;
	Netblock *netblock = NULL;

	char *ip_address = NULL;
	pthread_mutex_lock(&to_do_queue_mutex);
	for (int i = 0; i < netblocks.size(); i++) {
		netblock = new Netblock(netblocks[i].c_str());

		if (netblock && netblock->isValid()) {
			for (int j = 0; ip_address = (*netblock)[j]; j++) {
				to_do_queue.push(ip_address);
				SAFE_DELETE_ARRAY(ip_address);
			}
		} else {
			LMSG("Netblock Parse Error %s.", netblocks[i].c_str());

			result = FAILURE;
		}

		SAFE_DELETE(netblock);
	}
	pthread_mutex_unlock(&to_do_queue_mutex);

	return result;
}

RESULT SNMP::spawnThreads() {
	RESULT result = SUCCESS;
	pthread_t thread_id = (pthread_t) 0;
	unsigned long param_size = sizeof(this);

	for (unsigned i = 0; i < maxThreadCount; i++) {
		if (createThread(thread_id, (void* (*)(void*))SNMP::query_thread, (void*) this) == SUCCESS) {
			children.push_back(thread_id);
			usleep(USLEEP_NORMLOOP_TIME);
		} else {
			LMSG("Unable to spawn thread for SNMP Discovery Module (%i/%i spawned), throttling back target thread count (by %d) and going to sleep for a while before trying again.", children.size(), maxThreadCount, (maxThreadCount - i) / 2);
			i += ((maxThreadCount - i) / 4);
			usleep(USLEEP_LONGLOOP_TIME);

		}
	}
	if (children.size() < 1) {
		result = FAILURE;
	}

	return result;
}

RESULT SNMP::queryThreads() {
	RESULT result = SUCCESS;
	STATUS status;
	list<pthread_t>::iterator it;

	it = children.begin();
	while (it != children.end()) {
		getThreadStatus(*it, status);
		if (status == THREAD_COMPLETED || status == THREAD_KILLED) {
			getThreadResult(*it, result);
			it = children.erase(it);
		} else {
			++it;
		}
	}
}

RESULT SNMP::writeResults() {
	RESULT result = SUCCESS;
	string sql, ip;
	StorageData data;
	unsigned long count = 0, throttle = 5;

	while (!result_queue.empty() && throttle++ > 0 && !interrupt) {
		pthread_mutex_lock(&result_queue_mutex);

		ip = result_queue.front().first;
		while (!result_queue.front().second->empty() && count < commit_count) {
			count++;
			sql.append("('" + ip + "','" + result_queue.front().second->front().first + "'," + data.quoteString(result_queue.front().second->front().second) + "),");
			result_queue.front().second->pop();
		}

		if (result_queue.front().second->empty()) {
			SAFE_DELETE(result_queue.front().second);
			result_queue.pop();
		}

		pthread_mutex_unlock(&result_queue_mutex);

		if (!sql.empty()) {
			sql.insert(0, "INSERT INTO discovery_snmp_oid (ip, name, value) VALUES");
			sql.erase(sql.size() - 1, 1);

			if (storage.query(data, sql.c_str()) == FAILURE) {
				LMSG("Error: Failed to write to database.");
			}
			sql.erase();
			count = 0;
		}
	}
	interrupt = false;

	return result;
}

RESULT SNMP::sendRequest(pthread_t thread_id, const char* dst_address, query_result_t &ret_val) {
	RESULT result = FAILURE;
	int liberr, syserr, i, j, oid_count, community_count;
	char *errmsg, *buf, *buf2;
	bool running, correct_community;
	void *opaque;
	struct snmp_session session, *ss;
	struct snmp_pdu *pdu;
	struct snmp_pdu *response;
	struct variable_list *variables;
	pair<string, string> string_pair;
	oid root[MAX_OID_LEN];
	oid name[MAX_OID_LEN];
	size_t root_length = MAX_OID_LEN;
	size_t name_length = MAX_OID_LEN;

	buf = new char[BUF_SIZE];
	buf2 = new char[BUF_SIZE];

	/* Initialize our local session */
	snmp_sess_init(&session);
	session.version = SNMP_VERSION_1;
	session.retries = retry;
	session.timeout = timeout;
	session.peername = (char*) dst_address;

	community_count = communities.size();
	for (i = 0; i < community_count; i++) {
		correct_community = true;

		oid_count = real_oids.size();
		/* Open our connection to dst_address. Community doesn't matter yet. */
		opaque = snmp_sess_open(&session);
		if (opaque == NULL) {
			snmp_sess_error(opaque, &liberr, &syserr, &errmsg);
			DMSG(FLAG_ERROR, "SNMP Error: %s", errmsg);
			free(errmsg);
			LMSG("Failed to load SNMP library.");
			result = FAILURE;
		} else {
			/* Grab a pointer to the session being used internally for this thread. */
			ss = snmp_sess_session(opaque);
			free(ss->community);
			ss->community = (u_char*) strdup(communities[i].c_str());
			ss->community_len = communities[i].length();
		}

		/* Instruct the libary to leave the OID and Hints off sprint_value() output. */
#ifdef HAVE_NET_SNMP
		netsnmp_ds_set_boolean(NETSNMP_DS_LIBRARY_ID, NETSNMP_DS_LIB_PRINT_NUMERIC_OIDS, 1);
		netsnmp_ds_set_boolean(NETSNMP_DS_LIBRARY_ID, NETSNMP_DS_LIB_QUICK_PRINT, 1);
		netsnmp_ds_set_boolean(NETSNMP_DS_LIBRARY_ID, NETSNMP_DS_LIB_PRINT_BARE_VALUE, 1);
#else
		ds_set_boolean(DS_LIBRARY_ID, DS_LIB_PRINT_NUMERIC_OIDS, 1);
		ds_set_boolean(DS_LIBRARY_ID, DS_LIB_QUICK_PRINT, 1);
		ds_set_boolean(DS_LIBRARY_ID, DS_LIB_PRINT_BARE_VALUE, 1);
#endif

		for (j = 0; j < oid_count && correct_community; j++) {
			convOidToChar(buf, BUF_SIZE, real_oids[j].first, real_oids[j].second);
			if (dst_address == NULL) {
				DMSG(FLAG_ERROR, "Passed parameter, dst_address, is NULL.");

				result = FAILURE;
			} else { /* Everything is good.. do it. */
				/* Set root from the resolved variable vector */
				root_length = real_oids[j].second;
				memmove(root, real_oids[j].first, sizeof(oid) * root_length);
				/* Set name with root. */
				name_length = root_length;
				memmove(name, root, sizeof(oid) * root_length);

				running = true;
				while (running) {
					int status;
					DMSG(FLAG_INFO, "get_values: sending SNMP packet to host.");
					pdu = snmp_pdu_create(SNMP_MSG_GETNEXT);
					snmp_add_null_var(pdu, name, name_length);
			 
					DMSG(FLAG_INFO, "Sending and then waiting for response to SNMP query.");
		
					status = snmp_sess_synch_response(opaque, pdu, &response);
			 
					DMSG(FLAG_INFO, "Checking SNMP response.");
		
					switch (status) {
						case STAT_SUCCESS:
							DMSG(FLAG_INFO, "SNMP packet status is STAT_SUCCESS.");
							if (response->errstat == SNMP_ERR_NOERROR) {
								for (variables = response->variables; variables; variables = variables->next_variable) {
									if (memcmp(root, variables->name, root_length * sizeof(oid))) {
										running = false;
										continue;
									}
									result = SUCCESS;
#ifdef USE_NEW_CODE
									string str;
									if ((running = convOidValueToString(str, variables)) == false) {
										LMSG("SNMP response variable type is SNMP_ENDOFMIBVIEW, SNMP_NOSUCHOBJECT or SNMP_NOSUCHINSTANCE.");
										result = FAILURE;
									}
									snprintf(buf2, BUF_SIZE, "%s", str.c_str());
#else
									switch (variables->type) {
										case SNMP_ENDOFMIBVIEW:
											LMSG("SNMP response variable type is SNMP_ENDOFMIBVIEW.");
											result = FAILURE;
											running = false;
											break;
			 
										case SNMP_NOSUCHOBJECT:
											LMSG("SNMP response variable type is SNMP_NOSUCHOBJECT.");
											result = FAILURE;
											running = false;
											break;
			 
										case SNMP_NOSUCHINSTANCE:
											LMSG("SNMP response variable type is SNMP_NOSUCHINSTANCE.");
											result = FAILURE;
											running = false;
											break;

										case ASN_INTEGER:
											snprintf(buf2, BUF_SIZE, "%ld", *(long*)variables->val.integer);
											break;

										case ASN_OBJECT_ID:
											convOidToChar(buf2, BUF_SIZE, (oid*) variables->val.string, (int) (variables->val_len / sizeof(oid)));
											break;

										case ASN_GAUGE:
										case ASN_COUNTER:
										case ASN_TIMETICKS:
											snprintf(buf2, BUF_SIZE, "%lu", *(long*)variables->val.integer);
											break;

										case ASN_IPADDRESS:
											convVariableToIP(buf2, BUF_SIZE, variables);
											break;

										default:
											LMSG("DEBUG: SNMP::sendRequest() hit default for variables->type == %d", variables->type);
										case ASN_NULL:
										case ASN_BOOLEAN:
										case ASN_COUNTER64:
										case ASN_OCTET_STR:
											/* If val_len == 0, sprint_value() fills in the next oid. No idea why. */
											if (variables->val_len == 0) {
												buf2[0] = '\0';
											} else {
												snprint_value(buf2, BUF_SIZE, variables->name, variables->name_length, variables);
#if 1
												/* FIXME: Hack to remove quotes around output. There is some way to disable this in NET-SNMP. */
												int buf2_len = strlen(buf2) - 1;
												if (buf2[0] == '"' && buf2[buf2_len] == '"') {
													for (int x = 0; x < buf2_len; x++) {
														buf2[x] = buf2[x + 1];
													}
													buf2[buf2_len - 1] = '\0';
												}
#endif
											}
											break;
									}
#endif
									if (result == SUCCESS) {
										convOidToChar(buf, BUF_SIZE, variables->name, variables->name_length);
										string_pair.first = buf;
										string_pair.second = buf2;
										ret_val.push(string_pair);
									}
									/* Move us to the next OID */
									memmove((char *)name, (char *)variables->name, variables->name_length * sizeof(oid));
									name_length = variables->name_length;
								}
							} else {
								LMSG("Error in SNMP packet (%d).", snmp_errstring(response->errstat));
								DMSG(FLAG_ERROR, "Error in SNMP packet (%d).", snmp_errstring(response->errstat));
								running = false;
								result = FAILURE;
							}
							break;
			 
						case STAT_TIMEOUT:
							DMSG(FLAG_INFO, "SNMP packet status is STAT_TIMEOUT.");
							running = false;
							correct_community = false;
							result = FAILURE;
							break;
			 
						case STAT_ERROR:
							LMSG("SNMP packet status is STAT_ERROR.");
							DMSG(FLAG_ERROR, "SNMP packet status is STAT_ERROR.");
							running = false;
							result = FAILURE;
							break;
			 
						default:
							LMSG("SNMP packet status is unknown.");
							DMSG(FLAG_ERROR, "SNMP packet status is unknown.");
							running = false;
							result = FAILURE;
					}
		
					/* Release response before we allocate another one. */
					if (response) {
						snmp_free_pdu(response);
						response = NULL;
					}
				}
			}
		}
		DMSG(FLAG_INFO, "Closing SNMP connection to remote host.");
		snmp_sess_close(opaque);
	}

	/* Cleanup */
	SAFE_DELETE_ARRAY(buf);
	SAFE_DELETE_ARRAY(buf2);

	return result;
}

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

	pair<string, query_result_t *> result_pair;
	string ip_address;
	SNMP* instance = (SNMP*) args;

	if (instance->trace.thread) {
		LMSG("SNMP::query_thread() starting thread");
	}

	instance->setThreadStatus(thread_id, THREAD_RUNNING);

	while (instance->getThreadStatus(thread_id, status) == SUCCESS && status == THREAD_RUNNING) {
		pthread_mutex_lock(&(instance->to_do_queue_mutex));
		if (!instance->to_do_queue.empty()) {
			ip_address = instance->to_do_queue.top();
			instance->to_do_queue.pop();
			pthread_mutex_unlock(&(instance->to_do_queue_mutex));
		} else {
			pthread_mutex_unlock(&(instance->to_do_queue_mutex));
			instance->setThreadStatus(thread_id, THREAD_COMPLETED);
			break;
		}

		result_pair.second = new query_result_t;
		if (!result_pair.second) {
			LMSG("Failed to allocate result_pair.second");
		} else {
			if (instance->trace.debug) {
				LMSG("Polling [SNMP]: %s", ip_address.c_str());
			}
			instance->sendRequest(thread_id, ip_address.c_str(), *(result_pair.second)); 

			/* Insert data from SNMP query into result_queue, if needed. */
			if (!result_pair.second->empty()) {
				result_pair.first = ip_address;
				pthread_mutex_lock(&(instance->result_queue_mutex));
				instance->result_queue.push(result_pair);
				pthread_mutex_unlock(&(instance->result_queue_mutex));
			} else {
				SAFE_DELETE(result_pair.second);
			}
		}
	}

	if (instance->trace.thread) {
		LMSG("SNMP::query_thread() ending thread");
	}


	THREAD_RETURN(result);
}

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