#!/usr/bin/env python

"""
This is a rewrite of the original bash script called check_bgp_neighbors and
provides a better checking mechanism with the ability to check the number of
BGP prefixes in memory per router. The command line arguments passed to this
script follow the standards set by other nagios plugins, which are not the same
as many unix power tools
"""

################################################################################
#
# Writen By: Scott McCarty
# Date: 3/2010
# Email: scott.mccarty@gmail.com
#
# Copyright (C) 2010 Scott McCarty
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
#
#################################################################################

# Standard imports
from optparse import OptionParser
from types import *
import sys
import os
import signal
import re
import syslog
import commands
import logging

# Process Signals
## Ignore problems when piping to head
signal.signal(signal.SIGPIPE, signal.SIG_DFL)

## Exit when control-C is pressed
def signal_handler(signal, frame):
	sys.exit(0)
signal.signal(signal.SIGINT, signal_handler)

class BGP:

	filename="check_bgp_neighbors"

	# Host Values
	host=""
	community=""

	# Neighbor Values
	neighbors=[]
	bgp_prefixes_threshold={}
	bgp_status={}
	bgp_prefixes={}
	current_rx={}
	current_tx={}
	current_rx_counter={}
	current_tx_counter={}
	last_rx_counter={}
	last_tx_counter={}

	# Thresholds
	bgp_prefixes_default_threshold=300000
	status_threshold=6
	rx_threshold=5
	tx_threshold=1

	# Others
	error_message_report=[]

	def __init__(self, host, community, nts):

		self.host = host
		self.community = community

		for value in nts:
			if re.match(".*:[0-9]*", value):
				(neighbor, threshold) = value.split(":")
			else:
				neighbor = value
				threshold = self.bgp_prefixes_default_threshold

			self.neighbors.append(neighbor)
			self.bgp_prefixes_threshold[neighbor]=int(threshold)

	def snmpget(self, oid):
		"""Uses snmpget command to get oid"""

		snmpgetcmd = "/usr/bin/snmpget -v1 -On -Oe"
		cmd = snmpgetcmd+" -c"+self.community+" "+self.host+" "+oid
		(status, output) = commands.getstatusoutput(cmd)
		if status == 0:
			return int(output.split()[3])
		else:
			print "SNMP: problem with snmpget subcommand"
			sys.exit(1)

	def get_status(self, neighbor):
		return self.snmpget(".1.3.6.1.2.1.15.3.1.2."+neighbor)

	def get_prefixes(self, neighbor):
		return self.snmpget("1.3.6.1.4.1.9.9.187.1.2.4.1.1."+neighbor+".1.1")

	def get_rx(self, neighbor):
		return self.snmpget(".1.3.6.1.2.1.15.3.1.12."+neighbor)

	def get_tx(self, neighbor):
		return self.snmpget(".1.3.6.1.2.1.15.3.1.13."+neighbor)

	def get_current(self, neighbor):

		# Status & Number of Prefixes don't need history to be determined correctly
		self.bgp_status[neighbor] = self.get_status(neighbor)
		self.bgp_prefixes[neighbor] = self.get_prefixes(neighbor)

		# Messages recieved and sent are counters, so they need last value to determine current difference
		self.current_rx_counter[neighbor]=self.get_rx(neighbor)
		self.current_tx_counter[neighbor]=self.get_tx(neighbor)

	def get_last(self, neighbor):

		if os.path.exists("/tmp/"+self.filename+"."+self.host):
			# File format on each line is: neighbor:172.16.0.2:sent:2470623:received:852786
			f = open("/tmp/"+self.filename+"."+self.host, 'rw')
			for line in f.readlines():
				(k_neighbor, v_neighbor, k_rx, v_rx, k_tx, v_tx) = line.split(":")
				if v_neighbor == neighbor:
					self.last_rx_counter[neighbor]=int(v_rx)
					self.last_tx_counter[neighbor]=int(v_tx)
			f.close();
		else:
			for neighbor in self.neighbors:
				# Messages recieved and sent are counters, so they need last value to determine current difference
				self.last_rx_counter[neighbor]=self.get_rx(neighbor)
				self.last_tx_counter[neighbor]=self.get_tx(neighbor)


	def calculate_values(self, neighbor):

		# Determine values
		self.current_rx[neighbor] = self.current_rx_counter[neighbor] - self.last_rx_counter[neighbor]
		self.current_tx[neighbor] = self.current_tx_counter[neighbor] - self.last_tx_counter[neighbor]

	def save_current(self):

		buffer = []

		# File format on each line is: neighbor:172.16.0.2:sent:2470623:received:852786
		f = open("/tmp/"+self.filename+"."+self.host, 'w')

		for neighbor in self.neighbors:
			buffer.append("neighbor:"+neighbor+":received:"+str(self.current_rx_counter[neighbor])+":sent:"+str(self.current_tx_counter[neighbor])+"\r\n")

		f.writelines(buffer)
		f.close()
	
	def poll(self):
		"""Iterate each neighbor and poll the router for each one"""
		for neighbor in self.neighbors:
			self.get_current(neighbor)
			self.get_last(neighbor)
			self.calculate_values(neighbor)

		self.save_current()

	def append_error_message(self, message):		
		if message not in self.error_message_report:
			self.error_message_report.append(message)

	def failed(self, neighbor):
		error_message={}

		# Calculate correct assertions values (4 checks * number of neighbors)
		assert_failure=4
		error_message["status"]="BGP status did not return correct value from one or more routers"
		error_message["prefix"]="Number of prefixes in memory too low"
		error_message["rx"]="Number of messages sent too low"
		error_message["tx"]="Number of messages received too low"

		# Status
		logging.info("Status Current: "+str(self.bgp_status[neighbor])+" Status Thresh: "+str(self.status_threshold))
		if self.bgp_status[neighbor] == self.status_threshold:
			assert_failure -= 1
		else:
			self.append_error_message(error_message["status"])

		# Prefixes
		logging.info("Prefixes Current: "+str(self.bgp_prefixes[neighbor])+" Prefixes Thresh: "+str(self.bgp_prefixes_threshold[neighbor]))
		if self.bgp_prefixes[neighbor] > self.bgp_prefixes_threshold[neighbor]:
			assert_failure -= 1
		else:
			self.append_error_message(error_message["prefix"])
			
		# RX	
		logging.info("RX Current: "+str(self.current_rx[neighbor])+" RX Thresh: "+str(self.rx_threshold))
		if self.current_rx[neighbor] > self.rx_threshold:
			assert_failure -= 1
		else:
			self.append_error_message(error_message["rx"])

		# TX
		logging.info("TX Current: "+str(self.current_tx[neighbor])+" TX Thresh: "+str(self.tx_threshold))
		if self.current_tx[neighbor] > self.tx_threshold:
			assert_failure -= 1
		else:
			self.append_error_message(error_message["tx"])

		logging.debug("Assert Failure: "+str(assert_failure))

		if assert_failure == 0:
			return False
		else:
			return True

	def report(self):

		report_failure = False

		for neighbor in self.neighbors:
			logging.debug("Neighbor: "+neighbor)
			logging.debug("Status: "+str(self.bgp_status[neighbor]))
			logging.debug("Prefixes: "+str(self.bgp_prefixes[neighbor]))
			logging.debug("Curr RX:"+str(self.current_rx_counter[neighbor]))
			logging.debug("Curr TX:"+str(self.current_tx_counter[neighbor]))
			logging.debug("Last RX:"+str(self.last_rx_counter[neighbor]))
			logging.debug("Last TX:"+str(self.last_tx_counter[neighbor]))
			logging.debug("RX:"+str(self.current_rx[neighbor]))
			logging.debug("TX:"+str(self.current_tx[neighbor]))

		for neighbor in self.neighbors:
			if self.failed(neighbor):
				print "Failed: Neighbor: " + neighbor + \
					" status: "+str(self.bgp_status[neighbor])+ \
					" prefixes: "+str(self.bgp_prefixes[neighbor])+ \
					" sent: "+str(self.current_tx[neighbor])+ \
					" received: "+str(self.current_rx[neighbor])+ \
					" Message: "+str(self.error_message_report)
				report_failure = True
			else:
				print "Success: Neighbor: " + neighbor + \
					" status: "+str(self.bgp_status[neighbor])+ \
					" prefixes: "+str(self.bgp_prefixes[neighbor])+ \
					" sent: "+str(self.current_tx[neighbor])+ \
					" received: "+str(self.current_rx[neighbor])

		if report_failure == True:
			sys.exit(2)
		else:
			sys.exit(0)

def version():
	print "Version: 0.9"
	sys.exit(0)

def init(argv):
	""" Used to capture all of the command line args and perform initializations"""

	# Use global namespace for flags and other major variables
	global options

	# Declarations & Variables
	usage="usage: %prog -C public -H router.example.com -n <neighbor> ..."
	parser = OptionParser(usage)

	# Handle flags
	parser.add_option("-V", "--version", dest="mode", action="store_const", const="version", help="show program's version number and exit")
	parser.add_option("-v", "--verbose", dest="verbose", action="count", help="Show verbose output")
	parser.add_option("-n", "--nts", dest="nts", action="append", help="Add neighbor/threshold pair")
	parser.add_option("-C", "", dest="community", action="store", type="string", help="Community string used for snmp")
	parser.add_option("-H", "", dest="host", action="store", type="string", help="Host to check")
	(options, args) = parser.parse_args()

	# Determine mode
	if options.mode == "version":
		version()

	# Check options
	check_options(parser)
	setup_logging()

def setup_logging():

	global logging

	# Set Verbosity
	log_level = logging.WARNING
	if options.verbose == 1:
		log_level = logging.INFO
		print "Info Logging"
	elif options.verbose >= 2:
		log_level = logging.DEBUG
		print "Debug Logging"

	# Set up basic configuration, out to stderr with a reasonable default format.
	logging.basicConfig(level=log_level)

def check_options(parser):
	if type(options.nts) == NoneType \
	or type(options.community) == NoneType \
	or type(options.host) == NoneType:
		parser.print_usage()
		sys.exit()

if __name__ == "__main__":
        init(sys.argv[1:])
	x = BGP(options.host, options.community, options.nts)
	x.poll()
	x.report()