Adds SSO auth to gsutil

third_party/gsutil/plugins/sso_auth.py

+# Copyright (c) 2013 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.
+
+"""AuthHandler plugin for gsutil's boto to support LOAS based auth."""
+
+import getpass
+import json
+import os
+import re
+import subprocess
+import time
+import urllib2
+
+from boto.auth_handler import AuthHandler
+from boto.auth_handler import NotReadyToAuthenticate
+
+CMD = ['stubby', '--proto2', 'call', 'blade:sso', 'CorpLogin.Exchange']
+
+STUBBY_CMD = """target: {
+  scope: GAIA_USER
+  name: "%s"
+}
+target_credential: {
+  type: OAUTH2_TOKEN
+  oauth2_attributes: {
+    scope: 'https://www.googleapis.com/auth/devstorage.read_only'
+  }
+}"""
+
+COOKIE_LOCATION = os.path.expanduser('~/.devstore_token')
+
+TOKEN_EXPIRY = 300
+
+
+class SSOAuthError(Exception):
+  pass
+
+
+class SSOAuth(AuthHandler):
+  """SSO based auth handler."""
+
+  capability = ['google-oauth2', 's3']
+
+  def __init__(self, path, config, provider):
+    if provider.name == 'google' and self.has_prodaccess():
+      # If we don't have a loas token, then bypass this auth handler.
+      if subprocess.call('loas_check',
+                         stdout=subprocess.PIPE,
+                         stderr=subprocess.PIPE):
+        raise NotReadyToAuthenticate()
+    else:
+      raise NotReadyToAuthenticate()
+    self.token = None
+    self.expire = 0
+
+  def GetAccessToken(self):
+    """Returns a valid devstore access token.
+
+    This will return from an in-memory cache if the token is there already,
+    then try a filesystem cache, and then runs a stubby call if none of the
+    caches have a valid token.
+    """
+    if self.token and self.expire > time.time():
+      return self.token
+
+    # Try to retrieve token from filesystem cache.
+    if os.path.exists(COOKIE_LOCATION):
+      last_modified = os.path.getmtime(COOKIE_LOCATION)
+      if time.time() - last_modified < TOKEN_EXPIRY:
+        with open(COOKIE_LOCATION, 'rb') as f:
+          self.token = f.read()
+        self.expire = last_modified + TOKEN_EXPIRY
+        return self.token
+
+    # If the token is not in either caches, or has expired, then fetch token.
+    username = '%s@google.com' % getpass.getuser()
+    proc = subprocess.Popen(CMD, stdin=subprocess.PIPE, stdout=subprocess.PIPE)
+    out, err = proc.communicate(STUBBY_CMD % username)
+    if proc.returncode:
+      raise SSOAuthError('Stubby returned %d\n%s' % (proc.returncode, err))
+    token_match = re.search(r'oauth2_token: "(.*)"$', out)
+
+    if not token_match:
+      raise SSOAuthError('Oauth2 token not found in %s' % out)
+
+    token = token_match.group(1)
+    self.token = token
+    self.expire = time.time() + TOKEN_EXPIRY
+    with os.fdopen(os.open(COOKIE_LOCATION,
+                           os.O_WRONLY | os.O_CREAT,
+                           0600), 'wb') as f:
+      f.write(token)
+    return token
+
+  def add_auth(self, http_request):
+    http_request.headers['Authorization'] = 'OAuth %s' % self.GetAccessToken()
+
+  @staticmethod
+  def has_prodaccess():
+    for path in os.environ['PATH'].split(os.pathsep):
+      exe_file = os.path.join(path, 'prodaccess')
+      if os.path.exists(exe_file) and os.access(exe_file, os.X_OK):
+        return True
+    return False