| #!/usr/bin/python |
| # 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()) |