Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion services/token-service/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
82 changes: 68 additions & 14 deletions src/aieng_platform_onboard/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)


Expand All @@ -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,
Expand All @@ -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.
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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.

Expand Down Expand Up @@ -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()

Expand All @@ -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
Expand All @@ -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):
Expand Down
185 changes: 171 additions & 14 deletions src/aieng_platform_onboard/utils.py
Original file line number Diff line number Diff line change
@@ -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

Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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]:
Expand Down