diff --git a/codecarbon/cli/main.py b/codecarbon/cli/main.py index 6f28c2309..516498f66 100644 --- a/codecarbon/cli/main.py +++ b/codecarbon/cli/main.py @@ -1,15 +1,20 @@ +import json import os import signal import sys import time +import webbrowser +from http.server import BaseHTTPRequestHandler, HTTPServer from pathlib import Path from typing import Optional +from urllib.parse import parse_qs, urlparse import questionary import requests import typer -from fief_client import Fief -from fief_client.integrations.cli import FiefAuth +from authlib.common.security import generate_token +from authlib.integrations.requests_client import OAuth2Session +from authlib.oauth2.rfc7636 import create_s256_code_challenge from rich import print from rich.prompt import Confirm from typing_extensions import Annotated @@ -22,7 +27,6 @@ get_existing_local_exp_id, overwrite_local_config, ) -from codecarbon.cli.monitor import run_and_monitor from codecarbon.core.api_client import ApiClient, get_datetime_with_timezone from codecarbon.core.schemas import ExperimentCreate, OrganizationCreate, ProjectCreate from codecarbon.emissions_tracker import EmissionsTracker, OfflineEmissionsTracker @@ -31,9 +35,16 @@ "AUTH_CLIENT_ID", "jsUPWIcUECQFE_ouanUuVhXx52TTjEVcVNNtNGeyAtU", ) -AUTH_SERVER_URL = os.environ.get( - "AUTH_SERVER_URL", "https://auth.codecarbon.io/codecarbon" +AUTH_SERVER_WELL_KNOWN = os.environ.get( + "AUTH_SERVER_WELL_KNOWN", + "https://auth.codecarbon.io/codecarbon/.well-known/openid-configuration", ) + +AUTH_CLIENT_ID = "codecarbon-cli" +AUTH_SERVER_WELL_KNOWN = ( + "https://kc.lloret.ovh/realms/codecarbon/.well-known/openid-configuration" +) + API_URL = os.environ.get("API_URL", "https://dashboard.codecarbon.io/api") DEFAULT_PROJECT_ID = "e60afa92-17b7-4720-91a0-1ae91e409ba1" @@ -115,17 +126,115 @@ def show_config(path: Path = Path("./.codecarbon.config")) -> None: ) -def get_fief_auth(): - fief = Fief(AUTH_SERVER_URL, AUTH_CLIENT_ID) - fief_auth = FiefAuth(fief, "./credentials.json") - return fief_auth +_REDIRECT_PORT = 8090 +_REDIRECT_URI = f"http://localhost:{_REDIRECT_PORT}/callback" +_CREDENTIALS_FILE = Path("./credentials.json") + + +class _CallbackHandler(BaseHTTPRequestHandler): + """HTTP handler that captures the OAuth2 authorization callback.""" + + callback_url = None + error = None + + def do_GET(self): + _CallbackHandler.callback_url = f"http://localhost:{_REDIRECT_PORT}{self.path}" + parsed = urlparse(self.path) + params = parse_qs(parsed.query) + + if "error" in params: + _CallbackHandler.error = params["error"][0] + self.send_response(400) + self.send_header("Content-Type", "text/html") + self.end_headers() + msg = params.get("error_description", [params["error"][0]])[0] + self.wfile.write( + f"

Login failed

{msg}

".encode() + ) + else: + self.send_response(200) + self.send_header("Content-Type", "text/html") + self.end_headers() + self.wfile.write( + b"

Login successful!

" + b"

You can close this window.

" + ) + + def log_message(self, format, *args): + pass # Suppress server logs + + +def _discover_endpoints(): + """Fetch OpenID Connect discovery document.""" + resp = requests.get(AUTH_SERVER_WELL_KNOWN) + resp.raise_for_status() + return resp.json() + + +def _authorize(): + """Run the OAuth2 Authorization Code flow with PKCE.""" + discovery = _discover_endpoints() + + session = OAuth2Session( + client_id=AUTH_CLIENT_ID, + redirect_uri=_REDIRECT_URI, + scope="openid offline_access", + token_endpoint_auth_method="none", + ) + + code_verifier = generate_token(48) + code_challenge = create_s256_code_challenge(code_verifier) + + uri, state = session.create_authorization_url( + discovery["authorization_endpoint"], + code_challenge=code_challenge, + code_challenge_method="S256", + ) + + # Reset handler state + _CallbackHandler.callback_url = None + _CallbackHandler.error = None + + server = HTTPServer(("localhost", _REDIRECT_PORT), _CallbackHandler) + + print("Opening browser for authentication...") + webbrowser.open(uri) + + server.handle_request() + server.server_close() + + if _CallbackHandler.error: + raise ValueError(f"Authorization failed: {_CallbackHandler.error}") + + if not _CallbackHandler.callback_url: + raise ValueError("Authorization failed: no callback received") + + token = session.fetch_token( + discovery["token_endpoint"], + authorization_response=_CallbackHandler.callback_url, + code_verifier=code_verifier, + ) + + _save_credentials(token) + return token + + +def _save_credentials(tokens): + """Save OAuth tokens to credentials file.""" + with open(_CREDENTIALS_FILE, "w") as f: + json.dump(tokens, f) + + +def _load_credentials(): + """Load OAuth tokens from credentials file.""" + with open(_CREDENTIALS_FILE, "r") as f: + return json.load(f) def _get_access_token(): try: - access_token_info = get_fief_auth().access_token_info() - access_token = access_token_info["access_token"] - return access_token + creds = _load_credentials() + return creds["access_token"] except Exception as e: raise ValueError( f"Not able to retrieve the access token, please run `codecarbon login` first! (error: {e})" @@ -133,8 +242,8 @@ def _get_access_token(): def _get_id_token(): - id_token = get_fief_auth()._tokens["id_token"] - return id_token + creds = _load_credentials() + return creds["id_token"] @codecarbon.command( @@ -152,7 +261,7 @@ def api_get(): @codecarbon.command("login", short_help="Login to CodeCarbon") def login(): - get_fief_auth().authorize() + _authorize() api = ApiClient(endpoint_url=API_URL) # TODO: get endpoint from config access_token = _get_access_token() api.set_access_token(access_token) @@ -346,12 +455,11 @@ def config(): context_settings={"allow_extra_args": True, "ignore_unknown_options": True}, ) def monitor( - ctx: typer.Context, measure_power_secs: Annotated[ - int, typer.Option(help="Interval between two measures.") + int, typer.Argument(help="Interval between two measures.") ] = 10, api_call_interval: Annotated[ - int, typer.Option(help="Number of measures between API calls.") + int, typer.Argument(help="Number of measures between API calls.") ] = 30, api: Annotated[ bool, typer.Option(help="Choose to call Code Carbon API or not") @@ -365,25 +473,17 @@ def monitor( ] = None, ): """Monitor your machine's carbon emissions.""" - - # Shared tracker args so monitor and run_and_monitor behave the same - tracker_args = { - "measure_power_secs": measure_power_secs, - "api_call_interval": api_call_interval, - } - # Set up the tracker arguments based on mode (offline vs online) and validate required args for each mode if offline: if not country_iso_code: print( "ERROR: country_iso_code is required for offline mode", file=sys.stderr ) raise typer.Exit(1) - - tracker_args = { - **tracker_args, - "country_iso_code": country_iso_code, - "region": region, - } + tracker = OfflineEmissionsTracker( + measure_power_secs=measure_power_secs, + country_iso_code=country_iso_code, + region=region, + ) else: experiment_id = get_existing_local_exp_id() if api and experiment_id is None: @@ -393,17 +493,11 @@ def monitor( ) raise typer.Exit(1) - tracker_args = {**tracker_args, "save_to_api": api} - - # If extra args are provided (e.g. `codecarbon monitor -- my_script.py`), delegate to `run_and_monitor` - if getattr(ctx, "args", None): - return run_and_monitor(ctx, **tracker_args) - - # Instantiate the tracker - if offline: - tracker = OfflineEmissionsTracker(**tracker_args) - else: - tracker = EmissionsTracker(**tracker_args) + tracker = EmissionsTracker( + measure_power_secs=measure_power_secs, + api_call_interval=api_call_interval, + save_to_api=api, + ) def signal_handler(signum, frame): print("\nReceived signal to stop. Saving emissions data...")