# Copyright 2012 Google Inc. All Rights Reserved.
# Author: mrdmnd@ (Matt Redmond)
"""A client to pull data from Bartlett.

Inspired by //depot/google3/experimental/mobile_gwp/database/app_engine_pull.py

The server houses perf.data.gz, board, chrome version for each upload.
This script first authenticates with a proper @google.com account, then
downloads a sample (if it's not already cached) and unzips perf.data

  Authenticate(): Gets login info and returns an auth token
  DownloadSamples(): Download and unzip samples.
  _GetServePage(): Pulls /serve page from the app engine server
  _DownloadSampleFromServer(): Downloads a local compressed copy of a sample
  _UncompressSample(): Decompresses a sample, deleting the compressed version.
"""
import cookielib
import getpass
import gzip
import optparse
import os
import urllib
import urllib2

SERVER_NAME = 'http://chromeoswideprofiling.appspot.com'
APP_NAME = 'chromeoswideprofiling'
DELIMITER = '~'


def Authenticate(server_name):
  """Gets credentials from user and attempts to retrieve auth token.
     TODO: Accept OAuth2 instead of password.
  Args:
    server_name: (string) URL that the app engine code is living on.
  Returns:
    authtoken: (string) The authorization token that can be used
                        to grab other pages.
  """

  if server_name.endswith('/'):
    server_name = server_name.rstrip('/')
  # Grab username and password from user through stdin.
  username = raw_input('Email (must be @google.com account): ')
  password = getpass.getpass('Password: ')
  # Use a cookie to authenticate with GAE.
  cookiejar = cookielib.LWPCookieJar()
  opener = urllib2.build_opener(urllib2.HTTPCookieProcessor(cookiejar))
  urllib2.install_opener(opener)
  # Get an AuthToken from Google accounts service.
  auth_uri = 'https://www.google.com/accounts/ClientLogin'
  authreq_data = urllib.urlencode({'Email': username,
                                   'Passwd': password,
                                   'service': 'ah',
                                   'source': APP_NAME,
                                   'accountType': 'HOSTED_OR_GOOGLE'})
  auth_req = urllib2.Request(auth_uri, data=authreq_data)
  try:
    auth_resp = urllib2.urlopen(auth_req)
  except urllib2.URLError:
    print 'Error logging in to Google accounts service.'
    return None
  body = auth_resp.read()
  # Auth response contains several fields.
  # We care about the part after Auth=
  auth_resp_dict = dict(x.split('=') for x in body.split('\n') if x)
  authtoken = auth_resp_dict['Auth']
  return authtoken


def DownloadSamples(server_name, authtoken, output_dir, start, stop):
  """Download every sample and write unzipped version
     to output directory.
  Args:
    server_name: (string) URL that the app engine code is living on.
    authtoken:   (string) Authorization token.
    output_dir   (string) Filepath to write output to.
    start:       (int)    Index to start downloading from, starting at top.
    stop:        (int)    Index to stop downloading, non-inclusive. -1 for end.
  Returns:
    None
  """

  if server_name.endswith('/'):
    server_name = server_name.rstrip('/')

  serve_page_string = _GetServePage(server_name, authtoken)
  if serve_page_string is None:
    print 'Error getting /serve page.'
    return

  sample_list = serve_page_string.split('</br>')
  print 'Will download:'
  sample_list_subset = sample_list[start:stop]
  for sample in sample_list_subset:
    print sample
  for sample in sample_list_subset:
    assert sample, 'Sample should be valid.'
    sample_info = [s.strip() for s in sample.split(DELIMITER)]
    key = sample_info[0]
    time = sample_info[1]
    time = time.replace(' ', '_')  # No space between date and time.
    # sample_md5 = sample_info[2]
    board = sample_info[3]
    version = sample_info[4]

    # Put a compressed copy of the samples in output directory.
    _DownloadSampleFromServer(server_name, authtoken, key, time, board, version,
                              output_dir)
    _UncompressSample(key, time, board, version, output_dir)


def _BuildFilenameFromParams(key, time, board, version):
  """Return the filename for our sample.
  Args:
    key:  (string) Key indexing our sample in the datastore.
    time: (string) Date that the sample was uploaded.
    board: (string) Board that the sample was taken on.
    version: (string) Version string from /etc/lsb-release
  Returns:
    filename (string)
  """
  filename = DELIMITER.join([key, time, board, version])
  return filename


def _DownloadSampleFromServer(server_name, authtoken, key, time, board, version,
                              output_dir):
  """Downloads sample_$(samplekey).gz to current dir.
  Args:
    server_name: (string) URL that the app engine code is living on.
    authtoken:   (string) Authorization token.
    key:  (string) Key indexing our sample in the datastore
    time: (string) Date that the sample was uploaded.
    board: (string) Board that the sample was taken on.
    version: (string) Version string from /etc/lsb-release
    output_dir:  (string) Filepath to write to output to.
  Returns:
    None
  """
  filename = _BuildFilenameFromParams(key, time, board, version)
  compressed_filename = filename + '.gz'

  if os.path.exists(os.path.join(output_dir, filename)):
    print 'Already downloaded %s, skipping.' % filename
    return

  serv_uri = server_name + '/serve/' + key
  serv_args = {'continue': serv_uri, 'auth': authtoken}
  full_serv_uri = server_name + '/_ah/login?%s' % urllib.urlencode(serv_args)
  serv_req = urllib2.Request(full_serv_uri)
  serv_resp = urllib2.urlopen(serv_req)
  f = open(os.path.join(output_dir, compressed_filename), 'w+')
  f.write(serv_resp.read())
  f.close()


def _UncompressSample(key, time, board, version, output_dir):
  """Uncompresses a given sample.gz file and deletes the compressed version.
  Args:
    key: (string) Sample key to uncompress.
    time: (string) Date that the sample was uploaded.
    board: (string) Board that the sample was taken on.
    version: (string) Version string from /etc/lsb-release
    output_dir: (string) Filepath to find sample key in.
  Returns:
    None
  """
  filename = _BuildFilenameFromParams(key, time, board, version)
  compressed_filename = filename + '.gz'

  if os.path.exists(os.path.join(output_dir, filename)):
    print 'Already decompressed %s, skipping.' % filename
    return

  out_file = open(os.path.join(output_dir, filename), 'wb')
  in_file = gzip.open(os.path.join(output_dir, compressed_filename), 'rb')
  out_file.write(in_file.read())
  in_file.close()
  out_file.close()
  os.remove(os.path.join(output_dir, compressed_filename))


def _DeleteSampleFromServer(server_name, authtoken, key):
  """Opens the /delete page with the specified key
     to delete the sample off the datastore.
    Args:
      server_name: (string) URL that the app engine code is living on.
      authtoken:   (string) Authorization token.
      key:  (string) Key to delete.
    Returns:
      None
  """

  serv_uri = server_name + '/del/' + key
  serv_args = {'continue': serv_uri, 'auth': authtoken}
  full_serv_uri = server_name + '/_ah/login?%s' % urllib.urlencode(serv_args)
  serv_req = urllib2.Request(full_serv_uri)
  urllib2.urlopen(serv_req)


def _GetServePage(server_name, authtoken):
  """Opens the /serve page and lists all keys.
  Args:
    server_name: (string) URL the app engine code is living on.
    authtoken:   (string) Authorization token.
  Returns:
    The text of the /serve page (including HTML tags)
  """

  serv_uri = server_name + '/serve'
  serv_args = {'continue': serv_uri, 'auth': authtoken}
  full_serv_uri = server_name + '/_ah/login?%s' % urllib.urlencode(serv_args)
  serv_req = urllib2.Request(full_serv_uri)
  serv_resp = urllib2.urlopen(serv_req)
  return serv_resp.read()


def main():
  parser = optparse.OptionParser()
  parser.add_option('--output_dir',
                    dest='output_dir',
                    action='store',
                    help='Path to output perf data files.')
  parser.add_option('--start',
                    dest='start_ind',
                    action='store',
                    default=0,
                    help='Start index.')
  parser.add_option('--stop',
                    dest='stop_ind',
                    action='store',
                    default=-1,
                    help='Stop index.')
  options = parser.parse_args()[0]
  if not options.output_dir:
    print 'Must specify --output_dir.'
    return 1
  if not os.path.exists(options.output_dir):
    print 'Specified output_dir does not exist.'
    return 1

  authtoken = Authenticate(SERVER_NAME)
  if not authtoken:
    print 'Could not obtain authtoken, exiting.'
    return 1
  DownloadSamples(SERVER_NAME, authtoken, options.output_dir, options.start_ind,
                  options.stop_ind)
  print 'Downloaded samples.'
  return 0


if __name__ == '__main__':
  exit(main())