From f1a110cebdd9407be5c464bac21936080a37a50c Mon Sep 17 00:00:00 2001 From: George Pollard Date: Thu, 10 Nov 2022 01:42:16 +0000 Subject: [PATCH] Use browser login instead of device flow --- src/cli/onefuzz/backend.py | 54 ++++++++++++++++++++++---------------- 1 file changed, 31 insertions(+), 23 deletions(-) diff --git a/src/cli/onefuzz/backend.py b/src/cli/onefuzz/backend.py index 41700560e42..96f971e21b2 100644 --- a/src/cli/onefuzz/backend.py +++ b/src/cli/onefuzz/backend.py @@ -114,7 +114,7 @@ def __init__( self.config = config self.token_cache: Optional[msal.SerializableTokenCache] = None self.init_cache() - self.app: Optional[Any] = None + self.app: Optional[msal.ClientApplication] = None self.token_expires = 0 self.load_config() self.session = requests.Session() @@ -197,7 +197,7 @@ def get_access_token(self) -> Any: if self.client_secret: return self.access_token_from_client_secret(scopes) - return self.device_login(scopes) + return self.do_login(scopes) def access_token_from_client_secret(self, scopes: List[str]) -> Any: if not self.app: @@ -229,7 +229,7 @@ def access_token_from_client_secret(self, scopes: List[str]) -> Any: ) return result - def device_login(self, scopes: List[str]) -> Any: + def do_login(self, scopes: List[str]) -> Any: if not self.app: self.app = msal.PublicClientApplication( self.config.client_id, @@ -237,6 +237,7 @@ def device_login(self, scopes: List[str]) -> Any: token_cache=self.token_cache, ) + access_token = None for scope in scopes: accounts = self.app.get_accounts() if accounts: @@ -248,26 +249,33 @@ def device_login(self, scopes: List[str]) -> Any: for scope in scopes: LOGGER.info("Attempting interactive device login") - print("Please login", flush=True) - - flow = self.app.initiate_device_flow(scopes=[scope]) - - check_msal_error(flow, ["user_code", "message"]) - # setting the expiration time to allow us to retry the interactive login with a new scope - flow["expires_at"] = int(time.time()) + 90 # 90 seconds from now - print(flow["message"], flush=True) - - access_token = self.app.acquire_token_by_device_flow(flow) - # AADSTS70016: OAuth 2.0 device flow error. Authorization is pending - # this happens when the intractive login request times out. This heppens when the login - # fails because of a scope mismatch. - if ( - "error" in access_token - and "AADSTS70016" in access_token["error_description"] - ): - LOGGER.warning(f"failed to get access token with scope {scope}") - continue - check_msal_error(access_token, ["access_token"]) + try: + access_token = self.app.acquire_token_interactive(scopes=[scope]) + check_msal_error(access_token, ["access_token"]) + except KeyboardInterrupt: + result = input("\nInteractive login cancelled. Use device login (Y/n)? ") + if result == "" or result.startswith("y") or result.startswith("Y"): + print("Falling back to device flow, please sign in:", flush=True) + flow = self.app.initiate_device_flow(scopes=[scope]) + + check_msal_error(flow, ["user_code", "message"]) + # setting the expiration time to allow us to retry the interactive login with a new scope + flow["expires_at"] = int(time.time()) + 90 # 90 seconds from now + print(flow["message"], flush=True) + + access_token = self.app.acquire_token_by_device_flow(flow) + # AADSTS70016: OAuth 2.0 device flow error. Authorization is pending + # this happens when the intractive login request times out. This heppens when the login + # fails because of a scope mismatch. + if ( + "error" in access_token + and "AADSTS70016" in access_token["error_description"] + ): + LOGGER.warning(f"failed to get access token with scope {scope}") + continue + check_msal_error(access_token, ["access_token"]) + else: + continue LOGGER.info("Interactive device authentication succeeded") print("Login succeeded", flush=True)