authentication.py

#!/usr/bin/env python3
"""Example of connecting to MCP servers with authentication.

This script demonstrates:
1. Bearer token authentication with configurable environment variables
2. Direct token passing
3. Using the MaggClient for automatic authentication
"""
import argparse
import asyncio
import logging
import os
import sys
from typing import Any

from fastmcp import Client
from fastmcp.client import BearerAuth
from magg.client import MaggClient


async def bearer_auth(args: argparse.Namespace) -> None:
    """Test bearer token authentication with configurable options."""
    jwt = args.token
    if not jwt:
        jwt = os.environ.get(args.env_var)
        if not jwt:
            print(f"Error: No JWT token provided. Set {args.env_var} or use --token", file=sys.stderr)
            sys.exit(1)

    print(f"Connecting to {args.url} with bearer token authentication...")

    if args.magg:
        # Use MaggClient
        print(f"Using MaggClient{' with provided token' if args.token else f' (loading from {args.env_var})'}")

        # If using custom env var, set MAGG_JWT for MaggClient
        if args.env_var != "MAGG_JWT" and not args.token:
            os.environ["MAGG_JWT"] = jwt

        if args.token:
            auth = BearerAuth(args.token)
            client = MaggClient(args.url, auth=auth)
        else:
            # Let MaggClient handle auth from MAGG_JWT env var
            client = MaggClient(args.url)

        async with client:
            print(f"Transparent proxy mode: {client._transparent}")

            # Test the connection
            await check_client_capabilities(client)

            # Test proxy-specific functionality if available
            try:
                print("\nTesting proxy functionality...")
                tools = await client.proxy("tool", "list")
                print(f"Proxy list returned {len(tools)} items")
            except Exception as e:
                print(f"Proxy functionality not available: {e}")
    else:
        # Use regular FastMCP Client
        print(f"Using FastMCP Client")
        auth = BearerAuth(jwt)

        async with Client(args.url, auth=auth) as client:
            # Test the connection
            await check_client_capabilities(client)


async def check_client_capabilities(client: Any) -> None:
    """Test basic MCP client capabilities."""
    # List available tools
    tools = await client.list_tools()
    print(f"\nFound {len(tools)} tools:")
    for tool in tools[:5]:  # Show first 5 tools
        print(f"  - {tool.name}: {tool.description}")
    if len(tools) > 5:
        print(f"  ... and {len(tools) - 5} more")

    # List resources
    try:
        resources = await client.list_resources()
        print(f"\nFound {len(resources)} resources")
    except Exception as e:
        print(f"\nResource listing not available: {e}")

    # List prompts
    try:
        prompts = await client.list_prompts()
        print(f"\nFound {len(prompts)} prompts")
    except Exception as e:
        print(f"\nPrompt listing not available: {e}")

    # Call a simple tool if available
    if tools:
        # Look for a simple info/status tool
        info_tools = [t for t in tools if "status" in t.name.lower() or "info" in t.name.lower()]
        if info_tools:
            tool = info_tools[0]
            print(f"\nCalling tool: {tool.name}")
            try:
                result = await client.call_tool(tool.name)
                print(f"Result: {result[:200]}..." if len(str(result)) > 200 else f"Result: {result}")
            except Exception as e:
                print(f"Tool call failed: {e}")


def create_parser() -> argparse.ArgumentParser:
    """Create command line parser."""
    parser = argparse.ArgumentParser(
        description="Test MCP server authentication",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Test with bearer token from environment
  %(prog)s bearer

  # Test with custom URL
  %(prog)s bearer http://localhost:9000

  # Test with custom environment variable
  %(prog)s bearer --env-var MY_TOKEN

  # Test with direct token
  %(prog)s bearer --token "eyJ..."

  # Test using MaggClient (auto-loads from MAGG_JWT)
  %(prog)s bearer --magg

  # Test MaggClient with direct token
  %(prog)s bearer --magg --token "eyJ..."
"""
    )

    parser.add_argument(
        "--debug",
        action="store_true",
        help="Enable debug output"
    )

    subparsers = parser.add_subparsers(dest="auth_type", help="Authentication type")

    # Bearer authentication subcommand
    bearer_parser = subparsers.add_parser(
        "bearer",
        help="Test bearer token authentication"
    )
    bearer_parser.add_argument(
        "url",
        nargs="?",
        default="http://localhost:8000/mcp",
        help="MCP server URL (default: http://localhost:8000/mcp)"
    )
    bearer_parser.add_argument(
        "--env-var",
        default="MAGG_JWT",
        help="Environment variable name for JWT (default: MAGG_JWT)"
    )
    bearer_parser.add_argument(
        "--token",
        help="JWT token (overrides environment variable)"
    )
    bearer_parser.add_argument(
        "--magg",
        action="store_true",
        help="Use MaggClient instead of regular FastMCP Client"
    )

    return parser


async def main():
    """Main entry point."""
    parser = create_parser()
    args = parser.parse_args()

    if args.debug:
        logging.basicConfig(level=logging.DEBUG)

    # Check if subcommand was provided
    if not args.auth_type:
        parser.print_help()
        sys.exit(1)

    try:
        if args.auth_type == "bearer":
            await bearer_auth(args)
    except KeyboardInterrupt:
        print("\nInterrupted by user")
        sys.exit(130)
    except Exception as e:
        print(f"\nError: {e}", file=sys.stderr)
        if args.debug:
            import traceback
            traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    asyncio.run(main())