crewAI/src/crewai/cli/authentication/main.py

import time
import webbrowser
from typing import Any, Dict, Optional

import requests
from rich.console import Console
from pydantic import BaseModel, Field


from .utils import validate_jwt_token
from crewai.cli.shared.token_manager import TokenManager
from crewai.cli.config import Settings

console = Console()


class Oauth2Settings(BaseModel):
    provider: str = Field(
        description="OAuth2 provider used for authentication (e.g., workos, okta, auth0)."
    )
    client_id: str = Field(
        description="OAuth2 client ID issued by the provider, used during authentication requests."
    )
    domain: str = Field(
        description="OAuth2 provider's domain (e.g., your-org.auth0.com) used for issuing tokens."
    )
    audience: Optional[str] = Field(
        description="OAuth2 audience value, typically used to identify the target API or resource.",
        default=None,
    )

    @classmethod
    def from_settings(cls):
        settings = Settings()

        return cls(
            provider=settings.oauth2_provider,
            domain=settings.oauth2_domain,
            client_id=settings.oauth2_client_id,
            audience=settings.oauth2_audience,
        )


class ProviderFactory:
    @classmethod
    def from_settings(cls, settings: Optional[Oauth2Settings] = None):
        settings = settings or Oauth2Settings.from_settings()

        import importlib

        module = importlib.import_module(
            f"crewai.cli.authentication.providers.{settings.provider.lower()}"
        )
        provider = getattr(module, f"{settings.provider.capitalize()}Provider")

        return provider(settings)


class AuthenticationCommand:
    def __init__(self):
        self.token_manager = TokenManager()
        self.oauth2_provider = ProviderFactory.from_settings()

    def login(self) -> None:
        """Sign up to CrewAI+"""
        console.print("Signing in to CrewAI Enterprise...\n", style="bold blue")

        device_code_data = self._get_device_code()
        self._display_auth_instructions(device_code_data)

        return self._poll_for_token(device_code_data)

    def _get_device_code(self) -> Dict[str, Any]:
        """Get the device code to authenticate the user."""

        device_code_payload = {
            "client_id": self.oauth2_provider.get_client_id(),
            "scope": "openid",
            "audience": self.oauth2_provider.get_audience(),
        }
        response = requests.post(
            url=self.oauth2_provider.get_authorize_url(),
            data=device_code_payload,
            timeout=20,
        )
        response.raise_for_status()
        return response.json()

    def _display_auth_instructions(self, device_code_data: Dict[str, str]) -> None:
        """Display the authentication instructions to the user."""
        console.print("1. Navigate to: ", device_code_data["verification_uri_complete"])
        console.print("2. Enter the following code: ", device_code_data["user_code"])
        webbrowser.open(device_code_data["verification_uri_complete"])

    def _poll_for_token(self, device_code_data: Dict[str, Any]) -> None:
        """Polls the server for the token until it is received, or max attempts are reached."""

        token_payload = {
            "grant_type": "urn:ietf:params:oauth:grant-type:device_code",
            "device_code": device_code_data["device_code"],
            "client_id": self.oauth2_provider.get_client_id(),
        }

        console.print("\nWaiting for authentication... ", style="bold blue", end="")

        attempts = 0
        while True and attempts < 10:
            response = requests.post(
                self.oauth2_provider.get_token_url(), data=token_payload, timeout=30
            )
            token_data = response.json()

            if response.status_code == 200:
                self._validate_and_save_token(token_data)

                console.print(
                    "Success!",
                    style="bold green",
                )

                self._login_to_tool_repository()

                console.print(
                    "\n[bold green]Welcome to CrewAI Enterprise![/bold green]\n"
                )
                return

            if token_data["error"] not in ("authorization_pending", "slow_down"):
                raise requests.HTTPError(token_data["error_description"])

            time.sleep(device_code_data["interval"])
            attempts += 1

        console.print(
            "Timeout: Failed to get the token. Please try again.", style="bold red"
        )

    def _validate_and_save_token(self, token_data: Dict[str, Any]) -> None:
        """Validates the JWT token and saves the token to the token manager."""

        jwt_token = token_data["access_token"]
        issuer = self.oauth2_provider.get_issuer()
        jwt_token_data = {
            "jwt_token": jwt_token,
            "jwks_url": self.oauth2_provider.get_jwks_url(),
            "issuer": issuer,
            "audience": self.oauth2_provider.get_audience(),
        }

        decoded_token = validate_jwt_token(**jwt_token_data)

        expires_at = decoded_token.get("exp", 0)
        self.token_manager.save_tokens(jwt_token, expires_at)

    def _login_to_tool_repository(self) -> None:
        """Login to the tool repository."""

        from crewai.cli.tools.main import ToolCommand

        try:
            console.print(
                "Now logging you in to the Tool Repository... ",
                style="bold blue",
                end="",
            )

            ToolCommand().login()

            console.print(
                "Success!\n",
                style="bold green",
            )

            settings = Settings()
            console.print(
                f"You are authenticated to the tool repository as [bold cyan]'{settings.org_name}'[/bold cyan] ({settings.org_uuid})",
                style="green",
            )
        except Exception:
            console.print(
                "\n[bold yellow]Warning:[/bold yellow] Authentication with the Tool Repository failed.",
                style="yellow",
            )
            console.print(
                "Other features will work normally, but you may experience limitations "
                "with downloading and publishing tools."
                "\nRun [bold]crewai login[/bold] to try logging in again.\n",
                style="yellow",
            )