| # 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()) |