SAML Login can now handle step up authentication and uses TOTP if ava…

…ilable and configured
SUNET · Jan 11, 2024 · 824da44 · 824da44
1 parent 09ad05c
commit 824da44
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 3 deletions.
diff --git a/sunetdrive.py b/sunetdrive.py
@@ -346,6 +346,19 @@ def get_samluserpassword(self, userid):
         else:
             raise NotImplementedError
 
+    def get_samlusertotpsecret(self, userid):
+        if self.platform == 'win32':
+            pwdcmd = 'powershell -command "& { . ./NodeCredentials.ps1; Get-SamlUserTotpSecret ' + userid + ' ' + self.target + ' }"'
+            process = os.popen(pwdcmd)
+            pwd = process.read()
+            process.close()
+            return pwd
+        elif self.platform == 'linux':
+            env = "DRIVE_SELENIUM_SAML_MFA_SECRET_" + userid.upper() + "_" + self.target.upper()
+            return os.environ.get(env)
+        else:
+            raise NotImplementedError        
+
     def save_samlusercredentials(self, userid):
         if self.platform == 'win32':
             cmd = 'powershell -command "& { . ./NodeCredentials.ps1; Save-SamlUserCredentials ' + userid + ' ' + self.target + ' }"'

diff --git a/test_login_selenium.py b/test_login_selenium.py
@@ -6,6 +6,7 @@
 import unittest
 import sunetdrive
 from webdav3.client import Client
+import pyotp
 
 import time
 
@@ -251,16 +252,17 @@ def test_saml_eduid_nomfa(self):
     def test_saml_su_nomfa(self):
         delay = 30 # seconds
         drv = sunetdrive.TestTarget()
+        nodeName = 'su'
         if len(drv.allnodes) == 1:
-            if drv.allnodes[0] != 'su':
+            if drv.allnodes[0] != nodeName:
                 self.logger.info(f'Only testing {drv.allnodes[0]}, not testing su saml')
                 return
 
         loginurl = drv.get_gss_url()
         self.logger.info(f'Login url: {loginurl}')
 
-        samluser=drv.get_samlusername("su")
-        samlpassword=drv.get_samluserpassword("su")
+        samluser=drv.get_samlusername(nodeName)
+        samlpassword=drv.get_samluserpassword(nodeName)
 
         try:
             options = Options()
@@ -292,6 +294,22 @@ def test_saml_su_nomfa(self):
         wait.until(EC.presence_of_element_located((By.ID, 'password'))).send_keys(samlpassword + Keys.ENTER)
         # wait.until(EC.presence_of_element_located((By.CLASS_NAME, 'login-form-button'))).click()
 
+        # Wait for TOTP screen
+        requireTotp = False
+        try:
+            self.logger.info(f'Check if TOTP selection dialogue is visible')
+            totpselect = driver.find_element(By.XPATH, '//a[@href="'+ '/index.php/login/challenge/totp?redirect_url=/index.php/apps/dashboard/' +'"]')
+            self.logger.warning(f'Found TOTP selection dialogue')
+            requireTotp = True
+            totpselect.click()
+        except:
+            self.logger.info(f'No need to select TOTP provider')
+
+        if requireTotp:
+            nodetotpsecret = drv.get_samlusertotpsecret(nodeName)
+            totp = pyotp.TOTP(nodetotpsecret)
+            wait.until(EC.presence_of_element_located((By.XPATH, '//*[@id="body-login"]/div[1]/div/main/div/form/input'))).send_keys(totp.now() + Keys.ENTER)
+
         try:
             wait.until(EC.presence_of_element_located((By.CLASS_NAME, 'app-menu')))
             self.logger.info(f'App menu is ready!')