From bfe2f916f2e249ebd4b054d6fa604fdafb06cb4e Mon Sep 17 00:00:00 2001
From: Dirk-jan Mollema <dirkjan@outsidersecurity.nl>
Date: Tue, 10 Oct 2023 18:31:34 +0200
Subject: [PATCH] roadtx: support for device code auth in selenium scenarios
 (primarily for testing scenarios)

---
 roadtx/roadtools/roadtx/main.py     |  6 +++-
 roadtx/roadtools/roadtx/selenium.py | 46 +++++++++++++++++++----------
 2 files changed, 36 insertions(+), 16 deletions(-)

diff --git a/roadtx/roadtools/roadtx/main.py b/roadtx/roadtools/roadtx/main.py
index f6dcbed..a57b045 100644
--- a/roadtx/roadtools/roadtx/main.py
+++ b/roadtx/roadtools/roadtx/main.py
@@ -337,6 +337,10 @@ def main():
     kdbauth_parser.add_argument('--federated',
                                 action='store_true',
                                 help='Fill in password on Federation server login page (assumes AD FS)')
+    kdbauth_parser.add_argument('--device-code',
+                                action='store',
+                                help='Authenticate with the given device code')
+
 
 
     # Interactive auth using Selenium - inject PRT
@@ -705,7 +709,7 @@ def main():
         if not service:
             return
         selauth.driver = selauth.get_webdriver(service)
-        result = selauth.selenium_login(url, args.username, password, otpseed, keep=args.keep_open, capture=args.capture_code, federated=args.federated)
+        result = selauth.selenium_login(url, args.username, password, otpseed, keep=args.keep_open, capture=args.capture_code, federated=args.federated, devicecode=args.device_code)
         if args.capture_code:
             if result:
                 print(f'Captured auth code: {result}')
diff --git a/roadtx/roadtools/roadtx/selenium.py b/roadtx/roadtools/roadtx/selenium.py
index 68adea9..99e777b 100644
--- a/roadtx/roadtools/roadtx/selenium.py
+++ b/roadtx/roadtools/roadtx/selenium.py
@@ -97,12 +97,19 @@ class SeleniumAuthentication():
             otpseed = None
         return userpassword, otpseed
 
-    def selenium_login(self, url, identity=None, password=None, otpseed=None, keep=False, capture=False, federated=False):
+    def selenium_login(self, url, identity=None, password=None, otpseed=None, keep=False, capture=False, federated=False, devicecode=None):
         '''
         Selenium based login with optional autofill of whatever is provided
         '''
         driver = self.driver
+        # Change if using device code auth
+        if devicecode:
+            url = 'https://login.microsoftonline.com/common/oauth2/deviceauth'
         driver.get(url)
+        # Enter code first if device code flow
+        if devicecode:
+            el = WebDriverWait(driver, 3000).until(lambda d: d.find_element(By.ID, "otc"))
+            el.send_keys(devicecode + Keys.ENTER)
         if identity and not 'login_hint' in url:
             el = WebDriverWait(driver, 3000).until(lambda d: d.find_element(By.ID, "i0116"))
             el.send_keys(identity + Keys.ENTER)
@@ -157,20 +164,29 @@ class SeleniumAuthentication():
                 # No MFA?
                 pass
 
-        try:
-            WebDriverWait(driver, 120).until(lambda d: '?code=' in d.current_url)
-            res = urlparse(driver.current_url)
-            params = parse_qs(res.query)
-            code = params['code'][0]
-            if not keep:
-                driver.close()
-            if capture:
-                return code
-            return self.auth.authenticate_with_code_native(code, self.redirurl)
-        except TimeoutException:
-            if not keep:
-                driver.close()
-                raise AuthenticationException('Authentication did not complete within time limit')
+        if devicecode:
+            try:
+                els = WebDriverWait(driver, 10).until(lambda d: d.find_element(By.ID, "idSIButton9"))
+                els.click()
+            except TimeoutException:
+                if not keep:
+                    driver.close()
+                    raise AuthenticationException('Could not complete device code auth within the time limit (button not found: idSIButton9)')
+        else:
+            try:
+                WebDriverWait(driver, 120).until(lambda d: '?code=' in d.current_url)
+                res = urlparse(driver.current_url)
+                params = parse_qs(res.query)
+                code = params['code'][0]
+                if not keep:
+                    driver.close()
+                if capture:
+                    return code
+                return self.auth.authenticate_with_code_native(code, self.redirurl)
+            except TimeoutException:
+                if not keep:
+                    driver.close()
+                    raise AuthenticationException('Authentication did not complete within time limit')
         return False
 
     def selenium_login_with_prt(self, url, identity=None, password=None, otpseed=None, keep=False, prtcookie=None, capture=False):