Add support to tools/ffdb.py to take a screenshot of the FFOS device screen.

1 files changed, 37 insertions, 1 deletions

diff --git a/tools/ffdb.py b/tools/ffdb.py
index 15c986a1..1f1aac29 100755
--- a/tools/ffdb.py
+++ b/tools/ffdb.py
@@ -1,6 +1,6 @@
 #!/usr/bin/env python
 
-import socket, json, sys, uuid, datetime, time, logging, cgi, zipfile, os, tempfile, atexit, subprocess
+import socket, json, sys, uuid, datetime, time, logging, cgi, zipfile, os, tempfile, atexit, subprocess, re, base64
 
 LOG_VERBOSE = False  # Verbose printing enabled with --verbose
 HOST = 'localhost'   # The remote host to connect to the B2G device
@@ -168,6 +168,9 @@ def main():
     log <app> [--clear]: Starts a persistent log listener that reads web console messages from the given application.
                          If --clear is passed, the message log for that application is cleared instead.
     navigate <url>: Opens the given web page in the B2G browser.
+    screenshot [filename.png]: Takes a screenshot of the current contents displayed on the device. If an optional
+                               filename is specified, the screenshot is saved to that file. Otherwise the filename
+                               will be autogenerated.
 
   In the above, whenever a command requires an <app> to be specified, either the human-readable name, 
   localId or manifestURL of the application can be used.'''
@@ -357,6 +360,39 @@ def main():
         log_b2g_message(msg)
     else:
       print 'Application "' + sys.argv[2] + '" is not running!'
+  elif sys.argv[1] == 'screenshot':
+    if len(sys.argv) >= 3:
+      filename = sys.argv[2]
+      if not filename.endswith('.png'):
+        print >> sys.stderr, "Writing screenshots only to .png files are supported!"
+        sys.exit(1)
+    else:
+      filename = time.strftime("screen_%Y%m%d_%H%M%S.png", time.gmtime())
+
+    data_reply = send_b2g_cmd(deviceActorName, 'screenshotToDataURL')
+    data = data_reply['value']
+    data_get_actor = data['actor']
+    data_len = int(data['length'])
+    data_str = data['initial']
+    delim = re.search(",", data_str).start()
+    data_format = data_str[:delim]
+    if data_format != "data:image/png;base64":
+      print >> sys.stderr, "Error: Received screenshot from device in an unexpected format '" + data_format + "'!"
+      sys.exit(1)
+    data = data_str[delim+1:]
+    chunk_size = 65000
+    pos = len(data_str)
+    while pos < data_len:
+      bytes_to_read = min(data_len - pos, chunk_size)
+      data_reply = send_b2g_cmd(data_get_actor, 'substring', { 'start': str(pos), 'end': str(pos + bytes_to_read) })
+      if len(data_reply['substring']) != bytes_to_read:
+        print >> sys.stderr, 'Error! Expected to receive ' + str(bytes_to_read) + ' bytes of image data, but got ' + str(len(data_reply['substring'])) + ' bytes instead!'
+        sys.exit(1)
+      data += data_reply['substring']
+      pos += bytes_to_read
+    binary_data = base64.b64decode(data)
+    open(filename, 'wb').write(binary_data)
+    print "Wrote " + sizeof_fmt(len(binary_data)) + " to file '" + filename + "'."
   else:
     print "Unknown command '" + sys.argv[1] + "'! Pass --help for instructions."