basicmachines-co-basic-memory/src/basic_memory/cli/auth.py

"""WorkOS OAuth Device Authorization for CLI."""

import base64
import hashlib
import json
import os
import secrets
import time
import webbrowser

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"
        # PKCE parameters
        self.code_verifier = None
        self.code_challenge = None

    def generate_pkce_pair(self) -> tuple[str, str]:
        """Generate PKCE code verifier and challenge."""
        # Generate code verifier (43-128 characters)
        code_verifier = base64.urlsafe_b64encode(secrets.token_bytes(32)).decode("utf-8")
        code_verifier = code_verifier.rstrip("=")

        # Generate code challenge (SHA256 hash of verifier)
        challenge_bytes = hashlib.sha256(code_verifier.encode("utf-8")).digest()
        code_challenge = base64.urlsafe_b64encode(challenge_bytes).decode("utf-8")
        code_challenge = code_challenge.rstrip("=")

        return code_verifier, code_challenge

    async def request_device_authorization(self) -> dict | None:
        """Request device authorization from WorkOS with PKCE."""
        device_auth_url = f"{self.authkit_domain}/oauth2/device_authorization"

        # Generate PKCE pair
        self.code_verifier, self.code_challenge = self.generate_pkce_pair()

        data = {
            "client_id": self.client_id,
            "scope": "openid profile email offline_access",
            "code_challenge": self.code_challenge,
            "code_challenge_method": "S256",
        }

        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",
            "code_verifier": self.code_verifier,
        }

        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 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 Basic Memory Cloud![/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]")