codekingpro/portable-devtools

1from requests import Session2from requests.auth import AuthBase3from authlib.oauth2.client import OAuth2Client4from authlib.oauth2.auth import ClientAuth, TokenAuth5from ..base_client import (6    OAuthError,7    InvalidTokenError,8    MissingTokenError,9    UnsupportedTokenTypeError,10)11from .utils import update_session_configure12 13__all__ = ['OAuth2Session', 'OAuth2Auth']14 15 16class OAuth2Auth(AuthBase, TokenAuth):17    """Sign requests for OAuth 2.0, currently only bearer token is supported."""18 19    def ensure_active_token(self):20        if self.client and not self.client.ensure_active_token(self.token):21            raise InvalidTokenError()22 23    def __call__(self, req):24        self.ensure_active_token()25        try:26            req.url, req.headers, req.body = self.prepare(27                req.url, req.headers, req.body)28        except KeyError as error:29            description = f'Unsupported token_type: {str(error)}'30            raise UnsupportedTokenTypeError(description=description)31        return req32 33 34class OAuth2ClientAuth(AuthBase, ClientAuth):35    """Attaches OAuth Client Authentication to the given Request object.36    """37    def __call__(self, req):38        req.url, req.headers, req.body = self.prepare(39            req.method, req.url, req.headers, req.body40        )41        return req42 43 44class OAuth2Session(OAuth2Client, Session):45    """Construct a new OAuth 2 client requests session.46 47    :param client_id: Client ID, which you get from client registration.48    :param client_secret: Client Secret, which you get from registration.49    :param authorization_endpoint: URL of the authorization server's50        authorization endpoint.51    :param token_endpoint: URL of the authorization server's token endpoint.52    :param token_endpoint_auth_method: client authentication method for53        token endpoint.54    :param revocation_endpoint: URL of the authorization server's OAuth 2.055        revocation endpoint.56    :param revocation_endpoint_auth_method: client authentication method for57        revocation endpoint.58    :param scope: Scope that you needed to access user resources.59    :param state: Shared secret to prevent CSRF attack.60    :param redirect_uri: Redirect URI you registered as callback.61    :param token: A dict of token attributes such as ``access_token``,62        ``token_type`` and ``expires_at``.63    :param token_placement: The place to put token in HTTP request. Available64        values: "header", "body", "uri".65    :param update_token: A function for you to update token. It accept a66        :class:`OAuth2Token` as parameter.67    :param default_timeout: If settled, every requests will have a default timeout.68    """69    client_auth_class = OAuth2ClientAuth70    token_auth_class = OAuth2Auth71    oauth_error_class = OAuthError72    SESSION_REQUEST_PARAMS = (73        'allow_redirects', 'timeout', 'cookies', 'files',74        'proxies', 'hooks', 'stream', 'verify', 'cert', 'json'75    )76 77    def __init__(self, client_id=None, client_secret=None,78                 token_endpoint_auth_method=None,79                 revocation_endpoint_auth_method=None,80                 scope=None, state=None, redirect_uri=None,81                 token=None, token_placement='header',82                 update_token=None, default_timeout=None, **kwargs):83        Session.__init__(self)84        self.default_timeout = default_timeout85        update_session_configure(self, kwargs)86 87        OAuth2Client.__init__(88            self, session=self,89            client_id=client_id, client_secret=client_secret,90            token_endpoint_auth_method=token_endpoint_auth_method,91            revocation_endpoint_auth_method=revocation_endpoint_auth_method,92            scope=scope, state=state, redirect_uri=redirect_uri,93            token=token, token_placement=token_placement,94            update_token=update_token, **kwargs95        )96 97    def fetch_access_token(self, url=None, **kwargs):98        """Alias for fetch_token."""99        return self.fetch_token(url, **kwargs)100 101    def request(self, method, url, withhold_token=False, auth=None, **kwargs):102        """Send request with auto refresh token feature (if available)."""103        if self.default_timeout:104            kwargs.setdefault('timeout', self.default_timeout)105        if not withhold_token and auth is None:106            if not self.token:107                raise MissingTokenError()108            auth = self.token_auth109        return super().request(110            method, url, auth=auth, **kwargs)111