v1.1.7

atwoodjw · Feb 23, 2021 · d499681 · d499681
1 parent d5657d0
commit d499681
Show file tree

Hide file tree

Showing 3 changed files with 55 additions and 44 deletions.
diff --git a/README.md b/README.md
@@ -17,7 +17,7 @@ Please note that the script is called `aws2-wrap` to show that it works with AWS
 
 <https://pypi.org/project/aws2-wrap>
 
-`pip3 install aws2-wrap==1.1.6`
+`pip3 install aws2-wrap==1.1.7`
 
 ## Run a command using AWS SSO credentials
 
@@ -106,4 +106,4 @@ Note that because the profile is being specified via `AWS_PROFILE`, it is someti
 
 ## Credits
 
-Thanks to @damian-bisignano, @flyinprogrammer, @abeluck, @topu, @bigwheel, @krabbit, @jscook2345, @hieki, @blazdivjak, @fukushun1994 and @johann8384 for their contributions.
+Thanks to @damian-bisignano, @flyinprogrammer, @abeluck, @topu, @bigwheel, @krabbit, @jscook2345, @hieki, @blazdivjak, @fukushun1994, @johann8384 and @ppezoldt for their contributions.
diff --git a/aws2wrap/__init__.py b/aws2wrap/__init__.py
@@ -237,6 +237,56 @@ def process_cred_generation(
     print("The credentials will expire at %s" % expiration)
 
 
+def run_command(access_key, secret_access_key, session_token, profile, args):
+    """ Run the specified command with the credentials set up """
+    os.environ["AWS_ACCESS_KEY_ID"] = access_key
+    os.environ["AWS_SECRET_ACCESS_KEY"] = secret_access_key
+    os.environ["AWS_SESSION_TOKEN"] = session_token
+    status = None # ensure this is initialised
+    # If region is specified in profile, also set AWS_DEFAULT_REGION
+    if "AWS_DEFAULT_REGION" not in os.environ and "region" in profile:
+        os.environ["AWS_DEFAULT_REGION"] = retrieve_attribute(profile, "region")
+    if args.exec is not None:
+        status = os.system(args.exec)
+    elif args.command is not None:
+        status = os.system(" ".join(args.command))
+    # The return value of os.system is not simply the exit code of the process
+    # see: https://mail.python.org/pipermail/python-list/2003-May/207712.html
+    # noinspection PyUnboundLocalVariable
+    if status is None:
+        sys.exit(0)
+    # noinspection PyUnboundLocalVariable
+    if status % 256 == 0:
+        sys.exit(status//256)
+    sys.exit(status % 256)
+
+
+def export_credentials(access_key, secret_access_key, session_token, profile):
+    """ Export the AWS credentials to environment variables """
+    # On Windows, parent process is aws2-wrap.exe, in unix it's the shell
+    if os.name == "nt":
+        shell_name = psutil.Process().parent().parent().name()
+    else:
+        shell_name = psutil.Process().parent().name()
+
+    is_powershell = bool(re.fullmatch('pwsh|pwsh.exe|powershell.exe', shell_name))
+
+    if is_powershell:
+        print("$ENV:AWS_ACCESS_KEY_ID=\"%s\"" % access_key)
+        print("$ENV:AWS_SECRET_ACCESS_KEY=\"%s\"" % secret_access_key)
+        print("$ENV:AWS_SESSION_TOKEN=\"%s\"" % session_token)
+        # If region is specified in profile, also export AWS_DEFAULT_REGION
+        if "AWS_DEFAULT_REGION" not in os.environ and "region" in profile:
+            print("$ENV:AWS_DEFAULT_REGION=\"%s\"" % retrieve_attribute(profile, "region"))
+    else:
+        print("export AWS_ACCESS_KEY_ID=%s" % access_key)
+        print("export AWS_SECRET_ACCESS_KEY=%s" % secret_access_key)
+        print("export AWS_SESSION_TOKEN=%s" % session_token)
+        # If region is specified in profile, also export AWS_DEFAULT_REGION
+        if "AWS_DEFAULT_REGION" not in os.environ and "region" in profile:
+            print("export AWS_DEFAULT_REGION=%s" % retrieve_attribute(profile, "region"))
+
+
 def main():
     """ Main! """
     args = process_arguments()
@@ -257,27 +307,7 @@ def main():
     expiration = grc_structure["roleCredentials"]["expiration"]
     if args.export:
         # On Windows, parent process is aws2-wrap.exe, in unix it's the shell
-        if os.name == "nt":
-            shell_name = psutil.Process().parent().parent().name()
-        else:
-            shell_name = psutil.Process().parent().name()
-
-        is_powershell = bool(re.fullmatch('pwsh|pwsh.exe|powershell.exe', shell_name))
-
-        if is_powershell:
-            print("$ENV:AWS_ACCESS_KEY_ID=\"%s\"" % access_key)
-            print("$ENV:AWS_SECRET_ACCESS_KEY=\"%s\"" % secret_access_key)
-            print("$ENV:AWS_SESSION_TOKEN=\"%s\"" % session_token)
-            # If region is specified in profile, also export AWS_DEFAULT_REGION
-            if "AWS_DEFAULT_REGION" not in os.environ and "region" in profile:
-                print("$ENV:AWS_DEFAULT_REGION=\"%s\"" % retrieve_attribute(profile, "region"))
-        else:
-            print("export AWS_ACCESS_KEY_ID=%s" % access_key)
-            print("export AWS_SECRET_ACCESS_KEY=%s" % secret_access_key)
-            print("export AWS_SESSION_TOKEN=%s" % session_token)
-            # If region is specified in profile, also export AWS_DEFAULT_REGION
-            if "AWS_DEFAULT_REGION" not in os.environ and "region" in profile:
-                print("export AWS_DEFAULT_REGION=%s" % retrieve_attribute(profile, "region"))
+        export_credentials(access_key, secret_access_key, session_token, profile)
     elif args.generate:
         if args.outprofile is not None:
             process_cred_generation(
@@ -293,26 +323,7 @@ def main():
         }
         print(json.dumps(output))
     else:
-        os.environ["AWS_ACCESS_KEY_ID"] = access_key
-        os.environ["AWS_SECRET_ACCESS_KEY"] = secret_access_key
-        os.environ["AWS_SESSION_TOKEN"] = session_token
-        status = None # ensure this is initialised
-        # If region is specified in profile, also set AWS_DEFAULT_REGION
-        if "AWS_DEFAULT_REGION" not in os.environ and "region" in profile:
-            os.environ["AWS_DEFAULT_REGION"] = retrieve_attribute(profile, "region")
-        if args.exec is not None:
-            status = os.system(args.exec)
-        elif args.command is not None:
-            status = os.system(" ".join(args.command))
-        # The return value of os.system is not simply the exit code of the process
-        # see: https://mail.python.org/pipermail/python-list/2003-May/207712.html
-        # noinspection PyUnboundLocalVariable
-        if status is None:
-            sys.exit(0)
-        # noinspection PyUnboundLocalVariable
-        if status % 256 == 0:
-            sys.exit(status//256)
-        sys.exit(status % 256)
+        run_command(access_key, secret_access_key, session_token, profile, args)
 
 
 if __name__ == '__main__':

diff --git a/setup.py b/setup.py
@@ -5,7 +5,7 @@
 
 setup(
     name="aws2-wrap",
-    version="1.1.6",
+    version="1.1.7",
     description="A wrapper for executing a command with AWS CLI v2 and SSO",
     long_description=long_description,
     long_description_content_type="text/markdown",