Add MFA session keeper module (weblogin.mfa)

doc/contents.tex

@@ -3,4 +3,5 @@
 \input{../src/weblogin/kth.tex}
 \input{../src/weblogin/seamlessaccess.tex}
 \input{../src/weblogin/ladok.tex}
+\input{../src/weblogin/mfa.tex}
 

pyproject.toml

@@ -28,6 +28,10 @@ exclude = [
 python = "^3.8"
 requests = "^2.28.1"
 lxml = "^5.3.0"
+selenium = {version = "^4.0", optional = true}
+
+[tool.poetry.extras]
+mfa = ["selenium"]
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^7.2.0"

src/weblogin/Makefile

@@ -2,6 +2,7 @@ MODULES+=	__init__.py
 MODULES+=	seamlessaccess.py
 MODULES+=	kth.py
 MODULES+=	ladok.py
+MODULES+=	mfa.py
 
 .PHONY: all
 all: ${MODULES}

src/weblogin/mfa.nw

@@ -0,0 +1,311 @@
+\section{Handling MFA at KTH}\label{MFAsessionkeeper}
+
+Many universities, including KTH, have introduced Multi-Factor
+Authentication (MFA) using Microsoft Authenticator with number matching.
+This breaks the standard [[weblogin]] approach: the login page is now
+heavily JavaScript-driven and requires a real browser to render the MFA
+challenge.
+
+The solution is a \emph{session keeper}: a script that uses a headless
+Chrome browser (via Selenium) to perform the login once, handle the MFA
+challenge interactively, and save the resulting cookies to a file.
+Other scripts can then load these cookies and make requests without
+going through the login flow again.
+
+Since Selenium is not needed for institutions without MFA, it is an
+optional dependency.
+Install it with:
+\begin{verbatim}
+pip install weblogin[mfa]
+\end{verbatim}
+The module can then be used from the command line:
+\begin{verbatim}
+python -m weblogin.mfa <trigger_url> <cookie_file>
+\end{verbatim}
+or imported to use the individual functions:
+\begin{verbatim}
+from weblogin.mfa import setup_browser, do_login, collect_cookies, save_cookies
+\end{verbatim}
+
+The module looks like this.
+<<[[mfa.py]]>>=
+import getpass
+import http.cookiejar
+import os
+import pickle
+import re
+import sys
+
+try:
+  from selenium import webdriver
+  from selenium.webdriver.chrome.options import Options
+  from selenium.webdriver.common.by import By
+  from selenium.webdriver.support import expected_conditions as EC
+  from selenium.webdriver.support.ui import WebDriverWait
+except ImportError:
+  raise ImportError(
+    "weblogin.mfa requires Selenium. Install with: pip install weblogin[mfa]")
+
+import requests.cookies
+
+<<set up headless browser>>
+<<perform login>>
+<<find MFA number>>
+<<collect cookies>>
+<<save cookies>>
+<<command-line interface>>
+@
+
+\subsection{Setting up the headless browser}
+
+We use Chrome in headless mode.
+The [[--no-sandbox]] and [[--disable-dev-shm-usage]] flags are needed
+when running on a server without a display.
+<<set up headless browser>>=
+def setup_browser():
+  """Start a headless Chrome browser and return the driver."""
+  options = Options()
+  options.add_argument("--headless=new")
+  options.add_argument("--no-sandbox")
+  options.add_argument("--disable-dev-shm-usage")
+  options.add_argument("--disable-gpu")
+  options.add_argument("--window-size=1280,900")
+  return webdriver.Chrome(options=options)
+@
+
+\subsection{Performing the login}
+
+The login flow is as follows.
+\begin{enumerate}
+\item Navigate to a URL that triggers a redirect to the KTH login page.
+\item Fill in credentials using JavaScript injection.
+  We cannot use [[send_keys]] because it fails with
+  \enquote{element not interactable} in headless Chrome on this page.
+\item Wait for either the MFA number to appear or for the login to
+  complete directly (some services skip MFA on trusted networks).
+\item If MFA is required, print the number and wait for the user to
+  approve on their phone.
+\end{enumerate}
+<<perform login>>=
+def do_login(driver, username, password, trigger_url):
+  """
+  Navigate to trigger_url, fill credentials, handle MFA if needed.
+  Raises RuntimeError on timeout.
+  """
+  driver.get(trigger_url)
+
+  <<wait for login page>>
+  <<fill credentials>>
+  <<wait for MFA or completion>>
+@
+
+We wait until we land on [[login.ug.kth.se]], which is KTH's
+authentication server.
+<<wait for login page>>=
+WebDriverWait(driver, 20).until(
+  lambda d: "login.ug.kth.se" in d.current_url)
+@
+
+We fill in credentials via JavaScript because [[send_keys]] does not
+work reliably in headless Chrome on this particular page.
+We append [[@ug.kth.se]] if needed, since the login form expects the
+full UG username.
+<<fill credentials>>=
+uname = username if "@ug.kth.se" in username \
+                 else username + "@ug.kth.se"
+driver.execute_script(
+  "document.getElementById('userNameInput').value = arguments[0];",
+  uname)
+driver.execute_script(
+  "document.getElementById('passwordInput').value = arguments[0];",
+  password)
+driver.execute_script(
+  "document.getElementById('submitButton').click();")
+@
+
+After submitting, we may land on an MFA challenge or proceed directly
+to the target page.
+We wait for either condition.
+We check whether we have left the login page \emph{before} scanning for
+an MFA number, to avoid false positives from two-digit numbers that may
+appear on the destination page.
+<<wait for MFA or completion>>=
+def mfa_or_done(d):
+  if not _on_login_page(d.current_url):
+    return ("done", None)
+  number = find_mfa_number(d)
+  if number:
+    return ("mfa", number)
+  return None
+
+result = WebDriverWait(driver, 30).until(mfa_or_done)
+
+if result[0] == "mfa":
+  <<handle MFA challenge>>
+@
+
+When MFA is required, we print the number and wait up to two minutes
+for the user to approve on their phone.
+<<handle MFA challenge>>=
+number = result[1]
+print("\n" + "=" * 50)
+print("  Enter  {}  in Microsoft Authenticator".format(number))
+print("=" * 50 + "\n")
+WebDriverWait(driver, 120).until(
+  lambda d: not _on_login_page(d.current_url))
+@
+
+We use a small helper to check if we are still on the login page.
+<<perform login>>=
+def _on_login_page(url):
+  return "login.ug.kth.se" in url or "microsoftonline.com" in url
+@
+
+\subsection{Finding the MFA number}
+
+Microsoft Authenticator number matching displays a two-digit number
+that the user must enter on their phone.
+The number appears in a specific element, but the element ID has varied
+across Microsoft updates.
+We therefore try several known IDs and fall back to scanning the page
+body for any standalone two-digit number.
+<<find MFA number>>=
+_MFA_ELEMENT_IDS = [
+  "idRichContext_DisplaySign",
+  "displaySign",
+  "idDiv_DisplaySign",
+]
+
+def find_mfa_number(driver):
+  """
+  Return the MFA number shown on the current page, or None if not present.
+  """
+  for eid in _MFA_ELEMENT_IDS:
+    elems = driver.find_elements(By.ID, eid)
+    if elems:
+      txt = elems[0].text.strip()
+      if txt:
+        return txt
+  <<fall back to body scan>>
+@
+
+If no known element is found, we scan the visible body text.
+This is a heuristic, but in practice the MFA number is the only
+prominent two-digit number on the page.
+<<fall back to body scan>>=
+body = driver.find_element(By.TAG_NAME, "body").text
+matches = re.findall(r'\b(\d{2})\b', body)
+return matches[0] if matches else None
+@
+
+\subsection{Collecting cookies}
+
+After a successful login, we collect all cookies from the browser and
+convert them into a [[requests.cookies.RequestsCookieJar]].
+We use [[http.cookiejar.Cookie]] objects directly rather than
+[[jar.set()]], because [[jar.set()]] does not correctly handle domain
+cookies with a leading dot (e.g.\ [[.kth.se]]).
+<<collect cookies>>=
+def collect_cookies(driver, extra_urls=None):
+  """
+  Collect cookies from the browser and return a RequestsCookieJar.
+  If extra_urls is given, also visit those URLs to pick up their cookies
+  (useful when SSO spans multiple domains).
+  """
+  jar = requests.cookies.RequestsCookieJar()
+  _add_driver_cookies(jar, driver)
+
+  for url in (extra_urls or []):
+    driver.get(url)
+    WebDriverWait(driver, 30).until(
+      lambda d: not _on_login_page(d.current_url))
+    _add_driver_cookies(jar, driver)
+
+  return jar
+
+def _add_driver_cookies(jar, driver):
+  for c in driver.get_cookies():
+    domain = c.get("domain", "")
+    cookie = http.cookiejar.Cookie(
+      version=0,
+      name=c["name"],
+      value=c["value"],
+      port=None, port_specified=False,
+      domain=domain,
+      domain_specified=bool(domain),
+      domain_initial_dot=domain.startswith("."),
+      path=c.get("path", "/"),
+      path_specified=bool(c.get("path")),
+      secure=c.get("secure", False),
+      expires=c.get("expiry"),
+      discard=False,
+      comment=None, comment_url=None,
+      rest={"HttpOnly": ""} if c.get("httpOnly") else {},
+    )
+    jar.set_cookie(cookie)
+@
+
+\subsection{Saving cookies}
+
+We save the cookie jar to a file using [[pickle]], so that other
+scripts can load it with [[pickle.load]] and pass the jar directly to
+a [[requests.Session]].
+<<save cookies>>=
+def save_cookies(jar, cookie_file):
+  """Save a RequestsCookieJar to cookie_file using pickle."""
+  os.makedirs(os.path.dirname(os.path.abspath(cookie_file)), exist_ok=True)
+  with open(cookie_file, "wb") as f:
+    pickle.dump(jar, f)
+@
+
+\subsection{Command-line interface}
+
+When run as a script ([[python -m weblogin.mfa]]), the module asks for
+credentials and a target URL, performs the login, and saves the cookies.
+If no cookie file is specified, it defaults to
+[[~/.cache/weblogin/<hostname>.pkl]], where [[<hostname>]] is derived
+from the trigger URL.
+This ensures that cookies for different services do not overwrite each
+other.
+<<command-line interface>>=
+def _default_cookie_file(trigger_url):
+  import urllib.parse
+  hostname = urllib.parse.urlparse(trigger_url).hostname or "weblogin"
+  cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "weblogin")
+  return os.path.join(cache_dir, hostname + ".pkl")
+
+def main():
+  """
+  Perform an MFA login and save cookies to a file.
+
+  Usage: python -m weblogin.mfa <trigger_url> [cookie_file]
+
+  If cookie_file is omitted, defaults to ~/.cache/weblogin/<hostname>.pkl.
+  """
+  if len(sys.argv) < 2:
+    print("Usage: python -m weblogin.mfa <trigger_url> [cookie_file]",
+          file=sys.stderr)
+    sys.exit(1)
+
+  trigger_url = sys.argv[1]
+  cookie_file = sys.argv[2] if len(sys.argv) >= 3 \
+                             else _default_cookie_file(trigger_url)
+
+  username = input("KTH username: ")
+  password = getpass.getpass("KTH password: ")
+
+  driver = setup_browser()
+  try:
+    do_login(driver, username, password, trigger_url)
+    jar = collect_cookies(driver)
+    save_cookies(jar, cookie_file)
+    print("Cookies saved to {}".format(cookie_file))
+  except Exception as e:
+    print("ERROR: {}".format(e), file=sys.stderr)
+    sys.exit(1)
+  finally:
+    driver.quit()
+
+if __name__ == "__main__":
+  main()
+@