from requests.auth import AuthBase


class TokenAuthentication(AuthBase):
    def __init__(self, token, scheme="Bearer"):
        """
        * Use an unauthenticated client, and make a request to obtain a token.
        * Create an authenticated client using eg. `TokenAuthentication(token="<token>")`
        """
        self.token = token
        self.scheme = scheme

    def __call__(self, request):
        request.headers["Authorization"] = "%s %s" % (self.scheme, self.token)
        return request


class SessionAuthentication(AuthBase):
    """
    Enables session based login.

    * Make an initial request to obtain a CSRF token.
    * Make a login request.
    """

    safe_methods = ("GET", "HEAD", "OPTIONS", "TRACE")

    def __init__(self, csrf_cookie_name=None, csrf_header_name=None):
        self.csrf_cookie_name = csrf_cookie_name
        self.csrf_header_name = csrf_header_name
        self.csrf_token = None

    def store_csrf_token(self, response, **kwargs):
        if self.csrf_cookie_name in response.cookies:
            self.csrf_token = response.cookies[self.csrf_cookie_name]

    def __call__(self, request):
        if (
            self.csrf_token
            and self.csrf_header_name is not None
            and (request.method not in self.safe_methods)
        ):
            request.headers[self.csrf_header_name] = self.csrf_token
        if self.csrf_cookie_name is not None:
            request.register_hook("response", self.store_csrf_token)
        return request
