def _make_token_request(self, data):
headers = {
'Content-Type':
'application/x-www-form-urlencoded',
'Authorization':
'Basic ' +
base64.b64encode(b'codalab_cli_client:').decode('utf-8'),
'X-Requested-With':
'XMLHttpRequest',
}
headers.update(self._extra_headers)
request = urllib.request.Request(
self._address + '/rest/oauth2/token',
headers=headers,
data=urllib.parse.urlencode(data).encode('utf-8'),
)
try:
response = urlopen_with_retry(request)
result = json.loads(response.read().decode())
return result
except urllib.error.HTTPError as e:
if e.code == 401:
return None
elif e.code == 404:
raise LoginPermissionError("Invalid username or password.")
raise e
def download_url(source_url, target_path, print_status=False):
"""
Download the file at |source_url| and write it to |target_path|.
"""
in_file = urlopen_with_retry(source_url)
total_bytes = in_file.info().get('Content-Length')
if total_bytes:
total_bytes = int(total_bytes)
num_bytes = 0
out_file = open(target_path, 'wb')
def status_str():
if total_bytes:
return 'Downloaded %s/%s (%d%%)' % (
formatting.size_str(num_bytes),
formatting.size_str(total_bytes),
100.0 * num_bytes / total_bytes,
)
else:
return 'Downloaded %s' % (formatting.size_str(num_bytes))
while True:
s = in_file.read(BUFFER_SIZE)
if not s:
break
out_file.write(s)
num_bytes += len(s)
if print_status:
print('\r' + status_str(), end=' ', file=sys.stderr)
sys.stderr.flush()
if print_status:
print('\r' + status_str() + ' [done]', file=sys.stderr)
def _authorize(self):
request_data = {
'grant_type': 'password',
'username': self._username,
'password': self._password,
}
headers = {
'Authorization': 'Basic ' + base64.b64encode(b'codalab_worker_client:').decode('utf-8'),
'Content-Type': 'application/x-www-form-urlencoded',
'X-Requested-With': 'XMLHttpRequest',
}
request_to_send = urllib.request.Request(
self._base_url + '/oauth2/token',
data=urllib.parse.urlencode(request_data).encode('utf-8'),
)
for k, v in headers.items():
request_to_send.add_unredirected_header(k, v)
with closing(urlopen_with_retry(request_to_send)) as response:
response_data = response.read().decode()
try:
token = json.loads(response_data)
except ValueError:
raise BundleServiceException('Invalid JSON: ' + response_data, False)
if token['token_type'] != 'Bearer':
raise BundleServiceException(
'Unknown authorization token type: ' + token['token_type'], True
)
self._access_token = token['access_token']
self._token_expiration_time = time.time() + token['expires_in']
def upload_to_bundle_store(self, bundle: Bundle, source: Source, git: bool,
unpack: bool):
"""Uploads the given source to the bundle store.
Given arguments are the same as UploadManager.upload_to_bundle_store().
Used when uploading from rest server."""
try:
# bundle_path = self._bundle_store.get_bundle_location(bundle.uuid)
is_url, is_fileobj, filename = self._interpret_source(source)
if is_url:
assert isinstance(source, str)
if git:
bundle_path = self._update_and_get_bundle_location(
bundle, is_directory=True)
self.write_git_repo(source, bundle_path)
else:
# If downloading from a URL, convert the source to a file object.
is_fileobj = True
source = (filename, urlopen_with_retry(source))
if is_fileobj:
source_filename, source_fileobj = cast(Tuple[str, IO[bytes]],
source)
source_ext = zip_util.get_archive_ext(source_filename)
if unpack and zip_util.path_is_archive(filename):
bundle_path = self._update_and_get_bundle_location(
bundle, is_directory=source_ext in ARCHIVE_EXTS_DIR)
self.write_fileobj(source_ext,
source_fileobj,
bundle_path,
unpack_archive=True)
else:
bundle_path = self._update_and_get_bundle_location(
bundle, is_directory=False)
self.write_fileobj(source_ext,
source_fileobj,
bundle_path,
unpack_archive=False)
except UsageError:
if FileSystems.exists(bundle_path):
path_util.remove(bundle_path)
raise
def _make_request(
self,
method,
path,
query_params=None,
headers=None,
data=None,
return_response=False,
authorized=True,
timeout_seconds=URLOPEN_TIMEOUT_SECONDS,
):
"""
`data` can be one of the following:
- bytes
- string (text/plain)
- dict (application/json)
"""
# Set headers
if headers is None:
headers = {}
headers['X-Requested-With'] = 'XMLHttpRequest'
access_token = self._get_access_token()
if authorized and access_token:
headers['Authorization'] = 'Bearer ' + self._get_access_token()
if isinstance(data, dict):
headers['Content-Type'] = 'application/json'
data = json.dumps(data) # Turn dict into string
if isinstance(data, str):
data = data.encode() # Turn string into bytes
# Emphasize utf-8 for non-bytes data.
if headers.get('Content-Type') in ('text/plain', 'application/json'):
headers['Content-Type'] += '; charset=utf-8'
headers.update(self._extra_headers)
# Set path
if query_params is not None:
path = path + '?' + urllib.parse.urlencode(query_params)
request_url = self._base_url + path
# Make the actual request
request = urllib.request.Request(request_url, data=data, headers=headers)
request.get_method = lambda: method
if return_response:
# Return a file-like object containing the contents of the response
# body, transparently decoding gzip streams if indicated by the
# Content-Encoding header.
response = urlopen_with_retry(request, timeout=timeout_seconds)
encoding = response.headers.get('Content-Encoding')
if not encoding or encoding == 'identity':
return response
elif encoding == 'gzip':
return un_gzip_stream(response)
else:
raise RestClientException('Unsupported Content-Encoding: ' + encoding, False)
with closing(urlopen_with_retry(request, timeout=timeout_seconds)) as response:
# If the response is a JSON document, as indicated by the
# Content-Type header, try to deserialize it and return the result.
# Otherwise, just ignore the response body and return None.
if response.headers.get('Content-Type') == 'application/json':
response_data = response.read().decode()
try:
return json.loads(response_data)
except ValueError:
raise RestClientException('Invalid JSON: ' + response_data, False)