From ac34fa43c52787e3b174e0db49a8a7ac93501508 Mon Sep 17 00:00:00 2001 From: Nick Venenga Date: Fri, 18 Sep 2026 18:17:32 -0400 Subject: [PATCH] Serialize concurrent runs and atomic-write credential/config dotfiles Wrap the writes to ~/.aws/credentials (write_sts_token / set_default_profile) and ~/.okta-aws (write_role_to_profile / write_applink_to_profile) with a cross-process advisory lock (filelock) and an atomic temp-file + os.replace, so parallel okta-awscli invocations can no longer corrupt or clobber each other's config. The ~/.okta-aws writers now re-read under the lock and merge, preserving keys written by a peer. A lock-acquisition timeout surfaces as a friendly CLI error instead of a traceback. Adds a small oktaawscli/_locking module (locked(), atomic_write()) and a filelock runtime dependency. Adapted by hand from amplify-education/okta-awscli (their base is 0.4.x, so the change was reimplemented on our code rather than cherry-picked): 02474e0 Add _locking module with locked() and atomic_write() helpers 14897cc Lock and atomic-write write_sts_token 0d27063 Lock and atomic-write copy_to_default 144f2ef Lock and merge-on-write _save_config_value 33f9046 Handle filelock.Timeout with a friendly CLI error --- CHANGELOG.md | 1 + oktaawscli/_locking.py | 40 ++++++++++++++++++++++++++ oktaawscli/aws_auth.py | 52 ++++++++++++++++++++-------------- oktaawscli/okta_auth_config.py | 41 ++++++++++++++++----------- oktaawscli/okta_awscli.py | 22 +++++++++----- pyproject.toml | 1 + uv.lock | 11 +++++++ 7 files changed, 123 insertions(+), 45 deletions(-) create mode 100644 oktaawscli/_locking.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 26d0290..1d17729 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,7 @@ - Add support for the AWS China (aws-cn) partition (upstream #182) - Add --no-default-profile flag to opt out of writing the [default] profile (from kvinck) - Show the device name behind each MFA factor (from kvinck) +- Serialize concurrent runs and write `~/.aws/credentials` and `~/.okta-aws` atomically under a cross-process file lock (`filelock`), so parallel okta-awscli invocations no longer corrupt or clobber each other's config. Adapted by hand from [amplify-education/okta-awscli](https://github.com/amplify-education/okta-awscli) (SHAs `02474e0`, `14897cc`, `0d27063`, `144f2ef`, `33f9046`). ## [0.5.5] 2024-03-19 - Bugfix [#199](https://github.com/okta-awscli/okta-awscli/issues/199) duplicates of data inside config file diff --git a/oktaawscli/_locking.py b/oktaawscli/_locking.py new file mode 100644 index 0000000..70840d8 --- /dev/null +++ b/oktaawscli/_locking.py @@ -0,0 +1,40 @@ +"""Cross-process advisory locking and atomic writes for dotfiles.""" + +import os +import tempfile +from contextlib import contextmanager + +from filelock import FileLock + +LOCK_TIMEOUT_SECONDS = 60 + + +def locked(path, timeout=LOCK_TIMEOUT_SECONDS): + """Return a FileLock guarding `path`, using `.lock` as the lock file.""" + return FileLock(f"{path}.lock", timeout=timeout) + + +@contextmanager +def atomic_write(path): + """Yield a write file handle that replaces `path` atomically on clean exit. + + The temp file is created next to `path` so the final rename stays on the + same filesystem (POSIX guarantees same-FS rename is atomic). If the with + block raises, the temp file is removed and `path` is left untouched. + """ + parent = os.path.dirname(path) or "." + fd, tmp_path = tempfile.mkstemp( + dir=parent, + prefix=os.path.basename(path) + ".", + suffix=".tmp", + ) + try: + with os.fdopen(fd, "w") as tmp: + yield tmp + os.replace(tmp_path, path) + except BaseException: + try: + os.unlink(tmp_path) + except OSError: + pass + raise diff --git a/oktaawscli/aws_auth.py b/oktaawscli/aws_auth.py index f443287..e6808eb 100644 --- a/oktaawscli/aws_auth.py +++ b/oktaawscli/aws_auth.py @@ -10,6 +10,7 @@ import boto3 from botocore.exceptions import ClientError, NoCredentialsError from subprocess import call +from oktaawscli._locking import atomic_write, locked class AwsPartition(Enum): @@ -55,12 +56,14 @@ def __init__(self, profile, okta_profile, lookup, verbose, logger, self.logger.debug("Setting AWS profile to %s" % self.profile) def set_default_profile(self, parser: RawConfigParser): + # Callers hold the lock on self.creds_file so this write does not need + # to (and must not) re-acquire it. if not parser.has_section('default'): parser.add_section('default') for key, value in parser.items(self.profile): parser.set('default', key, value) self.logger.info("Setting default profile.") - with open(self.creds_file, 'w+') as configfile: + with atomic_write(self.creds_file) as configfile: parser.write(configfile) def choose_aws_role(self, assertion, refresh_role): @@ -171,7 +174,10 @@ def check_sts_token(self): self.logger.info("STS credentials are valid. Nothing to do.") if self.should_set_default_profile: - AwsAuth.set_default_profile(self, parser) + with locked(self.creds_file): + parser = RawConfigParser() + parser.read(self.creds_file) + AwsAuth.set_default_profile(self, parser) return True @@ -179,26 +185,28 @@ def write_sts_token(self, access_key_id, secret_access_key, session_token_expiry """ Writes STS auth information to credentials file """ if not os.path.exists(self.creds_dir): os.makedirs(self.creds_dir) - config = RawConfigParser() - - if os.path.isfile(self.creds_file): - config.read(self.creds_file) - - if not config.has_section(self.profile): - config.add_section(self.profile) - - config.set(self.profile, 'aws_access_key_id', access_key_id) - config.set(self.profile, 'aws_secret_access_key', secret_access_key) - config.set(self.profile, 'aws_session_expiration', session_token_expiry) - config.set(self.profile, 'aws_session_token', session_token) - - with open(self.creds_file, 'w+') as configfile: - config.write(configfile) - self.logger.info("Temporary credentials written to profile: %s" % self.profile) - self.logger.info("Invoke using: aws --profile %s " % self.profile) - - if self.profile != 'default' and self.should_set_default_profile: - AwsAuth.set_default_profile(self, config) + + with locked(self.creds_file): + config = RawConfigParser() + + if os.path.isfile(self.creds_file): + config.read(self.creds_file) + + if not config.has_section(self.profile): + config.add_section(self.profile) + + config.set(self.profile, 'aws_access_key_id', access_key_id) + config.set(self.profile, 'aws_secret_access_key', secret_access_key) + config.set(self.profile, 'aws_session_expiration', session_token_expiry) + config.set(self.profile, 'aws_session_token', session_token) + + with atomic_write(self.creds_file) as configfile: + config.write(configfile) + self.logger.info("Temporary credentials written to profile: %s" % self.profile) + self.logger.info("Invoke using: aws --profile %s " % self.profile) + + if self.profile != 'default' and self.should_set_default_profile: + AwsAuth.set_default_profile(self, config) @staticmethod def __extract_available_roles_from(assertion): diff --git a/oktaawscli/okta_auth_config.py b/oktaawscli/okta_auth_config.py index 9f1d4f7..6711de4 100644 --- a/oktaawscli/okta_auth_config.py +++ b/oktaawscli/okta_auth_config.py @@ -6,6 +6,7 @@ from configparser import RawConfigParser from getpass import getpass import validators +from oktaawscli._locking import atomic_write, locked class OktaAuthConfig(): @@ -159,27 +160,35 @@ def duration_for(self, okta_profile): def write_role_to_profile(self, okta_profile, role_arn): """ Saves role to profile in config """ - if not self._value.has_section(okta_profile): - self._value.add_section(okta_profile) - base_url = self.base_url_for(okta_profile) - self._value.set(okta_profile, 'base-url', base_url) - self._value.set(okta_profile, 'role', role_arn) - - with open(self.config_path, 'w+') as configfile: - self._value.write(configfile) + with locked(self.config_path): + # Re-read inside the lock so concurrent saves merge instead of clobbering. + fresh = RawConfigParser() + fresh.read(self.config_path) + if not fresh.has_section(okta_profile): + fresh.add_section(okta_profile) + fresh.set(okta_profile, 'base-url', base_url) + fresh.set(okta_profile, 'role', role_arn) + + with atomic_write(self.config_path) as configfile: + fresh.write(configfile) + self._value = fresh def write_applink_to_profile(self, okta_profile, app_link): """ Saves app link to profile in config """ - if not self._value.has_section(okta_profile): - self._value.add_section(okta_profile) - base_url = self.base_url_for(okta_profile) - self._value.set(okta_profile, 'base-url', base_url) - self._value.set(okta_profile, 'app-link', app_link) - - with open(self.config_path, 'w+') as configfile: - self._value.write(configfile) + with locked(self.config_path): + # Re-read inside the lock so concurrent saves merge instead of clobbering. + fresh = RawConfigParser() + fresh.read(self.config_path) + if not fresh.has_section(okta_profile): + fresh.add_section(okta_profile) + fresh.set(okta_profile, 'base-url', base_url) + fresh.set(okta_profile, 'app-link', app_link) + + with atomic_write(self.config_path) as configfile: + fresh.write(configfile) + self._value = fresh @staticmethod def get_okta_profiles(): diff --git a/oktaawscli/okta_awscli.py b/oktaawscli/okta_awscli.py index 913d761..0c010be 100644 --- a/oktaawscli/okta_awscli.py +++ b/oktaawscli/okta_awscli.py @@ -5,6 +5,7 @@ import sys import logging import click +from filelock import Timeout from oktaawscli.version import __version__ from oktaawscli.okta_auth import OktaAuth from oktaawscli.okta_auth_config import OktaAuthConfig @@ -138,14 +139,21 @@ def main(okta_profile, profile, verbose, version, aws_auth = AwsAuth(profile, okta_profile, lookup, verbose, logger, set_default_profile=not no_default_profile) profile = aws_auth.profile - if force or not aws_auth.check_sts_token(): - if force and profile: - - logger.info("Force option selected, \ - getting new credentials anyway.") - get_credentials( - aws_auth, okta_profile, profile, verbose, logger, token, cache, refresh_role, okta_username, okta_password + try: + if force or not aws_auth.check_sts_token(): + if force and profile: + + logger.info("Force option selected, \ + getting new credentials anyway.") + get_credentials( + aws_auth, okta_profile, profile, verbose, logger, token, cache, refresh_role, okta_username, okta_password + ) + except Timeout as exc: + logger.error( + "Could not acquire lock on %s - another okta-awscli process is " + "holding it. Try again." % exc.lock_file ) + sys.exit(1) if awscli_args: aws_auth.execute_aws_args(awscli_args, logger) diff --git a/pyproject.toml b/pyproject.toml index a0e025b..fd136be 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,6 +30,7 @@ dependencies = [ "boto3", "ConfigParser", "validators", + "filelock", ] [project.optional-dependencies] diff --git a/uv.lock b/uv.lock index 1abe3e1..d79d6fc 100644 --- a/uv.lock +++ b/uv.lock @@ -228,6 +228,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/1e/77/dc8c558f7593132cf8fefec57c4f60c83b16941c574ac5f619abb3ae7933/dill-0.4.1-py3-none-any.whl", hash = "sha256:1e1ce33e978ae97fcfcff5638477032b801c46c7c65cf717f95fbc2248f79a9d", size = 120019, upload-time = "2026-01-19T02:36:55.663Z" }, ] +[[package]] +name = "filelock" +version = "4.0.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/5e/ac/8c98b17ee3900147b38ef7884a1e590b070b63f0d59fd8b46ef0205f4576/filelock-4.0.0.tar.gz", hash = "sha256:3611eca5d818ca9b00ec3cc7db1dcfe1e2aafc8d44fb4920d8cf60ad1f6bfda6", size = 237935, upload-time = "2026-09-17T03:59:02.967Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/01/fb/0e4489505dac46a0487ba925ea4b046612f1a40f432d63a000421de78549/filelock-4.0.0-py3-none-any.whl", hash = "sha256:a850aa9ec2acba8db9ca2e9fcf8a327fbc2f85e432725715b37bd76b7dd1f798", size = 106036, upload-time = "2026-09-17T03:59:01.336Z" }, +] + [[package]] name = "h11" version = "0.16.0" @@ -436,6 +445,7 @@ dependencies = [ { name = "boto3" }, { name = "click" }, { name = "configparser" }, + { name = "filelock" }, { name = "niquests" }, { name = "validators" }, ] @@ -456,6 +466,7 @@ requires-dist = [ { name = "boto3" }, { name = "click" }, { name = "configparser" }, + { name = "filelock" }, { name = "niquests" }, { name = "python-u2flib-host", marker = "extra == 'u2f'" }, { name = "validators" },