diff --git a/pilab/auth.py b/pilab/auth.py index 0c1cc28a8a667ab19e4227cf79df147eb3028141..e3811c617c282a0a870dd5e0765fb606ee7ea5f6 100644 --- a/pilab/auth.py +++ b/pilab/auth.py @@ -1,18 +1,14 @@ import logging -from collections import Counter -from typing import Optional, List, Dict, Callable +from typing import Optional, List + from fastapi import Header, HTTPException + from pilab.events import reservation -from datetime import datetime logger = logging.getLogger(__name__) ADMIN_GROUPS = ["/admin"] -cube_reservations: Dict[int: reservation.Reservation] = {} -res_thread: reservation.ReservationConsumer = None -custom_callback: Callable = None - def is_admin(x_forwarded_groups: str): admin = False @@ -30,22 +26,6 @@ def get_username(usernames: List[str]): return None -def get_active_users(cube_id: int): - res: reservation.Reservation = cube_reservations.get(cube_id) - if res and res.endtime > datetime.utcnow(): - return [*res.extraUsers, res.owner] - else: - return [] - - -def get_active_reservation(cube_id: int): - res: reservation.Reservation = cube_reservations.get(cube_id) - if res and res.endtime > datetime.utcnow(): - return res - else: - return None - - async def get_user(x_forwarded_user: Optional[str] = Header(None), x_forwarded_preferred_username: Optional[str] = Header(None), x_forwarded_groups: Optional[str] = Header(None)): @@ -69,7 +49,7 @@ async def verify_user(cube_id: int, x_forwarded_preferred_username: Optional[str username = get_username([x_forwarded_preferred_username, x_forwarded_user]) if admin: return True - if username is not None and username in get_active_users(cube_id): + if username is not None and username in reservation.get_active_users(cube_id): return True raise HTTPException(status_code=401, detail="Unauthorized") @@ -78,32 +58,7 @@ async def verify_user(cube_id: int, x_forwarded_preferred_username: Optional[str def verify_user_pi(cube_id: int, username: str, admin: bool): if admin: return True - if username is not None and username in get_active_users(cube_id): + if username is not None and username in reservation.get_active_users(cube_id): return True return False - - -def restart_reservation_consumer(ids: List[int], callback: Callable = None): - global res_thread - global custom_callback - - custom_callback = callback - - def res_event_callback(res: reservation.Reservation): - logger.info(f'Received Reservation for cube {res.cube_id}; {res}') - if custom_callback: - custom_callback(res) - cube_reservations[res.cube_id] = res - - if not res_thread: - res_thread = reservation.ReservationConsumer(cubes=ids, callback=res_event_callback) - res_thread.start() - elif Counter(ids) != Counter(res_thread.cubes): - res_thread.stop() - res_thread.join() - res_thread = reservation.ReservationConsumer(cubes=ids, callback=res_event_callback) - res_thread.start() - - return res_thread - diff --git a/pilab/events/reservation.py b/pilab/events/reservation.py index f297fb8f644ea0ef915c2010486d1941522d8d33..4b220b6e106d2ddff1a82d7c045b70a310b92789 100644 --- a/pilab/events/reservation.py +++ b/pilab/events/reservation.py @@ -3,6 +3,7 @@ import json import logging import threading import traceback +from collections import Counter from datetime import datetime from typing import List, Callable, Optional, Dict @@ -270,3 +271,54 @@ class ReservationConsumer(threading.Thread): else: self._connection.ioloop.stop() logger.info('Stopped') + + +cube_reservations: Dict[int: Reservation] = {} +res_thread: ReservationConsumer = None +custom_callback: Callable = None + + +def get_active_users(cube_id: int): + res: Reservation = cube_reservations.get(cube_id) + if not res_thread: + raise RuntimeError("Reservation consumer is not initialized") + + if res and res.endtime > datetime.utcnow(): + return [*res.extraUsers, res.owner] + else: + return [] + + +def get_active_reservation(cube_id: int): + res: Reservation = cube_reservations.get(cube_id) + if not res_thread: + raise RuntimeError("Reservation consumer is not initialized") + + if res and res.endtime > datetime.utcnow(): + return res + else: + return None + + +def restart_reservation_consumer(ids: List[int], callback: Callable = None): + global res_thread + global custom_callback + + custom_callback = callback + + def res_event_callback(res: Reservation): + logger.info(f'Received Reservation for cube {res.cube_id}; {res}') + if custom_callback: + custom_callback(res) + cube_reservations[res.cube_id] = res + + if not res_thread: + res_thread = ReservationConsumer(cubes=ids, callback=res_event_callback) + res_thread.start() + elif Counter(ids) != Counter(res_thread.cubes): + res_thread.stop() + res_thread.join() + res_thread = ReservationConsumer(cubes=ids, callback=res_event_callback) + res_thread.start() + + return res_thread diff --git a/setup.py b/setup.py index 12d1af7a2b7185a6dfcb1e6233c011e831e4651f..5b085279c9dd7d1c690c96bf00c0e8286e309312 100644 --- a/setup.py +++ b/setup.py @@ -2,7 +2,7 @@ from setuptools import setup, find_packages setup( name='pilab', - version='1.4.5', + version='1.5.5', description='Shared-Libs for the pi-lab microservices', url='https://code.fbi.h-da.de/api/v4/projects/27896/packages/pypi/pilab', author='Max Reinheimer',