Merge pull request #17 from VectorInstitute/use_firestore_token

Update cli to use firestore token service

Author: amrit110

services/token-service/main.py

@@ -14,7 +14,7 @@
 import jwt
 from firebase_admin import auth, credentials
 from flask import Flask, request
-from google.cloud import firestore  # type: ignore[attr-defined]
+from google.cloud import firestore
 
 
 # Initialize Flask app

src/aieng_platform_onboard/cli.py

@@ -17,15 +17,17 @@
 from rich.panel import Panel
 
 from aieng_platform_onboard.utils import (
+    check_onboarded_status,
     console,
     create_env_file,
-    fetch_token_from_secret_manager,
+    fetch_token_from_service,
     get_github_user,
     get_global_keys,
     get_participant_data,
     get_team_data,
     initialize_firestore_with_token,
     update_onboarded_status,
+    validate_env_file,
 )
 
 
@@ -45,7 +47,7 @@ def run_integration_test(test_script: Path) -> tuple[bool, str]:
     """
     try:
         result = subprocess.run(
-            [sys.executable, str(test_script)],
+            [sys.executable, "-m", "pytest", str(test_script)],
             check=False,
             capture_output=True,
             text=True,
@@ -63,7 +65,9 @@ def run_integration_test(test_script: Path) -> tuple[bool, str]:
 
 
 def _authenticate_and_connect(
-    bootcamp_name: str, gcp_project: str, firebase_api_key: str | None
+    bootcamp_name: str,
+    gcp_project: str,
+    firebase_api_key: str | None,
 ) -> tuple[str, Any] | tuple[None, None]:
     """
     Authenticate participant and connect to Firestore.
@@ -97,20 +101,17 @@ def _authenticate_and_connect(
 
     # Step 2: Fetch authentication token
     console.print("[bold]Step 2: Fetch Authentication Token[/bold]")
-    console.print("[cyan]Fetching token from Secret Manager...[/cyan]")
-
-    success, token, error = fetch_token_from_secret_manager(
-        github_user, bootcamp_name, gcp_project
-    )
+    console.print("[cyan]Fetching fresh token from service...[/cyan]")
+    success, token, error = fetch_token_from_service(github_user)
 
     if not success or not token:
         console.print(
             f"[red]✗ Failed to fetch authentication token:[/red]\n"
             f"  {error}\n\n"
             "[yellow]Possible reasons:[/yellow]\n"
-            "  • Tokens not yet generated by admin\n"
-            "  • Incorrect bootcamp name\n"
-            "  • Missing Secret Manager permissions\n\n"
+            "  • Token service not deployed or misconfigured\n"
+            "  • Missing permissions\n"
+            "  • Participant not found in Firestore\n\n"
             "[dim]Contact bootcamp admin for assistance[/dim]"
         )
         return None, None
@@ -323,7 +324,7 @@ def _run_tests_and_finalize(
     return True
 
 
-def main() -> int:
+def main() -> int:  # noqa: PLR0911
     """
     Onboard bootcamp participants with team-specific API keys.
 
@@ -370,6 +371,11 @@ def main() -> int:
         type=str,
         help="Firebase Web API key for token exchange (or set FIREBASE_WEB_API_KEY env var)",
     )
+    parser.add_argument(
+        "--force",
+        action="store_true",
+        help="Force re-onboarding even if already onboarded",
+    )
 
     args = parser.parse_args()
 
@@ -382,13 +388,60 @@ def main() -> int:
         )
     )
 
+    # Check if .env file already exists and is complete
+    output_path = Path(args.output_dir) / ".env"
+    if output_path.exists():
+        console.print("\n[bold]Checking existing .env file...[/bold]")
+        is_complete, missing = validate_env_file(output_path)
+
+        if is_complete:
+            console.print(
+                "[green]✓ .env file is already complete with all required keys[/green]\n"
+                f"[dim]Location: {output_path}[/dim]\n"
+            )
+            console.print(
+                Panel.fit(
+                    "[green bold]✓ ALREADY ONBOARDED[/green bold]\n\n"
+                    "Your environment is already set up. No action needed.\n\n"
+                    "[dim]To re-onboard, delete the .env file and run again.[/dim]",
+                    border_style="green",
+                    title="Success",
+                )
+            )
+            return 0
+        console.print(
+            f"[yellow]⚠ .env file exists but is incomplete[/yellow]\n"
+            f"[dim]Missing keys: {', '.join(missing)}[/dim]\n"
+            "[cyan]Continuing with onboarding to regenerate .env file...[/cyan]\n"
+        )
+
     # Authenticate and connect
     github_user, db = _authenticate_and_connect(
         args.bootcamp_name, args.gcp_project, args.firebase_api_key
     )
     if not github_user or not db:
         return 1
 
+    # Check if already onboarded in Firestore
+    console.print("[bold]Checking onboarded status...[/bold]")
+    success_check, is_onboarded = check_onboarded_status(db, github_user)
+
+    skip_onboarding = (
+        success_check and is_onboarded and not args.force and output_path.exists()
+    )
+    if skip_onboarding:
+        console.print(
+            "[green]✓ You are already marked as onboarded in Firestore[/green]\n"
+            "[yellow]Use --force to re-onboard anyway[/yellow]\n"
+        )
+        return 0
+
+    if success_check and is_onboarded and not args.force:
+        console.print(
+            "[green]✓ You are already marked as onboarded in Firestore[/green]\n"
+            "[cyan]But .env file is missing, continuing to create it...[/cyan]\n"
+        )
+
     # Fetch data
     participant_data, team_data, global_keys = _fetch_participant_and_team_data(
         db, github_user
@@ -397,9 +450,10 @@ def main() -> int:
         return 1
 
     # Setup environment
-    output_path = _setup_environment(args.output_dir, team_data, global_keys)
-    if not output_path:
+    env_output_path = _setup_environment(args.output_dir, team_data, global_keys)
+    if not env_output_path:
         return 1
+    output_path = env_output_path
 
     # Run tests and finalize
     if not _run_tests_and_finalize(db, github_user, args.skip_test, args.test_script):

src/aieng_platform_onboard/utils.py

@@ -1,12 +1,16 @@
 """Shared utilities for participant onboarding scripts."""
 
 import os
+import subprocess
+import time
 from datetime import datetime, timezone
 from pathlib import Path
 from typing import Any
 
+import google.auth
 import requests
-from google.cloud import firestore, secretmanager  # type: ignore[attr-defined]
+from google.auth import jwt as google_jwt
+from google.cloud import firestore
 from google.oauth2 import credentials as oauth2_credentials
 from rich.console import Console
 
@@ -50,38 +54,111 @@ def get_github_user() -> str | None:
     return github_user
 
 
-def fetch_token_from_secret_manager(
-    github_handle: str, bootcamp_name: str, project_id: str
+def fetch_token_from_service(  # noqa: PLR0911
+    github_handle: str, token_service_url: str | None = None
 ) -> tuple[bool, str | None, str | None]:
     """
-    Fetch Firebase token from GCP Secret Manager.
+    Fetch fresh Firebase token from Cloud Run token service.
+
+    This generates tokens on-demand using the workspace service account identity.
+    Tokens are always fresh (< 1 hour old).
 
     Parameters
     ----------
     github_handle : str
         GitHub handle of the participant.
-    bootcamp_name : str
-        Name of the bootcamp.
-    project_id : str
-        GCP project ID.
+    token_service_url : str | None, optional
+        Token service URL. If not provided, reads from TOKEN_SERVICE_URL env var
+        or .token-service-url file.
 
     Returns
     -------
     tuple[bool, str | None, str | None]
         Tuple of (success, token_value, error_message).
     """
     try:
-        client = secretmanager.SecretManagerServiceClient()
-        secret_name = f"{bootcamp_name}-token-{github_handle}"
-        name = f"projects/{project_id}/secrets/{secret_name}/versions/latest"
+        # Get token service URL
+        if not token_service_url:
+            token_service_url = os.environ.get("TOKEN_SERVICE_URL")
+
+        if not token_service_url:
+            # Try reading from config file
+            config_file = Path.home() / ".token-service-url"
+            if config_file.exists():
+                token_service_url = config_file.read_text().strip()
+
+        if not token_service_url:
+            return (
+                False,
+                None,
+                (
+                    "Token service URL not found. Set TOKEN_SERVICE_URL environment "
+                    "variable or ensure .token-service-url file exists."
+                ),
+            )
+
+        # Get identity token for authentication
+        # Try to get identity token
+        # In workspaces, this will use the compute service account
+        credentials, project_id = google.auth.default()  # type: ignore[no-untyped-call]
+
+        # Request an identity token instead of an access token
+        # This is required for Cloud Run authentication
+        if hasattr(credentials, "signer"):
+            # Service account - can create identity tokens
+            now = int(time.time())
+            payload = {
+                "iss": credentials.service_account_email,
+                "sub": credentials.service_account_email,
+                "aud": token_service_url,
+                "iat": now,
+                "exp": now + 3600,
+            }
+            id_token = google_jwt.encode(credentials.signer, payload)  # type: ignore[no-untyped-call]
+        else:
+            # User credentials or other - use gcloud
+            try:
+                result = subprocess.run(
+                    ["gcloud", "auth", "print-identity-token"],
+                    check=False,
+                    capture_output=True,
+                    text=True,
+                    timeout=10,
+                )
+                if result.returncode != 0:
+                    return False, None, f"Failed to get identity token: {result.stderr}"
+                id_token = result.stdout.strip()
+            except Exception as e:
+                return False, None, f"Failed to run gcloud: {str(e)}"
+
+        if not id_token:
+            return False, None, "Failed to get identity token for authentication"
+
+        # Call token service
+        url = f"{token_service_url.rstrip('/')}/generate-token"
+        headers = {
+            "Authorization": f"Bearer {id_token}",
+            "Content-Type": "application/json",
+        }
+        payload = {"github_handle": github_handle}
+
+        response = requests.post(url, json=payload, headers=headers, timeout=30)
+
+        if response.status_code != 200:
+            error_data = response.json() if response.content else {}
+            error_msg = error_data.get("message", f"HTTP {response.status_code}")
+            return False, None, f"Token service error: {error_msg}"
+
+        data = response.json()
+        token = data.get("token")
 
-        response = client.access_secret_version(request={"name": name})
-        token = response.payload.data.decode("UTF-8")
+        if not token:
+            return False, None, "No token in service response"
 
         return True, token, None
 
     except Exception as e:
-        return False, None, str(e)
+        return False, None, f"Failed to fetch token from service: {str(e)}"
 
 
 def exchange_custom_token_for_id_token(
@@ -359,6 +436,86 @@ def create_env_file(
         return False
 
 
+def check_onboarded_status(
+    db: firestore.Client, github_handle: str
+) -> tuple[bool, bool]:
+    """
+    Check if participant is already onboarded.
+
+    Parameters
+    ----------
+    db : firestore.Client
+        Firestore client instance.
+    github_handle : str
+        GitHub handle of the participant.
+
+    Returns
+    -------
+    tuple[bool, bool]
+        Tuple of (success, is_onboarded).
+    """
+    try:
+        doc_ref = db.collection("participants").document(github_handle)
+        doc = doc_ref.get()
+
+        if not doc.exists:
+            return True, False
+
+        data = doc.to_dict()
+        is_onboarded = data.get("onboarded", False) if data else False
+        return True, is_onboarded
+
+    except Exception as e:
+        console.print(f"[yellow]⚠ Failed to check onboarded status:[/yellow] {e}")
+        return False, False
+
+
+def validate_env_file(env_path: Path) -> tuple[bool, list[str]]:
+    """
+    Validate if .env file exists and contains all required keys.
+
+    Parameters
+    ----------
+    env_path : Path
+        Path to the .env file.
+
+    Returns
+    -------
+    tuple[bool, list[str]]
+        Tuple of (is_complete, missing_keys).
+    """
+    if not env_path.exists():
+        return False, ["File does not exist"]
+
+    required_keys = [
+        "OPENAI_API_KEY",
+        "EMBEDDING_BASE_URL",
+        "EMBEDDING_API_KEY",
+        "LANGFUSE_SECRET_KEY",
+        "LANGFUSE_PUBLIC_KEY",
+        "LANGFUSE_HOST",
+        "WEAVIATE_HTTP_HOST",
+        "WEAVIATE_GRPC_HOST",
+        "WEAVIATE_API_KEY",
+    ]
+
+    try:
+        with open(env_path) as f:
+            content = f.read()
+
+        missing_keys = []
+        for key in required_keys:
+            # Check if key exists and has a non-empty value
+            if f'{key}=""' in content or key not in content:
+                missing_keys.append(key)
+
+        return len(missing_keys) == 0, missing_keys
+
+    except Exception as e:
+        console.print(f"[yellow]⚠ Failed to validate .env file:[/yellow] {e}")
+        return False, [str(e)]
+
+
 def update_onboarded_status(
     db: firestore.Client, github_handle: str
 ) -> tuple[bool, str | None]: