mtlib/feedback.py

# Copyright (c) 2013 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import cookielib
import datetime
import time
import getpass
import imghdr
import os
import os.path
import re
import subprocess
import sys
import urllib
import urllib2
import StringIO
import multiprocessing
import random
import errno

# path to current script directory
script_dir = os.path.dirname(os.path.realpath(__file__))
cache_dir = os.path.realpath(os.path.join(script_dir, "..", "cache"))
if not os.path.exists(cache_dir):
  os.mkdir(cache_dir)

# path to the cookies file used for storing the login cookies.
rpc_cred_file = os.path.join(cache_dir, "rpc_cred")

# path to folder where downloaded reports are cached
log_cache_dir = os.path.join(cache_dir, "reports")
if not os.path.exists(log_cache_dir):
  os.mkdir(log_cache_dir)

class FeedbackDownloader():
  # Stubby command to download a feedback log:
  #   field_id - determines what data to download:
  #              1: feedback log ID #
  #              7: screenshot
  #              8: system log
  #   resource_id - feedback log ID #
  STUBBY_FILE_CMD = 'stubby --rpc_creds_file=' + rpc_cred_file + \
                    ' call blade:feedback-export-api-prod ' \
                    'ReportService.Get ' \
                    '\'result_mask <field <id:{field_id}>>, ' \
                    'resource_id: "{resource_id}"\''
  # Stubby command to download a list of feedback log IDs:
  #   product_id - 208 refers to ChromeOS
  #   submission_time_[start|end]_time_ms - feedback submission time range
  #   max_results - number of IDs to download
  #   token - page token
  STUBBY_LIST_CMD = 'stubby --rpc_creds_file=' + rpc_cred_file + \
                    ' call blade:feedback-export-api-prod ' \
                    'ReportService.List --proto2 ' \
                    '\'result_mask <field <id:1>>, ' \
                    'product_id: "208", ' \
                    'submission_time_start_ms: 0, ' \
                    'submission_time_end_ms: {end_time}, ' \
                    'page_selection {{ max_results: {max_results} ' \
                    'token: "{token}" }}\''
  # Refer to go/touch-feedback-download for information on getting permission
  # to download feedback logs.
  GAIA_CMD = '/google/data/ro/projects/gaiamint/bin/get_mint --type=loas ' \
             '--text --scopes=40700 --endusercreds > ' + rpc_cred_file
  SYSTEM_LOG_FIELD = 7
  SCREENSHOT_FIELD = 8
  sleep_sec = 4.0
  auth_lock = multiprocessing.Lock()

  def __init__(self, force_authenticate=False):
    if force_authenticate:
      self._Authenticate()

  def _OctetStreamToBinary(self, octetStream):
    """ The zip files are returned in an octet-stream format that must
    be decoded back into a binary.  This function scans through the stream
    and unescapes the special characters
    """
    binary = ''
    i = 0
    while i < len(octetStream):
      if ord(octetStream[i]) is ord('\\'):
        if re.match('\d\d\d', octetStream[i + 1:i + 4]):
          binary += chr(int(octetStream[i + 1:i + 4], 8))
          i += 4
        else:
          binary += octetStream[i:i + 2].decode("string-escape")
          i += 2
      else:
        binary += octetStream[i]
        i += 1
    return binary

  def _AuthenticateWithLock(self):
    is_owner = self.auth_lock.acquire(False)
    if is_owner:
      self._Authenticate()
      self.auth_lock.release()
    else:
      self.auth_lock.acquire()
      self.auth_lock.release()

  def _StubbyCall(self, cmd):
    if not os.path.exists(rpc_cred_file):
      try:
        os.mkdir(cache_dir)
      except OSError as ex:
        if ex.errno != errno.EEXIST:
          raise
        pass
      self._AuthenticateWithLock()
    while True:
      process = subprocess.Popen(cmd, shell=True,
                                 stdout=subprocess.PIPE,
                                 stderr=subprocess.PIPE)

      output, errors = process.communicate()
      errorcode = process.returncode

      if ('AUTH_FAIL' in errors or
          'CredentialsExpiredException' in output or
          'FORBIDDEN' in output):
        self._AuthenticateWithLock()
        continue
      elif ('exceeds rate limit' in errors or
            'WaitUntilNonEmpty' in errors):
        self._AuthenticateWithLock()

        sleep_time = self.sleep_sec + random.uniform(0.0, self.sleep_sec * 2)
        time.sleep(sleep_time)
        continue
      elif errorcode != 0:
        print errors
        print 'default error'
        print "An error (%d) occurred while downloading" % errorcode
        sys.exit(errorcode)
      return output

  def _DownloadAttachedFile(self, id, field):
    cmd = FeedbackDownloader.STUBBY_FILE_CMD.format(
        field_id=field, resource_id=id)
    output = self._StubbyCall(cmd)
    return output

  def _Authenticate(self):
    cmd = FeedbackDownloader.GAIA_CMD
    process = subprocess.Popen(cmd, shell=True,
                               stdout=subprocess.PIPE,
                               stderr=subprocess.PIPE)
    print "Authenticating..."
    output, errors = process.communicate()
    errorcode = process.returncode

    if errorcode != 0:
      print errors
      print "An error (%d) occurred while authenticating" % errorcode
      if errorcode == 126:
        print "You may need to run prodaccess"
        sys.exit(errorcode)
      return None
    print "Done Authenticating"

  def DownloadIDs(self, num, end_time=None, page_token=''):
    if not end_time:
      dt = datetime.datetime.utcnow() - datetime.datetime(1970, 1, 1)
      end_time = (((dt.days * 24 * 60 * 60 + dt.seconds) * 1000) +
                  (dt.microseconds / 10))

    cmd = FeedbackDownloader.STUBBY_LIST_CMD.format(
        end_time=end_time, max_results=num, token=page_token)
    output = self._StubbyCall(cmd)

    page_token = filter(lambda x: 'next_page_token' in x, output.split('\n'))[0]
    page_token = page_token.split(" ")[-1][1:-1]
    ids = filter(lambda x: 'id' in x, output.split('\n'))
    ids = [x[7:-1] for x in ids]
    return page_token, ids

  def DownloadSystemLog(self, id):
    report = self._DownloadAttachedFile(id,
                                        FeedbackDownloader.SYSTEM_LOG_FIELD)
    sleep_time = self.sleep_sec + random.uniform(0.0, self.sleep_sec * 2)
    time.sleep(sleep_time)
    data_line = None
    system_log = None
    for count, line in enumerate(StringIO.StringIO(report)):
      if 'name: "system_logs.zip"' in line:
        data_line = count + 2
      elif data_line and data_line == count:
        system_log = re.search('data: "(.*)"\s*', line).group(1)

    if not system_log or (system_log[0:2] != "BZ" and system_log[0:2] != "PK"):
      print "Report " + id + " does not seem to include include log files..."
      return None

    return self._OctetStreamToBinary(system_log)

  def DownloadScreenshot(self, id):
    print "Downloading screenshot from %s..." % id
    report = self._DownloadAttachedFile(id,
                                        FeedbackDownloader.SCREENSHOT_FIELD)
    data_line = None
    screenshot = None
    for count, line in enumerate(StringIO.StringIO(report)):
      if 'screenshot <' in line:
        data_line = count + 2
      elif data_line and data_line == count:
        screenshot = re.search('content: "(.*)"\s*', line).group(1)

    if not screenshot:
      print "Report does not seem to include include a screenshot..."
      return None

    return self._OctetStreamToBinary(screenshot)