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

"""OAuth management commands."""

import typer
from typing import Optional
from pydantic import AnyHttpUrl

from basic_memory.cli.app import app
from basic_memory.mcp.auth_provider import BasicMemoryOAuthProvider
from mcp.shared.auth import OAuthClientInformationFull


auth_app = typer.Typer(help="OAuth client management commands")
app.add_typer(auth_app, name="auth")


@auth_app.command()
def register_client(
    client_id: Optional[str] = typer.Option(
        None, help="Client ID (auto-generated if not provided)"
    ),
    client_secret: Optional[str] = typer.Option(
        None, help="Client secret (auto-generated if not provided)"
    ),
    issuer_url: str = typer.Option("http://localhost:8000", help="OAuth issuer URL"),
):
    """Register a new OAuth client for Basic Memory MCP server."""

    # Create provider instance
    provider = BasicMemoryOAuthProvider(issuer_url=issuer_url)

    # Create client info with required redirect_uris
    client_info = OAuthClientInformationFull(
        client_id=client_id or "",  # Provider will generate if empty
        client_secret=client_secret or "",  # Provider will generate if empty
        redirect_uris=[AnyHttpUrl("http://localhost:8000/callback")],  # Default redirect URI
        client_name="Basic Memory OAuth Client",
        grant_types=["authorization_code", "refresh_token"],
    )

    # Register the client
    import asyncio

    asyncio.run(provider.register_client(client_info))

    typer.echo("Client registered successfully!")
    typer.echo(f"Client ID: {client_info.client_id}")
    typer.echo(f"Client Secret: {client_info.client_secret}")
    typer.echo("\nSave these credentials securely - the client secret cannot be retrieved later.")


@auth_app.command()
def test_auth(
    issuer_url: str = typer.Option("http://localhost:8000", help="OAuth issuer URL"),
):
    """Test OAuth authentication flow.

    IMPORTANT: Use the same FASTMCP_AUTH_SECRET_KEY environment variable
    as your MCP server for tokens to validate correctly.
    """

    import asyncio
    import secrets
    from mcp.server.auth.provider import AuthorizationParams
    from pydantic import AnyHttpUrl

    async def test_flow():
        # Create provider with same secret key as server
        provider = BasicMemoryOAuthProvider(issuer_url=issuer_url)

        # Register a test client
        client_info = OAuthClientInformationFull(
            client_id=secrets.token_urlsafe(16),
            client_secret=secrets.token_urlsafe(32),
            redirect_uris=[AnyHttpUrl("http://localhost:8000/callback")],
            client_name="Test OAuth Client",
            grant_types=["authorization_code", "refresh_token"],
        )
        await provider.register_client(client_info)
        typer.echo(f"Registered test client: {client_info.client_id}")

        # Get the client
        client = await provider.get_client(client_info.client_id)
        if not client:
            typer.echo("Error: Client not found after registration", err=True)
            return

        # Create authorization request
        auth_params = AuthorizationParams(
            state="test-state",
            scopes=["read", "write"],
            code_challenge="test-challenge",
            redirect_uri=AnyHttpUrl("http://localhost:8000/callback"),
            redirect_uri_provided_explicitly=True,
        )

        # Get authorization URL
        auth_url = await provider.authorize(client, auth_params)
        typer.echo(f"Authorization URL: {auth_url}")

        # Extract auth code from URL
        from urllib.parse import urlparse, parse_qs

        parsed = urlparse(auth_url)
        params = parse_qs(parsed.query)
        auth_code = params.get("code", [None])[0]

        if not auth_code:
            typer.echo("Error: No authorization code in URL", err=True)
            return

        # Load the authorization code
        code_obj = await provider.load_authorization_code(client, auth_code)
        if not code_obj:
            typer.echo("Error: Invalid authorization code", err=True)
            return

        # Exchange for tokens
        token = await provider.exchange_authorization_code(client, code_obj)
        typer.echo(f"Access token: {token.access_token}")
        typer.echo(f"Refresh token: {token.refresh_token}")
        typer.echo(f"Expires in: {token.expires_in} seconds")

        # Validate access token
        access_token_obj = await provider.load_access_token(token.access_token)
        if access_token_obj:
            typer.echo("Access token validated successfully!")
            typer.echo(f"Client ID: {access_token_obj.client_id}")
            typer.echo(f"Scopes: {access_token_obj.scopes}")
        else:
            typer.echo("Error: Invalid access token", err=True)

    asyncio.run(test_flow())


if __name__ == "__main__":
    auth_app()