From bd1f0689cbd4a1e7d1df009bc6a4f41fd6c8676e Mon Sep 17 00:00:00 2001
From: Schlueter <b@schlueter.blue>
Date: Wed, 7 Feb 2024 17:10:56 -0500
Subject: [PATCH] feature:Allow configuration of a preferred_mfa_provider

---
 README.md                       |  4 ++++
 gimme_aws_creds/main.py         |  3 +++
 gimme_aws_creds/okta_classic.py | 17 +++++++++++++++++
 3 files changed, 24 insertions(+)

diff --git a/README.md b/README.md
index 79d8dd5..6c1e44f 100644
--- a/README.md
+++ b/README.md
@@ -228,6 +228,10 @@ A configuration wizard will prompt you to enter the necessary configuration para
   - web - DUO uses localhost webbrowser to support push|call|passcode
   - passcode - DUO uses `OKTA_MFA_CODE` or `--mfa-code` if set, or prompts user for passcode(OTP).
   - claims_provider - DUO Universal Prompt
+- preferred_mfa_provider - (optional) automatically select a particular provider when prompted for MFA:
+  - GOOGLE
+  - OKTA
+  - DUO
 - duo_universal_factor - (optional) Configure which type of factor to use with Duo Universal Prompt. Must be one of (case-sensitive):
   - `Duo Push` (default)
   - `Passcode`
diff --git a/gimme_aws_creds/main.py b/gimme_aws_creds/main.py
index c33627c..f4e7e7b 100644
--- a/gimme_aws_creds/main.py
+++ b/gimme_aws_creds/main.py
@@ -583,6 +583,9 @@ def okta(self):
             if self.conf_dict.get('preferred_mfa_type'):
                 okta.set_preferred_mfa_type(self.conf_dict['preferred_mfa_type'])
 
+            if self.conf_dict.get('preferred_mfa_provider'):
+                okta.set_preferred_mfa_provider(self.conf_dict['preferred_mfa_provider'])
+
             if self.conf_dict.get('duo_universal_factor'):
                 okta.set_duo_universal_factor(self.conf_dict.get('duo_universal_factor'))
 
diff --git a/gimme_aws_creds/okta_classic.py b/gimme_aws_creds/okta_classic.py
index 21cfc41..af308be 100644
--- a/gimme_aws_creds/okta_classic.py
+++ b/gimme_aws_creds/okta_classic.py
@@ -65,6 +65,7 @@ def __init__(self, gac_ui, okta_org_url, verify_ssl_certs=True, device_token=Non
         self._username = None
         self._password = None
         self._preferred_mfa_type = None
+        self._preferred_mfa_provider = None
         self._duo_universal_factor = 'Duo Push'
         self._mfa_code = None
         self._remember_device = None
@@ -105,6 +106,9 @@ def set_password(self, password):
     def set_preferred_mfa_type(self, preferred_mfa_type):
         self._preferred_mfa_type = preferred_mfa_type
 
+    def set_preferred_mfa_provider(self, preferred_mfa_provider):
+        self._preferred_mfa_provider = preferred_mfa_provider
+
     def set_mfa_code(self, mfa_code):
         self._mfa_code = mfa_code
 
@@ -837,6 +841,19 @@ def _choose_factor(self, factors):
             if not preferred_factors:
                 self.ui.notify('Preferred factor type of {} not available.'.format(self._preferred_mfa_type))
 
+        if self._preferred_mfa_provider is not None:
+            preferred_factors_with_preferred_provider = list(
+                filter(lambda item: item['provider'] == self._preferred_mfa_provider, preferred_factors)
+            )
+            # If filtering for the preferred provider yields no results, announce it,
+            # but don't update the list of preferred factors.
+            if preferred_factors and not preferred_factors_with_preferred_provider:
+                self.ui.notify('Preferred factor provider of {} not available. Will use available factors.'.format(
+                    self._preferred_mfa_provider
+                ))
+            else:
+                preferred_factors = preferred_factors_with_preferred_provider
+
         if len(preferred_factors) == 1:
             factor_name = self._build_factor_name(preferred_factors[0])
             self.ui.info(factor_name + ' selected')