Add support for externally provided refresh tokens.

auth.py

 # Copyright 2015 The Chromium Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
 """Google OAuth2 related functions."""
 
 import BaseHTTPServer
 import collections
 import datetime
 import functools
+import hashlib
 import json
 import logging
 import optparse
 import os
 import socket
 import sys
 import threading
 import urllib
 import urlparse
 import webbrowser
@@ skipping 34 matching lines @@
     os.path.expanduser('~'), '.depot_tools_oauth2_tokens')
 
 
 # Authentication configuration extracted from command line options.
 # See doc string for 'make_auth_config' for meaning of fields.
 AuthConfig = collections.namedtuple('AuthConfig', [
     'use_oauth2', # deprecated, will be always True
     'save_cookies', # deprecated, will be removed
     'use_local_webserver',
     'webserver_port',
+    'refresh_token_json',
 ])
 
 
 # OAuth access token with its expiration time (UTC datetime or None if unknown).
 AccessToken = collections.namedtuple('AccessToken', [
     'token',
     'expires_at',
 ])
 
 
+# Refresh token passed via --auth-refresh-token-json.
+RefreshToken = collections.namedtuple('RefreshToken', [
+    'client_id',
+    'client_secret',
+    'refresh_token',
+])
+
+
 class AuthenticationError(Exception):
   """Raised on errors related to authentication."""
 
 
 class LoginRequiredError(AuthenticationError):
   """Interaction with the user is required to authenticate."""
 
   def __init__(self, token_cache_key):
     # HACK(vadimsh): It is assumed here that the token cache key is a hostname.
     msg = (
         'You are not logged in. Please login first by running:\n'
         '  depot-tools-auth login %s' % token_cache_key)
     super(LoginRequiredError, self).__init__(msg)
 
 
 def make_auth_config(
     use_oauth2=None,
     save_cookies=None,
     use_local_webserver=None,
-    webserver_port=None):
+    webserver_port=None,
+    refresh_token_json=None):
   """Returns new instance of AuthConfig.
 
   If some config option is None, it will be set to a reasonable default value.
   This function also acts as an authoritative place for default values of
   corresponding command line options.
   """
   default = lambda val, d: val if val is not None else d
   return AuthConfig(
       default(use_oauth2, _should_use_oauth2()),
       default(save_cookies, True),
       default(use_local_webserver, not _is_headless()),
-      default(webserver_port, 8090))
+      default(webserver_port, 8090),
+      default(refresh_token_json, ''))
 
 
 def add_auth_options(parser, default_config=None):
   """Appends OAuth related options to OptionParser."""
   default_config = default_config or make_auth_config()
   parser.auth_group = optparse.OptionGroup(parser, 'Auth options')
   parser.add_option_group(parser.auth_group)
 
   # OAuth2 vs password switch.
   auth_default = 'use OAuth2' if default_config.use_oauth2 else 'use password'
@@ skipping 24 matching lines @@
       action='store_false',
       dest='use_local_webserver',
       default=default_config.use_local_webserver,
       help='Do not run a local web server when performing OAuth2 login flow.')
   parser.auth_group.add_option(
       '--auth-host-port',
       type=int,
       default=default_config.webserver_port,
       help='Port a local web server should listen on. Used only if '
           '--auth-no-local-webserver is not set. [default: %default]')
+  parser.auth_group.add_option(
+      '--auth-refresh-token-json',
+      default=default_config.refresh_token_json,
+      help='Path to a JSON file with role account refresh token to use.')
 
 
 def extract_auth_config_from_options(options):
   """Given OptionParser parsed options, extracts AuthConfig from it.
 
   OptionParser should be populated with auth options by 'add_auth_options'.
   """
   return make_auth_config(
       use_oauth2=options.use_oauth2,
       save_cookies=False if options.use_oauth2 else options.save_cookies,
       use_local_webserver=options.use_local_webserver,
-      webserver_port=options.auth_host_port)
+      webserver_port=options.auth_host_port,
+      refresh_token_json=options.auth_refresh_token_json)
 
 
 def auth_config_to_command_options(auth_config):
   """AuthConfig -> list of strings with command line options.
 
   Omits options that are set to default values.
   """
   if not auth_config:
     return []
   defaults = make_auth_config()
   opts = []
   if auth_config.use_oauth2 != defaults.use_oauth2:
     opts.append('--oauth2' if auth_config.use_oauth2 else '--no-oauth2')
   if auth_config.save_cookies != auth_config.save_cookies:
     if not auth_config.save_cookies:
       opts.append('--no-cookies')
   if auth_config.use_local_webserver != defaults.use_local_webserver:
     if not auth_config.use_local_webserver:
       opts.append('--auth-no-local-webserver')
   if auth_config.webserver_port != defaults.webserver_port:
     opts.extend(['--auth-host-port', str(auth_config.webserver_port)])
+  if auth_config.refresh_token_json != defaults.refresh_token_json:
+    opts.extend([
+        '--auth-refresh-token-json', str(auth_config.refresh_token_json)])
   return opts
 
 
 def get_authenticator_for_host(hostname, config):
   """Returns Authenticator instance to access given host.
 
   Args:
     hostname: a naked hostname or http(s)://<hostname>[/] URL. Used to derive
         a cache key for token cache.
     config: AuthConfig instance.
@@ skipping 21 matching lines @@
     config: AuthConfig object that holds authentication configuration.
   """
 
   def __init__(self, token_cache_key, config):
     assert isinstance(config, AuthConfig)
     assert config.use_oauth2
     self._access_token = None
     self._config = config
     self._lock = threading.Lock()
     self._token_cache_key = token_cache_key
+    self._external_token = None
+    if config.refresh_token_json:
+      self._external_token = _read_refresh_token_json(config.refresh_token_json)
 
   def login(self):
     """Performs interactive login flow if necessary.
 
     Raises:
       AuthenticationError on error or if interrupted.
     """
+    if self._external_token:
+      raise AuthenticationError(
+          'Can\'t run login flow when using --auth-refresh-token-json.')
     return self.get_access_token(
         force_refresh=True, allow_user_interaction=True)
 
   def logout(self):
     """Revokes the refresh token and deletes it from the cache.
 
-    Returns True if actually revoked a token.
+    Returns True if had some credentials cached.
     """
-    revoked = False
     with self._lock:
       self._access_token = None
       storage = self._get_storage()
       credentials = storage.get()
-      if credentials:
-        credentials.revoke(httplib2.Http())
-        revoked = True
+      if credentials and credentials.refresh_token and credentials.revoke_uri:
+        try:
+          credentials.revoke(httplib2.Http())
+        except client.TokenRevokeError as e:
+          logging.warning('Failed to revoke refresh token: %s', e)

[M-A Ruel, 2015/04/17 01:49:50]

you don't set credentials to None in else / except?

[Vadim Sh., 2015/04/17 02:38:07]

If cached refresh token is invalid (for example revoked), revoke raises
client.TokenRevokeError. In such case "depot-tools-auth logout" would still show
message "Done" (and a warning), since logout() returns True (because it did have
some credentials in the cache, and it did remove them from there).

       storage.delete()
-    return revoked
+    return bool(credentials)
 
   def has_cached_credentials(self):
     """Returns True if long term credentials (refresh token) are in cache.
 
     Doesn't make network calls.
 
     If returns False, get_access_token() later will ask for interactive login by
     raising LoginRequiredError.
 
     If returns True, most probably get_access_token() won't ask for interactive
     login, though it is not guaranteed, since cached token can be already
     revoked and there's no way to figure this out without actually trying to use
     it.
     """
     with self._lock:
-      credentials = self._get_storage().get()
-      return credentials and not credentials.invalid
+      return bool(self._get_cached_credentials())
 
   def get_access_token(self, force_refresh=False, allow_user_interaction=False):
     """Returns AccessToken, refreshing it if necessary.
 
     Args:
       force_refresh: forcefully refresh access token even if it is not expired.
       allow_user_interaction: True to enable blocking for user input if needed.
 
     Raises:
       AuthenticationError on error or if authentication flow was interrupted.
@@ skipping 64 matching lines @@
       else:
         return (resp, content)
 
     http.request = new_request
     return http
 
   ## Private methods.
 
   def _get_storage(self):
     """Returns oauth2client.Storage with cached tokens."""
+    # Do not mix cache keys for different externally provided tokens.
+    if self._external_token:
+      token_hash = hashlib.sha1(self._external_token.refresh_token).hexdigest()
+      cache_key = '%s:refresh_tok:%s' % (self._token_cache_key, token_hash)
+    else:
+      cache_key = self._token_cache_key
     return multistore_file.get_credential_storage_custom_string_key(
-        OAUTH_TOKENS_CACHE, self._token_cache_key)
+        OAUTH_TOKENS_CACHE, cache_key)
+
+  def _get_cached_credentials(self):
+    """Returns oauth2client.Credentials loaded from storage."""
+    storage = self._get_storage()
+    credentials = storage.get()
+
+    # Is using --auth-refresh-token-json?
+    if self._external_token:
+      # Cached credentials are valid and match external token -> use them. It is
+      # important to reuse credentials from the storage because they contain
+      # cached access token.
+      valid = (
+          credentials and not credentials.invalid and
+          credentials.refresh_token == self._external_token.refresh_token and
+          credentials.client_id == self._external_token.client_id and
+          credentials.client_secret == self._external_token.client_secret)
+      if valid:
+        return credentials
+      # Construct new credentials from externally provided refresh token,
+      # associate them with cache storage (so that access_token will be placed
+      # in the cache later too).
+      credentials = client.OAuth2Credentials(
+          access_token=None,
+          client_id=self._external_token.client_id,
+          client_secret=self._external_token.client_secret,
+          refresh_token=self._external_token.refresh_token,
+          token_expiry=None,
+          token_uri='https://accounts.google.com/o/oauth2/token',
+          user_agent=None,
+          revoke_uri=None)
+      credentials.set_store(storage)
+      storage.put(credentials)
+      return credentials
+
+    # Not using external refresh token -> return whatever is cached.
+    return credentials if (credentials and not credentials.invalid) else None
 
   def _load_access_token(self):
     """Returns cached AccessToken if it is not expired yet."""
-    credentials = self._get_storage().get()
-    if not credentials or credentials.invalid:
+    creds = self._get_cached_credentials()
+    if not creds or not creds.access_token or creds.access_token_expired:
       return None
-    if not credentials.access_token or credentials.access_token_expired:
-      return None
-    return AccessToken(str(credentials.access_token), credentials.token_expiry)
+    return AccessToken(str(creds.access_token), creds.token_expiry)
 
   def _create_access_token(self, allow_user_interaction=False):
     """Mints and caches a new access token, launching OAuth2 dance if necessary.
 
     Uses cached refresh token, if present. In that case user interaction is not
     required and function will finish quietly. Otherwise it will launch 3-legged
     OAuth2 flow, that needs user interaction.
 
     Args:
       allow_user_interaction: if True, allow interaction with the user (e.g.
           reading standard input, or launching a browser).
 
     Returns:
       AccessToken.
 
     Raises:
       AuthenticationError on error or if authentication flow was interrupted.
       LoginRequiredError if user interaction is required, but
           allow_user_interaction is False.
     """
-    storage = self._get_storage()
-    credentials = None
+    credentials = self._get_cached_credentials()
 
     # 3-legged flow with (perhaps cached) refresh token.
-    credentials = storage.get()
     refreshed = False
     if credentials and not credentials.invalid:
       try:
         credentials.refresh(httplib2.Http())
         refreshed = True
       except client.Error as err:
         logging.warning(
-            'OAuth error during access token refresh: %s. '
+            'OAuth error during access token refresh (%s). '
             'Attempting a full authentication flow.', err)
 
     # Refresh token is missing or invalid, go through the full flow.
     if not refreshed:
+      # Can't refresh externally provided token.
+      if self._external_token:
+        raise AuthenticationError(
+            'Token provided via --auth-refresh-token-json is no longer valid.')
       if not allow_user_interaction:
         raise LoginRequiredError(self._token_cache_key)
       credentials = _run_oauth_dance(self._config)
 
     logging.info(
         'OAuth access_token refreshed. Expires in %s.',
         credentials.token_expiry - datetime.datetime.utcnow())
+    storage = self._get_storage()
     credentials.set_store(storage)
     storage.put(credentials)
     return AccessToken(str(credentials.access_token), credentials.token_expiry)
 
 
 ## Private functions.
 
 
 def _should_use_oauth2():
   """Default value for use_oauth2 config option.
 
   Used to selectively enable OAuth2 by default.
   """
   return os.path.exists(os.path.join(DEPOT_TOOLS_DIR, 'USE_OAUTH2'))
 
 
 def _is_headless():
   """True if machine doesn't seem to have a display."""
   return sys.platform == 'linux2' and not os.environ.get('DISPLAY')
 
 
+def _read_refresh_token_json(path):
+  """Returns RefreshToken by reading it from the JSON file."""
+  try:
+    with open(path, 'r') as f:
+      data = json.load(f)
+      return RefreshToken(
+          client_id=str(data.get('client_id', OAUTH_CLIENT_ID)),
+          client_secret=str(data.get('client_secret', OAUTH_CLIENT_SECRET)),
+          refresh_token=str(data['refresh_token']))
+  except (IOError, ValueError) as e:
+    raise AuthenticationError(
+        'Failed to read refresh token from %s: %s' % (path, e))
+  except KeyError as e:
+    raise AuthenticationError(
+        'Failed to read refresh token from %s: missing key %s' % (path, e))
+
+
 def _needs_refresh(access_token):
   """True if AccessToken should be refreshed."""
   if access_token.expires_at is not None:
     # Allow 5 min of clock skew between client and backend.
     now = datetime.datetime.utcnow() + datetime.timedelta(seconds=300)
     return now >= access_token.expires_at
   # Token without expiration time never expires.
   return False
 
 
@@ skipping 103 matching lines @@
     self.end_headers()
     query = self.path.split('?', 1)[-1]
     query = dict(urlparse.parse_qsl(query))
     self.server.query_params = query
     self.wfile.write('<html><head><title>Authentication Status</title></head>')
     self.wfile.write('<body><p>The authentication flow has completed.</p>')
     self.wfile.write('</body></html>')
 
   def log_message(self, _format, *args):
     """Do not log messages to stdout while running as command line program."""