Accept ASTA_TOKEN env var (#28)

Author: rodneykinney

src/asta/auth/storage.py

@@ -2,9 +2,13 @@
 Secure token storage using platform-specific secure storage.
 
 Falls back to file-based storage with restrictive permissions.
+
+Both keyring and file are kept in sync to prevent token loss when
+refresh token rotation is enabled (rotated tokens are single-use).
 """
 
 import json
+import logging
 import os
 from pathlib import Path
 
@@ -14,6 +18,8 @@
 APP_NAME = "asta-cli"
 TOKEN_FILE_NAME = "tokens.json"
 
+logger = logging.getLogger(__name__)
+
 
 class TokenStorage:
     """Manages secure storage of authentication tokens."""
@@ -26,36 +32,49 @@ def __init__(self, use_keyring: bool = True):
         # Ensure config directory exists
         self.config_dir.mkdir(parents=True, exist_ok=True)
 
+    def _save_to_file(self, tokens: dict[str, str]) -> None:
+        """Save tokens to file with restrictive permissions."""
+        with open(self.token_file, "w") as f:
+            json.dump(tokens, f, indent=2)
+        os.chmod(self.token_file, 0o600)
+
     def save_tokens(self, tokens: dict[str, str]) -> None:
         """
         Save tokens securely.
 
-        Tries system keyring first, falls back to file with restrictive permissions.
+        Writes to both keyring and file to prevent token loss when
+        refresh token rotation invalidates old tokens. If keyring
+        fails, falls back to file-only storage.
         """
+        keyring_ok = False
         if self.use_keyring:
             try:
                 keyring.set_password(APP_NAME, "tokens", json.dumps(tokens))
-                return
-            except Exception:
-                # Keyring failed, fall back to file
-                pass
+                keyring_ok = True
+            except Exception as e:
+                logger.debug("Keyring save failed, using file storage: %s", e)
 
-        # File-based storage
-        with open(self.token_file, "w") as f:
-            json.dump(tokens, f, indent=2)
+        # Always save to file as well — ensures consistency if keyring
+        # becomes unavailable on a future read after token rotation.
+        self._save_to_file(tokens)
 
-        # Set restrictive permissions (owner read/write only)
-        os.chmod(self.token_file, 0o600)
+        if not keyring_ok and self.use_keyring:
+            logger.debug("Tokens saved to file only (keyring unavailable)")
 
     def load_tokens(self) -> dict[str, str] | None:
-        """Load tokens from storage."""
+        """
+        Load tokens from storage.
+
+        Prefers keyring, falls back to file. Both should be in sync
+        since save_tokens writes to both.
+        """
         if self.use_keyring:
             try:
                 token_json = keyring.get_password(APP_NAME, "tokens")
                 if token_json:
                     return json.loads(token_json)
-            except Exception:
-                pass
+            except Exception as e:
+                logger.debug("Keyring load failed, trying file: %s", e)
 
         # Try file-based storage
         if self.token_file.exists():
@@ -68,7 +87,7 @@ def load_tokens(self) -> dict[str, str] | None:
         return None
 
     def delete_tokens(self) -> None:
-        """Delete stored tokens."""
+        """Delete stored tokens from all backends."""
         if self.use_keyring:
             try:
                 keyring.delete_password(APP_NAME, "tokens")

src/asta/auth/token_manager.py

@@ -2,7 +2,9 @@
 High-level token management with automatic refresh.
 """
 
+import asyncio
 import json
+import logging
 import time
 import urllib.error
 import urllib.request
@@ -16,6 +18,12 @@
 from .exceptions import AuthenticationError
 from .storage import TokenStorage
 
+logger = logging.getLogger(__name__)
+
+# Retry settings for token refresh
+REFRESH_MAX_RETRIES = 3
+REFRESH_RETRY_DELAY = 2  # seconds, doubles each retry
+
 
 class TokenManager:
     """Manages authentication tokens with automatic refresh."""
@@ -121,26 +129,58 @@ async def get_valid_access_token(self) -> str:
                     "Please re-authenticate with 'asta auth login'."
                 )
 
-            try:
-                token_response = await self.flow.refresh_token(refresh_token)
-
-                # Save new tokens
-                self.storage.save_tokens(
-                    {
-                        "access_token": token_response.access_token,
-                        "refresh_token": token_response.refresh_token or refresh_token,
-                        "id_token": token_response.id_token,
-                        "expires_at": int(time.time()) + token_response.expires_in,
-                    }
-                )
-
-                return token_response.access_token
+            # Retry refresh with exponential backoff to handle transient errors
+            last_error = None
+            for attempt in range(REFRESH_MAX_RETRIES):
+                try:
+                    if attempt > 0:
+                        delay = REFRESH_RETRY_DELAY * (2 ** (attempt - 1))
+                        logger.info(
+                            "Retrying token refresh (attempt %d/%d) after %ds",
+                            attempt + 1,
+                            REFRESH_MAX_RETRIES,
+                            delay,
+                        )
+                        await asyncio.sleep(delay)
+
+                    token_response = await self.flow.refresh_token(refresh_token)
+
+                    # Save new tokens
+                    self.storage.save_tokens(
+                        {
+                            "access_token": token_response.access_token,
+                            "refresh_token": token_response.refresh_token
+                            or refresh_token,
+                            "id_token": token_response.id_token,
+                            "expires_at": int(time.time()) + token_response.expires_in,
+                        }
+                    )
+
+                    if attempt > 0:
+                        logger.info(
+                            "Token refresh succeeded on attempt %d", attempt + 1
+                        )
+
+                    return token_response.access_token
+
+                except AuthenticationError:
+                    # Auth errors (invalid grant, denied) won't be fixed by retry
+                    raise
+
+                except Exception as e:
+                    last_error = e
+                    logger.warning(
+                        "Token refresh attempt %d/%d failed: %s",
+                        attempt + 1,
+                        REFRESH_MAX_RETRIES,
+                        e,
+                    )
 
-            except Exception as e:
-                raise AuthenticationError(
-                    f"Token refresh failed: {e}. "
-                    f"Please re-authenticate with 'asta auth login'."
-                )
+            raise AuthenticationError(
+                f"Token refresh failed after {REFRESH_MAX_RETRIES} attempts: "
+                f"{last_error}. "
+                f"Please re-authenticate with 'asta auth login'."
+            )
 
         return access_token
 

src/asta/commands/auth.py

@@ -189,17 +189,37 @@ def status():
 
 @auth.command(name="print-token")
 @click.option("--raw", is_flag=True, help="Print raw base64-encoded token")
-def print_token(raw):
+@click.option(
+    "--refresh", is_flag=True, help="Refresh the token if expired before printing"
+)
+def print_token(raw, refresh):
     """Print the stored access token."""
-    storage = TokenStorage()
-    tokens = storage.load_tokens()
+    if refresh:
+        # Go through TokenManager to trigger auto-refresh if needed
+        from asta.auth.exceptions import AuthenticationError
 
-    if not tokens or not tokens.get("access_token"):
-        console.print("❌ [red]No token found[/red]")
-        console.print("   Run [cyan]asta auth login[/cyan] to authenticate")
-        raise click.Abort()
+        try:
+            settings = get_auth_settings()
+            manager = TokenManager(
+                auth0_domain=settings.auth0_domain,
+                client_id=settings.auth0_client_id,
+                audience=settings.auth0_audience,
+                gateway_url=settings.gateway_url,
+            )
+            access_token = asyncio.run(manager.get_valid_access_token())
+        except AuthenticationError as e:
+            console.print(f"❌ [red]{e}[/red]")
+            raise click.Abort()
+    else:
+        storage = TokenStorage()
+        tokens = storage.load_tokens()
+
+        if not tokens or not tokens.get("access_token"):
+            console.print("❌ [red]No token found[/red]")
+            console.print("   Run [cyan]asta auth login[/cyan] to authenticate")
+            raise click.Abort()
 
-    access_token = tokens["access_token"]
+        access_token = tokens["access_token"]
 
     if raw:
         # Print raw base64-encoded token

src/asta/utils/auth_helper.py

@@ -1,6 +1,7 @@
 """Helper utilities for authentication in CLI commands."""
 
 import asyncio
+import os
 
 from asta.auth.exceptions import AuthenticationError
 from asta.auth.token_manager import TokenManager
@@ -22,6 +23,9 @@ def get_access_token() -> str:
         This function handles token refresh automatically. If the token
         is expired and a refresh token is available, it will be refreshed.
     """
+    if token := os.environ.get("ASTA_TOKEN"):
+        return token
+
     try:
         # Get auth settings from config
         settings = get_auth_settings()
@@ -39,10 +43,8 @@ def get_access_token() -> str:
         return asyncio.run(manager.get_valid_access_token())
 
     except AuthenticationError:
-        # Re-raise with user-friendly message
-        raise AuthenticationError(
-            "Not authenticated. Please run 'asta auth login' to authenticate."
-        )
+        # Re-raise with original message preserved (includes refresh failure details)
+        raise
     except Exception as e:
         # Convert any other errors to AuthenticationError with helpful message
         raise AuthenticationError(

tests/test_auth_helper.py

@@ -54,13 +54,11 @@ def test_get_access_token_without_token(self):
                 )
                 MockManager.return_value = mock_manager
 
-                with pytest.raises(
-                    AuthenticationError, match="Please run 'asta auth login'"
-                ):
+                with pytest.raises(AuthenticationError):
                     get_access_token()
 
     def test_get_access_token_refresh_failure(self):
-        """Test that token refresh failures are handled with helpful message."""
+        """Test that token refresh failures are re-raised with original message."""
         from asta.auth.exceptions import AuthenticationError
         from asta.utils.auth_helper import get_access_token
 
@@ -81,9 +79,7 @@ def test_get_access_token_refresh_failure(self):
                 )
                 MockManager.return_value = mock_manager
 
-                with pytest.raises(
-                    AuthenticationError, match="Please run 'asta auth login'"
-                ):
+                with pytest.raises(AuthenticationError):
                     get_access_token()
 
     def test_get_access_token_other_error(self):