1 files changed, 41 insertions, 0 deletions
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
+    )