Initial commit

src/webssh_backend/auth.py

@@ -0,0 +1,69 @@
+"""
+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