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.
+

[Vadim Sh., 2013/12/03 23:32:42]

nit: one new line here

[Ryan Tseng, 2013/12/03 23:55:25]

Done.

+
+"""AuthHandler plugin for gsutil's boto to support LOAS based auth."""
+

[Vadim Sh., 2013/12/03 23:32:42]

nit: and here

[Ryan Tseng, 2013/12/03 23:55:25]

Done.

+
+import json
+import getpass
+import os
+import re
+import urllib2

[Vadim Sh., 2013/12/03 23:32:42]

nit: sort imports

[Ryan Tseng, 2013/12/03 23:55:25]

Done.

+import subprocess
+import time
+
+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):
+    # Return from in-memory cache if its there already.
+    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 token_match:

[Vadim Sh., 2013/12/03 23:32:42]

nit: swap if\else sections, that will make larger 'if' body less indented:

if not token_match:
  raise SSOAuthError(...)

token = token.match.group(1)
...

[Ryan Tseng, 2013/12/03 23:55:25]

Done.

+      token = token_match.group(1)
+      self.token = token
+      self.expire = time.time() + TOKEN_EXPIRY
+      with open(COOKIE_LOCATION, 'wb') as f:
+        f.write(token)
+      os.chmod(COOKIE_LOCATION, 0600)  # Mark cookie read-only to user.

[Vadim Sh., 2013/12/03 23:32:42]

Strictly speaking, there's a chance that someone (running under different
account) can open and read a token between f.close() and os.chmod(...) calls.

More bullet proof solution would be to use proper umask or os.open(...) with
correct mode flags, so that file is created with right mode from the start.

Up to you, it's optional. Probably not the largest security hole we have :)

[Ryan Tseng, 2013/12/03 23:55:25]

Done.

+      return token
+    raise SSOAuthError('Oauth2 token not found in %s' % out)
+
+  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