add auth to cloud cli calls

Signed-off-by: phernandez <paul@basicmachines.co>

Author: phernandez

src/basic_memory/cli/auth.py

@@ -0,0 +1,250 @@
+"""WorkOS OAuth Device Authorization for CLI."""
+
+import json
+import os
+import time
+import webbrowser
+from pathlib import Path
+
+import httpx
+from rich.console import Console
+
+from basic_memory.config import ConfigManager
+
+console = Console()
+
+
+class CLIAuth:
+    """Handles WorkOS OAuth Device Authorization for CLI tools."""
+
+    def __init__(self, client_id: str, authkit_domain: str):
+        self.client_id = client_id
+        self.authkit_domain = authkit_domain
+        app_config = ConfigManager().config
+        # Store tokens in data dir
+        self.token_file = app_config.data_dir_path / "basic-memory-cloud.json"
+
+    async def request_device_authorization(self) -> dict | None:
+        """Request device authorization from WorkOS."""
+        device_auth_url = f"{self.authkit_domain}/oauth2/device_authorization"
+
+        data = {"client_id": self.client_id, "scope": "openid profile email offline_access"}
+
+        try:
+            async with httpx.AsyncClient() as client:
+                response = await client.post(device_auth_url, data=data)
+
+                if response.status_code == 200:
+                    return response.json()
+                else:
+                    console.print(
+                        f"[red]Device authorization failed: {response.status_code} - {response.text}[/red]"
+                    )
+                    return None
+        except Exception as e:
+            console.print(f"[red]Device authorization error: {e}[/red]")
+            return None
+
+    def display_user_instructions(self, device_response: dict) -> None:
+        """Display user instructions for device authorization."""
+        user_code = device_response["user_code"]
+        verification_uri = device_response["verification_uri"]
+        verification_uri_complete = device_response.get("verification_uri_complete")
+
+        console.print("\n[bold blue]🔐 Authentication Required[/bold blue]")
+        console.print("\nTo authenticate, please visit:")
+        console.print(f"[bold cyan]{verification_uri}[/bold cyan]")
+        console.print(f"\nAnd enter this code: [bold yellow]{user_code}[/bold yellow]")
+
+        if verification_uri_complete:
+            console.print("\nOr for one-click access, visit:")
+            console.print(f"[bold green]{verification_uri_complete}[/bold green]")
+
+            # Try to open browser automatically
+            try:
+                console.print("\n[dim]Opening browser automatically...[/dim]")
+                webbrowser.open(verification_uri_complete)
+            except Exception:
+                pass  # Silently fail if browser can't be opened
+
+        console.print("\n[dim]Waiting for you to complete authentication in your browser...[/dim]")
+
+    async def poll_for_token(self, device_code: str, interval: int = 5) -> dict | None:
+        """Poll the token endpoint until user completes authentication."""
+        token_url = f"{self.authkit_domain}/oauth2/token"
+
+        data = {
+            "client_id": self.client_id,
+            "device_code": device_code,
+            "grant_type": "urn:ietf:params:oauth:grant-type:device_code",
+        }
+
+        max_attempts = 60  # 5 minutes with 5-second intervals
+        current_interval = interval
+
+        for _attempt in range(max_attempts):
+            try:
+                async with httpx.AsyncClient() as client:
+                    response = await client.post(token_url, data=data)
+
+                    if response.status_code == 200:
+                        return response.json()
+
+                    # Parse error response
+                    try:
+                        error_data = response.json()
+                        error = error_data.get("error")
+                    except Exception:
+                        error = "unknown_error"
+
+                    if error == "authorization_pending":
+                        # User hasn't completed auth yet, keep polling
+                        pass
+                    elif error == "slow_down":
+                        # Increase polling interval
+                        current_interval += 5
+                        console.print("[yellow]Slowing down polling rate...[/yellow]")
+                    elif error == "access_denied":
+                        console.print("[red]Authentication was denied by user[/red]")
+                        return None
+                    elif error == "expired_token":
+                        console.print("[red]Device code has expired. Please try again.[/red]")
+                        return None
+                    else:
+                        console.print(f"[red]Token polling error: {error}[/red]")
+                        return None
+
+            except Exception as e:
+                console.print(f"[red]Token polling request error: {e}[/red]")
+
+            # Wait before next poll
+            await self._async_sleep(current_interval)
+
+        console.print("[red]Authentication timeout. Please try again.[/red]")
+        return None
+
+    async def _async_sleep(self, seconds: int) -> None:
+        """Async sleep utility."""
+        import asyncio
+
+        await asyncio.sleep(seconds)
+
+    def save_tokens(self, tokens: dict) -> None:
+        """Save tokens to project root as .bm-auth.json."""
+        token_data = {
+            "access_token": tokens["access_token"],
+            "refresh_token": tokens.get("refresh_token"),
+            "expires_at": int(time.time()) + tokens.get("expires_in", 3600),
+            "token_type": tokens.get("token_type", "Bearer"),
+        }
+
+        with open(self.token_file, "w") as f:
+            json.dump(token_data, f, indent=2)
+
+        # Secure the token file
+        os.chmod(self.token_file, 0o600)
+
+        console.print(f"[green]✓ Tokens saved to {self.token_file}[/green]")
+
+    def load_tokens(self) -> dict | None:
+        """Load tokens from .bm-auth.json file."""
+        if not self.token_file.exists():
+            return None
+
+        try:
+            with open(self.token_file) as f:
+                return json.load(f)
+        except (OSError, json.JSONDecodeError):
+            return None
+
+    def is_token_valid(self, tokens: dict) -> bool:
+        """Check if stored token is still valid."""
+        expires_at = tokens.get("expires_at", 0)
+        # Add 60 second buffer for clock skew
+        return time.time() < (expires_at - 60)
+
+    async def refresh_token(self, refresh_token: str) -> dict | None:
+        """Refresh access token using refresh token."""
+        token_url = f"{self.authkit_domain}/oauth2/token"
+
+        data = {
+            "client_id": self.client_id,
+            "grant_type": "refresh_token",
+            "refresh_token": refresh_token,
+        }
+
+        try:
+            async with httpx.AsyncClient() as client:
+                response = await client.post(token_url, data=data)
+
+                if response.status_code == 200:
+                    return response.json()
+                else:
+                    console.print(
+                        f"[red]Token refresh failed: {response.status_code} - {response.text}[/red]"
+                    )
+                    return None
+        except Exception as e:
+            console.print(f"[red]Token refresh error: {e}[/red]")
+            return None
+
+    async def get_valid_token(self) -> str | None:
+        """Get valid access token, refresh if needed."""
+        tokens = self.load_tokens()
+        if not tokens:
+            return None
+
+        if self.is_token_valid(tokens):
+            return tokens["access_token"]
+
+        # Token expired - try to refresh if we have a refresh token
+        refresh_token = tokens.get("refresh_token")
+        if refresh_token:
+            console.print("[yellow]Access token expired, refreshing...[/yellow]")
+
+            new_tokens = await self.refresh_token(refresh_token)
+            if new_tokens:
+                # Save new tokens (may include rotated refresh token)
+                self.save_tokens(new_tokens)
+                console.print("[green]✓ Token refreshed successfully[/green]")
+                return new_tokens["access_token"]
+            else:
+                console.print("[yellow]Token refresh failed. Please run 'login' again.[/yellow]")
+                return None
+        else:
+            console.print("[yellow]No refresh token available. Please run 'login' again.[/yellow]")
+            return None
+
+    async def login(self) -> bool:
+        """Perform OAuth Device Authorization login flow."""
+        console.print("[blue]Initiating WorkOS authentication...[/blue]")
+
+        # Step 1: Request device authorization
+        device_response = await self.request_device_authorization()
+        if not device_response:
+            return False
+
+        # Step 2: Display user instructions
+        self.display_user_instructions(device_response)
+
+        # Step 3: Poll for token
+        device_code = device_response["device_code"]
+        interval = device_response.get("interval", 5)
+
+        tokens = await self.poll_for_token(device_code, interval)
+        if not tokens:
+            return False
+
+        # Step 4: Save tokens
+        self.save_tokens(tokens)
+
+        console.print("\n[green]✅ Successfully authenticated with WorkOS![/green]")
+        return True
+
+    def logout(self) -> None:
+        """Remove stored authentication tokens."""
+        if self.token_file.exists():
+            self.token_file.unlink()
+            console.print("[green]✓ Logged out successfully[/green]")
+        else:
+            console.print("[yellow]No stored authentication found[/yellow]")