Team Ai
Datasetpublic

codekingpro/portable-devtools

sourceHugging Faceupdated 5mo agoView on Hugging Face
1likes14kdownloads
client.py439 linesDownload Raw Back to oauth2
1from authlib.common.security import generate_token2from authlib.common.urls import url_decode3from .rfc6749.parameters import (4    prepare_grant_uri,5    prepare_token_request,6    parse_authorization_code_response,7    parse_implicit_response,8)9from .rfc7009 import prepare_revoke_token_request10from .rfc7636 import create_s256_code_challenge11from .auth import TokenAuth, ClientAuth12from .base import OAuth2Error13 14DEFAULT_HEADERS = {15    'Accept': 'application/json',16    'Content-Type': 'application/x-www-form-urlencoded;charset=UTF-8'17}18 19 20class OAuth2Client:21    """Construct a new OAuth 2 protocol client.22 23    :param session: Requests session object to communicate with24                    authorization server.25    :param client_id: Client ID, which you get from client registration.26    :param client_secret: Client Secret, which you get from registration.27    :param token_endpoint_auth_method: client authentication method for28        token endpoint.29    :param revocation_endpoint_auth_method: client authentication method for30        revocation endpoint.31    :param scope: Scope that you needed to access user resources.32    :param state: Shared secret to prevent CSRF attack.33    :param redirect_uri: Redirect URI you registered as callback.34    :param code_challenge_method: PKCE method name, only S256 is supported.35    :param token: A dict of token attributes such as ``access_token``,36        ``token_type`` and ``expires_at``.37    :param token_placement: The place to put token in HTTP request. Available38        values: "header", "body", "uri".39    :param update_token: A function for you to update token. It accept a40        :class:`OAuth2Token` as parameter.41    """42    client_auth_class = ClientAuth43    token_auth_class = TokenAuth44    oauth_error_class = OAuth2Error45 46    EXTRA_AUTHORIZE_PARAMS = (47        'response_mode', 'nonce', 'prompt', 'login_hint'48    )49    SESSION_REQUEST_PARAMS = []50 51    def __init__(self, session, client_id=None, client_secret=None,52                 token_endpoint_auth_method=None,53                 revocation_endpoint_auth_method=None,54                 scope=None, state=None, redirect_uri=None, code_challenge_method=None,55                 token=None, token_placement='header', update_token=None, **metadata):56 57        self.session = session58        self.client_id = client_id59        self.client_secret = client_secret60        self.state = state61 62        if token_endpoint_auth_method is None:63            if client_secret:64                token_endpoint_auth_method = 'client_secret_basic'65            else:66                token_endpoint_auth_method = 'none'67 68        self.token_endpoint_auth_method = token_endpoint_auth_method69 70        if revocation_endpoint_auth_method is None:71            if client_secret:72                revocation_endpoint_auth_method = 'client_secret_basic'73            else:74                revocation_endpoint_auth_method = 'none'75 76        self.revocation_endpoint_auth_method = revocation_endpoint_auth_method77 78        self.scope = scope79        self.redirect_uri = redirect_uri80        self.code_challenge_method = code_challenge_method81 82        self.token_auth = self.token_auth_class(token, token_placement, self)83        self.update_token = update_token84 85        token_updater = metadata.pop('token_updater', None)86        if token_updater:87            raise ValueError('update token has been redesigned, checkout the documentation')88 89        self.metadata = metadata90 91        self.compliance_hook = {92            'access_token_response': set(),93            'refresh_token_request': set(),94            'refresh_token_response': set(),95            'revoke_token_request': set(),96            'introspect_token_request': set(),97        }98        self._auth_methods = {}99 100    def register_client_auth_method(self, auth):101        """Extend client authenticate for token endpoint.102 103        :param auth: an instance to sign the request104        """105        if isinstance(auth, tuple):106            self._auth_methods[auth[0]] = auth[1]107        else:108            self._auth_methods[auth.name] = auth109 110    def client_auth(self, auth_method):111        if isinstance(auth_method, str) and auth_method in self._auth_methods:112            auth_method = self._auth_methods[auth_method]113        return self.client_auth_class(114            client_id=self.client_id,115            client_secret=self.client_secret,116            auth_method=auth_method,117        )118 119    @property120    def token(self):121        return self.token_auth.token122 123    @token.setter124    def token(self, token):125        self.token_auth.set_token(token)126 127    def create_authorization_url(self, url, state=None, code_verifier=None, **kwargs):128        """Generate an authorization URL and state.129 130        :param url: Authorization endpoint url, must be HTTPS.131        :param state: An optional state string for CSRF protection. If not132                      given it will be generated for you.133        :param code_verifier: An optional code_verifier for code challenge.134        :param kwargs: Extra parameters to include.135        :return: authorization_url, state136        """137        if state is None:138            state = generate_token()139 140        response_type = self.metadata.get('response_type', 'code')141        response_type = kwargs.pop('response_type', response_type)142        if 'redirect_uri' not in kwargs:143            kwargs['redirect_uri'] = self.redirect_uri144        if 'scope' not in kwargs:145            kwargs['scope'] = self.scope146 147        if code_verifier and response_type == 'code' and self.code_challenge_method == 'S256':148            kwargs['code_challenge'] = create_s256_code_challenge(code_verifier)149            kwargs['code_challenge_method'] = self.code_challenge_method150 151        for k in self.EXTRA_AUTHORIZE_PARAMS:152            if k not in kwargs and k in self.metadata:153                kwargs[k] = self.metadata[k]154 155        uri = prepare_grant_uri(156            url, client_id=self.client_id, response_type=response_type,157            state=state, **kwargs)158        return uri, state159 160    def fetch_token(self, url=None, body='', method='POST', headers=None,161                    auth=None, grant_type=None, state=None, **kwargs):162        """Generic method for fetching an access token from the token endpoint.163 164        :param url: Access Token endpoint URL, if not configured,165                    ``authorization_response`` is used to extract token from166                    its fragment (implicit way).167        :param body: Optional application/x-www-form-urlencoded body to add the168                     include in the token request. Prefer kwargs over body.169        :param method: The HTTP method used to make the request. Defaults170                       to POST, but may also be GET. Other methods should171                       be added as needed.172        :param headers: Dict to default request headers with.173        :param auth: An auth tuple or method as accepted by requests.174        :param grant_type: Use specified grant_type to fetch token175        :return: A :class:`OAuth2Token` object (a dict too).176        """177        state = state or self.state178        # implicit  grant_type179        authorization_response = kwargs.pop('authorization_response', None)180        if authorization_response and '#' in authorization_response:181            return self.token_from_fragment(authorization_response, state)182 183        session_kwargs = self._extract_session_request_params(kwargs)184 185        if authorization_response and 'code=' in authorization_response:186            grant_type = 'authorization_code'187            params = parse_authorization_code_response(188                authorization_response,189                state=state,190            )191            kwargs['code'] = params['code']192 193        if grant_type is None:194            grant_type = self.metadata.get('grant_type')195 196        if grant_type is None:197            grant_type = _guess_grant_type(kwargs)198            self.metadata['grant_type'] = grant_type199 200        body = self._prepare_token_endpoint_body(body, grant_type, **kwargs)201 202        if auth is None:203            auth = self.client_auth(self.token_endpoint_auth_method)204 205        if headers is None:206            headers = DEFAULT_HEADERS207 208        if url is None:209            url = self.metadata.get('token_endpoint')210 211        return self._fetch_token(212            url, body=body, auth=auth, method=method,213            headers=headers, **session_kwargs214        )215 216    def token_from_fragment(self, authorization_response, state=None):217        token = parse_implicit_response(authorization_response, state)218        if 'error' in token:219            raise self.oauth_error_class(220                error=token['error'],221                description=token.get('error_description')222            )223        self.token = token224        return token225 226    def refresh_token(self, url, refresh_token=None, body='',227                      auth=None, headers=None, **kwargs):228        """Fetch a new access token using a refresh token.229 230        :param url: Refresh Token endpoint, must be HTTPS.231        :param refresh_token: The refresh_token to use.232        :param body: Optional application/x-www-form-urlencoded body to add the233                     include in the token request. Prefer kwargs over body.234        :param auth: An auth tuple or method as accepted by requests.235        :param headers: Dict to default request headers with.236        :return: A :class:`OAuth2Token` object (a dict too).237        """238        session_kwargs = self._extract_session_request_params(kwargs)239        refresh_token = refresh_token or self.token.get('refresh_token')240        if 'scope' not in kwargs and self.scope:241            kwargs['scope'] = self.scope242        body = prepare_token_request(243            'refresh_token', body,244            refresh_token=refresh_token, **kwargs245        )246 247        if headers is None:248            headers = DEFAULT_HEADERS.copy()249 250        for hook in self.compliance_hook['refresh_token_request']:251            url, headers, body = hook(url, headers, body)252 253        if auth is None:254            auth = self.client_auth(self.token_endpoint_auth_method)255 256        return self._refresh_token(257            url, refresh_token=refresh_token, body=body, headers=headers,258            auth=auth, **session_kwargs)259 260    def ensure_active_token(self, token):261        if not token.is_expired():262            return True263        refresh_token = token.get('refresh_token')264        url = self.metadata.get('token_endpoint')265        if refresh_token and url:266            self.refresh_token(url, refresh_token=refresh_token)267            return True268        elif self.metadata.get('grant_type') == 'client_credentials':269            access_token = token['access_token']270            new_token = self.fetch_token(url, grant_type='client_credentials')271            if self.update_token:272                self.update_token(new_token, access_token=access_token)273            return True274 275    def revoke_token(self, url, token=None, token_type_hint=None,276                     body=None, auth=None, headers=None, **kwargs):277        """Revoke token method defined via `RFC7009`_.278 279        :param url: Revoke Token endpoint, must be HTTPS.280        :param token: The token to be revoked.281        :param token_type_hint: The type of the token that to be revoked.282                                It can be "access_token" or "refresh_token".283        :param body: Optional application/x-www-form-urlencoded body to add the284                     include in the token request. Prefer kwargs over body.285        :param auth: An auth tuple or method as accepted by requests.286        :param headers: Dict to default request headers with.287        :return: Revocation Response288 289        .. _`RFC7009`: https://tools.ietf.org/html/rfc7009290        """291        return self._handle_token_hint(292            'revoke_token_request', url,293            token=token, token_type_hint=token_type_hint,294            body=body, auth=auth, headers=headers, **kwargs)295 296    def introspect_token(self, url, token=None, token_type_hint=None,297                         body=None, auth=None, headers=None, **kwargs):298        """Implementation of OAuth 2.0 Token Introspection defined via `RFC7662`_.299 300        :param url: Introspection Endpoint, must be HTTPS.301        :param token: The token to be introspected.302        :param token_type_hint: The type of the token that to be revoked.303                                It can be "access_token" or "refresh_token".304        :param body: Optional application/x-www-form-urlencoded body to add the305                     include in the token request. Prefer kwargs over body.306        :param auth: An auth tuple or method as accepted by requests.307        :param headers: Dict to default request headers with.308        :return: Introspection Response309 310        .. _`RFC7662`: https://tools.ietf.org/html/rfc7662311        """312        return self._handle_token_hint(313            'introspect_token_request', url,314            token=token, token_type_hint=token_type_hint,315            body=body, auth=auth, headers=headers, **kwargs)316 317    def register_compliance_hook(self, hook_type, hook):318        """Register a hook for request/response tweaking.319 320        Available hooks are:321 322        * access_token_response: invoked before token parsing.323        * refresh_token_request: invoked before refreshing token.324        * refresh_token_response: invoked before refresh token parsing.325        * protected_request: invoked before making a request.326        * revoke_token_request: invoked before revoking a token.327        * introspect_token_request: invoked before introspecting a token.328        """329        if hook_type == 'protected_request':330            self.token_auth.hooks.add(hook)331            return332 333        if hook_type not in self.compliance_hook:334            raise ValueError('Hook type %s is not in %s.',335                             hook_type, self.compliance_hook)336        self.compliance_hook[hook_type].add(hook)337 338    def parse_response_token(self, resp):339        if resp.status_code >= 500:340            resp.raise_for_status()341 342        token = resp.json()343        if 'error' in token:344            raise self.oauth_error_class(345                error=token['error'],346                description=token.get('error_description')347            )348        self.token = token349        return self.token350 351    def _fetch_token(self, url, body='', headers=None, auth=None,352                     method='POST', **kwargs):353 354        if method.upper() == 'POST':355            resp = self.session.post(356                url, data=dict(url_decode(body)),357                headers=headers, auth=auth, **kwargs)358        else:359            if '?' in url:360                url = '&'.join([url, body])361            else:362                url = '?'.join([url, body])363            resp = self.session.request(method, url, headers=headers, auth=auth, **kwargs)364 365        for hook in self.compliance_hook['access_token_response']:366            resp = hook(resp)367 368        return self.parse_response_token(resp)369 370    def _refresh_token(self, url, refresh_token=None, body='', headers=None,371                       auth=None, **kwargs):372        resp = self._http_post(url, body=body, auth=auth, headers=headers, **kwargs)373 374        for hook in self.compliance_hook['refresh_token_response']:375            resp = hook(resp)376 377        token = self.parse_response_token(resp)378        if 'refresh_token' not in token:379            self.token['refresh_token'] = refresh_token380 381        if callable(self.update_token):382            self.update_token(self.token, refresh_token=refresh_token)383 384        return self.token385 386    def _handle_token_hint(self, hook, url, token=None, token_type_hint=None,387                           body=None, auth=None, headers=None, **kwargs):388        if token is None and self.token:389            token = self.token.get('refresh_token') or self.token.get('access_token')390 391        if body is None:392            body = ''393 394        body, headers = prepare_revoke_token_request(395            token, token_type_hint, body, headers)396 397        for hook in self.compliance_hook[hook]:398            url, headers, body = hook(url, headers, body)399 400        if auth is None:401            auth = self.client_auth(self.revocation_endpoint_auth_method)402 403        session_kwargs = self._extract_session_request_params(kwargs)404        return self._http_post(405            url, body, auth=auth, headers=headers, **session_kwargs)406 407    def _prepare_token_endpoint_body(self, body, grant_type, **kwargs):408        if grant_type == 'authorization_code':409            if 'redirect_uri' not in kwargs:410                kwargs['redirect_uri'] = self.redirect_uri411            return prepare_token_request(grant_type, body, **kwargs)412 413        if 'scope' not in kwargs and self.scope:414            kwargs['scope'] = self.scope415        return prepare_token_request(grant_type, body, **kwargs)416 417    def _extract_session_request_params(self, kwargs):418        """Extract parameters for session object from the passing ``**kwargs``."""419        rv = {}420        for k in self.SESSION_REQUEST_PARAMS:421            if k in kwargs:422                rv[k] = kwargs.pop(k)423        return rv424 425    def _http_post(self, url, body=None, auth=None, headers=None, **kwargs):426        return self.session.post(427            url, data=dict(url_decode(body)),428            headers=headers, auth=auth, **kwargs)429 430 431def _guess_grant_type(kwargs):432    if 'code' in kwargs:433        grant_type = 'authorization_code'434    elif 'username' in kwargs and 'password' in kwargs:435        grant_type = 'password'436    else:437        grant_type = 'client_credentials'438    return grant_type439 
codekingpro/portable-devtools · Team Ai