Fix regression bug in MFA device recognition. (#145)

* Fix regression bug that removed trusted device recongnition.
* Fix User-Agent field to work with Okta.
* Clean up cookie passing and handling, to simplify development.

Author: pcmxgti

README.md

@@ -23,15 +23,18 @@ your AWS accounts, returning
 tokens into your local `~/.aws/credentials` file.
 
 ## What's new
+
 See [Releases](https://github.com/dowjones/tokendito/releases) for a detailed Changelog.
+
 ### Tokendito 2.3.0
+
 Version 2.3.0 of Tokendito introduces the following new features:
-- Basic OIE support while forcing Classic mode.    
+
+- Basic OIE support while forcing Classic mode.
 - Misc bug fixes
 
 Note: This feature currently works with locally enabled OIE organizations, but it does not for Organizations with chained Authentication in mixed OIE/Classic environments.
 
-
 ### Tokendito 2.2.0
 
 Version 2.2.0 of Tokendito introduces the following new features:
@@ -40,7 +43,6 @@ Version 2.2.0 of Tokendito introduces the following new features:
 - Support for Step-Up Authorization (by @ruhulio)
 - Misc bug fixes
 
-
 ### Tokendito 2.1.0
 
 Version 2.1.0 of Tokendito introduces the following new features:
@@ -51,9 +53,9 @@ Version 2.1.0 of Tokendito introduces the following new features:
 - Docker container signing to ensure you are on a 'certified' Tokendito container
 - Misc bug fixes
 
-
 ### Tokendito 2.0.0
-With the release of tokendito 2.0, many changes and fixes were introduced. **It is a breaking release**: your configuration needs to be updated, the command line arguments have changed, and support for Python < 3.7 has been removed.
+
+With the release of tokendito 2.0, many changes and fixes were introduced. **It is a breaking release**: your configuration needs to be updated, the command line arguments have changed, and support for Python \< 3.7 has been removed.
 The following changes are part of this release:
 
 - Set the config file to be platform dependent, and follow the XDG standard.
@@ -71,25 +73,24 @@ Consult [additional notes](https://github.com/dowjones/tokendito/blob/main/docs/
 
 ## Requirements
 
--   Python 3.7+, or a working Docker environment
--   AWS account(s) federated with Okta
+- Python 3.7+, or a working Docker environment
+- AWS account(s) federated with Okta
 
 Tokendito is compatible with Python 3 and can be installed with either
 pip or pip3.
 
 ## Getting started
 
-1.  Install (via PyPi): `pip install tokendito`
-2.  Run `tokendito --configure`.
-3.  Run `tokendito`.
+1. Install (via PyPi): `pip install tokendito`
+1. Run `tokendito --configure`.
+1. Run `tokendito`.
 
 **NOTE**: Advanced users may shorten the `tokendito` interaction to a [single
 command](https://github.com/dowjones/tokendito/blob/main/docs/README.md#single-command-usage).
 
 Have multiple Okta tiles to switch between? View our [multi-tile
 guide](https://github.com/dowjones/tokendito/blob/main/docs/README.md#multi-tile-guide).
 
-
 ## Docker
 
 Using Docker eliminates the need to install tokendito and its requirements. We are providing experimental Docker image support in [Dockerhub](https://hub.docker.com/r/tokendito/tokendito)
@@ -98,13 +99,13 @@ Using Docker eliminates the need to install tokendito and its requirements. We a
 
 Run tokendito with the `docker run` command. Tokendito supports [DCT](https://docs.docker.com/engine/security/trust/), and we encourage you to enforce image signature validation before running any containers.
 
-``` shell
+```shell
 export DOCKER_CONTENT_TRUST=1
 ```
 
 then
 
-``` shell
+```shell
 docker run --rm -it tokendito/tokendito  --version
 ```
 
@@ -118,27 +119,29 @@ These can be covered by mapping a single volume to both the host and container u
 Be sure to set the `-it` flags to enable an interactive terminal session.
 
 On Windows, you can do the following:
-``` powershell
+
+```powershell
 docker run --rm -it -v "%USERPROFILE%\.aws":/app/.aws  -v "%USERPROFILE%\.config":/app/.config tokendito/tokendito
 ```
 
 In a Mac OS system, you can run:
-``` shell
+
+```shell
 docker run --rm -it -v "$HOME/.aws":/app/.aws  -v "$HOME/.config":/app/.config tokendito/tokendito
 ```
 
 On a Linux system, however, you must specify the user and group IDs for the mount mappings to work as expected.
 Additionally the mount points within the container move to a different location:
 
-``` shell
+```shell
 docker run --user $(id -u):$(id -g) --rm -it -v "$HOME/.aws":/.aws  -v "$HOME/.config":/.config tokendito/tokendito
 ```
 
 Tokendito command line arguments are supported as well.
 
 **NOTE**: In the following examples the entire home directory is exported for simplicity. This is not recommended as it exposes too much data to the running container:
 
-``` shell
+```shell
 docker run --rm -it -v "$HOME":/ tokendito/tokendito \
   --okta-tile https://acme.okta.com/home/amazon_aws/000000000000000000x0/123 \
   --username username@example.com \
@@ -151,7 +154,7 @@ docker run --rm -it -v "$HOME":/ tokendito/tokendito \
 
 Tokendito profiles are supported while using containers provided the proper volume mapping exists.
 
-``` shell
+```shell
 docker run --rm -ti -v "$HOME":/app tokendito/tokendito \
   --profile my-profile-name
 ```

tests/functional/test_auth.py

@@ -73,6 +73,9 @@ def test_generate_credentials(custom_args, config_file):
         f"{config.okta['username']}",
         "--password",
         f"{config.okta['password']}",
+        "--config-file",
+        f"{config.user['config_file']}",
+        "--use-device-token",
         "--loglevel",
         "DEBUG",
     ]
@@ -87,6 +90,16 @@ def test_generate_credentials(custom_args, config_file):
     assert '"sessionToken": "*****"' in proc["stderr"]
     assert proc["exit_status"] == 0
 
+    # Ensure the device token is written to the config file, and is correct.
+    device_token = None
+    match = re.search(r"(?<=okta_device_token': ')[^']+", proc["stderr"])
+    if match:
+        device_token = match.group(0)
+    with open(config.user["config_file"]) as cfg:
+        assert f"okta_device_token = {device_token}" in cfg.read()
+
+    # print(f"stderr: {proc['stderr']}")
+
 
 @pytest.mark.run("second")
 def test_aws_credentials(custom_args):

tests/unit/test_aws.py

@@ -110,7 +110,6 @@ def test_authenticate_to_roles(status_code, monkeypatch):
             "org": "https://acme.okta.org/",
         }
     )
-    cookies = {"some_cookie": "some_value"}
 
     with pytest.raises(SystemExit):
-        authenticate_to_roles(pytest_config, [("http://test.url.com", "")], cookies)
+        authenticate_to_roles(pytest_config, [("http://test.url.com", "")])

tests/unit/test_http_client.py

@@ -17,14 +17,41 @@ def client():
     return HTTPClient()
 
 
+@pytest.mark.parametrize(
+    "base_os, expected",
+    [
+        ("Darwin", "Macintosh"),
+        ("Linux", "X11"),
+        ("Windows", "Windows"),
+        ("Unknown", "compatible"),
+    ],
+)
+def test_generate_user_agent(mocker, base_os, expected):
+    """Test the generate_user_agent function."""
+    import platform
+
+    from tokendito.http_client import generate_user_agent
+
+    mocker.patch("platform.uname", return_value=(base_os, "", "pytest", "", "", ""))
+    python_version = platform.python_version()
+
+    user_agent = generate_user_agent()
+    assert user_agent == (
+        f"{__title__}/{__version__} "
+        f"({expected}; {base_os}/pytest) "
+        f"Python/{python_version}; "
+        f"requests/{requests.__version__})"
+    )
+
+
 def test_init(client):
     """Test initialization of HTTPClient instance."""
     # Check if the session property of the client is an instance of requests.Session
     assert isinstance(client.session, requests.Session)
 
     # Check if the User-Agent header was set correctly during initialization
     expected_user_agent = f"{__title__}/{__version__}"
-    assert client.session.headers["User-Agent"] == expected_user_agent
+    assert str(expected_user_agent) in str(client.session.headers["User-Agent"])
 
 
 def test_set_cookies(client):
@@ -166,6 +193,10 @@ def test_get_device_token(client):
     # Check if the device token is set correctly in the session
     assert client.get_device_token() == device_token
 
+    # Check no device token when the cookie is not set
+    client.session.cookies.clear()
+    assert client.get_device_token() is None
+
 
 def test_set_device_token(client):
     """Test setting device token in the session."""
@@ -174,3 +205,7 @@ def test_set_device_token(client):
 
     # Check if the device token is set correctly in the session
     assert client.session.cookies.get("DT") == device_token
+
+    # Check no device token set when the cookie is not set
+    client.session.cookies.clear()
+    assert client.set_device_token("http://test.com", None) is None

tests/unit/test_okta.py

@@ -4,6 +4,7 @@
 from unittest.mock import Mock
 
 import pytest
+import requests.cookies
 from tokendito.config import Config
 from tokendito.http_client import HTTP_client
 
@@ -304,11 +305,11 @@ def test_push_approval(mocker, return_value, side_effect, expected):
         ({"type": "SAML2"}, False),
     ],
 )
-def test_is_local_auth(auth_properties, expected):
+def test_local_authentication_enabled(auth_properties, expected):
     """Test local auth method."""
     from tokendito import okta
 
-    assert okta.is_local_auth(auth_properties) == expected
+    assert okta.local_authentication_enabled(auth_properties) == expected
 
 
 @pytest.mark.parametrize(
@@ -486,8 +487,12 @@ def test_send_saml_response(mocker):
     from tokendito.config import Config
     from tokendito.http_client import HTTP_client
 
+    cookies = requests.cookies.RequestsCookieJar()
+    cookies.set("sid", "pytestcookie")
     mock_response = Mock()
-    mock_response.cookies = {"sid": "pytestcookie"}
+    mock_response.status_code = 201
+    mock_response.session = Mock()
+    mock_response.session.cookies = cookies
 
     saml_response = {
         "response": "pytestresponse",
@@ -501,7 +506,7 @@ def test_send_saml_response(mocker):
 
     pytest_config = Config()
 
-    assert okta.send_saml_response(pytest_config, saml_response) == mock_response.cookies
+    assert okta.send_saml_response(pytest_config, saml_response) is None
 
 
 def test_idp_auth(mocker):
@@ -522,10 +527,10 @@ def test_idp_auth(mocker):
     mocker.patch("tokendito.okta.saml2_authenticate", return_value=sid)
 
     mocker.patch("tokendito.okta.get_auth_properties", return_value={"type": "OKTA"})
-    assert okta.idp_auth(pytest_config) == sid
+    assert okta.idp_auth(pytest_config) is None
 
     mocker.patch("tokendito.okta.get_auth_properties", return_value={"type": "SAML2"})
-    assert okta.idp_auth(pytest_config) == sid
+    assert okta.idp_auth(pytest_config) is None
 
     mocker.patch("tokendito.okta.get_auth_properties", return_value={"type": "UNKNOWN"})
     with pytest.raises(SystemExit) as error:
@@ -585,7 +590,7 @@ def test_step_up_authenticate(mocker):
     assert okta.step_up_authenticate(pytest_config, state_token) is False
 
 
-def test_local_auth(mocker):
+def test_local_authenticate(mocker):
     """Test local auth method."""
     from tokendito import okta
     from tokendito.config import Config
@@ -606,7 +611,7 @@ def test_local_auth(mocker):
         }
     )
 
-    assert okta.local_auth(pytest_config) == "pytesttoken"
+    assert okta.local_authenticate(pytest_config) == "pytesttoken"
 
 
 def test_saml2_authenticate(mocker):
@@ -634,5 +639,5 @@ def test_saml2_authenticate(mocker):
     }
 
     mocker.patch("tokendito.okta.send_saml_request", return_value=saml_response)
-    mocker.patch("tokendito.okta.send_saml_response", return_value="pytestsessionid")
-    assert okta.saml2_authenticate(pytest_config, auth_properties) == "pytestsessionid"
+    mocker.patch("tokendito.okta.send_saml_response", return_value=None)
+    assert okta.saml2_authenticate(pytest_config, auth_properties) is None

tokendito/__init__.py

@@ -1,7 +1,7 @@
 # vim: set filetype=python ts=4 sw=4
 # -*- coding: utf-8 -*-
 """Tokendito module initialization."""
-__version__ = "2.3.0"
+__version__ = "2.3.1"
 __title__ = "tokendito"
 __description__ = "Get AWS STS tokens from Okta SSO"
 __long_description_content_type__ = "text/markdown"

tokendito/aws.py

@@ -47,11 +47,11 @@ def get_output_types():
     return ["json", "text", "csv", "yaml", "yaml-stream"]
 
 
-def authenticate_to_roles(config, urls, cookies):
+def authenticate_to_roles(config, urls):
     """Authenticate AWS user with saml.
 
-    :param urls: list of tuples or tuple, with tiles info
-    :param cookies: html cookies
+    :param config: configuration object
+    :param urls: list of tuples or tuple, with tiles information
     :return: response text
 
     """
@@ -63,7 +63,8 @@ def authenticate_to_roles(config, urls, cookies):
     logger.info(f"Discovering roles in {tile_count} tile{plural}.")
     for url, label in url_list:
         session_url = config.okta["org"] + "/login/sessionCookieRedirect"
-        params = {"token": cookies.get("sessionToken"), "redirectUrl": url}
+        token = HTTP_client.session.cookies.get("sessionToken", None)
+        params = {"token": token, "redirectUrl": url}
         response = HTTP_client.get(session_url, params=params)
 
         saml_response_string = response.text
@@ -74,7 +75,7 @@ def authenticate_to_roles(config, urls, cookies):
             if "Extra Verification" in saml_response_string and state_token:
                 logger.info(f"Step-Up authentication required for {url}.")
                 if okta.step_up_authenticate(config, state_token):
-                    return authenticate_to_roles(config, urls, cookies)
+                    return authenticate_to_roles(config, urls)
 
                 logger.error("Step-Up Authentication required, but not supported.")
             elif "App Access Locked" in saml_response_string: