codekingpro/portable-devtools

1from requests import Session2from authlib.oauth2.rfc7521 import AssertionClient3from authlib.oauth2.rfc7523 import JWTBearerGrant4from .oauth2_session import OAuth2Auth5from .utils import update_session_configure6 7 8class AssertionAuth(OAuth2Auth):9    def ensure_active_token(self):10        if not self.token or self.token.is_expired() and self.client:11            return self.client.refresh_token()12 13 14class AssertionSession(AssertionClient, Session):15    """Constructs a new Assertion Framework for OAuth 2.0 Authorization Grants16    per RFC7521_.17 18    .. _RFC7521: https://tools.ietf.org/html/rfc752119    """20    token_auth_class = AssertionAuth21    JWT_BEARER_GRANT_TYPE = JWTBearerGrant.GRANT_TYPE22    ASSERTION_METHODS = {23        JWT_BEARER_GRANT_TYPE: JWTBearerGrant.sign,24    }25    DEFAULT_GRANT_TYPE = JWT_BEARER_GRANT_TYPE26 27    def __init__(self, token_endpoint, issuer, subject, audience=None, grant_type=None,28                 claims=None, token_placement='header', scope=None, default_timeout=None, **kwargs):29        Session.__init__(self)30        self.default_timeout = default_timeout31        update_session_configure(self, kwargs)32        AssertionClient.__init__(33            self, session=self,34            token_endpoint=token_endpoint, issuer=issuer, subject=subject,35            audience=audience, grant_type=grant_type, claims=claims,36            token_placement=token_placement, scope=scope, **kwargs37        )38 39    def request(self, method, url, withhold_token=False, auth=None, **kwargs):40        """Send request with auto refresh token feature."""41        if self.default_timeout:42            kwargs.setdefault('timeout', self.default_timeout)43        if not withhold_token and auth is None:44            auth = self.token_auth45        return super().request(46            method, url, auth=auth, **kwargs)47