vlanmang/vlanmang.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

from pysnmp.hlapi import *
from pysnmp.smi.builder import MibBuilder
from pysnmp.smi.view import MibViewController

import itertools
import random
import unittest

_mbuilder = MibBuilder()
_mvc = MibViewController(_mbuilder)

#import data

# received packages
# pvid: dot1qPvid
#
# tx packets:
# dot1qVlanStaticEgressPorts
# dot1qVlanStaticUntaggedPorts
#
# vlans:
# dot1qVlanCurrentTable
#	lists ALL vlans, including baked in ones
#
# note that even though an snmpwalk of dot1qVlanStaticEgressPorts
# skips over other vlans (only shows statics), the other vlans (1,2,3)
# are still accessible via that oid
#
# LLDP:
# 1.0.8802.1.1.2.1.4.1.1 aka LLDP-MIB, lldpRemTable

def getpvidmapping(data, lookupfun):
	'''Return a mapping from vlan based table to a port: vlan
	dictionary.'''

	res = []
	for id in data:
		for i in data[id]['u']:
			if isinstance(i, str):
				i = lookupfun(i)
			res.append((i, id))

	return dict(res)

def getportlist(data, lookupfun):
	'''Return a set of all the ports indexes in data.'''

	res = set()

	for id in data:
		res.update(data[id]['u'])
		res.update(data[id].get('t', []))

	# filter out the strings
	strports = set(x for x in res if isinstance(x, str))

	res.update(lookupfun(x) for x in strports)
	res.difference_update(strports)

	return res

class SNMPSwitch(object):
	'''A class for manipulating switches via standard SNMP MIBs.'''

	def __init__(self, host, community):
		self._eng = SnmpEngine()
		self._cd = CommunityData(community, mpModel=0)
		self._targ = UdpTransportTarget((host, 161))

	def _get(self, oid):
		oid = ObjectIdentity(*oid)
		oid.resolveWithMib(_mvc)

		errorInd, errorStatus, errorIndex, varBinds = \
		    next(getCmd(self._eng, self._cd, self._targ, ContextData(), ObjectType(oid)))

		if errorInd: # pragma: no cover
			raise ValueError(errorIndication)
		elif errorStatus:
			raise ValueError('%s at %s' %
			    (errorStatus.prettyPrint(), errorIndex and
			    varBinds[int(errorIndex)-1][0] or '?'))
		else:
			if len(varBinds) != 1: # pragma: no cover
				raise ValueError('too many return values')

			varBind = varBinds[0]
			return varBind[1]

	def _set(self, oid, value):
		oid = ObjectIdentity(*oid)
		oid.resolveWithMib(_mvc)

		if isinstance(value, (int, long)):
			value = Integer(value)
		elif isinstance(value, str):
			value = OctetString(value)

		errorInd, errorStatus, errorIndex, varBinds = \
		    next(setCmd(self._eng, self._cd, self._targ, ContextData(), ObjectType(oid, value)))

		if errorInd: # pragma: no cover
			raise ValueError(errorIndication)
		elif errorStatus: # pragma: no cover
			raise ValueError('%s at %s' %
			    (errorStatus.prettyPrint(), errorIndex and
			    varBinds[int(errorIndex)-1][0] or '?'))
		else:
			for varBind in varBinds:
				if varBind[1] != value: # pragma: no cover
					raise RuntimeError('failed to set: %s' % ' = '.join([x.prettyPrint() for x in varBind]))

	def _walk(self, *oid):
		oid = ObjectIdentity(*oid)
		# XXX - keep these, this might stop working, no clue what managed to magically make things work
		# ref: http://snmplabs.com/pysnmp/examples/smi/manager/browsing-mib-tree.html#mib-objects-to-pdu-var-binds
		# mibdump.py --mib-source '/Users/jmg/Nextcloud/Documents/user manuals/netgear/gs7xxt-v6.3.1.19-mibs' --mib-source /usr/share/snmp/mibs  --rebuild rfc1212 pbridge vlan
		#oid.addAsn1MibSource('/usr/share/snmp/mibs', '/Users/jmg/Nextcloud/Documents/user manuals/netgear/gs7xxt-v6.3.1.19-mibs')

		oid.resolveWithMib(_mvc)

		for (errorInd, errorStatus, errorIndex, varBinds) in nextCmd(
		    self._eng, self._cd, self._targ, ContextData(),
		    ObjectType(oid),
		    lexicographicMode=False):
			if errorInd: # pragma: no cover
				raise ValueError(errorIndication)
			elif errorStatus: # pragma: no cover
				raise ValueError('%s at %s' % (errorStatus.prettyPrint(), errorIndex and varBinds[int(errorIndex)-1][0] or '?'))
			else:
				for varBind in varBinds:
					yield varBind

	def findport(self, name):
		'''Look up a port name and return it's port index.  This
		looks up via the ifName table in IF-MIB.'''

		return [ x[0][-1] for x in self._walk('IF-MIB', 'ifName') if str(x[1]) == name ][0]

	def getvlanname(self, vlan):
		'''Return the name for the vlan.'''

		v = self._get(('Q-BRIDGE-MIB', 'dot1qVlanStaticName', vlan))

		return str(v).decode('utf-8')

	def createvlan(self, vlan, name):
		# createAndGo(4)
		self._set(('Q-BRIDGE-MIB', 'dot1qVlanStaticRowStatus',
		    int(vlan)), 4)
		self._set(('Q-BRIDGE-MIB', 'dot1qVlanStaticName', int(vlan)),
		    name)

	def deletevlan(self, vlan):
		self._set(('Q-BRIDGE-MIB', 'dot1qVlanStaticRowStatus',
		    int(vlan)), 6) # destroy(6)

	def getvlans(self):
		'''Return all the vlans.'''

		return (x[0][-1] for x in self._walk('Q-BRIDGE-MIB', 'dot1qVlanStatus'))

	def staticvlans(self):
		'''Return the staticly defined/configured vlans.  This
		sometimes excludes special built in vlans, like vlan 1.'''

		return (x[0][-1] for x in self._walk('Q-BRIDGE-MIB', 'dot1qVlanStaticName'))

class _TestMisc(unittest.TestCase):
	def test_pvid(self):
		data = {
			1: {
				'u':	[ 1, 5, 10 ] + range(13, 20)
			},
			10: {
				'u':	[ 2, 3, 6, 7, 8, 'lag2' ],
			},
			13: {
				'u':	[ 4, 9 ],
			},
		}
		lookup = {
			'lag2': 30
		}

		check = dict(itertools.chain(enumerate([ 1, 10, 10, 13, 1, 10,
		    10, 10, 13, 1 ], 1), enumerate([ 1 ] * 7, 13),
		    [ (30, 10) ]))

		# That a pvid mapping
		res = getpvidmapping(data, lookup.__getitem__)

		# is correct
		self.assertEqual(res, check)

		self.assertEqual(getportlist(data, lookup.__getitem__),
		    set(xrange(1, 11)) | set(xrange(13, 20)) | set([30]))

_skipSwitchTests = True

class _TestSwitch(unittest.TestCase):
	def setUp(self):
		args = open('test.creds').read().split()
		self.switch = SNMPSwitch(*args)

	def test_unpacktable(self):
		pass

	@unittest.skipIf(_skipSwitchTests, 'slow')
	def test_misc(self):
		switch = self.switch

		self.assertEqual(switch.findport('g1'), 1)
		self.assertEqual(switch.findport('l1'), 14)

	@unittest.skipIf(_skipSwitchTests, 'slow')
	def test_vlan(self):
		switch = self.switch

		existingvlans = set(switch.getvlans())

		while True:
			testvlan = random.randint(1,4095)
			if testvlan not in existingvlans:
				break

		# Test that getting a non-existant vlans raises an exception
		self.assertRaises(ValueError, switch.getvlanname, testvlan)

		self.assertTrue(set(switch.staticvlans()).issubset(existingvlans))

		testname = 'Sometestname'

		# Create test vlan
		switch.createvlan(testvlan, testname)
		try:
			# make sure the test vlan was created
			self.assertIn(testvlan, set(switch.staticvlans()))

			self.assertEqual(testname, switch.getvlanname(testvlan))
		finally:
			switch.deletevlan(testvlan)