From 755e9aa0d379058d71210a8d7ea78aee05639164 Mon Sep 17 00:00:00 2001
From: I-Al-Istannen <i-al-istannen@users.noreply.github.com>
Date: Fri, 8 May 2020 21:51:33 +0200
Subject: [PATCH] Try to add support for Shibboleth TFA token

---
 PFERD/authenticators.py       | 24 ++++++++++++++++++++++++
 PFERD/ilias/authenticators.py | 30 +++++++++++++++++++++++++++++-
 2 files changed, 53 insertions(+), 1 deletion(-)

diff --git a/PFERD/authenticators.py b/PFERD/authenticators.py
index 7fff772..b8cfe28 100644
--- a/PFERD/authenticators.py
+++ b/PFERD/authenticators.py
@@ -6,6 +6,30 @@ import getpass
 from typing import Optional, Tuple
 
 
+class TfaAuthenticator:
+    # pylint: disable=too-few-public-methods
+    """
+    An authenticator for a TFA token. Always prompts the user, as the token can not be cached.
+    """
+
+    def __init__(self, reason: str):
+        """
+        Create a new tfa authenticator.
+
+        Arguments:
+            reason {str} -- the reason for obtaining the credentials
+        """
+        self._reason = reason
+
+    def get_token(self) -> str:
+        # pylint: disable=no-self-use
+        """
+        Prompts the user for the token and returns it.
+        """
+        print(f"Enter credentials ({self._reason})")
+        return getpass.getpass("TFA Token: ")
+
+
 class UserPassAuthenticator:
     """
     An authenticator for username-password combinations that prompts the user
diff --git a/PFERD/ilias/authenticators.py b/PFERD/ilias/authenticators.py
index 5443d99..68faf00 100644
--- a/PFERD/ilias/authenticators.py
+++ b/PFERD/ilias/authenticators.py
@@ -9,7 +9,7 @@ from typing import Optional
 import bs4
 import requests
 
-from ..authenticators import UserPassAuthenticator
+from ..authenticators import TfaAuthenticator, UserPassAuthenticator
 from ..utils import soupify
 
 LOGGER = logging.getLogger(__name__)
@@ -39,6 +39,7 @@ class KitShibbolethAuthenticator(IliasAuthenticator):
 
     def __init__(self, username: Optional[str] = None, password: Optional[str] = None) -> None:
         self._auth = UserPassAuthenticator("KIT ILIAS Shibboleth", username, password)
+        self._tfa_auth = TfaAuthenticator("KIT ILIAS Shibboleth")
 
     def authenticate(self, sess: requests.Session) -> None:
         """
@@ -80,6 +81,9 @@ class KitShibbolethAuthenticator(IliasAuthenticator):
             }
             soup = soupify(sess.post(url, data=data))
 
+            if self._tfa_required(soup):
+                soup = self._authenticate_tfa(sess, soup)
+
             if not self._login_successful(soup):
                 print("Incorrect credentials.")
                 self._auth.invalidate_credentials()
@@ -96,8 +100,32 @@ class KitShibbolethAuthenticator(IliasAuthenticator):
         }
         sess.post(url, data=data)
 
+    def _authenticate_tfa(
+            self,
+            session: requests.Session,
+            soup: bs4.BeautifulSoup
+    ) -> bs4.BeautifulSoup:
+        # Searching the form here so that this fails before asking for
+        # credentials rather than after asking.
+        form = soup.find("form", {"method": "post"})
+        action = form["action"]
+
+        # Equivalent: Enter token in
+        # https://idp.scc.kit.edu/idp/profile/SAML2/Redirect/SSO
+        LOGGER.debug("Attempt to log in to Shibboleth with TFA token")
+        url = "https://idp.scc.kit.edu" + action
+        data = {
+            "_eventId_proceed": "",
+            "j_tokenNumber": self._tfa_auth.get_token()
+        }
+        return soupify(session.post(url, data=data))
+
     @staticmethod
     def _login_successful(soup: bs4.BeautifulSoup) -> bool:
         relay_state = soup.find("input", {"name": "RelayState"})
         saml_response = soup.find("input", {"name": "SAMLResponse"})
         return relay_state is not None and saml_response is not None
+
+    @staticmethod
+    def _tfa_required(soup: bs4.BeautifulSoup) -> bool:
+        return soup.find(id="j_tokenNumber") is not None