From 987ae987af7a36e6bec55f10bcbabf17b53fe35c Mon Sep 17 00:00:00 2001
From: Andreas Olsson <andreas.olsson@yubico.com>
Date: Tue, 20 Oct 2020 10:40:44 +0200
Subject: [PATCH] Add support for one-time security codes (skotp)

---
 README.rst                |  3 +++
 aws_google_auth/google.py | 39 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 42 insertions(+)

diff --git a/README.rst b/README.rst
index 9c7fb86..d754ce1 100644
--- a/README.rst
+++ b/README.rst
@@ -282,6 +282,9 @@ by this tool are:
 | Dual prompt      | ``.../signin/challenge/dp/...``     |
 |  (Validate 2FA ) |                                     |
 +------------------+-------------------------------------+
+| One-time         | ``.../signin/challenge/skotp/...``  |
+| security codes   |                                     |
++------------------+-------------------------------------+
 | Backup code      | ``... (unknown yet) ...``           |
 |  (printed codes) |                                     |
 +------------------+-------------------------------------+
diff --git a/aws_google_auth/google.py b/aws_google_auth/google.py
index 398eff5..02a4d67 100644
--- a/aws_google_auth/google.py
+++ b/aws_google_auth/google.py
@@ -329,6 +329,8 @@ def do_login(self):
             sess = self.handle_iap(sess)
         elif "challenge/dp/" in sess.url:
             sess = self.handle_dp(sess)
+        elif "challenge/skotp/" in sess.url:
+            sess = self.handle_skotp(sess)
         elif "challenge/ootp/5" in sess.url:
             raise NotImplementedError(
                 'Offline Google App OOTP not implemented')
@@ -823,6 +825,41 @@ def handle_iap(self, sess):
         # Submit SMS/VOICE token
         return self.post(challenge_url, data=payload)
 
+    def handle_skotp(self, sess):
+        response_page = BeautifulSoup(sess.text, 'html.parser')
+        ltmpl = response_page.find('input', {'name': 'ltmpl'}).get('value')
+        flowname = response_page.find('input', {'name': 'flowName'}).get('value')
+        tl = response_page.find('input', {'name': 'TL'}).get('value')
+        gxf = response_page.find('input', {'name': 'gxf'}).get('value')
+        challenge_url = sess.url.split("?")[0]
+        challenge_id = challenge_url.split("skotp/")[1]
+
+        skotp_code = input("Enter https://g.co/sc one-time security code: ") or None
+
+        if not skotp_code:
+            raise ValueError(
+                "One-time security code required for {} but none supplied.".format(
+                    self.config.username))
+
+        payload = {
+            'challengeId': challenge_id,
+            'challengeType': 37,
+            'continue': self.cont,
+            'ltmpl': ltmpl,
+            'scc': 1,
+            'sarp': 1,
+            'oauth': 1,
+            'flowName': flowname,
+            'faa': 1,
+            'TL': tl,
+            'gxf': gxf,
+            'Pin': skotp_code,
+            'TrustDevice': 'on',
+        }
+
+        # Submit SKOTP
+        return self.post(challenge_url, data=payload)
+
     def handle_selectchallenge(self, sess):
         response_page = BeautifulSoup(sess.text, 'html.parser')
 
@@ -838,6 +875,8 @@ def handle_selectchallenge(self, sess):
                 challenges.append(['SMS other phone', i.attrs.get("data-challengeentry")])
             elif "challenge/sk/" in action:
                 challenges.append(['YubiKey', i.attrs.get("data-challengeentry")])
+            elif "challenge/skotp/" in action:
+                challenges.append(['One-time security code', i.attrs.get("data-challengeentry")])
             elif "challenge/az/" in action:
                 challenges.append(['Google Prompt', i.attrs.get("data-challengeentry")])