mtlib/log.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.

from gzip import GzipFile
from StringIO import StringIO
import base64
import bz2
import cookielib
import getpass
import imghdr
import os
import os.path
import re
import subprocess
import sys
import tarfile
import tempfile
import urllib2
import zipfile

# path to current script directory
script_dir = os.path.dirname(os.path.realpath(__file__))

# path to the cookies file used for storing the login cookies.
cookies_file = os.path.join(script_dir, "mtedit.cookies")

log_cache_dir = os.path.join(script_dir, "reports")

# path for temporary screenshot
screenshot_filepath = "extension/screenshot.jpg"

def Log(source, options):
  """
    Returns a log object containg activity, evdev and system logs.
    source can be either an ip address to a device from which to
    download, a url pointing to a feedback report to pull or a filename.
  """
  if os.path.exists(source):
    if source.endswith(".zip") or source.endswith(".bz2"):
      return FeedbackLog(source, options)
    return FileLog(source, options)
  else:
    match = re.search("Report/([0-9]+)", source)
    if match:
      return FeedbackLog(match.group(1), options)
    else:
      # todo add regex and error message
      return DeviceLog(source, options)


class AbstractLog(object):
  """
    A log file consists of the activity log, the evdev log, a system log, and
    possibly a screenshot, which can be saved to disk all together.
  """
  def SaveAs(self, filename):
    open(filename, "w").write(self.activity)
    if self.evdev:
      open(filename + ".evdev", "w").write(self.evdev)
    if self.system:
      open(filename + ".system", "w").write(self.system)
    if self.image:
      open(filename + ".jpg", "w").write(self.image)

  def CleanUp(self):
    if os.path.exists(screenshot_filepath):
      os.remove(screenshot_filepath)


class FileLog(AbstractLog):
  """
    Loads log from file. Does not contain evdev or system logs.
  """
  def __init__(self, filename, options):
    self.activity = open(filename).read()
    self.evdev = ""
    self.system = ""
    self.image = None
    if options.evdev:
      self.evdev = open(options.evdev).read()
    elif os.path.exists(filename + ".evdev"):
      self.evdev = open(filename + ".evdev").read()
    if os.path.exists(filename + ".system"):
      self.system = open(filename + ".system").read()
    if os.path.exists(filename + ".jpg"):
      self.image = open(filename + ".jpg").read()
      file(screenshot_filepath, "w").write(self.image)

class FeedbackDownloader():
  def __init__(self):
    # setup cookies file and url opener
    self.cookies = cookielib.MozillaCookieJar(cookies_file)
    if os.path.exists(cookies_file):
      self.cookies.load()
    cookie_processor = urllib2.HTTPCookieProcessor(self.cookies)
    self.opener = urllib2.build_opener(cookie_processor)
    self._Login()

  def _Login(self):
    username = getpass.getuser()

    # check if already logged in
    url = "https://login.corp.google.com/"
    data = self.opener.open(url).read()
    if username in data:
      return

    # ask for credentials
    print "Login to corp.google.com:"
    password = getpass.getpass()
    sys.stdout.write("OTP: ")
    otp = sys.stdin.readline().strip()

    # execute login by posting userdata to login.corp.google.com
    url = "https://login.corp.google.com/login?ssoformat=CORP_SSO"
    data = "u=%s&pw=%s&otp=%s" % (username, password, otp)
    result = self.opener.open(url, data).read()

    # check if the result displays error
    if result.find("error") >= 0:
      print "Login failed"
      exit(-1)

    # login was successful. save cookies to file to be reused later.
    self.cookies.save()

  def DownloadFile(self, url):
    try:
      return self.opener.open(url).read()
    except urllib2.URLError:
      return None

  def DownloadSystemLog(self, id):
    # First download the report.zip file
    logfile = "report-%s-system_logs.zip" % id
    url = ("https://feedback.corp.googleusercontent.com/binarydata/" +
           "%s?id=%s&logIndex=0") % (logfile, id)
    report = self.DownloadFile(url)
    if not report or (report[0:2] != "BZ" and report[0:2] != "PK"):
      print "Report does not include log files at %s" % url
      return None
    return report

  def DownloadScreenshot(self, id):
    url = "https://feedback.corp.googleusercontent.com/screenshot?id=%s" % id
    image = self.DownloadFile(url)
    if not image or imghdr.what("", image) != 'jpeg':
      print "No screenshots available at %s" % url
      return None
    return image

class FeedbackLog(AbstractLog):
  """
    Downloads logs and (possibly) screenshot from a feedback id or file name
  """
  def __init__(self, id_or_filename, options=None, force_latest=None):
    self.image = None
    self.force_latest = force_latest
    self.try_screenshot = options and options.screenshot

    if id_or_filename.endswith(".zip") or id_or_filename.endswith(".bz2"):
      self.report = open(id_or_filename).read()
      screenshot_filename = id_or_filename[:-4] + '.jpg'
      try:
        self.image = open(screenshot_filename).read()
      except IOError:
        # No corresponding screenshot available.
        pass
    else:
      self.downloader = FeedbackDownloader()
      def GetAndUpdateCachedFile(cached_filename, downloader):
        try:
          return open(cached_filename).read()
        except IOError:
          data = downloader(id_or_filename)
          # Cache the data on disk
          if not os.path.exists(log_cache_dir):
            os.makedirs(log_cache_dir)
          file(cached_filename, "w").write(data)
          return data
      self.report = GetAndUpdateCachedFile(
          os.path.join(log_cache_dir, id_or_filename + ".zip"),
          self.downloader.DownloadSystemLog)
      self._ExtractSystemLog()
      self._ExtractLogFiles()
      if self.try_screenshot:
        self.image = GetAndUpdateCachedFile(
            os.path.join(log_cache_dir, id_or_filename + ".jpg"),
            self.downloader.DownloadScreenshot)

    # Only write to screenshot.jpg if we will be viewing the screenshot
    if options and not options.download and self.image:
      file(screenshot_filepath, "w").write(self.image)


  def _GetLatestFile(self, match, tar):
    # find file name containing match with latest timestamp
    names = filter(lambda n: n.find(match) >= 0, tar.getnames())
    names.sort()
    if not names:
      print "Cannot find files named %s in tar file" % match
      print "Tar file contents:", tar.getnames()
      sys.exit(-1)
    return tar.extractfile(names[-1])


  def _ExtractSystemLog(self):
    if self.report[0:2] == "BZ":
      self.system = bz2.decompress(self.report)
    elif self.report[0:2] == "PK":
      io = StringIO(self.report)
      zip = zipfile.ZipFile(io, "r")
      self.system = zip.read(zip.namelist()[0])
      zip.extractall("./")
    else:
      print "Cannot download logs file"
      sys.exit(-1)

  def _ExtractLogFiles(self):
    # Find embedded and uuencoded activity.tar in system log

    def ExtractByInterface(interface):
      assert interface == 'pad' or interface == 'screen'

      log_index = [
        (("hack-33025-touch{0}_activity=\"\"\"\n" +
          "begin-base64 644 touch{0}_activity_log.tar\n").format(interface),
         '"""'),
        (("touch{0}_activity=\"\"\"\n" +
          "begin-base64 644 touch{0}_activity_log.tar\n").format(interface),
         '"""'),
        (("hack-33025-touch{0}_activity=<multiline>\n" +
          "---------- START ----------\n" +
          "begin-base64 644 touch{0}_activity_log.tar\n").format(interface),
         "---------- END ----------"),
      ]

      start_index = end_index = None
      for start, end in log_index:
        if start in self.system:
          start_index = self.system.index(start) + len(start)
          end_index = self.system.index(end, start_index)
          break

      if start_index is None:
        return []

      activity_tar_enc = self.system[start_index:end_index]

      # base64 decode
      activity_tar_data = base64.b64decode(activity_tar_enc)

      # untar
      activity_tar_file = tarfile.open(fileobj=StringIO(activity_tar_data))

      def ExtractPadFiles(name):
        # find matching evdev file
        evdev_name = name[0:name.index('touchpad_activity')]
        evdev_name = evdev_name + "cmt_input_events"

        for file_name in activity_tar_file.getnames():
          if file_name.startswith(evdev_name):
            evdev_name = file_name
            break

        activity_gz = activity_tar_file.extractfile(name)
        evdev_gz = activity_tar_file.extractfile(evdev_name)

        # decompress log files
        return (GzipFile(fileobj=activity_gz).read(),
            GzipFile(fileobj=evdev_gz).read())

      def ExtractScreenFiles(name):
        # Always try for a screenshot with touchscreen view
        self.try_screenshot = True

        evdev = activity_tar_file.extractfile(name)
        return (evdev.read(), None)

      extract_func = {
        'pad': ExtractPadFiles,
        'screen': ExtractScreenFiles,
      }

      # return latest log files, we don't include cmt files
      return [(filename, extract_func[interface])
          for filename in activity_tar_file.getnames()
              if filename.find('cmt') == -1]

    if self.force_latest == "pad":
      logs = ExtractByInterface('pad')
      idx = 0
    elif self.force_latest == "screen":
      logs = ExtractByInterface('screen')
      idx = 0
    else:
      logs = ExtractByInterface('pad') + ExtractByInterface('screen')
      if not logs:
        print ("Cannot find touchpad_activity_log.tar or " +
            "touchscreen_activity_log.tar in systems log file.")
        sys.exit(-1)

      if len(logs) == 1:
        idx = 0
      else:
        while True:
          print "Which log file would you like to use?"
          for i, (name, extract_func) in enumerate(logs):
            if name.startswith('evdev'):
              name = 'touchscreen log - ' + name
            print i, name
          print ">"
          selection = sys.stdin.readline()
          try:
            idx = int(selection)
            if idx < 0 or idx >= len(logs):
              print "Number out of range"
            else:
              break
          except:
            print "Not a number"

    name, extract_func = logs[idx]
    self.activity, self.evdev = extract_func(name)


identity_message = """\
In order to access devices in TPTool, you need to have password-less
auth for chromebooks set up.
Would you like tptool to run the following command for you?
$ %s
Yes/No? (Default: No)"""

class DeviceLog(AbstractLog):
  """
    Downloads logs from a running chromebook via scp.
  """
  def __init__(self, ip, options):
    self.id_path = os.path.expanduser("~/.ssh/identity")
    if not os.path.exists(self.id_path):
      self._SetupIdentity()

    if options.new:
      self._GenerateLogs(ip)

    self.activity = self._Download(ip, "touchpad_activity_log.txt")
    self.evdev = self._Download(ip, "cmt_input_events.dat")
    self.system = ""

  def _GenerateLogs(self, ip):
    cmd = ("ssh -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no " +
           "-q root@%s /opt/google/touchpad/tpcontrol log") % ip
    process = subprocess.Popen(cmd.split())
    process.wait()

  def _Download(self, ip, filename):
    temp = tempfile.NamedTemporaryFile("r")
    cmd = ("scp -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no " +
           "-q root@%s:/var/log/%s %s")
    cmd = cmd % (ip, filename, temp.name)
    process = subprocess.Popen(cmd.split(), stdout=subprocess.PIPE)
    process.wait()
    temp.seek(0)
    return temp.read()

  def _SetupIdentity(self):
    source = os.path.expanduser("~/trunk/chromite/ssh_keys/testing_rsa")
    command = "cp %s %s && chmod 600 %s" % (source, self.id_path, self.id_path)
    print identity_message % command
    response = sys.stdin.readline().strip()
    if response in ("Yes", "yes", "Y", "y"):
      print command
      os.system(command)