robosystems-client 0.1.17__py3-none-any.whl → 0.1.18__py3-none-any.whl

robosystems_client/extensions/tests/test_token_utils.py

@@ -0,0 +1,274 @@
+"""Tests for JWT token utilities"""
+
+import pytest
+from datetime import datetime, timedelta
+import json
+import base64
+import os
+
+from robosystems_client.extensions.token_utils import (
+  validate_jwt_format,
+  extract_jwt_from_header,
+  decode_jwt_payload,
+  is_jwt_expired,
+  get_jwt_claims,
+  get_jwt_expiration,
+  extract_token_from_environment,
+  extract_token_from_cookie,
+  find_valid_token,
+  TokenManager,
+)
+
+
+def create_test_jwt(payload: dict = None, exp_delta_seconds: int = 3600) -> str:
+  """Create a test JWT token"""
+  header = {"alg": "HS256", "typ": "JWT"}
+
+  if payload is None:
+    payload = {
+      "sub": "test_user",
+      "user_id": "123",
+      "exp": int((datetime.now() + timedelta(seconds=exp_delta_seconds)).timestamp()),
+    }
+
+  # Encode header and payload
+  header_b64 = (
+    base64.urlsafe_b64encode(json.dumps(header).encode()).decode().rstrip("=")
+  )
+
+  payload_b64 = (
+    base64.urlsafe_b64encode(json.dumps(payload).encode()).decode().rstrip("=")
+  )
+
+  # Create fake signature
+  signature = "test_signature_123"
+
+  return f"{header_b64}.{payload_b64}.{signature}"
+
+
+class TestJWTValidation:
+  """Test JWT format validation"""
+
+  def test_validate_jwt_format_valid(self):
+    """Test validation of valid JWT format"""
+    token = create_test_jwt()
+    assert validate_jwt_format(token) is True
+
+  def test_validate_jwt_format_invalid(self):
+    """Test validation of invalid JWT formats"""
+    # Missing parts
+    assert validate_jwt_format("header.payload") is False
+    # Wrong format
+    assert validate_jwt_format("not-a-jwt") is False
+    # Empty
+    assert validate_jwt_format("") is False
+    # None
+    assert validate_jwt_format(None) is False
+    # Not a string
+    assert validate_jwt_format(123) is False
+
+  def test_validate_jwt_with_padding(self):
+    """Test JWT validation handles padding correctly"""
+    # JWT with different padding requirements
+    token = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIn0.signature"
+    assert validate_jwt_format(token) is True
+
+
+class TestJWTExtraction:
+  """Test JWT extraction from various sources"""
+
+  def test_extract_jwt_from_header_bearer(self):
+    """Test extraction from Bearer authorization header"""
+    token = create_test_jwt()
+
+    # Standard Bearer format
+    assert extract_jwt_from_header(f"Bearer {token}") == token
+    # Case insensitive
+    assert extract_jwt_from_header(f"bearer {token}") == token
+    # Extra spaces
+    assert extract_jwt_from_header(f"Bearer  {token}  ") == token
+
+  def test_extract_jwt_from_header_dict(self):
+    """Test extraction from headers dictionary"""
+    token = create_test_jwt()
+
+    headers = {"Authorization": f"Bearer {token}"}
+    assert extract_jwt_from_header(headers) == token
+
+    # Case variation
+    headers = {"authorization": f"Bearer {token}"}
+    assert extract_jwt_from_header(headers) == token
+
+  def test_extract_jwt_from_header_invalid(self):
+    """Test extraction returns None for invalid inputs"""
+    assert extract_jwt_from_header(None) is None
+    assert extract_jwt_from_header("") is None
+    assert extract_jwt_from_header("NotBearer token") is None
+    assert extract_jwt_from_header({"Other": "header"}) is None
+
+
+class TestJWTDecoding:
+  """Test JWT payload decoding"""
+
+  def test_decode_jwt_payload(self):
+    """Test decoding JWT payload"""
+    payload = {"sub": "test_user", "user_id": "123", "roles": ["admin", "user"]}
+    token = create_test_jwt(payload)
+
+    decoded = decode_jwt_payload(token)
+    assert decoded["sub"] == "test_user"
+    assert decoded["user_id"] == "123"
+    assert decoded["roles"] == ["admin", "user"]
+
+  def test_decode_jwt_payload_invalid(self):
+    """Test decoding invalid JWT returns None"""
+    assert decode_jwt_payload("invalid.token") is None
+    assert decode_jwt_payload("") is None
+
+  def test_get_jwt_claims(self):
+    """Test getting all claims from JWT"""
+    payload = {"claim1": "value1", "claim2": "value2"}
+    token = create_test_jwt(payload)
+
+    claims = get_jwt_claims(token)
+    assert claims["claim1"] == "value1"
+    assert claims["claim2"] == "value2"
+
+
+class TestJWTExpiration:
+  """Test JWT expiration checking"""
+
+  def test_is_jwt_expired_not_expired(self):
+    """Test checking non-expired token"""
+    # Token expires in 1 hour
+    token = create_test_jwt(exp_delta_seconds=3600)
+    assert is_jwt_expired(token) is False
+
+  def test_is_jwt_expired_expired(self):
+    """Test checking expired token"""
+    # Token expired 1 hour ago
+    token = create_test_jwt(exp_delta_seconds=-3600)
+    assert is_jwt_expired(token) is True
+
+  def test_is_jwt_expired_with_buffer(self):
+    """Test expiration with buffer time"""
+    # Token expires in 30 seconds
+    token = create_test_jwt(exp_delta_seconds=30)
+    # With 60 second buffer, should be considered expired
+    assert is_jwt_expired(token, buffer_seconds=60) is True
+    # With no buffer, should not be expired
+    assert is_jwt_expired(token, buffer_seconds=0) is False
+
+  def test_get_jwt_expiration(self):
+    """Test getting expiration datetime"""
+    exp_time = datetime.now() + timedelta(hours=1)
+    payload = {"exp": int(exp_time.timestamp())}
+    token = create_test_jwt(payload)
+
+    exp = get_jwt_expiration(token)
+    assert exp is not None
+    # Allow 1 second difference for test execution
+    assert abs((exp - exp_time).total_seconds()) < 1
+
+
+class TestTokenExtraction:
+  """Test token extraction from various sources"""
+
+  def test_extract_token_from_environment(self):
+    """Test extracting token from environment variable"""
+    token = create_test_jwt()
+    os.environ["ROBOSYSTEMS_TOKEN"] = token
+
+    try:
+      assert extract_token_from_environment() == token
+      # Custom env var
+      os.environ["CUSTOM_TOKEN"] = token
+      assert extract_token_from_environment("CUSTOM_TOKEN") == token
+    finally:
+      # Clean up
+      os.environ.pop("ROBOSYSTEMS_TOKEN", None)
+      os.environ.pop("CUSTOM_TOKEN", None)
+
+  def test_extract_token_from_cookie(self):
+    """Test extracting token from cookies"""
+    token = create_test_jwt()
+    cookies = {"auth-token": token}
+
+    assert extract_token_from_cookie(cookies) == token
+    # Custom cookie name
+    cookies = {"session_token": token}
+    assert extract_token_from_cookie(cookies, "session_token") == token
+    # Missing cookie
+    assert extract_token_from_cookie({}) is None
+
+  def test_find_valid_token(self):
+    """Test finding first valid token from multiple sources"""
+    token = create_test_jwt()
+
+    # Found in second source
+    result = find_valid_token(None, "invalid-token", token, "another-invalid")
+    assert result == token
+
+    # Found in headers dict
+    headers = {"Authorization": f"Bearer {token}"}
+    result = find_valid_token(None, headers)
+    assert result == token
+
+    # Not found
+    result = find_valid_token(None, "", "invalid")
+    assert result is None
+
+
+class TestTokenManager:
+  """Test TokenManager class"""
+
+  def test_token_manager_basic(self):
+    """Test basic TokenManager functionality"""
+    token = create_test_jwt()
+    manager = TokenManager(token)
+
+    assert manager.token == token
+    assert manager.is_valid() is True
+
+    claims = manager.get_claims()
+    assert claims["sub"] == "test_user"
+
+  def test_token_manager_refresh(self):
+    """Test token refresh functionality"""
+    old_token = create_test_jwt(exp_delta_seconds=30)
+    new_token = create_test_jwt(exp_delta_seconds=3600)
+
+    def refresh_callback():
+      return new_token
+
+    manager = TokenManager(
+      old_token, refresh_callback=refresh_callback, auto_refresh=True, refresh_buffer=60
+    )
+
+    # Token should be refreshed automatically
+    assert manager.token == new_token
+
+  def test_token_manager_invalid_token(self):
+    """Test TokenManager with invalid token"""
+    manager = TokenManager()
+
+    with pytest.raises(ValueError):
+      manager.token = "invalid-token"
+
+    assert manager.is_valid() is False
+    assert manager.get_claims() is None
+    assert manager.get_expiration() is None
+
+  def test_token_manager_manual_refresh(self):
+    """Test manual token refresh"""
+    token = create_test_jwt()
+    new_token = create_test_jwt(exp_delta_seconds=7200)
+
+    manager = TokenManager(
+      token, refresh_callback=lambda: new_token, auto_refresh=False
+    )
+
+    assert manager.token == token
+    refreshed = manager.refresh()
+    assert refreshed == new_token
+    assert manager.token == new_token

robosystems_client/extensions/token_utils.py

@@ -0,0 +1,417 @@
+"""JWT Token validation and management utilities for RoboSystems SDK
+
+Provides comprehensive JWT handling, validation, and extraction utilities.
+"""
+
+import base64
+import json
+import os
+from datetime import datetime, timedelta
+from typing import Dict, Any, Optional, Union
+from enum import Enum
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class TokenSource(Enum):
+  """Sources where tokens can be extracted from"""
+
+  HEADER = "header"
+  COOKIE = "cookie"
+  ENVIRONMENT = "environment"
+  CONFIG = "config"
+
+
+def validate_jwt_format(token: Optional[str]) -> bool:
+  """Validate JWT token format (basic validation without cryptographic verification)
+
+  Args:
+      token: JWT token string to validate
+
+  Returns:
+      True if token appears to be valid JWT format
+
+  Example:
+      >>> validate_jwt_format("eyJhbGc.eyJzdWI.SflKxwRJSM")
+      True
+      >>> validate_jwt_format("invalid-token")
+      False
+  """
+  if not token or not isinstance(token, str):
+    return False
+
+  try:
+    # JWT should have exactly 3 parts: header.payload.signature
+    parts = token.split(".")
+    if len(parts) != 3:
+      return False
+
+    # Each part should be base64url encoded
+    for part in parts[:2]:  # Check header and payload only
+      # Add padding if needed
+      padding = 4 - (len(part) % 4)
+      if padding != 4:
+        part += "=" * padding
+
+      try:
+        # Try to decode base64
+        base64.urlsafe_b64decode(part)
+      except Exception:
+        return False
+
+    return True
+  except Exception:
+    return False
+
+
+def extract_jwt_from_header(
+  auth_header: Optional[Union[str, Dict[str, str]]],
+) -> Optional[str]:
+  """Extract JWT token from Authorization header
+
+  Args:
+      auth_header: Authorization header value (e.g., "Bearer token123") or headers dict
+
+  Returns:
+      JWT token if found, None otherwise
+
+  Example:
+      >>> extract_jwt_from_header("Bearer eyJhbGc.eyJzdWI.SflKxwRJSM")
+      "eyJhbGc.eyJzdWI.SflKxwRJSM"
+      >>> extract_jwt_from_header({"Authorization": "Bearer token123"})
+      "token123"
+  """
+  if not auth_header:
+    return None
+
+  # Handle dict of headers
+  if isinstance(auth_header, dict):
+    auth_value = auth_header.get("Authorization") or auth_header.get("authorization")
+    if not auth_value:
+      return None
+    auth_header = auth_value
+
+  # Extract token from Bearer scheme
+  if isinstance(auth_header, str):
+    auth_header = auth_header.strip()
+    if auth_header.startswith("Bearer "):
+      token = auth_header[7:].strip()
+      return token if token else None
+    elif auth_header.startswith("bearer "):  # Case insensitive
+      token = auth_header[7:].strip()
+      return token if token else None
+
+  return None
+
+
+def decode_jwt_payload(token: str, verify: bool = False) -> Optional[Dict[str, Any]]:
+  """Decode JWT payload without verification (for reading claims only)
+
+  Args:
+      token: JWT token to decode
+      verify: If True, will validate format first (default: False)
+
+  Returns:
+      Decoded payload as dictionary, None if invalid
+
+  Note:
+      This does NOT verify the signature. Use only for reading non-sensitive claims.
+
+  Example:
+      >>> payload = decode_jwt_payload("eyJhbGc.eyJzdWI.SflKxwRJSM")
+      >>> payload.get("sub")  # Get subject claim
+  """
+  if verify and not validate_jwt_format(token):
+    return None
+
+  try:
+    # Split token and get payload (second part)
+    parts = token.split(".")
+    if len(parts) != 3:
+      return None
+
+    payload_part = parts[1]
+
+    # Add padding if needed
+    padding = 4 - (len(payload_part) % 4)
+    if padding != 4:
+      payload_part += "=" * padding
+
+    # Decode base64url
+    payload_bytes = base64.urlsafe_b64decode(payload_part)
+    payload = json.loads(payload_bytes.decode("utf-8"))
+
+    return payload
+  except Exception as e:
+    logger.debug(f"Failed to decode JWT payload: {e}")
+    return None
+
+
+def is_jwt_expired(token: str, buffer_seconds: int = 60) -> bool:
+  """Check if JWT token is expired based on exp claim
+
+  Args:
+      token: JWT token to check
+      buffer_seconds: Consider expired if expiring within this many seconds (default: 60)
+
+  Returns:
+      True if token is expired or expiring soon
+
+  Example:
+      >>> is_jwt_expired("eyJhbGc.eyJleHAiOjE2MzA0MjU2MDB9.SflKxwRJSM")
+      True  # If current time is past exp claim
+  """
+  payload = decode_jwt_payload(token)
+  if not payload:
+    return True
+
+  exp = payload.get("exp")
+  if not exp:
+    # No expiration claim, consider as non-expiring
+    return False
+
+  try:
+    exp_datetime = datetime.fromtimestamp(exp)
+    buffer = timedelta(seconds=buffer_seconds)
+    return datetime.now() >= (exp_datetime - buffer)
+  except Exception:
+    return True
+
+
+def get_jwt_claims(token: str) -> Optional[Dict[str, Any]]:
+  """Get all claims from JWT token
+
+  Args:
+      token: JWT token
+
+  Returns:
+      Dictionary of all claims, None if invalid
+
+  Example:
+      >>> claims = get_jwt_claims(token)
+      >>> user_id = claims.get("user_id")
+      >>> roles = claims.get("roles", [])
+  """
+  return decode_jwt_payload(token)
+
+
+def get_jwt_expiration(token: str) -> Optional[datetime]:
+  """Get expiration datetime from JWT token
+
+  Args:
+      token: JWT token
+
+  Returns:
+      Expiration datetime, None if no exp claim or invalid
+
+  Example:
+      >>> exp = get_jwt_expiration(token)
+      >>> if exp and exp > datetime.now():
+      ...     print(f"Token valid until {exp}")
+  """
+  payload = decode_jwt_payload(token)
+  if not payload:
+    return None
+
+  exp = payload.get("exp")
+  if not exp:
+    return None
+
+  try:
+    return datetime.fromtimestamp(exp)
+  except Exception:
+    return None
+
+
+def extract_token_from_environment(env_var: str = "ROBOSYSTEMS_TOKEN") -> Optional[str]:
+  """Extract JWT token from environment variable
+
+  Args:
+      env_var: Environment variable name (default: ROBOSYSTEMS_TOKEN)
+
+  Returns:
+      JWT token if found and valid format, None otherwise
+
+  Example:
+      >>> os.environ["ROBOSYSTEMS_TOKEN"] = "eyJhbGc..."
+      >>> token = extract_token_from_environment()
+  """
+  token = os.environ.get(env_var)
+  if token and validate_jwt_format(token):
+    return token
+  return None
+
+
+def extract_token_from_cookie(
+  cookies: Dict[str, str], cookie_name: str = "auth-token"
+) -> Optional[str]:
+  """Extract JWT token from cookies
+
+  Args:
+      cookies: Dictionary of cookies
+      cookie_name: Name of cookie containing token (default: auth-token)
+
+  Returns:
+      JWT token if found, None otherwise
+
+  Example:
+      >>> cookies = {"auth-token": "eyJhbGc..."}
+      >>> token = extract_token_from_cookie(cookies)
+  """
+  token = cookies.get(cookie_name)
+  if token and validate_jwt_format(token):
+    return token
+  return None
+
+
+def find_valid_token(*sources: Union[str, Dict[str, str], None]) -> Optional[str]:
+  """Find first valid JWT token from multiple sources
+
+  Args:
+      *sources: Variable number of potential token sources
+                (strings, dicts with Authorization header, etc.)
+
+  Returns:
+      First valid JWT token found, None if none found
+
+  Example:
+      >>> token = find_valid_token(
+      ...     os.environ.get("TOKEN"),
+      ...     headers,
+      ...     cookies.get("auth-token"),
+      ...     config.get("token")
+      ... )
+  """
+  for source in sources:
+    if not source:
+      continue
+
+    # Direct token string
+    if isinstance(source, str):
+      if validate_jwt_format(source):
+        return source
+
+    # Headers dict
+    elif isinstance(source, dict):
+      # Try as headers
+      token = extract_jwt_from_header(source)
+      if token and validate_jwt_format(token):
+        return token
+
+      # Try as cookies
+      for key in ["auth-token", "auth_token", "token", "jwt"]:
+        token = source.get(key)
+        if token and validate_jwt_format(token):
+          return token
+
+  return None
+
+
+class TokenManager:
+  """Manages JWT tokens with automatic refresh and validation"""
+
+  def __init__(
+    self,
+    token: Optional[str] = None,
+    refresh_callback: Optional[callable] = None,
+    auto_refresh: bool = True,
+    refresh_buffer: int = 300,
+  ):
+    """Initialize token manager
+
+    Args:
+        token: Initial JWT token
+        refresh_callback: Callback to refresh token when expired
+        auto_refresh: Automatically refresh before expiration
+        refresh_buffer: Seconds before expiration to trigger refresh (default: 300)
+    """
+    self._token = token
+    self._refresh_callback = refresh_callback
+    self._auto_refresh = auto_refresh
+    self._refresh_buffer = refresh_buffer
+
+  @property
+  def token(self) -> Optional[str]:
+    """Get current token, refreshing if needed"""
+    if self._auto_refresh and self._token and self._refresh_callback:
+      if is_jwt_expired(self._token, self._refresh_buffer):
+        self.refresh()
+    return self._token
+
+  @token.setter
+  def token(self, value: Optional[str]):
+    """Set new token"""
+    if value and not validate_jwt_format(value):
+      raise ValueError("Invalid JWT token format")
+    self._token = value
+
+  def refresh(self) -> Optional[str]:
+    """Refresh token using callback"""
+    if not self._refresh_callback:
+      raise RuntimeError("No refresh callback configured")
+
+    try:
+      new_token = self._refresh_callback()
+      if new_token and validate_jwt_format(new_token):
+        self._token = new_token
+        logger.info("Token refreshed successfully")
+        return new_token
+    except Exception as e:
+      logger.error(f"Token refresh failed: {e}")
+
+    return None
+
+  def is_valid(self) -> bool:
+    """Check if current token is valid"""
+    return bool(
+      self._token
+      and validate_jwt_format(self._token)
+      and not is_jwt_expired(self._token, 0)
+    )
+
+  def get_claims(self) -> Optional[Dict[str, Any]]:
+    """Get claims from current token"""
+    if self._token:
+      return get_jwt_claims(self._token)
+    return None
+
+  def get_expiration(self) -> Optional[datetime]:
+    """Get expiration time of current token"""
+    if self._token:
+      return get_jwt_expiration(self._token)
+    return None
+
+
+# Convenience function for quick token extraction from client config
+def extract_token_from_client(client) -> Optional[str]:
+  """Extract JWT token from RoboSystems client configuration
+
+  Args:
+      client: RoboSystems client instance
+
+  Returns:
+      JWT token if found, None otherwise
+  """
+  # Try to get from authenticated client
+  if hasattr(client, "token"):
+    return client.token
+
+  # Try from headers
+  if hasattr(client, "_headers"):
+    token = extract_jwt_from_header(client._headers)
+    if token:
+      return token
+
+  # Try from config
+  if hasattr(client, "config"):
+    config = client.config
+    if isinstance(config, dict):
+      # Direct token
+      if config.get("token"):
+        return config["token"]
+      # From headers in config
+      if config.get("headers"):
+        return extract_jwt_from_header(config["headers"])
+
+  return None