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