headless: Implement screenshot capturing

headless/app/headless_shell.cc

Index: headless/app/headless_shell.cc
diff --git a/headless/app/headless_shell.cc b/headless/app/headless_shell.cc
index 25c823d541cb196bba2579c004e3722c57d7da3a..d206eb890f8127de691c9c26d3bb6f1122919e6d 100644
--- a/headless/app/headless_shell.cc
+++ b/headless/app/headless_shell.cc
@@ -135,6 +135,9 @@ class HeadlessShell : public HeadlessWebContents::Observer, page::Observer {
           << "Type a Javascript expression to evaluate or \"quit\" to exit."
           << std::endl;
       InputExpression();
+    } else if (base::CommandLine::ForCurrentProcess()->HasSwitch(
+                   headless::switches::kScreenshot)) {
+      CaptureScreenshot();
     } else {
       Shutdown();
     }
@@ -181,6 +184,22 @@ class HeadlessShell : public HeadlessWebContents::Observer, page::Observer {
     InputExpression();
   }
 
+  void CaptureScreenshot() {
+    devtools_client_->GetPage()->GetExperimental()->CaptureScreenshot(
+        page::CaptureScreenshotParams::Builder().Build(),
+        base::Bind(&HeadlessShell::OnScreenshotCaptured,
+                   base::Unretained(this)));
+  }
+
+  void OnScreenshotCaptured(
+      std::unique_ptr<page::CaptureScreenshotResult> result) {
+    std::unique_ptr<base::Value> value = result->Serialize();
+    std::string result_json;
+    base::JSONWriter::Write(*value, &result_json);
+    std::cout << result_json << std::endl;
+    Shutdown();
+  }
+
   bool RemoteDebuggingEnabled() const {
     const base::CommandLine& command_line =
         *base::CommandLine::ForCurrentProcess();