Skip to content

Commit 3199307

Browse files
michaelshinmichaelshin
authored
Update auth provider (#96)
1 parent b69e9c5 commit 3199307

File tree

4 files changed

+145
-29
lines changed

4 files changed

+145
-29
lines changed

‎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 = 57983
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}?isCliAuthenticated=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: 28 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -9,24 +9,36 @@
99

1010

1111
def refresh_centml_token(refresh_token):
12-
api_key = settings.CENTML_FIREBASE_API_KEY
13-
14-
cred = requests.post(
15-
f"https://securetoken.googleapis.com/v1/token?key={api_key}",
16-
headers={"content-type": "application/json; charset=UTF-8"},
17-
data=json.dumps({"grantType": "refresh_token", "refreshToken": refresh_token}),
12+
payload = {
13+
"client_id": settings.CENTML_WORKOS_CLIENT_ID,
14+
"grant_type": "refresh_token",
15+
"refresh_token": refresh_token,
16+
}
17+
18+
response = requests.post(
19+
"https://auth.centml.com/user_management/authenticate",
20+
headers={"Content-Type": "application/json; charset=UTF-8"},
21+
json=payload,
1822
timeout=3,
19-
).json()
20-
21-
with open(settings.CENTML_CRED_FILE_PATH, 'w') as f:
22-
json.dump(cred, f)
23+
)
24+
response_dict = response.json()
25+
26+
# If there is an error, we should remove the credentials and the user needs to sign in again.
27+
if "error" in response_dict:
28+
if os.path.exists(settings.CENTML_CRED_FILE_PATH):
29+
os.remove(settings.CENTML_CRED_FILE_PATH)
30+
cred = None
31+
else:
32+
cred = {key: response_dict[key] for key in ("access_token", "refresh_token") if key in response_dict}
33+
with open(settings.CENTML_CRED_FILE_PATH, "w") as f:
34+
json.dump(cred, f)
2335

2436
return cred
2537

2638

2739
def store_centml_cred(token_file):
2840
try:
29-
with open(token_file, 'r') as f:
41+
with open(token_file, "r") as f:
3042
os.makedirs(settings.CENTML_CONFIG_PATH, exist_ok=True)
3143
refresh_token = json.load(f)["refresh_token"]
3244

@@ -39,24 +51,24 @@ def load_centml_cred():
3951
cred = None
4052

4153
if os.path.exists(settings.CENTML_CRED_FILE_PATH):
42-
with open(settings.CENTML_CRED_FILE_PATH, 'r') as f:
54+
with open(settings.CENTML_CRED_FILE_PATH, "r") as f:
4355
cred = json.load(f)
4456

4557
return cred
4658

4759

4860
def get_centml_token():
4961
cred = load_centml_cred()
50-
5162
if not cred:
5263
sys.exit("CentML credentials not found. Please login...")
53-
54-
exp_time = int(jwt.decode(cred["id_token"], options={"verify_signature": False})["exp"])
64+
exp_time = int(jwt.decode(cred["access_token"], options={"verify_signature": False})["exp"])
5565

5666
if time.time() >= exp_time - 100:
5767
cred = refresh_centml_token(cred["refresh_token"])
68+
if cred is None:
69+
sys.exit("Could not refresh credentials. Please login and try again...")
5870

59-
return cred["id_token"]
71+
return cred["access_token"]
6072

6173

6274
def remove_centml_cred():

‎centml/sdk/config.py‎

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,8 @@
44

55

66
class Config(BaseSettings):
7-
87
# It is possible to override the default values by setting the environment variables
9-
model_config = SettingsConfigDict(env_file=Path('.env'))
8+
model_config = SettingsConfigDict(env_file=Path(".env"))
109

1110
CENTML_WEB_URL: str = os.getenv("CENTML_WEB_URL", default="https://app.centml.com/")
1211
CENTML_CONFIG_PATH: str = os.getenv("CENTML_CONFIG_PATH", default=os.path.expanduser("~/.centml"))
@@ -15,9 +14,7 @@ class Config(BaseSettings):
1514

1615
CENTML_PLATFORM_API_URL: str = os.getenv("CENTML_PLATFORM_API_URL", default="https://api.centml.com")
1716

18-
CENTML_FIREBASE_API_KEY: str = os.getenv(
19-
"CENTML_FIREBASE_API_KEY", default="AIzaSyChPXy41cIAxS_Nd8oaYKyP_oKkIucobtY"
20-
)
17+
CENTML_WORKOS_CLIENT_ID: str = os.getenv("CENTML_WORKOS_CLIENT_ID", default="client_01JP5TWW2997MF8AYQXHJEGYR0")
2118

2219

2320
settings = Config()

‎setup.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@
1111

1212
setup(
1313
name='centml',
14-
version='0.3.5',
14+
version='0.4.0',
1515
packages=find_packages(),
1616
python_requires=">=3.10",
1717
long_description=open('README.md').read(),

0 commit comments

Comments
 (0)