authenticator.py

import dpath
import pendulum
import requests
from requests.auth import AuthBase
from typing import Any, List, Mapping, MutableMapping, Tuple, Union
from abc import abstractmethod


class AbstractGithubIntunedAuthenticator(AuthBase):
    """
    Abstract class for an OAuth authenticators that implements the OAuth token refresh flow. The authenticator
    is designed to generically perform the refresh flow without regard to how config fields are get/set by
    delegating that behavior to the classes implementing the interface.
    """

    def __call__(self, request: requests.Request) -> requests.Request:
        """Attach the HTTP headers required to authenticate on the HTTP request"""
        request.headers.update(self.get_auth_header())
        return request

    def get_auth_header(self) -> Mapping[str, Any]:
        """HTTP header to set on the requests"""
        return {"Authorization": f"Bearer {self.get_access_token()}"}

    def get_access_token(self) -> str:
        """Returns the access token"""
        if self.token_has_expired():
            current_datetime = pendulum.now()
            accessToken, expires_at = self.refresh_access_token()
            self.access_token = accessToken
            self.set_token_expiry_date(current_datetime, expires_at)
            return self.access_token
        return self.access_token

    def token_has_expired(self) -> bool:
        """Returns True if the token is expired"""
        if not self.get_token_expiry_date():
            return True
        return pendulum.now() > self.get_token_expiry_date()

    def _get_refresh_access_token_response(self):
        response = requests.request(method="POST", url=self.get_token_refresh_endpoint(), data=self.build_refresh_request_body())
        response.raise_for_status()
        return response.json()

    def refresh_access_token(self) -> Tuple[str, int]:
        """
        Returns the refresh token and its lifespan in seconds

        :return: a tuple of (access_token, token_lifespan_in_seconds)
        """
        try:
            response_json = self._get_refresh_access_token_response()
            return (
                response_json["accessToken"],
                response_json["expiresAt"],
            )
        except Exception as e:
            raise Exception(f"Error while refreshing access token: {e}") from e

    def build_refresh_request_body(self) -> Mapping[str, Any]:
        """
        Returns the request body to set on the refresh request

        Override to define additional parameters
        """
        payload: MutableMapping[str, Any] = {
            "installationId": self.get_installation_id(),
            "secret": self.get_secret(),
        }
        return payload

    @abstractmethod
    def get_token_refresh_endpoint(self) -> str:
        """Returns the endpoint to refresh the access token"""

    @abstractmethod
    def get_installation_id(self) -> int:
        """The client id to authenticate"""

    @abstractmethod
    def get_secret(self) -> str:
        """The client secret to authenticate"""

    @abstractmethod
    def get_token_expiry_date(self) -> pendulum.DateTime:
        """Expiration date of the access token"""

    @abstractmethod
    def set_token_expiry_date(self, initial_time: pendulum.DateTime, value: Union[str, int]):
        """Setter for access token expiration date"""

    @property
    @abstractmethod
    def access_token(self) -> str:
        """Returns the access token"""

    @access_token.setter
    @abstractmethod
    def access_token(self, value: str) -> str:
        """Setter for the access token"""

class GithubIntunedAuthenticator(AbstractGithubIntunedAuthenticator):
    """
    Generates OAuth2.0 access tokens from an OAuth2.0 refresh token and client credentials.
    The generated access token is attached to each request via the Authorization header.
    If a connector_config is provided any mutation of it's value in the scope of this class will emit AirbyteControlConnectorConfigMessage.
    """

    def __init__(
        self,
        token_refresh_endpoint: str,
        installation_id: str,
        secret: str,
    ):
        self._token_refresh_endpoint = token_refresh_endpoint
        self.installation_id = installation_id
        self.secret = secret
        self._token_expiry_date = None
        self._access_token = None

    def get_token_refresh_endpoint(self) -> str:
        return self._token_refresh_endpoint

    def get_installation_id(self) -> str:
        return self.installation_id

    def get_secret(self) -> str:
        return self.secret

    def get_access_token_name(self) -> str:
        return self._access_token_name

    def get_token_expiry_date(self) -> pendulum.DateTime:
        return self._token_expiry_date

    def set_token_expiry_date(self, initial_time: pendulum.DateTime, value: Union[str, int]):
        self._token_expiry_date = pendulum.parse(value)
     

    @property
    def access_token(self) -> str:
        return self._access_token

    @access_token.setter
    def access_token(self, value: str):
        self._access_token = value