Serialize concurrent runs and atomic-write credential/config dotfiles

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

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 `<path>.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

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,34 +174,39 @@ 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
 
     def write_sts_token(self, access_key_id, secret_access_key, session_token_expiry, session_token):
         """ 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 <service> <command>" % 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 <service> <command>" % 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):

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():

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)

pyproject.toml

@@ -30,6 +30,7 @@ dependencies = [
     "boto3",
     "ConfigParser",
     "validators",
+    "filelock",
 ]
 
 [project.optional-dependencies]