From 7d5dbd1f72e4bcc704d35ba46da1244214626d44 Mon Sep 17 00:00:00 2001 From: arithmetic1728 Date: Wed, 7 Apr 2021 15:45:54 -0700 Subject: [PATCH] update --- google/oauth2/challenges.py | 22 +-------------- google/oauth2/reauth.py | 44 ++++++++++++++---------------- tests/oauth2/test_challenges.py | 14 ---------- tests/oauth2/test_reauth.py | 48 ++++++++++++++++++--------------- 4 files changed, 47 insertions(+), 81 deletions(-) diff --git a/google/oauth2/challenges.py b/google/oauth2/challenges.py index d41677027..891d2a876 100644 --- a/google/oauth2/challenges.py +++ b/google/oauth2/challenges.py @@ -133,27 +133,7 @@ def obtain_challenge_input(self, metadata): return None -class SamlChallenge(ReauthChallenge): - """Challenge that asks the users to browse to their ID Providers.""" - - @property - def name(self): - return "SAML" - - @property - def is_locally_eligible(self): - return True - - def obtain_challenge_input(self, metadata): - # Magic Arch has not fully supported returning a proper dedirect URL - # for programmatic SAML users today. So we error our here and request - # users to complete a web login. - raise exceptions.ReauthFailError( - "SAML login is required for the current account to complete reauthentication." - ) - - AVAILABLE_CHALLENGES = { challenge.name: challenge - for challenge in [SecurityKeyChallenge(), PasswordChallenge(), SamlChallenge()] + for challenge in [SecurityKeyChallenge(), PasswordChallenge()] } diff --git a/google/oauth2/reauth.py b/google/oauth2/reauth.py index 517a15994..ed60fd1bc 100644 --- a/google/oauth2/reauth.py +++ b/google/oauth2/reauth.py @@ -54,6 +54,11 @@ _CHALLENGE_PENDING = "CHALLENGE_PENDING" +# Override this global variable to set custom max number of rounds of reauth +# challenges should be run. +RUN_CHALLENGE_RETRY_LIMIT = 5 + + def _get_challenges( request, supported_challenge_types, access_token, requested_scopes=None ): @@ -161,7 +166,7 @@ def _run_next_challenge(msg, request, access_token): return None -def _obtain_rapt(request, access_token, requested_scopes, rounds_num=5): +def _obtain_rapt(request, access_token, requested_scopes): """Given an http request method and reauth access token, get rapt token. Args: @@ -169,10 +174,6 @@ def _obtain_rapt(request, access_token, requested_scopes, rounds_num=5): HTTP requests. access_token (str): reauth access token requested_scopes (Sequence[str]): scopes required by the client application - rounds_num (Optional(int)): max number of attempts to get a rapt after the next - challenge, before failing the reauth. This defines total number of - challenges + number of additional retries if the chalenge input - wasn't accepted. Returns: str: The rapt token. @@ -180,21 +181,17 @@ def _obtain_rapt(request, access_token, requested_scopes, rounds_num=5): Raises: google.auth.exceptions.ReauthError: if reauth failed """ - msg = None - - for _ in range(0, rounds_num): + msg = _get_challenges( + request, + list(challenges.AVAILABLE_CHALLENGES.keys()), + access_token, + requested_scopes, + ) - if not msg: - msg = _get_challenges( - request, - list(challenges.AVAILABLE_CHALLENGES.keys()), - access_token, - requested_scopes, - ) - - if msg["status"] == _AUTHENTICATED: - return msg["encodedProofOfReauthToken"] + if msg["status"] == _AUTHENTICATED: + return msg["encodedProofOfReauthToken"] + for _ in range(0, RUN_CHALLENGE_RETRY_LIMIT): if not ( msg["status"] == _CHALLENGE_REQUIRED or msg["status"] == _CHALLENGE_PENDING ): @@ -204,18 +201,17 @@ def _obtain_rapt(request, access_token, requested_scopes, rounds_num=5): ) ) - """Check if we are in an interractive environment. - - If the rapt token needs refreshing, the user needs to answer the - challenges. - """ if not _helpers.is_interactive(): raise exceptions.ReauthFailError( - "Reauthentication challenge could not be answered because you are not in an interactive session." + "Reauthentication challenge could not be answered because you are not" + " in an interactive session." ) msg = _run_next_challenge(msg, request, access_token) + if msg["status"] == _AUTHENTICATED: + return msg["encodedProofOfReauthToken"] + # If we got here it means we didn't get authenticated. raise exceptions.ReauthFailError() diff --git a/tests/oauth2/test_challenges.py b/tests/oauth2/test_challenges.py index 854266401..2cd48a94f 100644 --- a/tests/oauth2/test_challenges.py +++ b/tests/oauth2/test_challenges.py @@ -125,17 +125,3 @@ def test_password_challenge(getpass_mock): assert challenges.PasswordChallenge().obtain_challenge_input({}) == { "credential": " " } - - -def test_saml_challenge(): - metadata = { - "status": "READY", - "challengeId": 1, - "challengeType": "SAML", - "securityKey": {}, - } - challenge = challenges.SamlChallenge() - assert challenge.is_locally_eligible - assert challenge.name == "SAML" - with pytest.raises(exceptions.ReauthFailError): - challenge.obtain_challenge_input(metadata) diff --git a/tests/oauth2/test_reauth.py b/tests/oauth2/test_reauth.py index 4c036e1a0..3dde8c1ee 100644 --- a/tests/oauth2/test_reauth.py +++ b/tests/oauth2/test_reauth.py @@ -162,21 +162,30 @@ def test__run_next_challenge_success(): ) -def test__obtain_rapt_not_authenticated(): - with pytest.raises(exceptions.ReauthFailError) as excinfo: - reauth._obtain_rapt(MOCK_REQUEST, "token", None, rounds_num=0) - assert excinfo.match(r"Reauthentication failed. None") - - def test__obtain_rapt_authenticated(): with mock.patch( "google.oauth2.reauth._get_challenges", return_value=CHALLENGES_RESPONSE_AUTHENTICATED, ): - assert ( - reauth._obtain_rapt(MOCK_REQUEST, "token", None, rounds_num=1) - == "new_rapt_token" - ) + assert reauth._obtain_rapt(MOCK_REQUEST, "token", None) == "new_rapt_token" + + +def test__obtain_rapt_authenticated_after_run_next_challenge(): + with mock.patch( + "google.oauth2.reauth._get_challenges", + return_value=CHALLENGES_RESPONSE_TEMPLATE, + ): + with mock.patch( + "google.oauth2.reauth._run_next_challenge", + side_effect=[ + CHALLENGES_RESPONSE_TEMPLATE, + CHALLENGES_RESPONSE_AUTHENTICATED, + ], + ): + with mock.patch("google.auth._helpers.is_interactive", return_value=True): + assert ( + reauth._obtain_rapt(MOCK_REQUEST, "token", None) == "new_rapt_token" + ) def test__obtain_rapt_unsupported_status(): @@ -186,7 +195,7 @@ def test__obtain_rapt_unsupported_status(): "google.oauth2.reauth._get_challenges", return_value=challenges_response ): with pytest.raises(exceptions.ReauthFailError) as excinfo: - reauth._obtain_rapt(MOCK_REQUEST, "token", None, rounds_num=1) + reauth._obtain_rapt(MOCK_REQUEST, "token", None) assert excinfo.match(r"API error: STATUS_UNSPECIFIED") @@ -197,24 +206,19 @@ def test__obtain_rapt_not_interactive(): ): with mock.patch("google.auth._helpers.is_interactive", return_value=False): with pytest.raises(exceptions.ReauthFailError) as excinfo: - reauth._obtain_rapt(MOCK_REQUEST, "token", None, rounds_num=1) + reauth._obtain_rapt(MOCK_REQUEST, "token", None) assert excinfo.match(r"not in an interactive session") -def test__obtain_rapt_run_next_challenge(): +def test__obtain_rapt_not_authenticated(): with mock.patch( "google.oauth2.reauth._get_challenges", return_value=CHALLENGES_RESPONSE_TEMPLATE, ): - with mock.patch( - "google.oauth2.reauth._run_next_challenge", - return_value=CHALLENGES_RESPONSE_AUTHENTICATED, - ): - with mock.patch("google.auth._helpers.is_interactive", return_value=True): - assert ( - reauth._obtain_rapt(MOCK_REQUEST, "token", None, rounds_num=2) - == "new_rapt_token" - ) + with mock.patch("google.oauth2.reauth.RUN_CHALLENGE_RETRY_LIMIT", 0): + with pytest.raises(exceptions.ReauthFailError) as excinfo: + reauth._obtain_rapt(MOCK_REQUEST, "token", None) + assert excinfo.match(r"Reauthentication failed") def test_get_rapt_token():