headless: Implement screenshot capturing

headless/app/headless_shell.cc

 // 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.
 
 #include <iostream>
 #include <memory>
 
 #include "base/bind.h"
 #include "base/callback.h"
 #include "base/command_line.h"
@@ skipping 117 matching lines @@
 
     if (base::CommandLine::ForCurrentProcess()->HasSwitch(
             headless::switches::kDumpDom)) {
       FetchDom();
     } else if (base::CommandLine::ForCurrentProcess()->HasSwitch(
                    headless::switches::kRepl)) {
       std::cout
           << "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();
     }
   }
 
   void FetchDom() {
     devtools_client_->GetRuntime()->Evaluate(
         "document.body.innerHTML",
         base::Bind(&HeadlessShell::OnDomFetched, base::Unretained(this)));
   }
@@ skipping 26 matching lines @@
   }
 
   void OnExpressionResult(std::unique_ptr<runtime::EvaluateResult> 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;
     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;

[Sami, 2016/06/02 14:22:19]

Could you make this actually write out the png instead dumping the base64?

[Eric Seckler, 2016/06/03 11:30:40]

Done.

+    Shutdown();
+  }
+
   bool RemoteDebuggingEnabled() const {
     const base::CommandLine& command_line =
         *base::CommandLine::ForCurrentProcess();
     return command_line.HasSwitch(switches::kRemoteDebuggingPort);
   }
 
  private:
   GURL url_;
   HeadlessBrowser* browser_;  // Not owned.
   std::unique_ptr<HeadlessDevToolsClient> devtools_client_;
@@ skipping 49 matching lines @@
 
   if (command_line.HasSwitch(switches::kHostResolverRules)) {
     builder.SetHostResolverRules(
         command_line.GetSwitchValueASCII(switches::kHostResolverRules));
   }
 
   return HeadlessBrowserMain(
       builder.Build(),
       base::Bind(&HeadlessShell::OnStart, base::Unretained(&shell)));
 }