From 442853122bbaef6d197c91368e5ae79425ed496e Mon Sep 17 00:00:00 2001 From: Geoff Mishkin Date: Tue, 8 Oct 2024 12:54:34 -0400 Subject: [PATCH] refactor 401 retrying into both auth methods --- jira/client.py | 122 ++++++++++++++++++++++++++++++++----------------- 1 file changed, 80 insertions(+), 42 deletions(-) diff --git a/jira/client.py b/jira/client.py index 36dc2fea9..c33fda876 100644 --- a/jira/client.py +++ b/jira/client.py @@ -307,30 +307,24 @@ def _sort_and_quote_values(self, values): return [quote(value, safe="~") for value in ordered_values] -class JiraCookieAuth(AuthBase): - """Jira Cookie Authentication. - - Allows using cookie authentication as described by `jira api docs `_ - """ +class RetryingJiraAuth(AuthBase): + """Base class for Jira authentication handlers that need to retry requests on 401 responses.""" - def __init__( - self, session: ResilientSession, session_api_url: str, auth: tuple[str, str] - ): - """Cookie Based Authentication. - - Args: - session (ResilientSession): The Session object to communicate with the API. - session_api_url (str): The session api url to use. - auth (Tuple[str, str]): The username, password tuple. - """ + def __init__(self, session: ResilientSession | None = None): self._session = session - self._session_api_url = session_api_url # e.g ."/rest/auth/1/session" - self.__auth = auth self._retry_counter_401 = 0 self._max_allowed_401_retries = 1 # 401 aren't recoverable with retries really + def init_session(self): + """Auth mechanism specific code to re-initialize the Jira session.""" + raise NotImplementedError() + @property def cookies(self): + """Return the cookies from the session.""" + assert ( + self._session is not None + ) # handle_401 should've caught this before attempting retry return self._session.cookies def _increment_401_retry_counter(self): @@ -339,22 +333,6 @@ def _increment_401_retry_counter(self): def _reset_401_retry_counter(self): self._retry_counter_401 = 0 - def __call__(self, request: requests.PreparedRequest): - request.register_hook("response", self.handle_401) - return request - - def init_session(self): - """Initialise the Session object's cookies, so we can use the session cookie. - - Raises HTTPError if the post returns an erroring http response - """ - username, password = self.__auth - authentication_data = {"username": username, "password": password} - r = self._session.post( # this also goes through the handle_401() hook - self._session_api_url, data=json.dumps(authentication_data) - ) - r.raise_for_status() - def handle_401(self, response: requests.Response, **kwargs) -> requests.Response: """Refresh cookies if the session cookie has expired. Then retry the request. @@ -364,14 +342,19 @@ def handle_401(self, response: requests.Response, **kwargs) -> requests.Response Returns: requests.Response """ - if ( + is_retryable_401 = ( response.status_code == 401 and self._retry_counter_401 < self._max_allowed_401_retries - ): + ) + + if is_retryable_401 and self._session is not None: LOG.info("Trying to refresh the cookie auth session...") self._increment_401_retry_counter() self.init_session() response = self.process_original_request(response.request.copy()) + elif is_retryable_401 and self._session is None: + LOG.warning("No session was passed to constructor, can't refresh cookies.") + self._reset_401_retry_counter() return response @@ -379,21 +362,67 @@ def process_original_request(self, original_request: requests.PreparedRequest): self.update_cookies(original_request) return self.send_request(original_request) + def update_cookies(self, original_request: requests.PreparedRequest): + """Auth mechanism specific cookie handling prior to retrying.""" + raise NotImplementedError() + + def send_request(self, request: requests.PreparedRequest): + if self._session is not None: + request.prepare_cookies(self.cookies) # post-update re-prepare + return self._session.send(request) + + +class JiraCookieAuth(RetryingJiraAuth): + """Jira Cookie Authentication. + + Allows using cookie authentication as described by `jira api docs `_ + """ + + def __init__( + self, session: ResilientSession, session_api_url: str, auth: tuple[str, str] + ): + """Cookie Based Authentication. + + Args: + session (ResilientSession): The Session object to communicate with the API. + session_api_url (str): The session api url to use. + auth (Tuple[str, str]): The username, password tuple. + """ + super().__init__(session) + self._session_api_url = session_api_url # e.g ."/rest/auth/1/session" + self.__auth = auth + + def __call__(self, request: requests.PreparedRequest): + request.register_hook("response", self.handle_401) + return request + + def init_session(self): + """Initialise the Session object's cookies, so we can use the session cookie. + + Raises HTTPError if the post returns an erroring http response + """ + assert ( + self._session is not None + ) # Constructor for this subclass always takes a session + username, password = self.__auth + authentication_data = {"username": username, "password": password} + r = self._session.post( # this also goes through the handle_401() hook + self._session_api_url, data=json.dumps(authentication_data) + ) + r.raise_for_status() + def update_cookies(self, original_request: requests.PreparedRequest): # Cookie header needs first to be deleted for the header to be updated using the # prepare_cookies method. See request.PrepareRequest.prepare_cookies if "Cookie" in original_request.headers: del original_request.headers["Cookie"] - original_request.prepare_cookies(self.cookies) - - def send_request(self, request: requests.PreparedRequest): - return self._session.send(request) -class TokenAuth(AuthBase): +class TokenAuth(RetryingJiraAuth): """Bearer Token Authentication.""" - def __init__(self, token: str): + def __init__(self, token: str, session: ResilientSession | None = None): + super().__init__(session) # setup any auth-related data here self._token = token @@ -402,6 +431,15 @@ def __call__(self, r: requests.PreparedRequest): r.headers["authorization"] = f"Bearer {self._token}" return r + def init_session(self): + pass # token should still work, only thing needed is to clear session cookies which happens next + + def update_cookies(self, _): + assert ( + self._session is not None + ) # handle_401 on the superclass should've caught this before attempting retry + self._session.cookies.clear_session_cookies() + class JIRA: """User interface to Jira. @@ -4306,7 +4344,7 @@ def _create_token_session(self, token_auth: str): Header structure: "authorization": "Bearer ". """ - self._session.auth = TokenAuth(token_auth) + self._session.auth = TokenAuth(token_auth, session=self._session) def _set_avatar(self, params, url, avatar): data = {"id": avatar}