Merge pull request #84 from matin/return_on_mfa

* support return_on_mfa

* bump version

* ignore this exception in coverage

* use fewer asserts

.markdownlint.json

@@ -1,5 +1,6 @@
 {
     "MD033": {
         "allowed_elements": ["img"]
-    }
+    },
+    "MD046": false
 }

README.md

@@ -1,8 +1,12 @@
 # Garth
 
-[![CI](https://github.com/matin/garth/workflows/CI/badge.svg?event=push)](https://github.com/matin/garth/actions/workflows/ci.yml?query=event%3Apush+branch%3Amain+workflow%3ACI)
-[![codecov](https://codecov.io/gh/matin/garth/branch/main/graph/badge.svg?token=0EFFYJNFIL)](https://codecov.io/gh/matin/garth)
-[![PyPI - Downloads](https://img.shields.io/pypi/dm/garth)](https://pypistats.org/packages/garth)
+[![CI](https://github.com/matin/garth/workflows/CI/badge.svg?event=push)](
+    https://github.com/matin/garth/actions/workflows/ci.yml?query=event%3Apush+branch%3Amain+workflow%3ACI)
+[![codecov](
+    https://codecov.io/gh/matin/garth/branch/main/graph/badge.svg?token=0EFFYJNFIL)](
+    https://codecov.io/gh/matin/garth)
+[![PyPI - Downloads](https://img.shields.io/pypi/dm/garth)](
+    https://pypistats.org/packages/garth)
 
 Garmin SSO auth + Connect Python client
 
@@ -16,7 +20,8 @@ individual days and the trend. The Colab retrieves up to three years of daily
 data. If there's less than three years of data, it retrieves whatever is
 available.
 
-![Stress: Garph of 28-day rolling average](https://github.com/matin/garth/assets/98985/868ecf25-4644-4879-b28f-ed0706a9e7b9)
+![Stress: Garph of 28-day rolling average](
+    https://github.com/matin/garth/assets/98985/868ecf25-4644-4879-b28f-ed0706a9e7b9)
 
 ### [Sleep analysis over 90 days](https://colab.research.google.com/github/matin/garth/blob/main/colabs/sleep.ipynb)
 
@@ -27,13 +32,15 @@ daily sleep quality in 28-day pages, but that doesn't show details. Using
 day with enough detail to product a stacked bar graph of the daily sleep
 stages.
 
-![Sleep stages over 90 days](https://github.com/matin/garth/assets/98985/ba678baf-0c8a-4907-aa91-be43beec3090)
+![Sleep stages over 90 days](
+    https://github.com/matin/garth/assets/98985/ba678baf-0c8a-4907-aa91-be43beec3090)
 
 One specific graph that's useful but not available in the Connect app is
 sleep start and end times over an extended period. This provides context
 to the sleep hours and stages.
 
-![Sleep times over 90 days](https://github.com/matin/garth/assets/98985/c5583b9e-ab8a-4b5c-bfe6-1cb0ca95d1de)
+![Sleep times over 90 days](
+    https://github.com/matin/garth/assets/98985/c5583b9e-ab8a-4b5c-bfe6-1cb0ca95d1de)
 
 ### [ChatGPT analysis of Garmin stats](https://colab.research.google.com/github/matin/garth/blob/main/colabs/chatgpt_analysis_of_stats.ipynb)
 
@@ -123,14 +130,22 @@ garth.save("~/.garth")
 
 ### Custom MFA handler
 
-There's already a default MFA handler that prompts for the code in the
-terminal. You can provide your own handler. The handler should return the
-MFA code through your custom prompt.
+By default, MFA will prompt for the code in the terminal. You can provide your
+own handler:
 
 ```python
 garth.login(email, password, prompt_mfa=lambda: input("Enter MFA code: "))
 ```
 
+For advanced use cases (like async handling), MFA can be handled separately:
+
+```python
+result = garth.login(email, password, return_on_mfa=True)
+if isinstance(result, dict):  # MFA is required
+    mfa_code = "123456"  # Get this from your custom MFA flow
+    oauth1, oauth2 = garth.resume_login(result['client_state'], mfa_code)
+```
+
 ### Configure
 
 #### Set domain for China
@@ -188,7 +203,7 @@ list(sleep.keys())
 ### Stats
 
 ```python
-stress =  garth.connectapi("/usersummary-service/stats/stress/weekly/2023-07-05/52")
+stress = garth.connectapi("/usersummary-service/stats/stress/weekly/2023-07-05/52")
 ```
 
 ```json

garth/sso.py

@@ -41,13 +41,55 @@ def __init__(
             self.verify = parent.verify
 
 
+def _complete_login(
+    client: "http.Client", html: str
+) -> Tuple[OAuth1Token, OAuth2Token]:
+    """Complete the login process after successful authentication.
+
+    Args:
+        client: The HTTP client
+        html: The HTML response containing the ticket
+
+    Returns:
+        Tuple[OAuth1Token, OAuth2Token]: The OAuth tokens
+    """
+    # Parse ticket
+    m = re.search(r'embed\?ticket=([^"]+)"', html)
+    if not m:
+        raise GarthException(
+            "Couldn't find ticket in response"
+        )  # pragma: no cover
+    ticket = m.group(1)
+
+    oauth1 = get_oauth1_token(ticket, client)
+    oauth2 = exchange(oauth1, client)
+
+    return oauth1, oauth2
+
+
 def login(
     email: str,
     password: str,
     /,
     client: "http.Client | None" = None,
-    prompt_mfa: Callable = lambda: input("MFA code: "),
-) -> Tuple[OAuth1Token, OAuth2Token]:
+    prompt_mfa: Callable | None = lambda: input("MFA code: "),
+    return_on_mfa: bool = False,
+) -> Tuple[OAuth1Token, OAuth2Token] | dict:
+    """Login to Garmin Connect.
+
+    Args:
+        email: Garmin account email
+        password: Garmin account password
+        client: Optional HTTP client to use
+        prompt_mfa: Callable that prompts for MFA code. Returns on MFA if None.
+        return_on_mfa: If True, returns dict with MFA info instead of prompting
+
+    Returns:
+        If return_on_mfa=False (default):
+            Tuple[OAuth1Token, OAuth2Token]: OAuth tokens after login
+        If return_on_mfa=True and MFA required:
+            dict: Contains needs_mfa and client_state for resume_login()
+    """
     client = client or http.client
 
     # Define params based on domain
@@ -98,28 +140,34 @@ def login(
 
     # Handle MFA
     if "MFA" in title:
+        if return_on_mfa or prompt_mfa is None:
+            return {
+                "needs_mfa": True,
+                "client_state": {
+                    "csrf_token": csrf_token,
+                    "signin_params": SIGNIN_PARAMS,
+                    "client": client,
+                },
+            }
+
         handle_mfa(client, SIGNIN_PARAMS, prompt_mfa)
         title = get_title(client.last_resp.text)
 
-    assert title == "Success", f"Unexpected title: {title}"
-
-    # Parse ticket
-    m = re.search(r'embed\?ticket=([^"]+)"', client.last_resp.text)
-    assert m, "Couldn't find ticket in response"
-    ticket = m.group(1)
-
-    oauth1 = get_oauth1_token(ticket, client)
-    oauth2 = exchange(oauth1, client)
-
-    return oauth1, oauth2
+    if title != "Success":
+        raise GarthException(f"Unexpected title: {title}")  # pragma: no cover
+    return _complete_login(client, client.last_resp.text)
 
 
 def get_oauth1_token(ticket: str, client: "http.Client") -> OAuth1Token:
     sess = GarminOAuth1Session(parent=client.sess)
+    base_url = f"https://connectapi.{client.domain}/oauth-service/oauth/"
+    login_url = f"https://sso.{client.domain}/sso/embed"
+    url = (
+        f"{base_url}preauthorized?ticket={ticket}&login-url={login_url}"
+        "&accepts-mfa-tokens=true"
+    )
     resp = sess.get(
-        f"https://connectapi.{client.domain}/oauth-service/oauth/"
-        f"preauthorized?ticket={ticket}&login-url="
-        f"https://sso.{client.domain}/sso/embed&accepts-mfa-tokens=true",
+        url,
         headers=USER_AGENT,
         timeout=client.timeout,
     )
@@ -136,12 +184,15 @@ def exchange(oauth1: OAuth1Token, client: "http.Client") -> OAuth2Token:
         parent=client.sess,
     )
     data = dict(mfa_token=oauth1.mfa_token) if oauth1.mfa_token else {}
+    base_url = f"https://connectapi.{client.domain}/oauth-service/oauth/"
+    url = f"{base_url}exchange/user/2.0"
+    headers = {
+        **USER_AGENT,
+        **{"Content-Type": "application/x-www-form-urlencoded"},
+    }
     token = sess.post(
-        f"https://connectapi.{client.domain}/oauth-service/oauth/exchange/user/2.0",
-        headers={
-            **USER_AGENT,
-            **{"Content-Type": "application/x-www-form-urlencoded"},
-        },
+        url,
+        headers=headers,
         data=data,
         timeout=client.timeout,
     ).json()
@@ -190,3 +241,38 @@ def get_title(html: str) -> str:
     if not m:
         raise GarthException("Couldn't find title")
     return m.group(1)
+
+
+def resume_login(
+    client_state: dict, mfa_code: str
+) -> Tuple[OAuth1Token, OAuth2Token]:
+    """Complete login after MFA code is provided.
+
+    Args:
+        client_state: The client state from login() when MFA was needed
+        mfa_code: The MFA code provided by the user
+
+    Returns:
+        Tuple[OAuth1Token, OAuth2Token]: The OAuth tokens after login
+    """
+    client = client_state["client"]
+    signin_params = client_state["signin_params"]
+    csrf_token = client_state["csrf_token"]
+
+    client.post(
+        "sso",
+        "/sso/verifyMFA/loginEnterMfaCode",
+        params=signin_params,
+        referrer=True,
+        data={
+            "mfa-code": mfa_code,
+            "embed": "true",
+            "_csrf": csrf_token,
+            "fromPage": "setupEnterMfaCode",
+        },
+    )
+
+    title = get_title(client.last_resp.text)
+    if title != "Success":
+        raise GarthException(f"Unexpected title: {title}")  # pragma: no cover
+    return _complete_login(client, client.last_resp.text)

garth/version.py

@@ -1 +1 @@
-__version__ = "0.5.0"
+__version__ = "0.5.1"