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

static const char rcsid[] = "$Id: snmp.cxx,v 1.23 2006/01/25 07:18:46 sella Exp $";

#include "snmp.h"
#include "daemon.h"
#include "common/netblock.h"
#include "common/ip_util.h"
#include <errno.h>
#include <pthread.h>

#define BUF_SIZE 1024

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

SNMP::SNMP(): maxThreadCount(0), timeout(0), retry(0), commit_count(0) {
	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;

   ConfigNode node;

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

		return FAILURE;
	}
	
	node.assignValue("threads", maxThreadCount, DISCOVERY_SNMP_THREADCOUNT_DEF, DISCOVERY_SNMP_THREADCOUNT_MIN, DISCOVERY_SNMP_THREADCOUNT_MAX);

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

	/* Used to limit the size of INSERTs, both for delay and size of SQL query. 300 gives ~16500 characters. */
	node.assignValue("commit-count", commit_count, DISCOVERY_SNMP_COMMITCOUNT_DEF, DISCOVERY_SNMP_COMMITCOUNT_MIN, DISCOVERY_SNMP_COMMITCOUNT_MAX);

	node.assignValue("timeout", timeout, DISCOVERY_SNMP_TIMEOUT_DEF, DISCOVERY_SNMP_TIMEOUT_MIN, DISCOVERY_SNMP_TIMEOUT_MAX);

	node.assignValue("retry", retry, DISCOVERY_SNMP_RETRY_DEF, DISCOVERY_SNMP_RETRY_MIN, DISCOVERY_SNMP_RETRY_MAX);

	if (!node.assignValues("community", communities)) {
		communities.push_back(DISCOVERY_SNMP_COMMUNITY_DEF);
	}

	node.assignValues("oid", v);
	LMSG("DEBUG: We should insert a default set of oids here.");
	
	/* Resolve all our str oids into real oids. */
	pair<oid*, size_t> pair_ptr;
	oid root[MAX_OID_LEN];
	size_t root_length;

	for (vector<string>::iterator it = v.begin(); it != v.end(); ++it) {
		root_length = MAX_OID_LEN;
		TRACE(TRACE_INTERNAL, "reading object id: %s.", it->c_str());
		if (read_objid(it->c_str(), root, &root_length) == 0) {
			TRACE(TRACE_INTERNAL, "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);
	}

	node.assignValues("prefix", netblocks);

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

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

	printConfig();

	return result;
}

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

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

	TRACE(TRACE_CONFIG, "threads: %lu", this->maxThreadCount);
	TRACE(TRACE_CONFIG, "commit-count: %lu", this->commit_count);
	TRACE(TRACE_CONFIG, "timeout: %lu msec (%0.1f sec)", this->timeout, this->timeout / 1000000.0);
	TRACE(TRACE_CONFIG, "retry: %lu", this->retry);

	if (config->assignValues("discovery snmp 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 snmp 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 SNMP::execute() {
	RESULT result = SUCCESS;
	pthread_t thread_id;
	StorageData data;

	TRACE(TRACE_THREAD, "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") == FAILURE) {
		LMSG("Error: Failed to delete old info from discovery_snmp.");
	}

	if (pthread_create(&thread_id, NULL, (void* (*)(void*))SNMP::result_thread, this) != 0) {
		LMSG("Failed to create result_thread. Shutting down");
		interrupt = true;
		result = FAILURE;
	}

	if (this->queryThreads() == FAILURE) {
		LMSG("queryThreads() failed");

		result = FAILURE;
	}

	pthread_join(thread_id, NULL);

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

	TRACE(TRACE_THREAD, "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 (size_t 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;

	int error = 0;

	for (unsigned i = 0; i < maxThreadCount; i++) {
		if ((error = pthread_create(&thread_id, NULL, (void* (*)(void*))SNMP::query_thread, this)) == 0) {
			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 0
	if ((error = pthread_create(&thread_id, NULL, (void* (*)(void*))SNMP::result_thread, this)) == 0) {
		children.push_back(thread_id);
	}
#endif
	if (children.size() < 1) {
		result = FAILURE;
	}
	return result;
}

RESULT SNMP::result_thread(void* args) {
	SNMP* instance = (SNMP*) args;

	while (!interrupt && (!instance->children.empty() || !instance->result_queue.empty())) {
		if (instance->result_queue.empty()) {
			usleep(USLEEP_NORMLOOP_TIME);
		} else {
			if (instance->writeResults() == FAILURE) {
				LMSG("Dependant function, SNMP::writeResults(), failed.");
			}
		}
	}
	pthread_exit(0);
}

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

	it = children.begin();
	while (it != children.end()) {
		pthread_join(*it, NULL);
		children.erase(it);

		it = children.begin();
	}

	return result;
}

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 (ip, variable, 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, get_count = 0;
	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);
			LMSG("snmp_sess_open(): %s", errmsg);
			free(errmsg);
			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 && !interrupt; j++) {
			convOidToChar(buf, BUF_SIZE, real_oids[j].first, real_oids[j].second);
			if (dst_address == NULL) {
				LMSG("passed parameter, dst_address, is NULL.");

				result = FAILURE;
			} else { /* Everything is good.. do it. */
				/* Set root from the resolved oid 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 && !interrupt) {
					int status;
					pdu = snmp_pdu_create(SNMP_MSG_GETNEXT);
					snmp_add_null_var(pdu, name, name_length);
			 
					TRACE(TRACE_INTERNAL, "SNMPGET query to %s", dst_address);
		
					status = snmp_sess_synch_response(opaque, pdu, &response);
		
					switch (status) {
						case STAT_SUCCESS:
							TRACE(TRACE_INTERNAL, "SNMPGET status for %s is STAT_SUCCESS", dst_address);
							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;
									string str;
									if ((running = convOidValueToString(str, variables)) == false) {
										LMSG("SNMP response variable type is SNMP_ENDOFMIBVIEW, SNMP_NOSUCHOBJECT or SNMP_NOSUCHINSTANCE for host %s", dst_address);
										result = FAILURE;
									}
									snprintf(buf2, BUF_SIZE, "%s", str.c_str());
									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);

										TRACE(TRACE_INTERNAL, "SNMPGET response from %s: %s = %s", dst_address, buf, buf2);
									}
									/* Move us to the next OID */
									memmove((char *)name, (char *)variables->name, variables->name_length * sizeof(oid));
									name_length = variables->name_length;
									get_count++;
								}
							} else {
								LMSG("Error in SNMP packet (%d) for host %s", snmp_errstring(response->errstat), dst_address);
								running = false;
								result = FAILURE;
							}
							break;
			 
						case STAT_TIMEOUT:
							TRACE(TRACE_DEBUG, "SNMPGET for %s timed out", dst_address);
							running = false;
							correct_community = false;
							result = FAILURE;
							break;
			 
						case STAT_ERROR:
							TRACE(TRACE_DEBUG, "SNMPGET for %s returned in error", dst_address);
							running = false;
							result = FAILURE;
							break;
			 
						default:
							TRACE(TRACE_DEBUG, "SNMPGET for %s returned an unknown status", dst_address);
							running = false;
							result = FAILURE;
					}
		
					/* Release response before we allocate another one. */
					if (response) {
						snmp_free_pdu(response);
						response = NULL;
					}
				}
			}
		}
		TRACE(TRACE_DEBUG, "completed %d SNMPGETs to %s", get_count, dst_address);
		TRACE(TRACE_INTERNAL, "closing SNMP connection to host %s", dst_address);
		snmp_sess_close(opaque);
	}

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

	return result;
}

RESULT SNMP::query_thread(void* args) {
	pthread_t thread_id = pthread_self();

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


	TRACE(TRACE_THREAD, "SNMP::query_thread() starting thread");

	while (!interrupt) {
		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));
			break;
		}

		result_pair.second = new query_result_t;
		if (!result_pair.second) {
			LMSG("Failed to allocate result_pair.second");
		} else {
			TRACE(TRACE_DEBUG, "polling: %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);
			}
		}
	}

	TRACE(TRACE_THREAD, "SNMP::query_thread() ending thread");

	pthread_exit(0);
}

RESULT SNMP::readKnownAddress() {
	RESULT result = FAILURE;
	vector<string>::iterator v_it;

	variable_binding_t	policy_input;
	variable_binding_t	policy_output;

	StorageData data;

	if (storage.query(data, "SELECT ip from discovery_icmp") == SUCCESS) {
		while (data.nextRow()) {
			data.getString(policy_input["ip"], "ip");
			policy_input["ip-except"] = policy_input["ip"];

			if (this->inputPolicy.evaluate(policy_output, policy_input) != ACTION_REJECT) {
				netblocks.insert(netblocks.begin(), policy_input["ip"]);
			}
		}
		result = SUCCESS;
	} else {
		LMSG("failed to read IPs from discovery_icmp.");
	}

	/* Sort and Uniq the netblocks */
	sort(netblocks.begin(), netblocks.end());
	v_it = unique(netblocks.begin(), netblocks.end());
	netblocks.erase(v_it, netblocks.end());

	return result;
}

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