first commit

4 files changed, 127 insertions, 0 deletions

diff --git a/awscli_plugin_passtotp/__init__.py b/awscli_plugin_passtotp/__init__.py
new file mode 100644
index 0000000..aeb8606
--- /dev/null
+++ b/awscli_plugin_passtotp/__init__.py
@@ -0,0 +1,15 @@
+from .prompter import inject_pass_totp_prompter
+from .commands import SessionEnv
+
+
+def awscli_initialize(cli):
+    cli.register(
+        "session-initialized",
+        inject_pass_totp_prompter,
+        unique_id="inject_pass_totp_provider",
+    )
+    cli.register("building-command-table.main", awscli_register_commands)
+
+
+def awscli_register_commands(command_table, session, **kwargs):
+    command_table["session-env"] = SessionEnv(session)
diff --git a/awscli_plugin_passtotp/boto_plugin.py b/awscli_plugin_passtotp/boto_plugin.py
new file mode 100644
index 0000000..9a2f249
--- /dev/null
+++ b/awscli_plugin_passtotp/boto_plugin.py
@@ -0,0 +1,12 @@
+from . import inject_pass_totp_prompter
+import botocore.session
+
+old_init = botocore.session.Session.__init__
+
+
+def patched_session_init(self, *args, **kwargs):
+    old_init(self, *args, **kwargs)
+    inject_pass_totp_prompter(self)
+
+
+botocore.session.Session.__init__ = patched_session_init
diff --git a/awscli_plugin_passtotp/commands.py b/awscli_plugin_passtotp/commands.py
new file mode 100644
index 0000000..d924925
--- /dev/null
+++ b/awscli_plugin_passtotp/commands.py
@@ -0,0 +1,59 @@
+from awscli.customizations.commands import BasicCommand
+import sys
+from . import boto_plugin
+
+
+class SessionEnv(BasicCommand):
+    NAME = "session-env"
+    DESCRIPTION = (
+        "prints the current session's credentials in the form of "
+        "environment variables.\n"
+        "\n"
+        "You can use the ``--profile`` argument to select a different set "
+        "of credentials.\n"
+        "\n"
+        ".. note::\n"
+        "  If you have set the environment variables in your shell, subsequent "
+        "  calls to aws session-env will use those credentials. In order to use "
+        "  your default profile, you have to explicitly specify it::\n"
+        "\n"
+        "    $(aws session-env --profile default\n"
+    )
+    SYNOPSIS = "aws session-env"
+    EXAMPLES = (
+        "Print temporary session tokens for a given profile::\n"
+        "\n"
+        "  $ aws session-env --profile profile_name\n"
+        "  export AWS_ACCESS_KEY_ID=AKIAI44QH8DHBEXAMPLE\n"
+        "  export AWS_SECRET_ACCESS_KEY=wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY\n"
+        "  export AWS_SESSION_TOKEN=AQoDYXdzEJr...\n"
+        "  The acquired credentials will be valid for 60:00 minutes\n"
+        "\n"
+        "Directly set the environment variables for use in other applications::\n"
+        "\n"
+        "  $ $(aws session-env --profile profile_name\n"
+        "  The acquired credentials will be valid for 60:00 minutes\n"
+    )
+
+    ARG_TABLE = []
+
+    def _run_main(self, args, parsed_globals):
+        credentials = self._session.get_credentials()
+
+        frozen_credentials = credentials.get_frozen_credentials()
+        print("export AWS_ACCESS_KEY_ID={}".format(frozen_credentials.access_key))
+        print("export AWS_SECRET_ACCESS_KEY={}".format(frozen_credentials.secret_key))
+        if frozen_credentials.token is None:
+            print("unset AWS_SESSION_TOKEN")
+        else:
+            print("export AWS_SESSION_TOKEN={}".format(frozen_credentials.token))
+
+        if hasattr(credentials, "_seconds_remaining"):
+            seconds_to_expire = int(credentials._seconds_remaining())
+            print(
+                "The acquired credentials will be valid for {:.0f}:{:02.0f} minutes".format(
+                    seconds_to_expire // 60, seconds_to_expire % 60
+                ),
+                file=sys.stderr,
+            )
+        return 0
diff --git a/awscli_plugin_passtotp/prompter.py b/awscli_plugin_passtotp/prompter.py
new file mode 100644
index 0000000..59e56ad
--- /dev/null
+++ b/awscli_plugin_passtotp/prompter.py
@@ -0,0 +1,41 @@
+from botocore.exceptions import ProfileNotFound
+import subprocess
+import sys
+
+
+class PassTotpPrompter(object):
+    def __init__(self, mfa_path, original_prompter=None):
+        self.mfa_path = mfa_path
+        self._original_prompter = original_prompter
+
+    def __call__(self, prompt):
+        try:
+            pass_result = subprocess.run(
+                ["pass", "otp", self.mfa_path], capture_output=True
+            )
+            token = pass_result.stdout.decode("utf-8").strip()
+            return token
+        except subprocess.CalledProcessError as e:
+            print(e, file=sys.stderr)
+
+        if self._original_prompter:
+            return self._original_prompter(prompt)
+
+        return None
+
+def inject_pass_totp_prompter(session, **kwargs):
+    try:
+        providers = session.get_component("credential_provider")
+    except ProfileNotFound:
+        return
+
+    config = session.get_scoped_config()
+    mfa_path = config.get("mfa_path")
+    if mfa_path is None:
+        return
+
+    assume_role_provider = providers.get_provider("assume-role")
+    original_prompter = assume_role_provider._prompter
+    assume_role_provider._prompter = PassTotpPrompter(
+        mfa_path, original_prompter=original_prompter
+    )