Fix regression bug in MFA device recognition.

dowjones · Nov 16, 2023 · e9777d4 · e9777d4
1 parent eba165e
commit e9777d4
Show file tree

Hide file tree

Showing 6 changed files with 118 additions and 24 deletions.
diff --git a/README.md b/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 [email protected] \
@@ -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
 ```

diff --git a/tests/functional/test_auth.py b/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_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_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):

diff --git a/tests/unit/test_http_client.py b/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
diff --git a/tokendito/http_client.py b/tokendito/http_client.py
@@ -3,6 +3,7 @@
 """This module handles HTTP client operations."""
 
 import logging
+import platform
 import sys
 from urllib.parse import urlparse
 
@@ -13,12 +14,37 @@
 logger = logging.getLogger(__name__)
 
 
+def generate_user_agent():
+    """Generate a user agent string."""
+    python_version = platform.python_version()
+    (system, _, release, _, _, _) = platform.uname()
+
+    base_os = "compatible"
+    if system == "Darwin":
+        base_os = "Macintosh"
+    elif system == "Linux":
+        base_os = "X11"
+    elif system == "Windows":
+        base_os = "Windows"
+    else:
+        logger.warning(f"Unknown platform: {system}")
+
+    user_agent = (
+        f"{__title__}/{__version__} "
+        f"({base_os}; {system}/{release}) "
+        f"Python/{python_version}; "
+        f"requests/{requests.__version__})"
+    )
+    logger.debug(f"User agent: {user_agent}")
+    return user_agent
+
+
 class HTTPClient:
     """Handles HTTP client operations."""
 
     def __init__(self):
         """Initialize the HTTPClient with a session object."""
-        user_agent = f"{__title__}/{__version__}"
+        user_agent = generate_user_agent()
         self.session = requests.Session()
         self.session.headers.update({"User-Agent": user_agent})
 

diff --git a/tokendito/okta.py b/tokendito/okta.py
@@ -651,6 +651,7 @@ def create_authn_cookies(authn_org_url, session_token):
     domain = urlparse(url).netloc
     cookies.set("sid", session_id, domain=urlparse(url).netloc, path="/")
     cookies.set("sessionToken", session_token, domain=domain, path="/")
+    HTTP_client.set_cookies(cookies)
     return cookies
 
 
@@ -690,7 +691,7 @@ def idp_auth(config):
         # authentication sends us a token
         # which we then put in our session cookies
 
-        HTTP_client.session.cookies = create_authn_cookies(config.okta["org"], session_token)
+        create_authn_cookies(config.okta["org"], session_token)
         logger.debug(
             f"""
             authenticated via local_authenticate
@@ -714,7 +715,8 @@ def idp_auth(config):
                 session_cookies: {HTTP_client.session.cookies}
                 """
             )
-            HTTP_client.session.cookies = oauth2_authorize(config)
+            cookies = oauth2_authorize(config)
+            HTTP_client.set_cookies(cookies)
 
     logger.debug(f"Returning session cookies: {HTTP_client.session.cookies}")
     return HTTP_client.session.cookies

diff --git a/tokendito/user.py b/tokendito/user.py
@@ -64,6 +64,16 @@ def cmd_interface(args):
         )
         sys.exit(1)
 
+    if config.user["use_device_token"]:
+        device_token = config.okta["device_token"]
+        if device_token:
+            HTTP_client.set_device_token(config.okta["org"], device_token)
+        else:
+            logger.warning(
+                f"Device token unavailable for config profile {args.user_config_profile}. "
+                "May see multiple MFA requests this time."
+            )
+
     # get authentication and authorization cookies from okta
     session_cookies = okta.idp_auth(config)
     logger.debug(
@@ -99,6 +109,12 @@ def cmd_interface(args):
         output=config.aws["output"],
     )
 
+    device_token = HTTP_client.get_device_token()
+    if config.user["use_device_token"] and device_token:
+        logger.info(f"Saving device token to config profile {args.user_config_profile}")
+        config.okta["device_token"] = device_token
+        update_device_token(config)
+
     display_selected_role(profile_name=config.aws["profile"], role_response=role_response)
 
 
@@ -353,7 +369,7 @@ def select_role_arn(authenticated_tiles):
         if roles.count(config.aws["profile"]) > 1:
             logger.error(
                 "There are multiple matches for the profile selected, "
-                "please use the --role-arn option to select one"
+                "please use the --aws-role-arn option to select one"
             )
             sys.exit(2)
 
@@ -702,7 +718,6 @@ def process_arguments(args):
     pattern = re.compile(r"^(.*?)_(.*)")
 
     for key, val in vars(args).items():
-        logger.debug(f"key is {key} and val is {val}")
         match = re.search(pattern, key.lower())
         if match:
             if match.group(1) not in get_submodule_names():