diff --git a/CHANGES.rst b/CHANGES.rst index 0526581..c88f58e 100644 --- a/CHANGES.rst +++ b/CHANGES.rst @@ -9,6 +9,14 @@ prawcore follows `semantic versioning `_. **Added** - Support 202 "Accepted" HTTP responses. +- Calls to method ``ScriptAuthorizer.refresh`` are handled by a private method, + ``ScriptAuthorizer._refresh_with_retries``. The latter provides the option to retry + requests that result in invalid grants. + +**Changed** + +- Method ``ScriptAuthorizer.two_factor_callback`` can return either a string or a + tuple. **Fixed** @@ -21,7 +29,7 @@ prawcore follows `semantic versioning `_. **Added** - Add a ``URITooLarge`` exception. -- :class:`.ScriptAuthorizer` has a new parameter ``two_factor_callback `` that supplies +- :class:`.ScriptAuthorizer` has a new parameter ``two_factor_callback`` that supplies OTPs (One-Time Passcodes) when :meth:`.ScriptAuthorizer.refresh` is called. - Add a ``TooManyRequests`` exception. diff --git a/prawcore/auth.py b/prawcore/auth.py index 844594d..58a1907 100644 --- a/prawcore/auth.py +++ b/prawcore/auth.py @@ -357,28 +357,61 @@ class ScriptAuthorizer(Authorizer): AUTHENTICATOR_CLASS = TrustedAuthenticator def __init__( - self, authenticator, username, password, two_factor_callback=None + self, + authenticator, + username, + password, + two_factor_callback=None, ): """Represent a single personal-use authorization to Reddit's API. :param authenticator: An instance of :class:`TrustedAuthenticator`. :param username: The Reddit username of one of the application's developers. :param password: The password associated with ``username``. - :param two_factor_callback: A function that returns OTPs (One-Time - Passcodes), also known as 2FA auth codes. If this function is - provided, prawcore will call it when authenticating. + :param two_factor_callback: (Optional) A function that returns OTPs (One-Time + Passcodes), also known as 2FA auth codes. If provided, this function should + return either a string of six digits or a 3-tuple of the form + ``(, , )``, where ```` is a string of six + digits, ```` is an integer that represents the number of seconds + to sleep between invalid authorization attempts, and ```` is an + integer that represents the maximum number of authorization attempts to + make before an OAuthException is raised. """ - super(ScriptAuthorizer, self).__init__(authenticator) + super().__init__(authenticator) self._username = username self._password = password self._two_factor_callback = two_factor_callback + def _refresh_with_retries(self, count=1, delay=0, maxcount=1): + if delay > 0: + time.sleep(delay) + additional_kwargs = {} + otp = self._two_factor_callback and self._two_factor_callback() + if otp: + if isinstance(otp, tuple): + if otp[0]: + additional_kwargs["otp"] = otp[0] + else: + additional_kwargs["otp"] = otp + try: + self._request_token( + grant_type="password", + username=self._username, + password=self._password, + **additional_kwargs, + ) + except OAuthException: + if otp and isinstance(otp, tuple) and len(otp) == 3: + _, delay, maxcount = otp + if count >= min(maxcount, 10): + raise + self._refresh_with_retries( + count=count + 1, delay=delay, maxcount=maxcount + ) + else: + raise + def refresh(self): """Obtain a new personal-use script type access token.""" - self._request_token( - grant_type="password", - username=self._username, - password=self._password, - otp=self._two_factor_callback and self._two_factor_callback(), - ) + self._refresh_with_retries() diff --git a/tests/test_authorizer.py b/tests/test_authorizer.py index 813f321..777428d 100644 --- a/tests/test_authorizer.py +++ b/tests/test_authorizer.py @@ -2,6 +2,7 @@ import unittest from betamax import Betamax +from mock import Mock, patch import prawcore @@ -384,3 +385,19 @@ def test_refresh__with_invalid_username_or_password(self): ): self.assertRaises(prawcore.OAuthException, authorizer.refresh) self.assertFalse(authorizer.is_valid()) + + @patch("time.sleep", return_value=None) + @patch("prawcore.Requestor.request") + def test_refresh_with_retries(self, mock_post, _): + response = Mock( + json=lambda: {"error": "invalid grant"}, status_code=200 + ) + mock_post.side_effect = [response, response] + authorizer = prawcore.ScriptAuthorizer( + self.authentication, + "dummy", + "dummy", + two_factor_callback=lambda: ("123456", 31, 2), + ) + with self.assertRaises(prawcore.OAuthException): + authorizer.refresh()