diff --git a/bitelab/__init__.py b/bitelab/__init__.py index 8efb6ef..d4d9228 100644 --- a/bitelab/__init__.py +++ b/bitelab/__init__.py @@ -320,7 +320,7 @@ def get_authorized_board_parms(board_id, token: str = Depends(oauth2_scheme), return dict(board_id=board_id, token=token, data=data, brdmgr=brdmgr) @contextlib.asynccontextmanager -async def validate_board_params(board_id, token, data, brdmgr): +async def validate_board_params(board_id, data, brdmgr, user=None, token=None): '''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). @@ -329,7 +329,8 @@ async def validate_board_params(board_id, token, data, brdmgr): brd = brdmgr.boards[board_id] async with brd.lock: - user = await lookup_user(token, data) + if user is None: + user = await lookup_user(token, data) try: brduser = await data.BoardStatus.objects.get(board=board_id) @@ -448,11 +449,12 @@ async def reserve_board(board_id_or_class, return brd class HandleExec(WSFWDServer): - def __init__(self, *args, board_id, data, **kwargs): + def __init__(self, *args, board_id, data, brdmgr, **kwargs): super().__init__(*args, **kwargs) self._board_id = board_id self._data = data + self._brdmgr = brdmgr self._auth_user = None self._did_exec = False @@ -514,9 +516,16 @@ class HandleExec(WSFWDServer): if self._auth_user is None: raise RuntimeError('not authenticated') - self._proc = await asyncio.create_subprocess_exec( - 'jexec', self._board_id, *msg['args'], - stdin=subprocess.PIPE, stdout=subprocess.PIPE) + try: + async with validate_board_params(self._board_id, self._data, + self._brdmgr, user=self._auth_user) as brd: + self._proc = await \ + asyncio.create_subprocess_exec('jexec', + self._board_id, *msg['args'], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE) + except BITEError as e: + raise RuntimeError(e.errobj.error) self._did_exec = True @@ -547,7 +556,7 @@ async def board_exec_ws( try: async with HandleExec(websocket.receive_bytes, websocket.send_bytes, data=data, - board_id=board_id) as server: + board_id=board_id, brdmgr=brdmgr) as server: await server.get_finish_handler() finally: await websocket.close() @@ -773,7 +782,29 @@ class TestWebSocket(TestCommon): # that a valid auth token works await client.auth(dict(bearer='thisisanapikey')) - # XXX - enforce board reservation and correct user + # That since the board isn't reserved, it fails + with self.assertRaisesRegex(RuntimeError, + 'Board not reserved.'): + await client.exec([ 'sshd', '-i' ], stdin=1, + stdout=2) + + # that when the board is reserved by the wrong user + brd = self.brdmgr.boards['cora-1'] + obrdreq = await self.data.BoardStatus.objects.create( + board='cora-1', user='bar') + async with brd.lock: + await brd.reserve() + + # that it fails + with self.assertRaisesRegex(RuntimeError, 'Board reserved by \'bar\'.'): + await client.exec([ 'sshd', '-i' ], stdin=1, stdout=2) + + brduser = await self.data.BoardStatus.objects.get(board='cora-1') + obrdreq = await self.data.BoardStatus.delete(brduser) + + # that when the board is reserved by the correct user + obrdreq = await self.data.BoardStatus.objects.create( + board='cora-1', user='foo') echodata = b'somedata' wrap_subprocess_exec(cse, stdout=echodata, retcode=0)