netsnmptestenv.py

#!/usr/bin/env python
# encoding: utf-8
#
# python-netsnmpagent module
# Copyright (c) 2013 Pieter Hollants <pieter@hollants.com>
# Licensed under the GNU Public License (GPL) version 3
#
# net-snmp test environment module
#

""" Sets up net-snmp test environments.

This module allows to run net-snmp instances with user privileges that do not
interfere with any system-wide running net-snmp instance. """

import sys, os, tempfile, subprocess, re, inspect, signal, time, shutil

class netsnmpTestEnv(object):
	""" Implements a net-snmp test environment. """

	def __init__(self, **args):
		""" Initializes a new net-snmp test environment. """

		self.agentport  = 6555
		self.informport = 6556
		self.smuxport   = 6557

		# Create a temporary directory to hold the snmpd files
		self.tmpdir = tempfile.mkdtemp(os.path.basename(sys.argv[0]))

		# Compose paths to files inside the temp dir
		conffile          = os.path.join(self.tmpdir, "snmpd.conf")
		self.mastersocket = os.path.join(self.tmpdir, "snmpd-agentx.sock")
		self.statedir     = os.path.join(self.tmpdir, "state")
		self.pidfile      = os.path.join(self.tmpdir, "snmpd.pid")
		indexesfile       = os.path.join(self.tmpdir, "mib_indexes")

		# Create a minimal snmpd configuration file
		with open(conffile, "w") as f:
			f.write("[snmpd]\n")
			f.write("rocommunity public 127.0.0.1\n")
			f.write("rwcommunity simple 127.0.0.1\n")
			f.write("agentaddress localhost:{0}\n".format(self.agentport))
			f.write("informsink localhost:{0}\n".format(self.informport))
			f.write("smuxsocket localhost:{0}\n".format(self.smuxport))
			f.write("master agentx\n")
			f.write("agentXSocket {0}\n\n".format(self.mastersocket))
			f.write("[snmp]\n")
			f.write("persistentDir {0}\n".format(self.statedir))

		# Create an empty mib_indexes file
		open(indexesfile, "w").close()

		# Start the snmpd instance
		cmd = "/usr/sbin/snmpd -r -LE warning -C -c{0} -p{1}".format(
			conffile, self.pidfile
		)
		subprocess.check_call(cmd, shell=True)

	def shutdown(self):
		# Explicitly import used Python modules once more because they may
		# have been __del__'d before us and then we'd hit AttributeError
		# exceptions due to "os" being None
		import os, time

		def kill_process(pid):
			def is_running(pid):
				return os.path.exists("/proc/{0}".format(pid))

			if not is_running(pid):
				return

			starttime = time.clock()
			os.kill(pid, signal.SIGTERM)
			while time.clock() == starttime:
				time.sleep(0.25)

			if not is_running(pid):
				return

			os.kill(pid, signal.SIGTERM)
			while is_running(pid):
				time.sleep(0.25)

		# Check for existance of snmpd's PID file
		if os.access(self.pidfile, os.R_OK):
			# Read the PID
			with open(self.pidfile, "r") as f:
				pid = int(f.read())

			# And kill it
			kill_process(pid)

		# Recursively remove the temporary directory
		if os.access(self.tmpdir, os.R_OK):
			shutil.rmtree(self.tmpdir)

	def __del__(self):
		self.shutdown()

	class SNMPTimeoutError(Exception):
		pass

	class UnknownOIDError(Exception):
		pass

	class MIBUnavailableError(Exception):
		pass

	class NotWritableError(Exception):
		pass

	@staticmethod
	def snmpcmd(op, oid, data=None, datatype=None):
		""" Executes a SNMP client operation in the net-snmp test environment.
		    
		    "op" is either "get", "set", "walk" or "table".
		    "oid" is the OID to run the operation against.
		    "data" is the data to set in case of a "set" operation.
			"datatype" is the type of the data (as specified to "snmpset"). """

		# Compose the SNMP client command
		if op == "set":
			cmd = "/usr/bin/snmp{0} -M+. -r0 -v 2c -c simple localhost:6555 {1} {2} {3}"
			cmd = cmd.format(op, oid, datatype, data)
		else:
			cmd = "/usr/bin/snmp{0} -M+. -r0 -v 2c -c public localhost:6555 {1}"
			cmd = cmd.format(op, oid)

		# Python 2.6 (used eg. in SLES11SP2) does not yet know about
		# subprocess.check_output(), so we wrap subprocess.Popen() instead.
		#
		# Execute the command with stderr redirected to stdout and stdout
		# redirected to a pipe that we capture below
		proc = subprocess.Popen(
			cmd, shell=True, env={ "LANG": "C" },
			stdout=subprocess.PIPE, stderr=subprocess.STDOUT
		)
		output = proc.communicate()[0].strip()
		rc = proc.poll()
		if rc == 0:
			if re.search(" = No Such Object available on this agent at this OID", output):
				raise netsnmpTestEnv.MIBUnavailableError(oid)
			return output

		if re.search(": Unknown Object Identifier", output):
				raise netsnmpTestEnv.UnknownOIDError(oid)

		if re.search("Timeout: No Response from ", output):
			raise netsnmpTestEnv.SNMPTimeoutError("localhost:6555")

		if re.search("Reason: notWritable \(That object does not support modification\)", output):
			raise netsnmpTestEnv.NotWritableError(oid)

		# SLES11 SP2's Python 2.6 has a subprocess module whose
		# CalledProcessError exception does not yet know the third "output"
		# argument, so we monkey-patch support into it
		if len(inspect.getargspec(subprocess.CalledProcessError.__init__).args) == 3:
			def new_init(self, returncode, cmd, output=None):
				self.returncode = returncode
				self.cmd        = cmd
				self.output     = output
			subprocess.CalledProcessError.__init__ = new_init

		raise subprocess.CalledProcessError(rc, cmd, output)

	@classmethod
	def snmpget(self, oid):
		""" Executes a "snmpget" operation in the net-snmp test environment.

		    "oid" is the OID to run the operation against.

		    Returns a two-tuple (data, datatype). """

		data = self.snmpcmd("get", oid).split("=")[1]
		if ":" in data:
			(datatype, data) = data.split(":", 1)
			datatype = datatype.strip()
		else:
			datatype = "STRING"
		data = data.strip()
		if data.startswith('"') and data.endswith('"'):
			data = data[1:-1]
		return (data, datatype)

	@classmethod
	def snmpset(self, oid, data, datatype):
		""" Executes a "snmpset" operation in the net-snmp test environment.

		    "oid" is the OID to run the operation against.
		    "data" is the data to set.
		    "datatype" is the type of the data (as specified to "snmpset"). """

		return self.snmpcmd("set", oid, data, datatype)

	@classmethod
	def snmpwalk(self, oid):
		""" Executes a "snmpwalk" operation in the net-snmp test environment.

		    "oid" is the OID to run the operation against. """

		return self.snmpcmd("walk", oid)

	@classmethod
	def snmptable(self, oid):
		""" Executes a "snmpwalk" operation in the net-snmp test environment.

		    "oid" is the OID to run the operation against. """

		return self.snmpcmd("table", oid)