Skip to content

Commit b380635

Browse files
mmourafiqmmourafiq
committed
Add impersonate auth apis
1 parent e5788a7 commit b380635

File tree

4 files changed

+174
-10
lines changed

4 files changed

+174
-10
lines changed

‎polyaxon_client/api/auth.py

Lines changed: 72 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -47,12 +47,17 @@ def get_user(self, token=None):
4747

4848
return self.prepare_results(response_json=user_dict, config=UserConfig)
4949

50-
def _persist_token(self, token):
50+
def _persist_token(self, token, token_path):
5151
create_polyaxon_tmp()
52-
with open(settings.TMP_AUTH_TOKEN_PATH, "w") as config_file:
52+
with open(token_path, "w") as config_file:
5353
config_file.write(json.dumps({settings.SECRET_USER_TOKEN_KEY: token}))
5454

55-
def _process_token(self, request_url, response, set_token=True, persist_token=False):
55+
def _process_token(self,
56+
request_url,
57+
response,
58+
set_token=True,
59+
persist_token=False,
60+
token_path=None):
5661
try:
5762
token_dict = response.json()
5863
response.raise_for_status()
@@ -73,7 +78,7 @@ def _process_token(self, request_url, response, set_token=True, persist_token=Fa
7378
if set_token:
7479
self.config.token = token
7580
if persist_token:
76-
self._persist_token(token)
81+
self._persist_token(token, token_path)
7782
return token
7883

7984
def login(self, credentials, set_token=False):
@@ -118,10 +123,72 @@ def login_experiment_ephemeral_token(self,
118123
token = self._process_token(request_url=request_url,
119124
response=response,
120125
set_token=set_token,
121-
persist_token=persist_token)
126+
persist_token=persist_token,
127+
token_path=settings.TMP_AUTH_TOKEN_PATH)
122128
# Destroy ephemeral token
123129
if os.environ.get(settings.SECRET_EPHEMERAL_TOKEN_KEY):
124130
del os.environ[settings.SECRET_EPHEMERAL_TOKEN_KEY]
125131
if hasattr(settings, 'SECRET_EPHEMERAL_TOKEN') and settings.SECRET_EPHEMERAL_TOKEN:
126132
del settings.SECRET_EPHEMERAL_TOKEN
127133
return token
134+
135+
def _login_impersonate_token(self,
136+
request_url,
137+
internal_token,
138+
set_token=True,
139+
persist_token=True):
140+
try:
141+
response = self.transport.post(
142+
request_url,
143+
headers={
144+
'Authorization': '{} {}'.format(
145+
AuthenticationTypes.INTERNAL_TOKEN, internal_token)
146+
})
147+
except requests.ConnectionError:
148+
raise PolyaxonHTTPError(
149+
request_url,
150+
None,
151+
"Connection error.",
152+
None)
153+
token = self._process_token(request_url=request_url,
154+
response=response,
155+
set_token=set_token,
156+
persist_token=persist_token,
157+
token_path=settings.CONTEXT_AUTH_TOKEN_PATH)
158+
return token
159+
160+
def login_experiment_impersonate_token(self,
161+
username,
162+
project_name,
163+
experiment_id,
164+
internal_token,
165+
set_token=True,
166+
persist_token=True):
167+
request_url = self.build_url(self._get_http_url('/'),
168+
username,
169+
project_name,
170+
'experiments',
171+
experiment_id,
172+
'imporsonatetoken')
173+
return self._login_impersonate_token(request_url=request_url,
174+
internal_token=internal_token,
175+
set_token=set_token,
176+
persist_token=persist_token)
177+
178+
def login_job_impersonate_token(self,
179+
username,
180+
project_name,
181+
job_id,
182+
internal_token,
183+
set_token=True,
184+
persist_token=True):
185+
request_url = self.build_url(self._get_http_url('/'),
186+
username,
187+
project_name,
188+
'jobs',
189+
job_id,
190+
'imporsonatetoken')
191+
return self._login_impersonate_token(request_url=request_url,
192+
internal_token=internal_token,
193+
set_token=set_token,
194+
persist_token=persist_token)

‎polyaxon_client/settings.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
from rhea import RheaError # noqa
1111

1212
TMP_AUTH_TOKEN_PATH = '/tmp/.polyaxon/.authtoken'
13+
CONTEXT_AUTH_TOKEN_PATH = '/plx-context/.authtoken'
1314
CLIENT_CONFIG_PATH = os.path.join(polyaxon_user_path(), '.polyaxonclient')
1415
CONFIG_PATH = os.path.join(polyaxon_user_path(), '.polyaxonconfig')
1516
AUTH_PATH = os.path.join(polyaxon_user_path(), '.polyaxonauth')
@@ -25,7 +26,8 @@
2526
config = rhea.Rhea.read_configs([
2627
rhea.ConfigSpec(CLIENT_CONFIG_PATH, config_type='.json', check_if_exists=False),
2728
os.environ,
28-
rhea.ConfigSpec(TMP_AUTH_TOKEN_PATH, config_type='.json', check_if_exists=False)
29+
rhea.ConfigSpec(TMP_AUTH_TOKEN_PATH, config_type='.json', check_if_exists=False),
30+
rhea.ConfigSpec(CONTEXT_AUTH_TOKEN_PATH, config_type='.json', check_if_exists=False)
2931
])
3032

3133
IN_CLUSTER = config.get_boolean('POLYAXON_IN_CLUSTER',

‎tests/test_api/test_auth.py

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
# -*- coding: utf-8 -*-
22
from __future__ import absolute_import, division, print_function
33

4+
import tempfile
5+
46
import httpretty
57
import json
68
import os
@@ -18,6 +20,8 @@ class TestAuthApi(TestBaseApi):
1820

1921
def setUp(self):
2022
super(TestAuthApi, self).setUp()
23+
settings.CONTEXT_AUTH_TOKEN_PATH = '{}/{}'.format(tempfile.mkdtemp(), '.authtoken')
24+
settings.TMP_AUTH_TOKEN_PATH = '{}/{}'.format(tempfile.mkdtemp(), '.authtoken')
2125
self.api_handler = AuthApi(transport=self.transport, config=self.api_config)
2226

2327
@httpretty.activate
@@ -136,3 +140,93 @@ def test_login_experiment_ephemeral_token(self):
136140
ephemeral_token='foo',
137141
set_token=True,
138142
persist_token=True)
143+
144+
@httpretty.activate
145+
def test_login_experiment_impersonate_token(self):
146+
token = uuid.uuid4().hex
147+
httpretty.register_uri(
148+
httpretty.POST,
149+
BaseApiHandler.build_url(
150+
self.api_config.base_url,
151+
'/',
152+
'user',
153+
'project',
154+
'experiments',
155+
'1',
156+
'imporsonatetoken'
157+
),
158+
body=json.dumps({'token': token}),
159+
content_type='application/json', status=200)
160+
161+
# Login without updating the token and without persistence
162+
if os.path.exists(settings.CONTEXT_AUTH_TOKEN_PATH):
163+
os.remove(settings.CONTEXT_AUTH_TOKEN_PATH)
164+
assert self.api_config.token == 'token'
165+
assert token == self.api_handler.login_experiment_impersonate_token(
166+
username='user',
167+
project_name='project',
168+
experiment_id=1,
169+
internal_token='foo',
170+
set_token=False,
171+
persist_token=False)
172+
assert self.api_config.token == 'token'
173+
assert os.path.exists(settings.CONTEXT_AUTH_TOKEN_PATH) is False
174+
175+
# Login and update the token and persistence
176+
if os.path.exists(settings.CONTEXT_AUTH_TOKEN_PATH):
177+
os.remove(settings.CONTEXT_AUTH_TOKEN_PATH)
178+
assert self.api_config.token == 'token'
179+
assert token == self.api_handler.login_experiment_impersonate_token(
180+
username='user',
181+
project_name='project',
182+
experiment_id=1,
183+
internal_token='foo',
184+
set_token=True,
185+
persist_token=True)
186+
assert self.api_config.token == token
187+
assert os.path.exists(settings.CONTEXT_AUTH_TOKEN_PATH) is True
188+
189+
@httpretty.activate
190+
def test_login_job_impersonate_token(self):
191+
token = uuid.uuid4().hex
192+
httpretty.register_uri(
193+
httpretty.POST,
194+
BaseApiHandler.build_url(
195+
self.api_config.base_url,
196+
'/',
197+
'user',
198+
'project',
199+
'jobs',
200+
'1',
201+
'imporsonatetoken'
202+
),
203+
body=json.dumps({'token': token}),
204+
content_type='application/json', status=200)
205+
206+
# Login without updating the token and without persistence
207+
if os.path.exists(settings.CONTEXT_AUTH_TOKEN_PATH):
208+
os.remove(settings.CONTEXT_AUTH_TOKEN_PATH)
209+
assert self.api_config.token == 'token'
210+
assert token == self.api_handler.login_job_impersonate_token(
211+
username='user',
212+
project_name='project',
213+
job_id=1,
214+
internal_token='foo',
215+
set_token=False,
216+
persist_token=False)
217+
assert self.api_config.token == 'token'
218+
assert os.path.exists(settings.CONTEXT_AUTH_TOKEN_PATH) is False
219+
220+
# Login and update the token and persistence
221+
if os.path.exists(settings.CONTEXT_AUTH_TOKEN_PATH):
222+
os.remove(settings.CONTEXT_AUTH_TOKEN_PATH)
223+
assert self.api_config.token == 'token'
224+
assert token == self.api_handler.login_job_impersonate_token(
225+
username='user',
226+
project_name='project',
227+
job_id=1,
228+
internal_token='foo',
229+
set_token=True,
230+
persist_token=True)
231+
assert self.api_config.token == token
232+
assert os.path.exists(settings.CONTEXT_AUTH_TOKEN_PATH) is True

‎tests/test_polyaxon_client.py

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
from __future__ import absolute_import, division, print_function
33

44
import os
5+
import tempfile
56

67
from unittest import TestCase
78

@@ -24,10 +25,10 @@
2425

2526

2627
class TestPolyaxonClient(TestCase):
27-
def tearDown(self):
28-
super(TestPolyaxonClient, self).tearDown()
29-
if os.path.exists(settings.TMP_AUTH_TOKEN_PATH):
30-
os.remove(settings.TMP_AUTH_TOKEN_PATH)
28+
29+
def setUp(self):
30+
super(TestPolyaxonClient, self).setUp()
31+
settings.CONTEXT_AUTH_TOKEN_PATH = '{}/{}'.format(tempfile.mkdtemp(), '.authtoken')
3132

3233
def test_client(self):
3334
settings.SECRET_USER_TOKEN = None

0 commit comments

Comments
 (0)