diff --git a/google_keep_tasks/auth.py b/google_keep_tasks/auth.py index 238c2cd..518c672 100644 --- a/google_keep_tasks/auth.py +++ b/google_keep_tasks/auth.py @@ -21,8 +21,12 @@ class GoogleKeepFileAuth(object): def __init__(self, file=None): self.file = file or AUTH_FILE - def save_credentials(self, username, password): - data = {'username': username, 'password': password} + def save_credentials(self, username, password=None, master_token=None): + data = {'username': username} + if master_token: + data['master_token'] = master_token + elif password: + data['password'] = password directory = os.path.dirname(AUTH_FILE) if not os.path.exists(directory): os.makedirs(directory, 0o700) @@ -37,9 +41,18 @@ def get_credentials(self): raise LoginError('Invalid json from credentials file: {}. Error: {}'.format( self.file, e )) - if not isinstance(data, dict) or 'username' not in data or 'password' not in data: - raise LoginError('Invalid credentials format file. username and password are required.') - return data['username'], data['password'] + if not isinstance(data, dict) or 'username' not in data: + raise LoginError('Invalid credentials format file. username is required.') + + # Support both old (password) and new (master_token) format + username = data['username'] + master_token = data.get('master_token') + password = data.get('password') + + if not master_token and not password: + raise LoginError('Either password or master_token is required.') + + return username, password, master_token class GoogleKeep(object): @@ -50,9 +63,9 @@ def __init__(self): def login_or_input(self): auth_changed = False input_login = False - username = password = None + username = password = master_token = None try: - username, password = self.auth.get_credentials() + username, password, master_token = self.auth.get_credentials() except UnavailableLoginError: click.echo('Welcome to Google Keep. Enter your username and password below. ' 'If your account is protected, you need an application password: ' @@ -66,8 +79,12 @@ def login_or_input(self): # Request new credentials auth_changed = True username, password = self.get_credencials_assistant(username) + master_token = None try: - self.keep.login(username, password) + if master_token: + self.keep.resume(username, master_token) + else: + self.keep.login(username, password) except LoginException: choice = choices_prompt('Authentication failed, what do you want to do?', [ 'Enter new credentials', @@ -82,7 +99,7 @@ def login_or_input(self): raise Abort else: if auth_changed: - self.auth.save_credentials(username, password) + self.auth.save_credentials(username, password, master_token) break def get_credencials_assistant(self, default_username):