From 1ee8b25ca34fed62fc4bb39730ce64ad570a5551 Mon Sep 17 00:00:00 2001 From: pcmxgti <16561338+pcmxgti@users.noreply.github.com> Date: Wed, 8 Nov 2023 15:28:24 -0500 Subject: [PATCH] Adjust API calls and tests so that calls to DUO work --- tests/unit/test_duo.py | 29 +++++++++++++++++------------ tokendito/duo.py | 17 +++++++++-------- 2 files changed, 26 insertions(+), 20 deletions(-) diff --git a/tests/unit/test_duo.py b/tests/unit/test_duo.py index efb95b07..266beb6f 100644 --- a/tests/unit/test_duo.py +++ b/tests/unit/test_duo.py @@ -104,11 +104,10 @@ def test_get_mfa_response(): mfa_result = Mock() # Test if response is correct - mfa_result.json = Mock(return_value={"response": "test_value"}) - assert get_mfa_response(mfa_result) == "test_value" + assert get_mfa_response({"response": "test_value"}) == "test_value" # Test if response is incorrect - mfa_result.json = Mock(return_value={"badresponse": "FAIL"}) + mfa_result = Mock(return_value={"badresponse": "FAIL"}) with pytest.raises(SystemExit) as err: get_mfa_response(mfa_result) assert err.value.code == 1 @@ -175,8 +174,7 @@ def test_parse_mfa_challenge(): mfa_challenge = Mock() # Test successful challenge - mfa_challenge.json = Mock(return_value={"stat": "OK", "response": {"txid": "pytest"}}) - assert parse_mfa_challenge(mfa_challenge) == "pytest" + assert parse_mfa_challenge({"stat": "OK", "response": {"txid": "pytest"}}) == "pytest" # Test error mfa_challenge.json = Mock(return_value={"stat": "OK", "response": "error"}) @@ -220,10 +218,9 @@ def test_mfa_challenge(mocker): passcode = "pytest_passcode" mfa_option = {"factor": "pytest_factor", "device": "pytest_device - pytest_device_name"} - duo_api_response = mocker.Mock() - duo_api_response.json.return_value = {"stat": "OK", "response": {"txid": "pytest_txid"}} - - mocker.patch("tokendito.duo.api_post", return_value=duo_api_response) + mocker.patch( + "tokendito.duo.api_post", return_value={"stat": "OK", "response": {"txid": "pytest_txid"}} + ) txid = mfa_challenge(duo_info, mfa_option, passcode) assert txid == "pytest_txid" @@ -284,18 +281,26 @@ def test_factor_callback(mocker): verify_mfa = {"result_url": "/pytest_result_url"} duo_api_response = mocker.Mock() - duo_api_response.json.return_value = { + duo_api_response.return_value = { "stat": "OK", "response": {"txid": "pytest_txid", "cookie": "pytest_cookie"}, } - mocker.patch("tokendito.duo.api_post", return_value=duo_api_response) + + mocker.patch( + "tokendito.duo.api_post", + return_value={ + "stat": "OK", + "response": {"txid": "pytest_txid", "cookie": "pytest_cookie"}, + }, + ) # Test successful retrieval of the cookie sig_response = factor_callback(duo_info, verify_mfa) assert sig_response == "pytest_cookie:pytest_tile_sig" # Test failure to retrieve the cookie - duo_api_response.json.return_value = {"stat": "FAIL", "response": "pytest_error"} + duo_api_response.return_value = {"stat": "FAIL", "response": "pytest_error"} + mocker.patch("tokendito.duo.api_post", return_value=duo_api_response) with pytest.raises(SystemExit) as err: factor_callback(duo_info, verify_mfa) assert err.value.code == 2 diff --git a/tokendito/duo.py b/tokendito/duo.py index 3853dcbd..7e31c94e 100644 --- a/tokendito/duo.py +++ b/tokendito/duo.py @@ -45,7 +45,7 @@ def prepare_info(selected_okta_factor): return duo_info -def api_post(url, params=None, headers=None, payload=None): +def api_post(url, params=None, headers=None, payload=None, return_json=True): """Error handling and response parsing wrapper for Duo POSTs. :param url: The URL being connected to. @@ -54,7 +54,9 @@ def api_post(url, params=None, headers=None, payload=None): :param payload: Request body. :return response: Response to the API request. """ - response = HTTP_client.post(url, params=params, headers=headers, data=payload, return_json=True) + response = HTTP_client.post( + url, params=params, headers=headers, data=payload, return_json=return_json + ) return response @@ -71,7 +73,7 @@ def get_sid(duo_info): url = f"https://{duo_info['host']}/frame/web/v1/auth" logger.debug(f"Calling Duo {urlparse(url).path} with params {params.keys()}") - duo_auth_response = api_post(url, params=params) + duo_auth_response = api_post(url, params=params, return_json=False) try: duo_auth_redirect = urlparse(f"{unquote(duo_auth_response.url)}").query @@ -124,10 +126,9 @@ def parse_mfa_challenge(mfa_challenge): :return txid: Duo transaction ID. """ try: - mfa_challenge = mfa_challenge.json() mfa_status = mfa_challenge["stat"] txid = mfa_challenge["response"]["txid"] - except (TypeError, ValueError) as err: + except (TypeError, ValueError, AttributeError) as err: logger.error(f"The Duo API returned a non-json response: {err}") sys.exit(1) except KeyError as key_error: @@ -185,7 +186,7 @@ def get_mfa_response(mfa_result): :return verify_mfa: json response from mfa api """ try: - verify_mfa = mfa_result.json()["response"] + verify_mfa = mfa_result["response"] except KeyError as key_error: logger.error(f"The mfa challenge response is missing a required parameter: {key_error}") logger.debug(json.dumps(mfa_result.json())) @@ -266,7 +267,7 @@ def factor_callback(duo_info, verify_mfa): factor_callback = api_post(factor_callback_url, payload={"sid": duo_info["sid"]}) try: - sig_response = f"{factor_callback.json()['response']['cookie']}:{duo_info['tile_sig']}" + sig_response = f"{factor_callback['response']['cookie']}:{duo_info['tile_sig']}" except Exception as sig_error: logger.error("There was an error getting your application signature. Please try again.") logger.debug(f"from Duo: {sig_error}") @@ -331,5 +332,5 @@ def authenticate(selected_okta_factor): } # Send Okta callback and return payload - api_post(duo_info["okta_callback_url"], payload=payload) + api_post(duo_info["okta_callback_url"], payload=payload, return_json=False) return payload