Enable "get html from page" functionality for PyAuto....

chrome/browser/automation/automation_provider.cc

 // Copyright (c) 2010 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 "chrome/browser/automation/automation_provider.h"
 
 #include <set>
 
 #include "app/l10n_util.h"
 #include "app/message_box_flags.h"
 #include "base/callback.h"
+#include "base/file_path.h"
 #include "base/file_version_info.h"
 #include "base/json/json_reader.h"
 #include "base/json/json_writer.h"
 #include "base/keyboard_codes.h"
 #include "base/message_loop.h"
 #include "base/path_service.h"
 #include "base/process_util.h"
 #include "base/stl_util-inl.h"
 #include "base/string_util.h"
 #include "base/thread.h"
@@ skipping 1579 matching lines @@
   if (args->GetInteger(L"width", &width))
     rect.set_width(width);
   if (args->GetInteger(L"height", &height))
     rect.set_height(height);
   browser->window()->SetBounds(rect);
   AutomationMsg_SendJSONRequest::WriteReplyParams(
       reply_message, std::string("{}"), true);
   Send(reply_message);
 }
 
+std::string AutomationProvider::JSONErrorString(std::string err) {

[Paweł Hajdan Jr., 2010/06/24 06:21:48]

Could you move this from being AutomationProvider's member to an anonymous
namespace in this file? This would help to make the header file less cluttered.

+  std::string prefix = "{\"error\": \"";
+  std::string no_quote_err = err;
+  std::string suffix = "\"}";
+
+  // Don't allow input string to break JSON by embedding quotes.

[Paweł Hajdan Jr., 2010/06/24 06:21:48]

How about using functions from base/json/string_escape.h instead of this hack?

+  // Try and make sure the input string won't break json quoting rules.
+  if (no_quote_err.find("\"") != std::string::npos)
+    no_quote_err = "unhappy about embedded quote in error string";
+
+  return prefix + no_quote_err + suffix;
+}
+
 // Sample json input: { "command": "GetBrowserInfo" }
 // Refer to GetBrowserInfo() in chrome/test/pyautolib/pyauto.py for
 // sample json output.
 void AutomationProvider::GetBrowserInfo(Browser* browser,
                                         DictionaryValue* args,
                                         IPC::Message* reply_message) {
   std::string json_return;
   bool reply_return = true;
 
   DictionaryValue* properties = new DictionaryValue;
@@ skipping 168 matching lines @@
     hs->AddPage(gurl, time,
                 id_scope,
                 0,
                 GURL(),
                 PageTransition::LINK,
                 history::RedirectList(),
                 false);
     if (title.length())
       hs->SetPageTitle(gurl, title);
   } else {
-    json_return = "{\"error\": \"bad args (no URL in dict?).\"}";
+    json_return = JSONErrorString("bad args (no URL in dict?)");
     reply_return = false;
   }
 
   AutomationMsg_SendJSONRequest::WriteReplyParams(
       reply_message, json_return, reply_return);
   Send(reply_message);
 }
 
 // Sample json input: { "command": "GetDownloadsInfo" }
 // Refer chrome/test/pyautolib/download_info.py for sample json output.
 void AutomationProvider::GetDownloadsInfo(Browser* browser,
                                           DictionaryValue* args,
                                           IPC::Message* reply_message) {
   std::string json_return;
   bool reply_return = true;
   AutomationProviderDownloadManagerObserver observer;
   std::vector<DownloadItem*> downloads;
   scoped_ptr<DictionaryValue> return_value(new DictionaryValue);
 
   if (!profile_->HasCreatedDownloadManager()) {
-    json_return = "{\"error\": \"no download manager\"}";
+    json_return = JSONErrorString("no download manager");
     reply_return = false;
   } else {
     // Use DownloadManager's GetDownloads() method and not GetCurrentDownloads()
     // since that would be transient; a download might enter and empty out
     // the current download queue too soon to be noticed.
     profile_->GetDownloadManager()->GetDownloads(&observer, L"");
     downloads = observer.Downloads();
   }
 
   std::map<DownloadItem::DownloadState, std::string> state_to_string;
@@ skipping 45 matching lines @@
     Browser* browser,
     DictionaryValue* args,
     IPC::Message* reply_message) {
   std::string json_return;
   bool reply_return = true;
   AutomationProviderDownloadManagerObserver observer;
   std::vector<DownloadItem*> downloads;
 
   // Look for a quick return.
   if (!profile_->HasCreatedDownloadManager()) {
-    json_return = "{\"error\": \"no download manager\"}";
+    json_return = JSONErrorString("no download manager");
     reply_return = false;
   } else {
     profile_->GetDownloadManager()->GetCurrentDownloads(&observer,
                                                         FilePath());
     downloads = observer.Downloads();
     if (downloads.size() == 0) {
       json_return = "{}";
     }
   }
   if (!json_return.empty()) {
@@ skipping 44 matching lines @@
                                   IPC::Message* reply_message) {
   bool reply_return = true;
   std::string json_return = "{}";
   std::wstring path;
   Value* val;
   if (args->GetString(L"path", &path) && args->Get(L"value", &val)) {
     PrefService* pref_service = profile_->GetPrefs();
     const PrefService::Preference* pref =
         pref_service->FindPreference(path.c_str());
     if (!pref) {  // Not a registered pref.
-      json_return = "{\"error\": \"pref not registered.\"}";
+      json_return = JSONErrorString("pref not registered.");
       reply_return = false;
     } else if (pref->IsManaged()) {  // Do not attempt to change a managed pref.
-      json_return = "{\"error\": \"pref is managed. cannot be changed.\"}";
+      json_return = JSONErrorString("pref is managed. cannot be changed.");
       reply_return = false;
     } else {  // Set the pref.
       pref_service->Set(path.c_str(), *val);
     }
   } else {
-    json_return = "{\"error\": \"no pref path or value given.\"}";
+    json_return = JSONErrorString("no pref path or value given.");
     reply_return = false;
   }
 
   AutomationMsg_SendJSONRequest::WriteReplyParams(
       reply_message, json_return, reply_return);
   Send(reply_message);
 }
 
 // Sample json input: { "command": "GetOmniboxInfo" }
 // Refer chrome/test/pyautolib/omnibox_info.py for sample json output.
@@ skipping 41 matching lines @@
 // Sample json input: { "command": "SetOmniboxText",
 //                      "text": "goog" }
 void AutomationProvider::SetOmniboxText(Browser* browser,
                                         DictionaryValue* args,
                                         IPC::Message* reply_message) {
   std::string json_return = "{}";
   bool reply_return = true;
   std::wstring text;
 
   if (!args->GetString(L"text", &text)) {
-    json_return = "{\"error\": \"text missing\"}";
+    json_return = JSONErrorString("text missing");
     reply_return = false;
   } else {
     browser->FocusLocationBar();
     LocationBar* loc_bar = browser->window()->GetLocationBar();
     AutocompleteEditView* edit_view = loc_bar->location_entry();
     edit_view->model()->OnSetFocus(false);
     edit_view->SetUserText(text);
   }
 
   AutomationMsg_SendJSONRequest::WriteReplyParams(
       reply_message, json_return, reply_return);
   Send(reply_message);
 }
 
 // Sample json input: { "command": "OmniboxMovePopupSelection",
 //                      "count": 1 }
 // Negative count implies up, positive implies down. Count values will be
 // capped by the size of the popup list.
 void AutomationProvider::OmniboxMovePopupSelection(
     Browser* browser,
     DictionaryValue* args,
     IPC::Message* reply_message) {
   std::string json_return = "{}";
   bool reply_return = true;
   int count;
 
   if (!args->GetInteger(L"count", &count)) {
-    json_return = "{\"error\": \"count missing\"}";
+    json_return = JSONErrorString("count missing");
     reply_return = false;
   } else {
     LocationBar* loc_bar = browser->window()->GetLocationBar();
     AutocompleteEditModel* model = loc_bar->location_entry()->model();
     model->OnUpOrDownKeyPressed(count);
   }
 
   AutomationMsg_SendJSONRequest::WriteReplyParams(
       reply_message, json_return, reply_return);
   Send(reply_message);
@@ skipping 85 matching lines @@
 // Sample json input:
 //    { "command": "EnablePlugin",
 //      "path": "/Library/Internet Plug-Ins/Flash Player.plugin" }
 void AutomationProvider::EnablePlugin(Browser* browser,
                                       DictionaryValue* args,
                                       IPC::Message* reply_message) {
   std::string json_return = "{}";
   bool reply_return = true;
   FilePath::StringType path;
   if (!args->GetString(L"path", &path)) {
-    json_return = "{\"error\": \"path not specified.\"}";
+    json_return = JSONErrorString("path not specified.");
     reply_return = false;
   } else if (!NPAPI::PluginList::Singleton()->EnablePlugin(FilePath(path))) {
     json_return = StringPrintf("{\"error\": \"Could not enable plugin"
                                " for path %s.\"}", path.c_str());
     reply_return = false;
   }
 
   AutomationMsg_SendJSONRequest::WriteReplyParams(
       reply_message, json_return, reply_return);
   Send(reply_message);
 }
 
 // Sample json input:
 //    { "command": "DisablePlugin",
 //      "path": "/Library/Internet Plug-Ins/Flash Player.plugin" }
 void AutomationProvider::DisablePlugin(Browser* browser,
                                        DictionaryValue* args,
                                        IPC::Message* reply_message) {
   std::string json_return = "{}";
   bool reply_return = true;
   FilePath::StringType path;
   if (!args->GetString(L"path", &path)) {
-    json_return = "{\"error\": \"path not specified.\"}";
+    json_return = JSONErrorString("path not specified.");
     reply_return = false;
   } else if (!NPAPI::PluginList::Singleton()->DisablePlugin(FilePath(path))) {
     json_return = StringPrintf("{\"error\": \"Could not enable plugin"
                                " for path %s.\"}", path.c_str());
     reply_return = false;
   }
 
   AutomationMsg_SendJSONRequest::WriteReplyParams(
       reply_message, json_return, reply_return);
   Send(reply_message);
 }
 
+// Sample json input:
+//    { "command": "SaveTabContents",
+//      "tab_index": 0,
+//      "filename": <a full pathname> }
+// Sample json output:
+//    {}
+void AutomationProvider::SaveTabContents(Browser* browser,
+                                         DictionaryValue* args,
+                                         IPC::Message* reply_message) {
+  std::string json_return;
+  int tab_index = 0;
+  FilePath::StringType filename;
+  FilePath::StringType parent_directory;
+  TabContents* tab_contents = NULL;
+
+  if (!args->GetInteger(L"tab_index", &tab_index) ||
+      !args->GetString(L"filename", &filename)) {
+    json_return = JSONErrorString("tab_index or filename param missing");
+  } else {
+    tab_contents = browser->GetTabContentsAt(tab_index);
+    if (!tab_contents) {
+      json_return = JSONErrorString("no tab at tab_index");
+    }
+  }
+  if (tab_contents) {
+    // We're doing a SAVE_AS_ONLY_HTML so the the directory path isn't
+    // used.  Nevertheless, SavePackage requires it be valid.  Sigh.
+    parent_directory = FilePath(filename).DirName().value();
+    if (!tab_contents->SavePage(FilePath(filename), FilePath(parent_directory),
+                                SavePackage::SAVE_AS_ONLY_HTML)) {
+      json_return = JSONErrorString("Could not initiate SavePage");
+    } else {
+      // The observer will delete itself when done.
+      new SavePackageNotificationObserver(tab_contents->save_package(),
+                                          this, reply_message);
+      return;
+    }
+  }
+
+  // if we get here, error.

[Paweł Hajdan Jr., 2010/06/24 06:21:48]

nit: "if" -> "If".

+  DCHECK(!json_return.empty());
+  AutomationMsg_SendJSONRequest::WriteReplyParams(
+      reply_message, json_return, false);
+  Send(reply_message);
+}
+
 void AutomationProvider::SendJSONRequest(int handle,
                                          std::string json_request,
                                          IPC::Message* reply_message) {
   Browser* browser = NULL;
   std::string error_string;
   scoped_ptr<Value> values;
 
   // Basic error checking.
   if (browser_tracker_->ContainsHandle(handle)) {
     browser = browser_tracker_->GetResource(handle);
@@ skipping 45 matching lines @@
   handler_map["SetPrefs"] = &AutomationProvider::SetPrefs;
 
   handler_map["SetWindowDimensions"] = &AutomationProvider::SetWindowDimensions;
 
   handler_map["GetDownloadsInfo"] = &AutomationProvider::GetDownloadsInfo;
   handler_map["WaitForAllDownloadsToComplete"] =
       &AutomationProvider::WaitForDownloadsToComplete;
 
   handler_map["GetInitialLoadTimes"] = &AutomationProvider::GetInitialLoadTimes;
 
+  handler_map["SaveTabContents"] = &AutomationProvider::SaveTabContents;
+
   if (error_string.empty()) {
     if (handler_map.find(std::string(command)) != handler_map.end()) {
       (this->*handler_map[command])(browser, dict_value, reply_message);
       return;
     } else {
       error_string = "Unknown command. Options: ";
       for (std::map<std::string, JsonHandler>::const_iterator it =
            handler_map.begin(); it != handler_map.end(); ++it) {
         error_string += it->first + ", ";
       }
@@ skipping 1175 matching lines @@
 }
 
 void AutomationProvider::WaitForPopupMenuToOpen(IPC::Message* reply_message) {
   NOTIMPLEMENTED();
 }
 #endif  // !defined(TOOLKIT_VIEWS)
 
 void AutomationProvider::ResetToDefaultTheme() {
   profile_->ClearTheme();
 }