1 files changed, 90 insertions, 0 deletions
diff --git a/catapult/telemetry/telemetry/util/screenshot.py b/catapult/telemetry/telemetry/util/screenshot.py
new file mode 100644
index 00000000..bdd30142
--- /dev/null
+++ b/catapult/telemetry/telemetry/util/screenshot.py
@@ -0,0 +1,90 @@
+# Copyright 2016 The Chromium 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 datetime
+import logging
+import os
+import random
+import tempfile
+
+from py_utils import cloud_storage  # pylint: disable=import-error
+from telemetry.util import image_util
+from telemetry.internal.util import file_handle
+
+
+def TryCaptureScreenShot(platform, tab=None):
+  """ If the platform or tab supports screenshot, attempt to take a screenshot
+  of the current browser.
+
+  Args:
+    platform: current platform
+    tab: browser tab if available
+
+  Returns:
+    file handle of the tempoerary file path for the screenshot if
+    present, None otherwise.
+  """
+  try:
+    # TODO(nednguyen): once all platforms support taking screenshot,
+    # remove the tab checking logic and consider moving this to story_runner.
+    # (crbug.com/369490)
+    if platform.CanTakeScreenshot():
+      tf = tempfile.NamedTemporaryFile(delete=False, suffix='.png')
+      tf.close()
+      platform.TakeScreenshot(tf.name)
+      return file_handle.FromTempFile(tf)
+    elif tab and tab.IsAlive() and tab.screenshot_supported:
+      tf = tempfile.NamedTemporaryFile(delete=False, suffix='.png')
+      tf.close()
+      image = tab.Screenshot()
+      image_util.WritePngFile(image, tf.name)
+      return file_handle.FromTempFile(tf)
+    else:
+      logging.warning(
+          'Either tab has crashed or browser does not support taking tab '
+          'screenshot. Skip taking screenshot on failure.')
+      return None
+  except Exception as e:
+    logging.warning('Exception when trying to capture screenshot: %s', repr(e))
+    return None
+
+
+def TryCaptureScreenShotAndUploadToCloudStorage(platform, tab=None):
+  """ If the platform or tab supports screenshot, attempt to take a screenshot
+  of the current browser.  If present it uploads this local path to cloud
+  storage and returns the URL of the cloud storage path.
+
+  Args:
+    platform: current platform
+    tab: browser tab if available
+
+  Returns:
+    url of the cloud storage path if screenshot is present, None otherwise
+  """
+  fh = TryCaptureScreenShot(platform, tab)
+  if fh is not None:
+    return _UploadScreenShotToCloudStorage(fh)
+
+  return None
+
+def _GenerateRemotePath(fh):
+  return ('browser-screenshot_%s-%s%-d%s' % (
+          fh.id,
+          datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S'),
+          random.randint(1, 100000),
+          fh.extension))
+
+def _UploadScreenShotToCloudStorage(fh):
+  """ Upload the given screenshot image to cloud storage and return the
+    cloud storage url if successful.
+  """
+  try:
+    return cloud_storage.Insert(cloud_storage.TELEMETRY_OUTPUT,
+                                _GenerateRemotePath(fh), fh.GetAbsPath())
+  except cloud_storage.CloudStorageError as err:
+    logging.error('Cloud storage error while trying to upload screenshot: %s'
+                  % repr(err))
+    return '<Missing link>'
+  finally:  # Must clean up screenshot file if exists.
+    os.remove(fh.GetAbsPath())