def client(self, address, is_cli=True, use_rest=False):
'''
Return a client given the address. Note that this can either be called
by the CLI (is_cli=True) or the server (is_cli=False).
If called by the CLI, we don't need to authenticate.
Cache the Client if necessary.
'''
# FIXME(sckoo): temporary hack
# When BundleCLI is no longer dependent on RemoteBundleClient
# or LocalBundleClient, simplify everything to only using
# JsonApiClient, and remove all use_rest kwargs.
# Users should also then update their aliases to point at the
# correct addresses of the REST servers.
if use_rest:
address = self.derive_rest_address(address)
# Return cached client
# TODO(sckoo): Remove is_cli check when REST migration complete
if not is_cli and address in self.clients:
return self.clients[address]
# Create new client
if use_rest:
# Create RestOAuthHandler that authenticates directly with
# OAuth endpoints on the REST server
from codalab.server.auth import RestOAuthHandler
auth_handler = RestOAuthHandler(address, None)
# Create JsonApiClient with a callback to get access tokens
from codalab.client.json_api_client import JsonApiClient
client = JsonApiClient(
address, lambda: self._authenticate(address, auth_handler))
elif is_local_address(address):
# if local force mockauth or if local server use correct auth
bundle_store = self.bundle_store()
model = self.model()
worker_model = self.worker_model()
upload_manager = self.upload_manager()
download_manager = self.download_manager()
auth_handler = self.auth_handler(mock=is_cli)
from codalab.client.local_bundle_client import LocalBundleClient
client = LocalBundleClient(address, bundle_store, model, worker_model, upload_manager, download_manager, auth_handler, self.cli_verbose)
if is_cli:
# Set current user
access_token = self._authenticate(client.address, client.auth_handler)
auth_handler.validate_token(access_token)
else:
from codalab.client.remote_bundle_client import RemoteBundleClient
client = RemoteBundleClient(address, lambda a_client: self._authenticate(a_client.address, a_client), self.cli_verbose)
self._authenticate(client.address, client)
# Cache and return client
self.clients[address] = client
return client
class AuthHelper(object):
REFRESH_BUFFER_SECONDS = 5 * 60
def __init__(self, host, username, password):
self.username = username
self.password = password
self.auth_handler = RestOAuthHandler(host)
self.grant = None
self.expires_at = None
def get_access_token(self):
if not self.grant or time.time(
) > self.expires_at - self.REFRESH_BUFFER_SECONDS:
self.grant = self.auth_handler.generate_token(
'credentials', self.username, self.password)
if self.grant is None:
raise PermissionError('Invalid username or password.')
self.expires_at = time.time() + self.grant['expires_in']
return self.grant['access_token']
def client(self, address):
"""
Return a client given the address.
"""
# Return cached client
if address in self.clients:
return self.clients[address]
# Create RestOAuthHandler that authenticates directly with
# OAuth endpoints on the REST server
from codalab.server.auth import RestOAuthHandler
auth_handler = RestOAuthHandler(address)
# Create JsonApiClient with a callback to get access tokens
client = JsonApiClient(
address, lambda: self._authenticate(address, auth_handler),
self.check_version)
# Cache and return client
self.clients[address] = client
return client
def __init__(self, host, username, password):
self.username = username
self.password = password
self.auth_handler = RestOAuthHandler(host)
self.grant = None
self.expires_at = None
def rest_oauth_handler(self):
from codalab.server.auth import RestOAuthHandler
address = 'http://%s:%d' % (self.config['server']['rest_host'],
self.config['server']['rest_port'])
return RestOAuthHandler(address)
