def handle_token(self):
"""This function is called periodically to check the status of the current
token if there is one, and request a new one if needed.
If the token request fails, it will be retried according to the retry policy.
A token refresh will be attempted if the token will expire soon.
This function will return a tuple of two booleans. The first represents whether
the token authentication has not completed within it's given timeout window. The
second indicates whether the token negotiation is still in progress.
:raises: ~uamqp.errors.AuthenticationException if the token authentication fails.
:raises: ~uamqp.errors.TokenExpired if the token has expired and cannot be
refreshed.
:returns: tuple[bool, bool]
"""
timeout = False
in_progress = False
self._lock.acquire()
try:
auth_status = self._cbs_auth.get_status()
auth_status = constants.CBSAuthStatus(auth_status)
if auth_status == constants.CBSAuthStatus.Error:
if self.retries >= self._retry_policy.retries: # pylint: disable=no-member
_logger.warning("Authentication Put-Token failed. Retries exhausted.")
raise errors.TokenAuthFailure(*self._cbs_auth.get_failure_info())
else:
_logger.info("Authentication Put-Token failed. Retrying.")
self.retries += 1 # pylint: disable=no-member
time.sleep(self._retry_policy.backoff)
self._cbs_auth.authenticate()
in_progress = True
elif auth_status == constants.CBSAuthStatus.Failure:
errors.AuthenticationException("Failed to open CBS authentication link.")
elif auth_status == constants.CBSAuthStatus.Expired:
raise errors.TokenExpired("CBS Authentication Expired.")
elif auth_status == constants.CBSAuthStatus.Timeout:
timeout = True
elif auth_status == constants.CBSAuthStatus.InProgress:
in_progress = True
elif auth_status == constants.CBSAuthStatus.RefreshRequired:
_logger.info("Token will expire soon - attempting to refresh.")
self.update_token()
self._cbs_auth.refresh(self.token, int(self.expires_at))
elif auth_status == constants.CBSAuthStatus.Idle:
self._cbs_auth.authenticate()
in_progress = True
elif auth_status != constants.CBSAuthStatus.Ok:
raise errors.AuthenticationException("Invalid auth state.")
except ValueError as e:
raise errors.AuthenticationException(
"Token authentication failed: {}".format(e))
except:
raise
finally:
self._lock.release()
return timeout, in_progress
async def _handle_exception(self, exception: Exception) -> Exception:
if not self.running and isinstance(exception, compat.TimeoutException):
exception = errors.AuthenticationException(
"Authorization timeout.")
return await _handle_exception(exception, self)
return await _handle_exception(exception, self)
def _management_request(self, mgmt_msg, op_type):
# type: (Message, bytes) -> Any
retried_times = 0
last_exception = None
while retried_times <= self._config.max_retries:
mgmt_auth = self._create_auth()
mgmt_client = AMQPClient(self._mgmt_target)
try:
conn = self._conn_manager.get_connection(
self._address.hostname, mgmt_auth) # pylint:disable=assignment-from-none
mgmt_client.open(connection=conn)
response = mgmt_client.mgmt_request(
mgmt_msg,
constants.READ_OPERATION,
op_type=op_type,
status_code_field=b"status-code",
description_fields=b"status-description",
)
status_code = response.application_properties[b"status-code"]
if status_code < 400:
return response
raise errors.AuthenticationException(
"Management request error. Status code: {}".format(
status_code))
except Exception as exception: # pylint: disable=broad-except
last_exception = _handle_exception(exception, self)
self._backoff(retried_times=retried_times,
last_exception=last_exception)
retried_times += 1
if retried_times > self._config.max_retries:
_LOGGER.info("%r returns an exception %r",
self._container_id, last_exception)
raise last_exception
finally:
mgmt_client.close()
def _handle_exception(self, exception, retry_count, max_retries,
timeout_time):
if not self.running and isinstance(exception, compat.TimeoutException):
exception = errors.AuthenticationException(
"Authorization timeout.")
return _handle_exception(exception, retry_count, max_retries, self,
timeout_time)
return _handle_exception(exception, retry_count, max_retries, self,
timeout_time)
async def _management_request_async(self, mgmt_msg: Message,
op_type: bytes) -> Any:
retried_times = 0
last_exception = None
while retried_times <= self._config.max_retries:
mgmt_auth = await self._create_auth_async()
mgmt_client = AMQPClientAsync(self._mgmt_target,
auth=mgmt_auth,
debug=self._config.network_tracing)
try:
conn = await self._conn_manager_async.get_connection(
self._address.hostname, mgmt_auth)
mgmt_msg.application_properties[
"security_token"] = mgmt_auth.token
await mgmt_client.open_async(connection=conn)
response = await mgmt_client.mgmt_request_async(
mgmt_msg,
constants.READ_OPERATION,
op_type=op_type,
status_code_field=MGMT_STATUS_CODE,
description_fields=MGMT_STATUS_DESC,
)
status_code = int(
response.application_properties[MGMT_STATUS_CODE])
description = response.application_properties.get(
MGMT_STATUS_DESC) # type: Optional[Union[str, bytes]]
if description and isinstance(description, six.binary_type):
description = description.decode('utf-8')
if status_code < 400:
return response
if status_code in [401]:
raise errors.AuthenticationException(
"Management authentication failed. Status code: {}, Description: {!r}"
.format(status_code, description))
if status_code in [404]:
raise ConnectError(
"Management connection failed. Status code: {}, Description: {!r}"
.format(status_code, description))
raise errors.AMQPConnectionError(
"Management request error. Status code: {}, Description: {!r}"
.format(status_code, description))
except asyncio.CancelledError: # pylint: disable=try-except-raise
raise
except Exception as exception: # pylint:disable=broad-except
last_exception = await _handle_exception(exception, self)
await self._backoff_async(retried_times=retried_times,
last_exception=last_exception)
retried_times += 1
if retried_times > self._config.max_retries:
_LOGGER.info("%r returns an exception %r",
self._container_id, last_exception)
raise last_exception
finally:
await mgmt_client.close_async()
async def _management_request_async(self, mgmt_msg: Message,
op_type: bytes) -> Any:
retried_times = 0
last_exception = None
while retried_times <= self._config.max_retries:
mgmt_auth = await self._create_auth_async()
mgmt_client = AMQPClientAsync(self._mgmt_target,
auth=mgmt_auth,
debug=self._config.network_tracing)
try:
conn = await self._conn_manager_async.get_connection(
self._address.hostname, mgmt_auth)
await mgmt_client.open_async(connection=conn)
response = await mgmt_client.mgmt_request_async(
mgmt_msg,
constants.READ_OPERATION,
op_type=op_type,
status_code_field=b"status-code",
description_fields=b"status-description",
)
status_code = response.application_properties[b"status-code"]
if status_code < 400:
return response
raise errors.AuthenticationException(
"Management request error. Status code: {}".format(
status_code))
except Exception as exception: # pylint:disable=broad-except
last_exception = await _handle_exception(exception, self)
await self._backoff_async(retried_times=retried_times,
last_exception=last_exception)
retried_times += 1
if retried_times > self._config.max_retries:
_LOGGER.info("%r returns an exception %r",
self._container_id, last_exception)
raise last_exception
finally:
await mgmt_client.close_async()
async def handle_token_async(self):
"""This coroutine is called periodically to check the status of the current
token if there is one, and request a new one if needed.
If the token request fails, it will be retried according to the retry policy.
A token refresh will be attempted if the token will expire soon.
This function will return a tuple of two booleans. The first represents whether
the token authentication has not completed within it's given timeout window. The
second indicates whether the token negotiation is still in progress.
:raises: ~uamqp.errors.AuthenticationException if the token authentication fails.
:raises: ~uamqp.errors.TokenExpired if the token has expired and cannot be
refreshed.
:rtype: tuple[bool, bool]
"""
# pylint: disable=protected-access
timeout = False
in_progress = False
try:
await self._connection.lock_async()
if self._connection._closing or self._connection._error:
return timeout, in_progress
auth_status = self._cbs_auth.get_status()
auth_status = constants.CBSAuthStatus(auth_status)
if auth_status == constants.CBSAuthStatus.Error:
if self.retries >= self._retry_policy.retries: # pylint: disable=no-member
_logger.warning(
"Authentication Put-Token failed. Retries exhausted.")
raise errors.TokenAuthFailure(
*self._cbs_auth.get_failure_info())
error_code, error_description = self._cbs_auth.get_failure_info(
)
_logger.info("Authentication status: %r, description: %r",
error_code, error_description)
_logger.info("Authentication Put-Token failed. Retrying.")
self.retries += 1 # pylint: disable=no-member
await asyncio.sleep(self._retry_policy.backoff)
self._cbs_auth.authenticate()
in_progress = True
elif auth_status == constants.CBSAuthStatus.Failure:
raise errors.AuthenticationException(
"Failed to open CBS authentication link.")
elif auth_status == constants.CBSAuthStatus.Expired:
raise errors.TokenExpired("CBS Authentication Expired.")
elif auth_status == constants.CBSAuthStatus.Timeout:
timeout = True
elif auth_status == constants.CBSAuthStatus.InProgress:
in_progress = True
elif auth_status == constants.CBSAuthStatus.RefreshRequired:
_logger.info(
"Token on connection %r will expire soon - attempting to refresh.",
self._connection.container_id)
await self.update_token()
self._cbs_auth.refresh(self.token, int(self.expires_at))
elif auth_status == constants.CBSAuthStatus.Idle:
self._cbs_auth.authenticate()
in_progress = True
elif auth_status != constants.CBSAuthStatus.Ok:
raise ValueError("Invalid auth state.")
except asyncio.TimeoutError:
_logger.debug(
"CBS auth timed out while waiting for lock acquisition.")
return None, None
except ValueError as e:
raise errors.AuthenticationException(
"Token authentication failed: {}".format(e))
finally:
self._connection.release_async()
return timeout, in_progress
