src/third_party/chromite/scripts/check_gdata_token.py

   1 # Copyright (c) 2011 The Chromium OS Authors. All rights reserved.
   2 # Use of this source code is governed by a BSD-style license that can be
   3 # found in the LICENSE file.
   4
   5 """Validate or replace the standard gdata authorization token."""
   6
   7 from __future__ import print_function
   8
   9 import filecmp
  10 import optparse
  11 import os
  12 import shutil
  13
  14 from chromite.cbuildbot import constants
  15 from chromite.lib import cros_build_lib as build_lib
  16 from chromite.lib import operation
  17
  18 MODULE = os.path.splitext(os.path.basename(__file__))[0]
  19 oper = operation.Operation(MODULE)
  20
  21 TOKEN_FILE = os.path.join(os.environ['HOME'], '.gdata_token')
  22 CRED_FILE = os.path.join(os.environ['HOME'], '.gdata_cred.txt')
  23
  24
  25 def _ChrootPathToExternalPath(path):
  26   """Translate |path| inside chroot to external path to same location."""
  27   if path:
  28     return os.path.join(constants.SOURCE_ROOT,
  29                         constants.DEFAULT_CHROOT_DIR,
  30                         path.lstrip('/'))
  31
  32   return None
  33
  34
  35 class OutsideChroot(object):
  36   """Class for managing functionality when run outside chroot."""
  37
  38   def __init__(self, args):
  39     self.args = args
  40
  41   def Run(self):
  42     """Re-start |args| inside chroot and copy out auth file."""
  43
  44     # Note that enter_chroot (cros_sdk) will automatically copy both
  45     # the token file and the cred file into the chroot, so no need
  46     # to do that here.
  47
  48     # Rerun the same command that launched this run inside the chroot.
  49     cmd = [MODULE] + self.args
  50     result = build_lib.RunCommand(cmd, enter_chroot=True,
  51                                   print_cmd=False, error_code_ok=True)
  52     if result.returncode != 0:
  53       oper.Die('Token validation failed, exit code was %r.' %
  54                result.returncode)
  55
  56     # Copy the token file back from chroot if different.
  57     chroot_token_file = _ChrootPathToExternalPath(TOKEN_FILE)
  58     if not os.path.exists(chroot_token_file):
  59       oper.Die('No token file generated inside chroot.')
  60     elif (not os.path.exists(TOKEN_FILE) or not
  61           filecmp.cmp(TOKEN_FILE, chroot_token_file)):
  62       oper.Notice('Copying new token file from chroot to %r' % TOKEN_FILE)
  63       shutil.copy2(chroot_token_file, TOKEN_FILE)
  64     else:
  65       oper.Notice('No change in token file.')
  66
  67
  68 class InsideChroot(object):
  69   """Class for managing functionality when run inside chroot.
  70
  71   Note that some additional imports happen within code in this class
  72   because those imports are only available inside the chroot.
  73   """
  74
  75   def __init__(self):
  76     self.creds = None     # gdata_lib.Creds object.
  77     self.gd_client = None # For interacting with Google Docs.
  78     self.it_client = None # For interacting with Issue Tracker.
  79
  80   def _LoadTokenFile(self):
  81     """Load existing auth token file."""
  82     if not os.path.exists(TOKEN_FILE):
  83       oper.Warning('No current token file at %r.' % TOKEN_FILE)
  84       return False
  85
  86     # Load token file, if it exists.
  87     self.creds.LoadAuthToken(TOKEN_FILE)
  88     return True
  89
  90   def _SaveTokenFile(self):
  91     """Save to auth toke file if anything changed."""
  92     self.creds.StoreAuthTokenIfNeeded(TOKEN_FILE)
  93
  94   def _ValidateDocsToken(self):
  95     """Validate the existing Docs token."""
  96     # pylint: disable=W0404
  97     import gdata.service
  98
  99     if not self.creds.docs_auth_token:
 100       return False
 101
 102     oper.Notice('Attempting to log into Docs using auth token.')
 103     self.gd_client.source = 'Package Status'
 104     self.gd_client.SetClientLoginToken(self.creds.docs_auth_token)
 105
 106     try:
 107       # Try to access generic spreadsheets feed, which will check access.
 108       self.gd_client.GetSpreadsheetsFeed()
 109
 110       # Token accepted.  We're done here.
 111       oper.Notice('Docs token validated.')
 112       return True
 113     except gdata.service.RequestError as ex:
 114       reason = ex[0]['reason']
 115       if reason == 'Token expired':
 116         return False
 117
 118       raise
 119
 120   def _GenerateDocsToken(self):
 121     """Generate a new Docs token from credentials."""
 122     # pylint: disable=W0404
 123     import gdata.service
 124
 125     oper.Warning('Docs token not valid.  Will try to generate a new one.')
 126     self.creds.LoadCreds(CRED_FILE)
 127     self.gd_client.email = self.creds.user
 128     self.gd_client.password = self.creds.password
 129
 130     try:
 131       self.gd_client.ProgrammaticLogin()
 132       self.creds.SetDocsAuthToken(self.gd_client.GetClientLoginToken())
 133
 134       oper.Notice('New Docs token generated.')
 135       return True
 136     except gdata.service.BadAuthentication:
 137       oper.Error('Credentials from %r not accepted.'
 138                  '  Unable to generate new Docs token.' % CRED_FILE)
 139       return False
 140
 141   def _ValidateTrackerToken(self):
 142     """Validate the existing Tracker token."""
 143     # pylint: disable=W0404
 144     import gdata.gauth
 145     import gdata.projecthosting.client
 146
 147     if not self.creds.tracker_auth_token:
 148       return False
 149
 150     oper.Notice('Attempting to log into Tracker using auth token.')
 151     self.it_client.source = 'Package Status'
 152     self.it_client.auth_token = gdata.gauth.ClientLoginToken(
 153       self.creds.tracker_auth_token)
 154
 155     try:
 156       # Try to access Tracker Issue #1, which will check access.
 157       query = gdata.projecthosting.client.Query(issue_id='1')
 158       self.it_client.get_issues('chromium-os', query=query)
 159
 160       # Token accepted.  We're done here.
 161       oper.Notice('Tracker token validated.')
 162       return True
 163     except gdata.client.Error:
 164       # Exception is gdata.client.Unauthorized in the case of bad token, but
 165       # I do not know what the error is for an expired token so I do not
 166       # want to limit the catching here.  All the errors for gdata.client
 167       # functionality extend gdata.client.Error (I do not see one that is
 168       # obviously about an expired token).
 169       return False
 170
 171   def _GenerateTrackerToken(self):
 172     """Generate a new Tracker token from credentials."""
 173     # pylint: disable=W0404
 174     import gdata.client
 175
 176     oper.Warning('Tracker token not valid.  Will try to generate a new one.')
 177     self.creds.LoadCreds(CRED_FILE)
 178
 179     try:
 180       self.it_client.ClientLogin(self.creds.user, self.creds.password,
 181                                  source='Package Status', service='code',
 182                                  account_type='GOOGLE')
 183       self.creds.SetTrackerAuthToken(self.it_client.auth_token.token_string)
 184
 185       oper.Notice('New Tracker token generated.')
 186       return True
 187     except gdata.client.BadAuthentication:
 188       oper.Error('Credentials from %r not accepted.'
 189                  '  Unable to generate new Tracker token.' % CRED_FILE)
 190       return False
 191
 192   def Run(self):
 193     """Validate existing auth token or generate new one from credentials."""
 194     # pylint: disable=W0404
 195     import chromite.lib.gdata_lib as gdata_lib
 196     import gdata.spreadsheet.service
 197
 198     self.creds = gdata_lib.Creds()
 199     self.gd_client = gdata.spreadsheet.service.SpreadsheetsService()
 200     self.it_client = gdata.projecthosting.client.ProjectHostingClient()
 201
 202     self._LoadTokenFile()
 203
 204     if not self._ValidateTrackerToken():
 205       if not self._GenerateTrackerToken():
 206         oper.Die('Failed to validate or generate Tracker token.')
 207
 208     if not self._ValidateDocsToken():
 209       if not self._GenerateDocsToken():
 210         oper.Die('Failed to validate or generate Docs token.')
 211
 212     self._SaveTokenFile()
 213
 214
 215 def _CreateParser():
 216   usage = 'Usage: %prog'
 217   epilog = ('\n'
 218             'Run outside of chroot to validate the gdata '
 219             'token file at %r or update it if it has expired.\n'
 220             'To update the token file there must be a valid '
 221             'credentials file at %r.\n'
 222             'If run inside chroot the updated token file is '
 223             'still valid but will not be preserved if chroot\n'
 224             'is deleted.\n' %
 225             (TOKEN_FILE, CRED_FILE))
 226
 227   return optparse.OptionParser(usage=usage, epilog=epilog)
 228
 229
 230 def main(argv):
 231   """Main function."""
 232   # Create a copy of args just to be safe.
 233   argv = list(argv)
 234
 235   # No actual options used, but --help is still supported.
 236   parser = _CreateParser()
 237   (_options, args) = parser.parse_args(argv)
 238
 239   if args:
 240     parser.print_help()
 241     oper.Die('No arguments allowed.')
 242
 243   if build_lib.IsInsideChroot():
 244     InsideChroot().Run()
 245   else:
 246     OutsideChroot(args).Run()