diff --git a/README.md b/README.md index 6902409..0081ea0 100644 --- a/README.md +++ b/README.md @@ -121,6 +121,16 @@ garth.login(email, password) 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. + +```python +garth.login(email, password, prompt_mfa=lambda: input("Enter MFA code: ")) +``` + ### Configure #### Set domain for China diff --git a/garth/http.py b/garth/http.py index 53d3407..037a1ea 100644 --- a/garth/http.py +++ b/garth/http.py @@ -156,8 +156,10 @@ def delete(self, *args, **kwargs) -> Response: def put(self, *args, **kwargs) -> Response: return self.request("PUT", *args, **kwargs) - def login(self, *args): - self.oauth1_token, self.oauth2_token = sso.login(*args, client=self) + def login(self, *args, **kwargs): + self.oauth1_token, self.oauth2_token = sso.login( + *args, **kwargs, client=self + ) def refresh_oauth2(self): assert self.oauth1_token diff --git a/garth/sso.py b/garth/sso.py index cf7eb54..13583b1 100644 --- a/garth/sso.py +++ b/garth/sso.py @@ -1,6 +1,6 @@ import re import time -from typing import Dict, Optional, Tuple +from typing import Callable, Dict, Optional, Tuple from urllib.parse import parse_qs import requests @@ -40,7 +40,11 @@ def __init__( def login( - email: str, password: str, /, client: Optional["http.Client"] = None + email: str, + password: str, + /, + client: Optional["http.Client"] = None, + prompt_mfa: Callable = lambda: input("MFA code: "), ) -> Tuple[OAuth1Token, OAuth2Token]: client = client or http.client @@ -92,7 +96,7 @@ def login( # Handle MFA if "MFA" in title: - handle_mfa(client, SIGNIN_PARAMS) + handle_mfa(client, SIGNIN_PARAMS, prompt_mfa) title = get_title(client.last_resp.text) assert title == "Success" @@ -142,9 +146,11 @@ def exchange(oauth1: OAuth1Token, client: "http.Client") -> OAuth2Token: return OAuth2Token(**set_expirations(token)) -def handle_mfa(client: "http.Client", signin_params: dict) -> None: +def handle_mfa( + client: "http.Client", signin_params: dict, prompt_mfa: Callable +) -> None: csrf_token = get_csrf_token(client.last_resp.text) - mfa_code = input("Enter MFA code: ") + mfa_code = prompt_mfa() client.post( "sso", "/sso/verifyMFA/loginEnterMfaCode", diff --git a/garth/version.py b/garth/version.py index 96dcabd..98faf4f 100644 --- a/garth/version.py +++ b/garth/version.py @@ -1 +1 @@ -__version__ = "0.4.44" +__version__ = "0.4.45"