bitelab/bitelab/__init__.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 typing import Optional, Union, Dict, Any
from dataclasses import dataclass
from functools import lru_cache, wraps
from io import StringIO

from fastapi import APIRouter, Depends, FastAPI, HTTPException, Request
from fastapi.security import OAuth2PasswordBearer
from httpx import AsyncClient, Auth
from starlette.responses import JSONResponse
from starlette.status import HTTP_200_OK
from starlette.status import HTTP_400_BAD_REQUEST, HTTP_401_UNAUTHORIZED, \
    HTTP_403_FORBIDDEN, HTTP_404_NOT_FOUND, HTTP_409_CONFLICT
from starlette.status import HTTP_500_INTERNAL_SERVER_ERROR
from unittest.mock import patch, AsyncMock, Mock, PropertyMock

from . import config
from .data import *
from .abstract import *
from .snmp import *
from .mocks import *

import asyncio
import json
import orm
import os
import socket
import sqlite3
import subprocess
import sys
import tempfile
import unittest
import urllib

# fix up parse_socket_addr for hypercorn
from hypercorn.utils import parse_socket_addr
from hypercorn.asyncio import tcp_server
def new_parse_socket_addr(domain, addr):
	if domain == socket.AF_UNIX:
		return (addr, -1)

	return parse_socket_addr(domain, addr)

tcp_server.parse_socket_addr = new_parse_socket_addr

class BoardImpl:
	def __init__(self, name, brdclass, options):
		self.name = name
		self.brdclass = brdclass
		self.options = options
		self.reserved = False
		self.attrmap = {}
		self.lock = asyncio.Lock()
		for i in options:
			self.attrmap[i.defattrname] = i

		self.attrcache = {}

	def __repr__(self): #pragma: no cover
		return repr(Board.from_orm(self))

	async def reserve(self):
		assert self.lock.locked() and not self.reserved

		self.reserved = True

	async def release(self):
		assert self.lock.locked() and self.reserved

		self.reserved = False

	async def update(self):
		for i in self.attrmap:
			self.attrcache[i] = await self.attrmap[i].getvalue()

	def add_info(self, d):
		self.attrcache.update(d)

	def clean_info(self):
		# clean up attributes
		for i in set(self.attrcache) - set(self.attrmap):
			del self.attrcache[i]

	@property
	def attrs(self):
		return dict(self.attrcache)

@dataclass
class BITEError(Exception):
	errobj: Error
	status_code: int

class BoardManager(object):
	board_class_info = {
		'cora-z7s': {
			'clsname': 'cora-z7s',
			'arch': 'arm-armv7',
		},
	}

	# Naming scheme:
	#  <abbreviated class>-<num>
	#
	board_gen = [
		dict(name='cora-1', brdclass='cora-z7s', options=[
			SNMPPower(host='poe', port=2),
		]),
	]

	def __init__(self, settings):
		self._settings = settings
		self.boards = dict(**{ x.name: x for x in
		    (BoardImpl(**y) for y in self.board_gen)})

	def classes(self):
		return self.board_class_info

def unhashable_lru():
	def newwrapper(fun):
		cache = {}

		@wraps(fun)
		def wrapper(*args, **kwargs):
			idargs = tuple(id(x) for x in args)
			idkwargs = tuple(sorted((k, id(v)) for k, v in
			    kwargs.items()))
			k = (idargs, idkwargs)
			if k in cache:
				realargs, realkwargs, res = cache[k]
				if all(x is y for x, y in zip(args,
				    realargs)) and all(realkwargs[x] is
				    kwargs[x] for x in realkwargs):
					return res

			res = fun(*args, **kwargs)
			cache[k] = (args, kwargs, res)

			return res

		return wrapper

	return newwrapper

class BiteAuth(Auth):
	def __init__(self, token):
		self.token = token

	def __eq__(self, o):
		return self.token == o.token

	def auth_flow(self, request):
		request.headers['Authorization'] = 'Bearer ' + self.token
		yield request

@lru_cache()
def sync_get_board_lock():
	return asyncio.Lock()

async def get_board_lock():
	return sync_get_board_lock()

# how to get coverage for this?
@lru_cache()
def get_settings(): # pragma: no cover
	return config.Settings()

# how to get coverage for this?
@unhashable_lru()
def get_data(settings: config.Settings = Depends(get_settings)):
	#print(repr(settings))
	database = data.databases.Database('sqlite:///' + settings.db_file)
	d = make_orm(database)
	return d

@unhashable_lru()
def sync_get_boardmanager(settings):
	return BoardManager(settings)

async def get_boardmanager(settings: config.Settings = Depends(get_settings)):
	return sync_get_boardmanager(settings)

oauth2_scheme = OAuth2PasswordBearer(tokenUrl='/nonexistent')

async def lookup_user(token: str = Depends(oauth2_scheme),
    data: data.DataWrapper = Depends(get_data)):
	try:
		return (await data.APIKey.objects.get(key=token)).user
	except orm.exceptions.NoMatch:
		raise HTTPException(
		    status_code=HTTP_401_UNAUTHORIZED,
		    detail='Invalid authentication credentials',
		    headers={'WWW-Authenticate': 'Bearer'},
		)

router = APIRouter()

def board_priority(request: Request):
	# Get the board, if any, from the connection
	scope = request.scope
	return scope['server']

@router.get('/board/classes', response_model=Dict[str, BoardClassInfo])
async def get_board_classes(user: str = Depends(lookup_user),
    brdmgr: BoardManager = Depends(get_boardmanager)):
	return brdmgr.classes()

@router.get('/board/{board_id}', response_model=Board)
async def get_board_info(board_id, user: str = Depends(lookup_user),
    brdmgr: BoardManager = Depends(get_boardmanager)):
	brd = brdmgr.boards[board_id]
	await brd.update()

	return brd

@router.post('/board/{board_id_or_class}/reserve', response_model=Union[Board, Error])
async def reserve_board(board_id_or_class, user: str = Depends(lookup_user),
    brdmgr: BoardManager = Depends(get_boardmanager),
    brdlck: asyncio.Lock = Depends(get_board_lock),
    settings: config.Settings = Depends(get_settings),
    data: data.DataWrapper = Depends(get_data)):
	board_id = board_id_or_class
	brd = brdmgr.boards[board_id]

	async with brd.lock:
		try:
			obrdreq = await data.BoardStatus.objects.create(board=board_id,
			    user=user)
			# XXX - There is a bug in orm where the returned
			# object has an incorrect board value
			# see: https://github.com/encode/orm/issues/47
			#assert obrdreq.board == board_id and \
			#    obrdreq.user == user
			brdreq = await data.BoardStatus.objects.get(board=board_id,
			    user=user)
			await brd.reserve()
		# XXX - orm isn't doing it's job here
		except sqlite3.IntegrityError:
			raise BITEError(
			    status_code=HTTP_409_CONFLICT,
			    errobj=Error(error='Board currently reserved.',
			    board=Board.from_orm(brd)),
			)

		# Initialize board
		try:
			sub = await asyncio.create_subprocess_exec(
			    settings.setup_script, 'reserve', brd.name, user,
			    stdout=subprocess.PIPE, stderr=subprocess.PIPE)
			stdout, stderr = await sub.communicate()
			if sub.returncode:
				raise RuntimeError(sub.returncode, stderr)

			brd.add_info(json.loads(stdout))
		except Exception as e:
			await brdreq.delete()
			await brd.release()
			if isinstance(e, RuntimeError):
				retcode, stderr = e.args
				raise BITEError(
				    status_code=HTTP_500_INTERNAL_SERVER_ERROR,
				    errobj=Error(error=
					'Failed to init board, ret: %d, stderr: %s' %
					(retcode, repr(stderr)),
				    board=Board.from_orm(brd)),
				)
			raise

	await brd.update()

	return brd

@router.post('/board/{board_id}/release', response_model=Union[Board, Error])
async def release_board(board_id, user: str = Depends(lookup_user),
    brdmgr: BoardManager = Depends(get_boardmanager),
    brdlck: asyncio.Lock = Depends(get_board_lock),
    settings: config.Settings = Depends(get_settings),
    data: data.DataWrapper = Depends(get_data)):
	brd = brdmgr.boards[board_id]

	async with brd.lock:
		try:
			brduser = await data.BoardStatus.objects.get(board=board_id)
			if user != brduser.user:
				raise BITEError(
				    status_code=HTTP_403_FORBIDDEN,
				    errobj=Error(error='Board reserved by %s.' % repr(brduser.user),
				    board=Board.from_orm(brd)))

		except orm.exceptions.NoMatch:
			raise BITEError(
			    status_code=HTTP_400_BAD_REQUEST,
			    errobj=Error(error='Board not reserved.',
			    board=Board.from_orm(brd)),
			)

		sub = await asyncio.create_subprocess_exec(
		    settings.setup_script, 'release', brd.name, user,
		    stdout=subprocess.PIPE, stderr=subprocess.PIPE)
		stdout, stderr = await sub.communicate()
		if sub.returncode:
			raise RuntimeError(sub.returncode, stderr)

		await data.BoardStatus.delete(brduser)
		await brd.release()

		brd.clean_info()

	await brd.update()

	return brd

@router.post('/board/{board_id_or_class}/attrs', response_model=Union[Board, Error])
async def set_board_attrs():
	pass

@router.get('/board/',response_model=Dict[str, Board])
async def get_boards(user: str = Depends(lookup_user),
    brdmgr: BoardManager = Depends(get_boardmanager)):
	brds = brdmgr.boards
	for i in brds:
		await brds[i].update()

	return brds

@router.get('/')
async def root_test(board_prio: dict = Depends(board_priority),
    settings: config.Settings = Depends(get_settings)):
	return { 'foo': 'bar', 'board': board_prio }

def getApp():
	app = FastAPI()
	app.include_router(router)

	@app.exception_handler(BITEError)
	async def error_handler(request, exc):
		return JSONResponse(exc.errobj.dict(), status_code=exc.status_code)

	return app

# uvicorn can't call the above function, while hypercorn can
#app = getApp()

class TestUnhashLRU(unittest.TestCase):
	def test_unhashlru(self):
		lsta = []
		lstb = []

		# that a wrapped function
		cachefun = unhashable_lru()(lambda x: object())

		# handles unhashable objects
		resa = cachefun(lsta)
		resb = cachefun(lstb)

		# that they return the same object again
		self.assertIs(resa, cachefun(lsta))
		self.assertIs(resb, cachefun(lstb))

		# that the object returned is not the same
		self.assertIsNot(cachefun(lsta), cachefun(lstb))

		# that a second wrapped funcion
		cachefun2 = unhashable_lru()(lambda x: object())

		# does not return the same object as the first cache
		self.assertIsNot(cachefun(lsta), cachefun2(lsta))

async def _setup_data(data):
	await data.APIKey.objects.create(user='foo', key='thisisanapikey')
	await data.APIKey.objects.create(user='bar', key='anotherlongapikey')

# Per RFC 5737 (https://tools.ietf.org/html/rfc5737):
#   The blocks 192.0.2.0/24 (TEST-NET-1), 198.51.100.0/24 (TEST-NET-2),
#   and 203.0.113.0/24 (TEST-NET-3) are provided for use in
#   documentation.

# Note: this will not work under python before 3.8 before
# IsolatedAsyncioTestCase was added.  The tearDown has to happen
# with the event loop running, otherwise the task and other things
# do not get cleaned up properly.
class TestBiteLab(unittest.IsolatedAsyncioTestCase):
	def get_settings_override(self):
		return self.settings

	def get_data_override(self):
		return self.data

	async def asyncSetUp(self):
		self.app = getApp()

		# setup test database
		self.dbtempfile = tempfile.NamedTemporaryFile()
		self.database = data.databases.Database('sqlite:///' +
		    self.dbtempfile.name)
		self.data = make_orm(self.database)

		await _setup_data(self.data)

		# setup settings
		self.settings = config.Settings(db_file=self.dbtempfile.name,
		    setup_script='somesetupscript',
		)

		self.app.dependency_overrides[get_settings] = \
		    self.get_settings_override
		self.app.dependency_overrides[get_data] = self.get_data_override

		self.client = AsyncClient(app=self.app,
		    base_url='http://testserver')

	def tearDown(self):
		self.app = None
		asyncio.run(self.client.aclose())
		self.client = None

	async def test_basic(self):
		res = await self.client.get('/')
		self.assertNotEqual(res.status_code, HTTP_404_NOT_FOUND)

	async def test_notauth(self):
		# test that simple accesses are denied
		res = await self.client.get('/board/classes')
		self.assertEqual(res.status_code, HTTP_401_UNAUTHORIZED)

		res = await self.client.get('/board/')
		self.assertEqual(res.status_code, HTTP_401_UNAUTHORIZED)

		# test that invalid api keys are denied
		res = await self.client.get('/board/classes',
		    auth=BiteAuth('badapikey'))
		self.assertEqual(res.status_code, HTTP_401_UNAUTHORIZED)

	async def test_classes(self):
		# that when requesting the board classes
		res = await self.client.get('/board/classes',
		    auth=BiteAuth('thisisanapikey'))

		# it is successful
		self.assertEqual(res.status_code, HTTP_200_OK)

		# and returns the correct data
		self.assertEqual(res.json(),  { 'cora-z7s': BoardClassInfo(**{
		    'arch': 'arm-armv7', 'clsname': 'cora-z7s', }) })

	@patch('asyncio.create_subprocess_exec')
	@patch('bitelab.snmp.snmpget')
	async def test_board_reserve_release(self, sg, cse):
		# that when releasing a board that is not yet reserved
		res = await self.client.post('/board/cora-1/release',
		    auth=BiteAuth('anotherlongapikey'))

		# that it returns an error
		self.assertEqual(res.status_code, HTTP_400_BAD_REQUEST)

		# that when snmpget returns False
		sg.return_value = False

		# that when the setup script will fail
		wrap_subprocess_exec(cse, stderr=b'error', retcode=1)

		# that reserving the board
		res = await self.client.post('/board/cora-1/reserve',
		    auth=BiteAuth('thisisanapikey'))

		# that it is a failure
		self.assertEqual(res.status_code, HTTP_500_INTERNAL_SERVER_ERROR)

		# and returns the correct data
		info = Error(error='Failed to init board, ret: 1, stderr: b\'error\'',
			board=Board(name='cora-1',
				brdclass='cora-z7s',
				reserved=False,
			),
		).dict()
		self.assertEqual(res.json(), info)

		# and that it called the start script
		cse.assert_called_with(self.settings.setup_script, 'reserve',
		    'cora-1', 'foo', stdout=subprocess.PIPE,
		    stderr=subprocess.PIPE)

		# that when the setup script returns
		wrap_subprocess_exec(cse,
		    json.dumps(dict(ip='192.0.2.10')).encode('utf-8'))

		# that reserving the board
		res = await self.client.post('/board/cora-1/reserve',
		    auth=BiteAuth('thisisanapikey'))

		# that it is successful
		self.assertEqual(res.status_code, HTTP_200_OK)

		# and returns the correct data
		brdinfo = Board(name='cora-1',
			brdclass='cora-z7s',
			reserved=True,
			attrs=dict(power=False,
				ip='192.0.2.10',
			),
		).dict()
		self.assertEqual(res.json(), brdinfo)

		# and that it called the start script
		cse.assert_called_with(self.settings.setup_script, 'reserve',
		    'cora-1', 'foo', stdout=subprocess.PIPE, stderr=subprocess.PIPE)

		# that another user reserving the board
		res = await self.client.post('/board/cora-1/reserve',
		    auth=BiteAuth('anotherlongapikey'))

		# that the request is fails with a conflict
		self.assertEqual(res.status_code, HTTP_409_CONFLICT)

		# and returns the correct data
		info = {
			'error': 'Board currently reserved.',
			'board': brdinfo,
		}
		self.assertEqual(res.json(), info)

		# that another user releases the board
		res = await self.client.post('/board/cora-1/release',
		    auth=BiteAuth('anotherlongapikey'))

		# that it is denied
		self.assertEqual(res.status_code, HTTP_403_FORBIDDEN)

		# and returns the correct data
		info = {
			'error': 'Board reserved by \'foo\'.',
			'board': brdinfo,
		}
		self.assertEqual(res.json(), info)

		# that when the correct user releases the board
		res = await self.client.post('/board/cora-1/release',
		    auth=BiteAuth('thisisanapikey'))

		# it is allowed
		self.assertEqual(res.status_code, HTTP_200_OK)

		# and returns the correct data
		info = {
			'name': 'cora-1',
			'brdclass': 'cora-z7s',
			'reserved': False,
			'attrs': { 'power': False },
		}
		self.assertEqual(res.json(), info)

		# and that it called the release script
		cse.assert_called_with(self.settings.setup_script, 'release',
		    'cora-1', 'foo', stdout=subprocess.PIPE, stderr=subprocess.PIPE)

		# that it can be reserved by a different user
		res = await self.client.post('/board/cora-1/reserve',
		    auth=BiteAuth('anotherlongapikey'))

		# that it is successful
		self.assertEqual(res.status_code, HTTP_200_OK)

	@patch('bitelab.snmp.snmpget')
	async def test_board_info(self, sg):
		# that when snmpget returns False
		sg.return_value = False

		# that getting the board info
		res = await self.client.get('/board/',
		    auth=BiteAuth('thisisanapikey'))

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

		# that it is successful
		self.assertEqual(res.status_code, HTTP_200_OK)

		# and returns the correct data
		info = {
			'cora-1': {
				'name': 'cora-1',
				'brdclass': 'cora-z7s',
				'reserved': False,
				'attrs': { 'power': False },
			},
		}
		self.assertEqual(res.json(), info)

		# that when snmpget returns True
		sg.return_value = True

		# that getting the board info
		res = await self.client.get('/board/cora-1',
		    auth=BiteAuth('thisisanapikey'))

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

		# that it is successful
		self.assertEqual(res.status_code, HTTP_200_OK)

		# and returns the correct data
		info = {
			'name': 'cora-1',
			'brdclass': 'cora-z7s',
			'reserved': False,
			'attrs': { 'power': True },
		}
		self.assertEqual(res.json(), info)

class TestDatabase(unittest.IsolatedAsyncioTestCase):
	def setUp(self):
		# setup temporary directory
		self.dbtempfile = tempfile.NamedTemporaryFile()

		self.database = data.databases.Database('sqlite:///' +
		    self.dbtempfile.name)
		self.data = make_orm(self.database)

	def tearDown(self):
		self.data = None
		self.database = None
		self.dbtempfile = None

	async def test_apikey(self):
		data = self.data

		# that the test database starts empty
		self.assertEqual(await data.APIKey.objects.all(), [])

		# that when it is populated with test data
		await _setup_data(data)

		# the data can be accessed
		self.assertEqual((await data.APIKey.objects.get(
		    key='thisisanapikey')).user, 'foo')
		self.assertEqual((await data.APIKey.objects.get(
		    key='anotherlongapikey')).user, 'bar')