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 mtlib.feedback import FeedbackDownloader
from StringIO import StringIO
import base64
import bz2
import imghdr
import os
import os.path
import re
import subprocess
import sys
import tarfile
import tempfile
import zipfile

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

def Log(source=None, options=None, activity=None, evdev=None):
  """ Load log from feedback url, device ip or local file path.

  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 source:
    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)

  if activity or evdev:
    log = AbstractLog()
    if activity:
      if os.path.exists(activity):
        log.activity = open(activity).read()
      else:
        log.activity = activity
    if evdev:
      if os.path.exists(evdev):
        log.evdev = open(evdev).read()
      else:
        log.evdev = evdev
    return log


class AbstractLog(object):
  """ Common functions for log files

  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 __init__(self):
    self.activity = None
    self.evdev = None
    self.system = None
    self.image = None

  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.

  evdev or system logs are optional.
  """
  def __init__(self, filename, options):
    self.activity = open(filename).read()
    self.evdev = ''
    self.system = ''
    self.image = None
    if options and 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 FeedbackLog(AbstractLog):
  """ Download log from feedback url.

  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
    self.report = None
    self.system = None
    self.activity = None
    self.evdev = None

    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()
      self.report = self.downloader.DownloadSystemLog(id_or_filename)
      if self.report:
        self._ExtractSystemLog()
      if self.system:
        self._ExtractLogFiles()
      if self.try_screenshot:
        self.image = self.downloader.DownloadScreenshot(id_or_filename)

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