diff --git a/src/sumo/wrapper/_auth_provider.py b/src/sumo/wrapper/_auth_provider.py index 3fb1a07..1f92153 100644 --- a/src/sumo/wrapper/_auth_provider.py +++ b/src/sumo/wrapper/_auth_provider.py @@ -183,6 +183,45 @@ def login(self): pass +class AuthProviderDeviceCode(AuthProvider): + def __init__(self, client_id, authority, resource_id): + super().__init__(resource_id) + cache = get_token_cache(resource_id) + self._app = msal.PublicClientApplication( + client_id=client_id, authority=authority, token_cache=cache + ) + self._resource_id = resource_id + self._scope = scope_for_resource(resource_id) + if self.get_token() is None: + self.login() + pass + return + + def login(self): + flow = self._app.initiate_device_flow([self._scope]) + + if "error" in flow: + raise ValueError( + "Failed to create device flow. Err: %s" + % json.dumps(flow, indent=4) + ) + + print(flow["message"]) + result = self._app.acquire_token_by_device_flow(flow) + + if "error" in result: + raise ValueError( + "Failed to acquire token by device flow. Err: %s" + % json.dumps(result, indent=4) + ) + + protect_token_cache(self._resource_id) + + return + + pass + + class AuthProviderManaged(AuthProvider): def __init__(self, resource_id): super().__init__(resource_id) @@ -203,6 +242,7 @@ def get_auth_provider( interactive=False, access_token=None, refresh_token=None, + devicecode=False, ): if refresh_token: return AuthProviderRefreshToken( @@ -214,7 +254,11 @@ def get_auth_provider( # ELSE if interactive: return AuthProviderInteractive(client_id, authority, resource_id) - + # ELSE + if devicecode: + # Potential issues with device-code + # under Equinor compliant device policy + return AuthProviderDeviceCode(client_id, authority, resource_id) # ELSE if all( [ @@ -228,8 +272,5 @@ def get_auth_provider( ] ): return AuthProviderManaged(resource_id) - # ELSE - # Device code login does not work with Equinor compliant device policy - # return AuthProviderDeviceCode(client_id, authority, resource_id) return AuthProviderInteractive(client_id, authority, resource_id) diff --git a/src/sumo/wrapper/login.py b/src/sumo/wrapper/login.py index a5dd0cd..a267ae5 100644 --- a/src/sumo/wrapper/login.py +++ b/src/sumo/wrapper/login.py @@ -37,6 +37,15 @@ def get_parser() -> ArgumentParser: help="Login interactively", ) + parser.add_argument( + "-d", + "--devicecode", + dest="devicecode", + action="store_true", + default=False, + help="Login with device-code", + ) + parser.add_argument( "-p", "--print", @@ -57,7 +66,9 @@ def main(): print("Login to Sumo environment: " + env) - sumo = SumoClient(args.env, interactive=args.interactive) + sumo = SumoClient( + args.env, interactive=args.interactive, devicecode=args.devicecode + ) token = sumo.authenticate() if args.print_token: diff --git a/src/sumo/wrapper/sumo_client.py b/src/sumo/wrapper/sumo_client.py index 41afb60..a74a042 100644 --- a/src/sumo/wrapper/sumo_client.py +++ b/src/sumo/wrapper/sumo_client.py @@ -29,6 +29,7 @@ def __init__( env: str, token: str = None, interactive: bool = False, + devicecode: bool = False, verbosity: str = "CRITICAL", retry_strategy=RetryStrategy(), ): @@ -81,6 +82,7 @@ def __init__( interactive=interactive, refresh_token=refresh_token, access_token=access_token, + devicecode=devicecode, ) if env == "localhost":