55from datetime import datetime , timedelta
66from typing import Set
77
8- from fastapi import APIRouter , Header , HTTPException , WebSocket
8+ from auth_lib .aiomethods import AsyncAuthLib
9+ from fastapi import APIRouter , Header , WebSocket , WebSocketException
910from fastapi_sqlalchemy import db
1011from pydantic import Field
1112from redis import Redis
13+ from starlette .status import WS_1000_NORMAL_CLOSURE
1214from typing_extensions import Annotated
1315
14- from print_service .exceptions import FileNotFound , InvalidPageRequest , IsNotUploaded , TerminalQRNotFound
16+ from print_service .exceptions import TerminalQRNotFound
1517from print_service .schema import BaseModel
1618from print_service .settings import Settings , get_settings
1719from print_service .utils import get_file
2022logger = logging .getLogger (__name__ )
2123settings : Settings = get_settings ()
2224router = APIRouter ()
25+ auth = AsyncAuthLib (auth_url = settings .AUTH_URL , userdata_url = settings .USERDATA_URL )
2326
2427
2528class InstantPrintCreate (BaseModel ):
@@ -49,13 +52,14 @@ class InstantPrintFetcher:
4952 def __init__ (self , terminal_token : str , settings : Settings = None ) -> None :
5053 self .terminal_token = terminal_token
5154 settings = settings or get_settings ()
52- self .redis = Redis .from_url (str (settings .REDIS_DSN ))
55+ self .redis : Redis = Redis .from_url (str (settings .REDIS_DSN ))
5356 self .ttl = settings .QR_TOKEN_TTL
5457 self .delay = settings .QR_TOKEN_DELAY
5558 self .symbols = settings .QR_TOKEN_SYMBOLS
5659 self .length = settings .QR_TOKEN_LENGTH
5760
5861 def new_qr (self ):
62+ logger .debug ("Generating new QR token" )
5963 for _ in range (5 ):
6064 qr_token = '' .join (random .choice (self .symbols ) for _ in range (self .length ))
6165 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]]:
7781 return {}
7882 return json .loads (raw_value )
7983
84+ async def check_token (self ):
85+ """Check if token valid and not used"""
86+ logger .info ("Checking token" )
87+
88+ # Token should be valid
89+ me = await auth .check_token (self .terminal_token )
90+ if me is None :
91+ logger .error ("Not authenticated" )
92+ raise Exception ("Not authenticated" )
93+
94+ for scope in me ['session_scopes' ]:
95+ if scope ['name' ] == "print.qr_task.get" :
96+ break
97+ else :
98+ logger .error ("Unauthorized" )
99+ logger .debug (me )
100+ raise Exception ("Unauthorized" )
101+
102+ # Token shouldn't be used yet
103+ for key in self .redis .keys ():
104+ value = self .redis .get (key )
105+ if self .redis .get (key ) == self .terminal_token .encode ():
106+ logger .error ("Token already used" )
107+ raise Exception ("Token already used" )
108+
80109 def __aiter__ (self ):
81110 return self
82111
@@ -93,7 +122,7 @@ async def __anext__(self):
93122@router .post ("" )
94123async def instant_print (options : InstantPrintCreate ):
95124 options .qr_token = options .qr_token .removeprefix (str (settings .QR_TOKEN_PREFIX ))
96- if redis_conn .send (** options .dict ()):
125+ if redis_conn .send (** options .model_dump ()):
97126 return {'status' : 'ok' }
98127 raise TerminalQRNotFound ()
99128
@@ -104,7 +133,15 @@ async def instant_print_terminal_connection(
104133 authorization : str = Header (),
105134):
106135 await websocket .accept ()
136+ logger .debug ("Websocket connection started" )
107137 manager = InstantPrintFetcher (authorization .removeprefix ("token " ))
138+ try :
139+ await manager .check_token ()
140+ except Exception as e :
141+ await websocket .send_text (json .dumps ({"error" : e .args [0 ]}))
142+ raise WebSocketException (WS_1000_NORMAL_CLOSURE , "Auth error" )
143+ logger .debug ("Websocket token checked" )
144+
108145 await websocket .send_text (json .dumps ({"qr_token" : str (settings .QR_TOKEN_PREFIX ) + manager .new_qr ()}))
109146 async for task in manager :
110147 task ['qr_token' ] = str (settings .QR_TOKEN_PREFIX ) + task ['qr_token' ]
0 commit comments