diff --git a/bitelab/__init__.py b/bitelab/__init__.py index 1924003..f16dc0d 100644 --- a/bitelab/__init__.py +++ b/bitelab/__init__.py @@ -48,6 +48,7 @@ from .snmp import * from .mocks import * import asyncio +import contextlib import json import orm import os @@ -96,6 +97,10 @@ class BoardImpl: self.reserved = False + async def update_attrs(self, **attrs): + for i in attrs: + self.attrcache[i] = await self.attrmap[i].setvalue(attrs[i]) + async def update(self): for i in self.attrmap: self.attrcache[i] = await self.attrmap[i].getvalue() @@ -179,13 +184,6 @@ class BiteAuth(Auth): 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 @@ -208,6 +206,42 @@ async def get_boardmanager(settings: config.Settings = Depends(get_settings)): oauth2_scheme = OAuth2PasswordBearer(tokenUrl='/nonexistent') +def get_authorized_board_parms(board_id, token: str = Depends(oauth2_scheme), + data: data.DataWrapper = Depends(get_data), + brdmgr: BoardManager = Depends(get_boardmanager)): + '''This dependancy is used to collect the parameters needed for + the validate_board_params context manager.''' + + return dict(board_id=board_id, token=token, data=data, brdmgr=brdmgr) + +@contextlib.asynccontextmanager +async def validate_board_params(board_id, token, data, brdmgr): + '''This context manager checks to see if the request is authorized + for the board_id. This requires that the board is reserved by + the user, or the connection came from the board's jail (TBI). + ''' + + brd = brdmgr.boards[board_id] + + async with brd.lock: + user = await lookup_user(token, data) + + try: + brduser = await data.BoardStatus.objects.get(board=board_id) + except orm.exceptions.NoMatch: + raise BITEError( + status_code=HTTP_403_FORBIDDEN, + errobj=Error(error='Board not reserved.', + board=Board.from_orm(brd))) + + 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))) + + yield brd + async def lookup_user(token: str = Depends(oauth2_scheme), data: data.DataWrapper = Depends(get_data)): try: @@ -242,7 +276,6 @@ async def get_board_info(board_id, user: str = Depends(lookup_user), @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 @@ -299,7 +332,6 @@ async def reserve_board(board_id_or_class, user: str = Depends(lookup_user), @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] @@ -337,10 +369,14 @@ async def release_board(board_id, user: str = Depends(lookup_user), return brd @router.post('/board/{board_id}/attrs', response_model=Union[Board, Error]) -async def set_board_attrs(board_id, +async def set_board_attrs( attrs: Dict[str, Any], - brdmgr: BoardManager = Depends(get_boardmanager)): - pass + brdparams: dict = Depends(get_authorized_board_parms)): + + async with validate_board_params(**brdparams) as brd: + await brd.update_attrs(**attrs) + + return brd @router.get('/board/',response_model=Dict[str, Board]) async def get_boards(user: str = Depends(lookup_user), @@ -410,6 +446,9 @@ class TestBiteLab(unittest.IsolatedAsyncioTestCase): def get_data_override(self): return self.data + def get_boardmanager_override(self): + return self.brdmgr + async def asyncSetUp(self): self.app = getApp() @@ -426,9 +465,12 @@ class TestBiteLab(unittest.IsolatedAsyncioTestCase): setup_script='somesetupscript', ) + self.brdmgr = BoardManager(self.settings) + self.app.dependency_overrides[get_settings] = \ self.get_settings_override self.app.dependency_overrides[get_data] = self.get_data_override + self.app.dependency_overrides[get_boardmanager] = self.get_boardmanager_override self.client = AsyncClient(app=self.app, base_url='http://testserver') @@ -633,3 +675,82 @@ class TestBiteLab(unittest.IsolatedAsyncioTestCase): 'attrs': { 'power': True }, } self.assertEqual(res.json(), info) + + @patch('bitelab.snmp.snmpset') + async def test_board_attrs(self, ss): + data = self.data + + # that when snmpset returns False + ss.return_value = False + + attrs = dict(power=False) + + # that setting the board attributes requires auth + res = await self.client.post('/board/cora-1/attrs', + auth=BiteAuth('badapi'), + json=attrs) + + # that it fails auth + self.assertEqual(res.status_code, HTTP_401_UNAUTHORIZED) + + # that when properly authorized, but board is not reserved + res = await self.client.post('/board/cora-1/attrs', + auth=BiteAuth('thisisanapikey'), + json=attrs) + + # that it is forbidden + self.assertEqual(res.status_code, HTTP_403_FORBIDDEN) + + # that the cora-1 board is reserved + brd = self.brdmgr.boards['cora-1'] + async with brd.lock: + await brd.reserve() + obrdreq = await data.BoardStatus.objects.create( + board='cora-1', user='foo') + + # that setting the board attributes + res = await self.client.post('/board/cora-1/attrs', + auth=BiteAuth('thisisanapikey'), + json=attrs) + + # that it is successful + self.assertEqual(res.status_code, HTTP_200_OK) + + # calls snmpset w/ the correct args + ss.assert_called_with('poe', 'pethPsePortAdminEnable.1.2', + 'bool', False) + + # and returns the correct data + info = { + 'name': 'cora-1', + 'brdclass': 'cora-z7s', + 'reserved': True, + 'attrs': { 'power': False }, + } + self.assertEqual(res.json(), info) + + # that when snmpset returns True + ss.return_value = True + + attrs = dict(power=True) + + # that setting the board attributes + res = await self.client.post('/board/cora-1/attrs', + auth=BiteAuth('thisisanapikey'), + json=attrs) + + # calls snmpget w/ the correct args + ss.assert_called_with('poe', 'pethPsePortAdminEnable.1.2', + 'bool', True) + + # 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': True, + 'attrs': { 'power': True }, + } + self.assertEqual(res.json(), info)