Skip to content

Commit b886e13

Browse files
michaelshinmichaelshin
committed
add pkce flow
1 parent dfaae02 commit b886e13

2 files changed

Lines changed: 115 additions & 8 deletions

File tree

‎centml/cli/login.py‎

Lines changed: 114 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,25 +1,132 @@
1+
import base64
2+
import hashlib
3+
from http.server import BaseHTTPRequestHandler, HTTPServer
4+
import json
5+
import secrets
6+
import urllib.parse
7+
import webbrowser
8+
19
import click
10+
import requests
11+
212

313
from centml.sdk import auth
414
from centml.sdk.config import settings
515

616

17+
CLIENT_ID = settings.CENTML_WORKOS_CLIENT_ID
18+
SERVER_HOST = "127.0.0.1"
19+
SERVER_PORT = 6789
20+
REDIRECT_URI = f"http://{SERVER_HOST}:{SERVER_PORT}/callback"
21+
AUTHORIZE_URL = "https://auth.centml.com/user_management/authorize"
22+
AUTHENTICATE_URL = "https://auth.centml.com/user_management/authenticate"
23+
PROVIDER = "authkit"
24+
25+
26+
def generate_pkce_pair():
27+
verifier = secrets.token_urlsafe(64)
28+
challenge = base64.urlsafe_b64encode(hashlib.sha256(verifier.encode()).digest()).decode().rstrip("=")
29+
return verifier, challenge
30+
31+
32+
def build_auth_url(client_id, redirect_uri, challenge):
33+
params = {
34+
"response_type": "code",
35+
"client_id": client_id,
36+
"redirect_uri": redirect_uri,
37+
"code_challenge": challenge,
38+
"code_challenge_method": "S256",
39+
"provider": PROVIDER,
40+
}
41+
return f"{AUTHORIZE_URL}?{urllib.parse.urlencode(params)}"
42+
43+
44+
class OAuthHandler(BaseHTTPRequestHandler):
45+
def do_GET(self):
46+
query = urllib.parse.urlparse(self.path).query
47+
params = urllib.parse.parse_qs(query)
48+
self.server.auth_code = params.get("code", [None])[0]
49+
50+
self.send_response(200)
51+
self.send_header("Content-type", "text/html")
52+
self.end_headers()
53+
self.wfile.write(
54+
"""
55+
<html>
56+
<body>
57+
<h1>Succesfully logged into CentML CLI</h1>
58+
<p>You can now close this tab and continue in the CLI.</p>
59+
</body>
60+
</html>
61+
""".encode(
62+
"utf-8"
63+
)
64+
)
65+
66+
def log_message(self, format, *args):
67+
# Override this to suppress logging
68+
pass
69+
70+
71+
def get_auth_code():
72+
server = HTTPServer((SERVER_HOST, SERVER_PORT), OAuthHandler)
73+
server.handle_request()
74+
return server.auth_code
75+
76+
77+
def exchange_code_for_token(code, code_verifier):
78+
data = {
79+
"grant_type": "authorization_code",
80+
"client_id": CLIENT_ID,
81+
"code": code,
82+
"redirect_uri": REDIRECT_URI,
83+
"code_verifier": code_verifier,
84+
}
85+
headers = {"Content-Type": "application/x-www-form-urlencoded"}
86+
response = requests.post(AUTHENTICATE_URL, data=data, headers=headers, timeout=3)
87+
response.raise_for_status()
88+
return response.json()
89+
90+
791
@click.command(help="Login to CentML")
892
@click.argument("token_file", required=False)
993
def login(token_file):
1094
if token_file:
1195
auth.store_centml_cred(token_file)
1296

13-
if auth.load_centml_cred():
14-
click.echo(f"Authenticating with credentials from {settings.CENTML_CRED_FILE_PATH}\n")
15-
click.echo("Login successful")
97+
cred = auth.load_centml_cred()
98+
if cred is not None and auth.refresh_centml_token(cred.get("refresh_token")):
99+
click.echo("Authenticating with stored credentials...\n")
100+
click.echo("✅ Login successful")
16101
else:
17-
click.echo("Login with CentML authentication token")
18-
click.echo("Usage: centml login TOKEN_FILE\n")
19-
choice = click.confirm("Do you want to download the token?")
102+
click.echo("Logging into CentML...")
20103

104+
choice = click.confirm("Do you want to log in with your browser now?", default=True)
21105
if choice:
22-
click.launch(f"{settings.CENTML_WEB_URL}?isCliMode=true")
106+
try:
107+
# PKCE Flow
108+
code_verifier, code_challenge = generate_pkce_pair()
109+
auth_url = build_auth_url(CLIENT_ID, REDIRECT_URI, code_challenge)
110+
click.echo("A browser window will open for you to authenticate.")
111+
click.echo("If it doesn't open automatically, you can copy and paste this URL:")
112+
click.echo(f" {auth_url}\n")
113+
webbrowser.open(auth_url)
114+
click.echo("Waiting for authentication...")
115+
116+
code = get_auth_code()
117+
response_dict = exchange_code_for_token(code, code_verifier)
118+
# If there is an error, we should remove the credentials and the user needs to sign in again.
119+
if "error" in response_dict:
120+
click.echo("Login failed. Please try again.")
121+
else:
122+
cred = {
123+
key: response_dict[key] for key in ("access_token", "refresh_token") if key in response_dict
124+
}
125+
with open(settings.CENTML_CRED_FILE_PATH, "w") as f:
126+
json.dump(cred, f)
127+
click.echo("✅ Login successful")
128+
except Exception as e:
129+
click.echo(f"Login failed: {e}")
23130
else:
24131
click.echo("Login unsuccessful")
25132

‎centml/sdk/auth.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ def refresh_centml_token(refresh_token):
1616
}
1717

1818
response = requests.post(
19-
"https://api.workos.com/user_management/authenticate",
19+
"https://auth.centml.com/user_management/authenticate",
2020
headers={"Content-Type": "application/json; charset=UTF-8"},
2121
json=payload,
2222
timeout=3,

0 commit comments

Comments
 (0)