70 lines
2.4 KiB
Python
70 lines
2.4 KiB
Python
"""
|
|
Auth module - JWT authentication and user management.
|
|
"""
|
|
from datetime import datetime, timedelta, timezone
|
|
from typing import Optional
|
|
from urllib.parse import parse_qs, urlparse
|
|
|
|
import jwt
|
|
from websockets.server import WebSocketServerProtocol
|
|
|
|
from .config import JWT_ALGORITHM, JWT_EXPIRE_HOURS, JWT_SECRET, USERS, logger
|
|
|
|
__all__ = ['verify_token', 'verify_token_ws', 'verify_token_http', 'generate_token', 'JWT_ALGORITHM', 'JWT_EXPIRE_HOURS', 'JWT_SECRET', 'USERS']
|
|
from .utils import get_client_ip
|
|
|
|
|
|
async def verify_token_ws(ws: WebSocketServerProtocol) -> Optional[str]:
|
|
"""Verify JWT token from WebSocket query parameters. Returns username or None."""
|
|
try:
|
|
# Get token from query string
|
|
path = ws.request.path if hasattr(ws.request, 'path') else str(ws.request)
|
|
parsed = urlparse(path)
|
|
params = parse_qs(parsed.query)
|
|
token = params.get('token', [''])[0]
|
|
|
|
if not token:
|
|
logger.warning(f"No token provided from {get_client_ip(ws)}")
|
|
return None
|
|
|
|
payload = jwt.decode(token, JWT_SECRET, algorithms=[JWT_ALGORITHM])
|
|
return payload['sub']
|
|
except jwt.ExpiredSignatureError:
|
|
logger.warning(f"Expired token from {get_client_ip(ws)}")
|
|
return None
|
|
except jwt.InvalidTokenError as e:
|
|
logger.warning(f"Invalid token from {get_client_ip(ws)}: {e}")
|
|
return None
|
|
except Exception as e:
|
|
logger.error(f"Token verification error: {e}")
|
|
return None
|
|
|
|
|
|
def generate_token(username: str) -> str:
|
|
"""Generate a new JWT token for the user."""
|
|
expire = datetime.now(timezone.utc) + timedelta(hours=JWT_EXPIRE_HOURS)
|
|
token = jwt.encode(
|
|
{'sub': username, 'exp': expire, 'iat': datetime.now(timezone.utc)},
|
|
JWT_SECRET,
|
|
algorithm=JWT_ALGORITHM
|
|
)
|
|
# Handle both PyJWT v1 (bytes) and v2 (str)
|
|
if isinstance(token, bytes):
|
|
token = token.decode('utf-8')
|
|
return token
|
|
|
|
|
|
def verify_token_http(token: str) -> Optional[str]:
|
|
"""Verify JWT token from HTTP request. Returns username or None."""
|
|
try:
|
|
payload = jwt.decode(token, JWT_SECRET, algorithms=[JWT_ALGORITHM])
|
|
return payload.get('sub')
|
|
except jwt.ExpiredSignatureError:
|
|
return None
|
|
except jwt.InvalidTokenError:
|
|
return None
|
|
|
|
|
|
# For backward compatibility, keep the original function name
|
|
verify_token = verify_token_ws
|