97 lines
3.3 KiB
Python
97 lines
3.3 KiB
Python
"""Authentication lifecycle and short-lived export download tokens."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import secrets
|
|
import threading
|
|
import time
|
|
from datetime import datetime, timezone
|
|
|
|
from quantux_client import QuantUXError
|
|
|
|
|
|
class QuantUXSessionManager:
|
|
"""Keeps the shared Quant-UX client authenticated without login storms."""
|
|
|
|
def __init__(self, client, email="", password="", refresh_leeway=300, clock=time.time):
|
|
self.client = client
|
|
self.email = email
|
|
self.password = password
|
|
self.refresh_leeway = max(0, int(refresh_leeway))
|
|
self.clock = clock
|
|
self._lock = threading.RLock()
|
|
|
|
def ensure_authenticated(self, force=False, failed_token=None):
|
|
now = self.clock()
|
|
if not force and self.client.token_is_valid(self.refresh_leeway, now):
|
|
return self.client
|
|
with self._lock:
|
|
now = self.clock()
|
|
if (
|
|
failed_token is not None
|
|
and self.client.token != failed_token
|
|
and self.client.token_is_valid(self.refresh_leeway, now)
|
|
):
|
|
return self.client
|
|
if not force and self.client.token_is_valid(self.refresh_leeway, now):
|
|
return self.client
|
|
if not self.email or not self.password:
|
|
raise QuantUXError(
|
|
"Quant-UX login expired and server credentials are not configured"
|
|
)
|
|
self.client.login(self.email, self.password)
|
|
if not self.client.token_is_valid(0, self.clock()):
|
|
raise QuantUXError("Quant-UX login returned an invalid or expired token")
|
|
return self.client
|
|
|
|
def refresh_after_auth_failure(self, failed_token):
|
|
return self.ensure_authenticated(force=True, failed_token=failed_token)
|
|
|
|
def status(self):
|
|
now = self.clock()
|
|
expires_at = self.client.token_expires_at()
|
|
valid = self.client.token_is_valid(0, now)
|
|
return {
|
|
"logged_in": valid,
|
|
"token_valid": valid,
|
|
"token_expires_at": (
|
|
datetime.fromtimestamp(expires_at, timezone.utc).isoformat()
|
|
if expires_at
|
|
else None
|
|
),
|
|
"token_refresh_due": bool(
|
|
valid and expires_at <= now + self.refresh_leeway
|
|
),
|
|
}
|
|
|
|
|
|
class DownloadTokenStore:
|
|
"""Issues scoped, short-lived tokens for one exported filename."""
|
|
|
|
def __init__(self, ttl_seconds=600, clock=time.time):
|
|
self.ttl_seconds = min(3600, max(60, int(ttl_seconds)))
|
|
self.clock = clock
|
|
self._tokens = {}
|
|
self._lock = threading.Lock()
|
|
|
|
def issue(self, filename):
|
|
token = secrets.token_urlsafe(32)
|
|
with self._lock:
|
|
self._purge_locked()
|
|
self._tokens[token] = (filename, self.clock() + self.ttl_seconds)
|
|
return token
|
|
|
|
def validate(self, token, filename):
|
|
if not token:
|
|
return False
|
|
with self._lock:
|
|
self._purge_locked()
|
|
value = self._tokens.get(token)
|
|
return bool(value and value[0] == filename and value[1] > self.clock())
|
|
|
|
def _purge_locked(self):
|
|
now = self.clock()
|
|
expired = [token for token, (_, expiry) in self._tokens.items() if expiry <= now]
|
|
for token in expired:
|
|
self._tokens.pop(token, None)
|