Skip to content

Commit 9d0b0f5

Browse files
prmths128prmths128
committed
feat: add missing oidc auth on cli
1 parent cdc791d commit 9d0b0f5

2 files changed

Lines changed: 365 additions & 15 deletions

File tree

‎codecarbon/cli/main.py‎

Lines changed: 118 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,20 @@
1+
import json
12
import os
23
import signal
34
import sys
45
import time
6+
import webbrowser
7+
from http.server import BaseHTTPRequestHandler, HTTPServer
58
from pathlib import Path
69
from typing import Optional
10+
from urllib.parse import parse_qs, urlparse
711

812
import questionary
913
import requests
1014
import typer
11-
from fief_client import Fief
12-
from fief_client.integrations.cli import FiefAuth
15+
from authlib.common.security import generate_token
16+
from authlib.integrations.requests_client import OAuth2Session
17+
from authlib.oauth2.rfc7636 import create_s256_code_challenge
1318
from rich import print
1419
from rich.prompt import Confirm
1520
from typing_extensions import Annotated
@@ -22,7 +27,6 @@
2227
get_existing_local_exp_id,
2328
overwrite_local_config,
2429
)
25-
from codecarbon.cli.monitor import run_and_monitor
2630
from codecarbon.core.api_client import ApiClient, get_datetime_with_timezone
2731
from codecarbon.core.schemas import ExperimentCreate, OrganizationCreate, ProjectCreate
2832
from codecarbon.emissions_tracker import EmissionsTracker, OfflineEmissionsTracker
@@ -31,8 +35,9 @@
3135
"AUTH_CLIENT_ID",
3236
"jsUPWIcUECQFE_ouanUuVhXx52TTjEVcVNNtNGeyAtU",
3337
)
34-
AUTH_SERVER_URL = os.environ.get(
35-
"AUTH_SERVER_URL", "https://auth.codecarbon.io/codecarbon"
38+
AUTH_SERVER_WELL_KNOWN = os.environ.get(
39+
"AUTH_SERVER_WELL_KNOWN",
40+
"https://auth.codecarbon.io/codecarbon/.well-known/openid-configuration",
3641
)
3742
API_URL = os.environ.get("API_URL", "https://dashboard.codecarbon.io/api")
3843

@@ -115,26 +120,124 @@ def show_config(path: Path = Path("./.codecarbon.config")) -> None:
115120
)
116121

117122

118-
def get_fief_auth():
119-
fief = Fief(AUTH_SERVER_URL, AUTH_CLIENT_ID)
120-
fief_auth = FiefAuth(fief, "./credentials.json")
121-
return fief_auth
123+
_REDIRECT_PORT = 8090
124+
_REDIRECT_URI = f"http://localhost:{_REDIRECT_PORT}/callback"
125+
_CREDENTIALS_FILE = Path("./credentials.json")
126+
127+
128+
class _CallbackHandler(BaseHTTPRequestHandler):
129+
"""HTTP handler that captures the OAuth2 authorization callback."""
130+
131+
callback_url = None
132+
error = None
133+
134+
def do_GET(self):
135+
_CallbackHandler.callback_url = f"http://localhost:{_REDIRECT_PORT}{self.path}"
136+
parsed = urlparse(self.path)
137+
params = parse_qs(parsed.query)
138+
139+
if "error" in params:
140+
_CallbackHandler.error = params["error"][0]
141+
self.send_response(400)
142+
self.send_header("Content-Type", "text/html")
143+
self.end_headers()
144+
msg = params.get("error_description", [params["error"][0]])[0]
145+
self.wfile.write(
146+
f"<html><body><h1>Login failed</h1><p>{msg}</p></body></html>".encode()
147+
)
148+
else:
149+
self.send_response(200)
150+
self.send_header("Content-Type", "text/html")
151+
self.end_headers()
152+
self.wfile.write(
153+
b"<html><body><h1>Login successful!</h1>"
154+
b"<p>You can close this window.</p></body></html>"
155+
)
156+
157+
def log_message(self, format, *args):
158+
pass # Suppress server logs
159+
160+
161+
def _discover_endpoints():
162+
"""Fetch OpenID Connect discovery document."""
163+
resp = requests.get(AUTH_SERVER_WELL_KNOWN)
164+
resp.raise_for_status()
165+
return resp.json()
166+
167+
168+
def _authorize():
169+
"""Run the OAuth2 Authorization Code flow with PKCE."""
170+
discovery = _discover_endpoints()
171+
172+
session = OAuth2Session(
173+
client_id=AUTH_CLIENT_ID,
174+
redirect_uri=_REDIRECT_URI,
175+
scope="openid offline_access",
176+
token_endpoint_auth_method="none",
177+
)
178+
179+
code_verifier = generate_token(48)
180+
code_challenge = create_s256_code_challenge(code_verifier)
181+
182+
uri, state = session.create_authorization_url(
183+
discovery["authorization_endpoint"],
184+
code_challenge=code_challenge,
185+
code_challenge_method="S256",
186+
)
187+
188+
# Reset handler state
189+
_CallbackHandler.callback_url = None
190+
_CallbackHandler.error = None
191+
192+
server = HTTPServer(("localhost", _REDIRECT_PORT), _CallbackHandler)
193+
194+
print("Opening browser for authentication...")
195+
webbrowser.open(uri)
196+
197+
server.handle_request()
198+
server.server_close()
199+
200+
if _CallbackHandler.error:
201+
raise ValueError(f"Authorization failed: {_CallbackHandler.error}")
202+
203+
if not _CallbackHandler.callback_url:
204+
raise ValueError("Authorization failed: no callback received")
205+
206+
token = session.fetch_token(
207+
discovery["token_endpoint"],
208+
authorization_response=_CallbackHandler.callback_url,
209+
code_verifier=code_verifier,
210+
)
211+
212+
_save_credentials(token)
213+
return token
214+
215+
216+
def _save_credentials(tokens):
217+
"""Save OAuth tokens to credentials file."""
218+
with open(_CREDENTIALS_FILE, "w") as f:
219+
json.dump(tokens, f)
220+
221+
222+
def _load_credentials():
223+
"""Load OAuth tokens from credentials file."""
224+
with open(_CREDENTIALS_FILE, "r") as f:
225+
return json.load(f)
122226

123227

124228
def _get_access_token():
125229
try:
126-
access_token_info = get_fief_auth().access_token_info()
127-
access_token = access_token_info["access_token"]
128-
return access_token
230+
creds = _load_credentials()
231+
return creds["access_token"]
129232
except Exception as e:
130233
raise ValueError(
131234
f"Not able to retrieve the access token, please run `codecarbon login` first! (error: {e})"
132235
)
133236

134237

135238
def _get_id_token():
136-
id_token = get_fief_auth()._tokens["id_token"]
137-
return id_token
239+
creds = _load_credentials()
240+
return creds["id_token"]
138241

139242

140243
@codecarbon.command(
@@ -152,7 +255,7 @@ def api_get():
152255

153256
@codecarbon.command("login", short_help="Login to CodeCarbon")
154257
def login():
155-
get_fief_auth().authorize()
258+
_authorize()
156259
api = ApiClient(endpoint_url=API_URL) # TODO: get endpoint from config
157260
access_token = _get_access_token()
158261
api.set_access_token(access_token)

0 commit comments

Comments
 (0)