diff --git a/Makefile b/Makefile index b7946a3..d10e3fe 100644 --- a/Makefile +++ b/Makefile @@ -1,5 +1,5 @@ run: - source ./venv/bin/activate && uvicorn --reload --log-level debug print_service.routes:app + source ./venv/bin/activate && uvicorn --reload --log-config logging_dev.conf print_service.routes:app configure: venv source ./venv/bin/activate && pip install -r requirements.dev.txt -r requirements.txt diff --git a/logging_dev.conf b/logging_dev.conf new file mode 100644 index 0000000..7837272 --- /dev/null +++ b/logging_dev.conf @@ -0,0 +1,21 @@ +[loggers] +keys=root + +[handlers] +keys=all + +[formatters] +keys=main + +[logger_root] +level=DEBUG +handlers=all + +[handler_all] +class=StreamHandler +formatter=main +level=DEBUG +args=(sys.stdout,) + +[formatter_main] +format=%(asctime)s %(levelname)-8s %(name)-15s %(message)s diff --git a/print_service/routes/qrprint.py b/print_service/routes/qrprint.py index c1d0782..fe6dfab 100644 --- a/print_service/routes/qrprint.py +++ b/print_service/routes/qrprint.py @@ -5,13 +5,15 @@ from datetime import datetime, timedelta from typing import Set -from fastapi import APIRouter, Header, HTTPException, WebSocket +from auth_lib.aiomethods import AsyncAuthLib +from fastapi import APIRouter, Header, WebSocket, WebSocketException from fastapi_sqlalchemy import db from pydantic import Field from redis import Redis +from starlette.status import WS_1000_NORMAL_CLOSURE from typing_extensions import Annotated -from print_service.exceptions import FileNotFound, InvalidPageRequest, IsNotUploaded, TerminalQRNotFound +from print_service.exceptions import TerminalQRNotFound from print_service.schema import BaseModel from print_service.settings import Settings, get_settings from print_service.utils import get_file @@ -20,6 +22,7 @@ logger = logging.getLogger(__name__) settings: Settings = get_settings() router = APIRouter() +auth = AsyncAuthLib(auth_url=settings.AUTH_URL, userdata_url=settings.USERDATA_URL) class InstantPrintCreate(BaseModel): @@ -49,13 +52,14 @@ class InstantPrintFetcher: def __init__(self, terminal_token: str, settings: Settings = None) -> None: self.terminal_token = terminal_token settings = settings or get_settings() - self.redis = Redis.from_url(str(settings.REDIS_DSN)) + self.redis: Redis = Redis.from_url(str(settings.REDIS_DSN)) self.ttl = settings.QR_TOKEN_TTL self.delay = settings.QR_TOKEN_DELAY self.symbols = settings.QR_TOKEN_SYMBOLS self.length = settings.QR_TOKEN_LENGTH def new_qr(self): + logger.debug("Generating new QR token") for _ in range(5): qr_token = ''.join(random.choice(self.symbols) for _ in range(self.length)) if not self.redis.get(qr_token): # If this qr already exists, generate new @@ -77,6 +81,31 @@ async def get_tasks(self) -> dict[str, list[str]]: return {} return json.loads(raw_value) + async def check_token(self): + """Check if token valid and not used""" + logger.info("Checking token") + + # Token should be valid + me = await auth.check_token(self.terminal_token) + if me is None: + logger.error("Not authenticated") + raise Exception("Not authenticated") + + for scope in me['session_scopes']: + if scope['name'] == "print.qr_task.get": + break + else: + logger.error("Unauthorized") + logger.debug(me) + raise Exception("Unauthorized") + + # Token shouldn't be used yet + for key in self.redis.keys(): + value = self.redis.get(key) + if self.redis.get(key) == self.terminal_token.encode(): + logger.error("Token already used") + raise Exception("Token already used") + def __aiter__(self): return self @@ -93,7 +122,7 @@ async def __anext__(self): @router.post("") async def instant_print(options: InstantPrintCreate): options.qr_token = options.qr_token.removeprefix(str(settings.QR_TOKEN_PREFIX)) - if redis_conn.send(**options.dict()): + if redis_conn.send(**options.model_dump()): return {'status': 'ok'} raise TerminalQRNotFound() @@ -104,7 +133,15 @@ async def instant_print_terminal_connection( authorization: str = Header(), ): await websocket.accept() + logger.debug("Websocket connection started") manager = InstantPrintFetcher(authorization.removeprefix("token ")) + try: + await manager.check_token() + except Exception as e: + await websocket.send_text(json.dumps({"error": e.args[0]})) + raise WebSocketException(WS_1000_NORMAL_CLOSURE, "Auth error") + logger.debug("Websocket token checked") + await websocket.send_text(json.dumps({"qr_token": str(settings.QR_TOKEN_PREFIX) + manager.new_qr()})) async for task in manager: task['qr_token'] = str(settings.QR_TOKEN_PREFIX) + task['qr_token']