bitelab/bitelab/snmp.py

#
# Copyright (c) 2020 The FreeBSD Foundation
#
# This software1 was developed by John-Mark Gurney under sponsorship
# from the FreeBSD Foundation.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
# 1.  Redistributions of source code must retain the above copyright
#     notice, this list of conditions and the following disclaimer.
# 2.  Redistributions in binary form must reproduce the above copyright
#     notice, this list of conditions and the following disclaimer in the
#     documentation and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
# OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
# OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
# SUCH DAMAGE.
#

from unittest.mock import patch

from .abstract import *
from .mocks import *

import asyncio
import subprocess
import unittest

__all__ = [ 'snmpget', 'snmpset', 'SNMPPower' ]

def _tosnmp(typ, value):
	if typ == 'bool':
		if value:
			outv = 'true'
		else:
			outv = 'false'

		return ('i', outv)

	raise RuntimeError('unknown type: %s' % repr(typ))

def _fromsnmp(typ, value):
	if typ == 'bool':
		if value == b'true':
			return True
		elif value == b'false':
			return False

		raise RuntimeError('unknown results for bool: %s' % repr(value))

	raise RuntimeError('unknown type: %s' % repr(typ))

async def snmpget(host, oid, typ):
	p = await asyncio.create_subprocess_exec('snmpget', '-Oqv', host, oid,
	    stdout=subprocess.PIPE)

	res = (await p.communicate())[0].strip()

	return _fromsnmp(typ, res)

async def snmpset(host, oid, typ, value):
	p = await asyncio.create_subprocess_exec('snmpset', '-Oqv', host, oid,
	    *_tosnmp(typ, value), stdout=subprocess.PIPE)

	res = (await p.communicate())[0].strip()

	return _fromsnmp(typ, res)

class SNMPPower(Power):
	def __init__(self, host, port):
		self.host = host
		self.port = port

	# Future - add caching + invalidation on set
	async def getvalue(self):
		return await snmpget(self.host,
		    'pethPsePortAdminEnable.1.%d' % self.port, 'bool')

	async def setvalue(self, v):
		return await snmpset(self.host,
		    'pethPsePortAdminEnable.1.%d' % self.port, 'bool', v)

	async def deactivate(self, brd):
		return await self.setvalue(False)

class TestSNMPWrapper(unittest.IsolatedAsyncioTestCase):
	@patch('asyncio.create_subprocess_exec')
	async def test_snmpset(self, cse):
		# that when snmpset returns false
		wrap_subprocess_exec(cse, b'false\n')

		# when being set to false
		r = await snmpset('somehost', 'snmpoid', 'bool', False)

		# that it returns false
		self.assertEqual(r, False)

		# and is called with the correct parameters
		cse.assert_called_with('snmpset', '-Oqv', 'somehost',
		    'snmpoid', 'i', 'false', stdout=subprocess.PIPE)

		# that when snmpset returns true
		wrap_subprocess_exec(cse, b'true\n')

		# when being set to true
		r = await snmpset('somehost', 'snmpoid', 'bool', True)

		# that it returns true
		self.assertEqual(r, True)

		# and is called with the correct parameters
		cse.assert_called_with('snmpset', '-Oqv', 'somehost',
		    'snmpoid', 'i', 'true', stdout=subprocess.PIPE)

	async def test_snmpset_wrongtype(self):
		with self.assertRaises(RuntimeError):
			await snmpset('somehost', 'snmpoid', 'boi', None)

	@patch('asyncio.create_subprocess_exec')
	async def test_snmpget(self, cse):
		wrap_subprocess_exec(cse, b'false\n')

		r = await snmpget('somehost', 'snmpoid', 'bool')

		self.assertEqual(r, False)

		cse.assert_called_with('snmpget', '-Oqv', 'somehost',
		    'snmpoid', stdout=subprocess.PIPE)

		wrap_subprocess_exec(cse, b'true\n')
		r = await snmpget('somehost', 'snmpoid', 'bool')

		self.assertEqual(r, True)

		# that a bogus return value
		wrap_subprocess_exec(cse, b'bogus\n')

		# raises an error
		with self.assertRaises(RuntimeError):
			await snmpget('somehost', 'snmpoid', 'bool')

		# that an unknown type, raises an error
		with self.assertRaises(RuntimeError):
			await snmpget('somehost', 'snmpoid', 'randomtype')

class TestSNMPPower(unittest.IsolatedAsyncioTestCase):
	@patch('bitelab.snmp.snmpset')
	@patch('bitelab.snmp.snmpget')
	async def test_snmppower(self, sg, ss):
		sp = SNMPPower('host', 5)

		# that when snmpget returns False
		sg.return_value = False

		self.assertFalse(await sp.getvalue())

		# calls snmpget w/ the correct args
		sg.assert_called_with('host', 'pethPsePortAdminEnable.1.5',
		    'bool')

		# that when setvalue is called
		await sp.setvalue(True)

		# calls snmpset w/ the correct args
		ss.assert_called_with('host', 'pethPsePortAdminEnable.1.5',
		    'bool', True)

		ss.reset_mock()

		# that when deactivate is called
		await sp.deactivate(None)

		# calls snmpset w/ the correct args
		ss.assert_called_with('host', 'pethPsePortAdminEnable.1.5',
		    'bool', False)