Merge pull request #3 from swsphn/mfa

Support PMHC MFA login

README.md

@@ -50,6 +50,7 @@ variables if they are set:
 ```
 PMHC_USERNAME
 PMHC_PASSWORD
+PMHC_TOTP_SECRET
 ```
 
 Otherwise, you will be prompted for credentials interactively.
@@ -60,15 +61,40 @@ follows:
 ``` ps1
 $env:PMHC_USERNAME='your_username_here'
 $env:PMHC_PASSWORD=python -c 'import getpass; print(getpass.getpass())'
+$env:PMHC_TOTP_SECRET=python -c 'import getpass; print(getpass.getpass("TOTP Secret: "))'
 ```
 
 In a Unix shell (Mac, Linux), you can do:
 
 ``` bash
 export PMHC_USERNAME='your_username_here'
 read -rs PMHC_PASSWORD && export PMHC_PASSWORD
+read -rs PMHC_TOTP_SECRET && export PMHC_TOTP_SECRET
 ```
 
+NOTE: `PMHC_TOTP_SECRET` is the unchanging base32-encoded TOTP secret,
+not the time-based six-digit code. You can likely find this secret in
+the 'advanced' section of your TOTP app. It will be a long string of
+upper-case letters and digits. See below for a list of TOTP apps which
+support viewing the TOTP secret. It is also possible to get the secret
+by scanning the setup QR code, or by clicking the button on the website
+to manually configure the TOTP app. The six-digit code will be
+automatically calculated based on the current time as required if
+`PMHC_TOTP_SECRET` is specified. Otherwise, the user will be prompted to
+enter the current six-digit code.
+
+Not all TOTP apps support viewing the secret. The following are known
+to support this:
+
+- [Aegis Authenticator](https://getaegis.app/) (Android only)
+- [Bitwarden
+  Authenticator](https://bitwarden.com/products/authenticator/)
+- [Ente Auth](https://github.com/ente-io/ente/tree/main/auth#readme)
+- [2FA Authenticator (2FAS)](https://2fas.com/)
+
+For more details, see the [list of recommended authenticator
+apps][mfa-apps] on our Data Wiki.
+
 ## Documentation
 
 See the [online documentation][docs].
@@ -106,3 +132,4 @@ The generated documentation can be viewed at `docs/_build/html/index.html`.
 [Playwright]: https://playwright.dev/python/
 [Sphinx]: https://www.sphinx-doc.org/
 [docs]: https://swsphn.github.io/pmhclib/
+[mfa-apps]: https://datawiki.swsphn.com.au/software/gui-tools/multi-factor-authentication-apps/

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pmhclib"
-version = "0.6.2"
+version = "0.7.0"
 description = "Python wrapper for unofficial PMHC MDS portal API"
 authors = [
     "David Wales <[email protected]>",
@@ -12,6 +12,7 @@ readme = "README.md"
 python = "^3.8"
 rich = "^13.7.0"
 playwright = "^1.40.0"
+pyotp = "^2.9.0"
 
 [tool.poetry.group.docs.dependencies]
 sphinx = "^7.0.0"

src/pmhclib/pmhc.py

@@ -4,12 +4,15 @@
 
 The script uses Python Playwright to do this
 Tested under Ubuntu WSL and PowerShell.
---no-headless runs best under PowerShell (it's slower under Ubuntu WSL)
+headless=True runs best under PowerShell (it's slower under Ubuntu WSL)
 
 No login details are saved anywhere
 To speed up usage when doing repeated calls, create the following local env variables:
 PMHC_USERNAME
 PMHC_PASSWORD
+PMHC_TOTP_SECRET
+
+Note: See PMHC.login() documentation for details about `PMHC_TOTP_SECRET`.
 """
 
 import logging
@@ -23,6 +26,7 @@
 from pathlib import Path
 from typing import Optional
 
+import pyotp
 import playwright.sync_api
 from playwright.sync_api import sync_playwright
 from rich.progress import Progress, TimeElapsedColumn
@@ -136,11 +140,20 @@ def login(self):
 
         - `PMHC_USERNAME`
         - `PMHC_PASSWORD`
+        - `PMHC_TOTP_SECRET`
+
+        NOTE: `PMHC_TOTP_SECRET` is _not_ the 6 digit time-dependent TOTP code,
+        but rather the long base32 encoded random secret. You might find this in
+        the 'advanced' section when editing the record in your TOTP app. It will
+        likely be a long string containing uppercase letters and numbers. This
+        will be automatically combined with the current time to derive the
+        correct 6 digit code.
         """
 
         # Prompt user for credentials if not set in env.
         username = os.getenv("PMHC_USERNAME")
         password = SecureString(os.getenv("PMHC_PASSWORD") or "")
+        totp_secret = SecureString(os.getenv("PMHC_TOTP_SECRET") or "")
 
         while not username:
             username = input("Enter PMHC username: ")
@@ -166,9 +179,39 @@ def login(self):
         password_field = self.page.locator('input[id="password"]')
         password_field.fill(password)
         password_field.press("Enter")
+        self.page.wait_for_load_state()
+
+        # Note: We get the code _after_ loading the page and entering
+        # the username and password, to ensure that it is still valid
+        # when we submit it.
+        logging.info("Entering TOTP MFA code")
+        if totp_secret:
+            totp = pyotp.TOTP(totp_secret)
+            totp_code = totp.now()
+        else:
+            totp_code = None
+            while not totp_code:
+                totp_code = SecureString(
+                    getpass(
+                        "Enter six-digit MFA code (keyboard input will be hidden): "
+                    )
+                )
 
+        mfa_field = self.page.locator('input[id="code"]')
+        mfa_field.fill(totp_code)
+        mfa_field.press("Enter")
         self.page.wait_for_load_state()
 
+        # Skip fingerprint/face recognition enrollment
+        # TODO: Only do this if we are actually on this page.
+        if self.page.url.startswith(
+            "https://login.logicly.com.au/u/mfa-webauthn-platform-enrollment"
+        ):
+            logging.info("Skipping fingerprint/face recognition enrollment")
+            skip_button = self.page.locator('button[value="refuse-add-device"]')
+            skip_button.click()
+            self.page.wait_for_load_state()
+
         # confirm login was successful
         user_query = self.page.request.get("https://pmhc-mds.net/api/current-user")
         self.user_info = user_query.json()
@@ -354,17 +397,17 @@ def wait_for_extract(self, uuid: str, max_retries: int) -> bool:
         For this reason, it's not sufficient to simply try the
         download URL until we get a success code. If there is
         a PMHC server error, we will end up retrying forever.
-        
+
         Instead, we need to fetch the list of extracts and filter
         for one with the required uuid. We can then check the
         extract status explicitly, which should be one of the
         following values:
-        
+
         - Completed
         - Processing
         - Queued
         - Error
-        
+
         If Completed, we can download the extract.
         If Processing or Queued, keep looping and waiting.
         If Error, the extract has failed. Exit.