Bug 754220: adding cookie support for Marionette, desktop only; r=mdas

testing/marionette/client/marionette/marionette.py

@@ -495,6 +495,55 @@ class Marionette(object):
             js = f.read()
         return self._send_message('importScript', 'ok', script=js)
 
+    def add_cookie(self, cookie):
+        """
+           Adds a cookie to your current session.
+
+           :Args:
+           - cookie_dict: A dictionary object, with required keys - "name" and "value";
+           optional keys - "path", "domain", "secure", "expiry"
+
+           Usage:
+              driver.add_cookie({'name' : 'foo', 'value' : 'bar'})
+              driver.add_cookie({'name' : 'foo', 'value' : 'bar', 'path' : '/'})
+              driver.add_cookie({'name' : 'foo', 'value' : 'bar', 'path' : '/',
+                                 'secure':True})
+        """
+        return self._send_message('addCookie', 'ok', cookie=cookie)
+
+    def delete_all_cookies(self):
+        """
+            Delete all cookies in the scope of the session.
+            :Usage:
+                driver.delete_all_cookies()
+        """
+        return self._send_message('deleteAllCookies', 'ok')
+
+    def delete_cookie(self, name):
+        """
+            Delete a cookie by its name
+            :Usage:
+                driver.delete_cookie('foo')
+
+        """
+        return self._send_message('deleteCookie', 'ok', name=name);
+
+    def get_cookie(self, name):
+        """
+            Get a single cookie by name. Returns the cookie if found, None if not.
+
+            :Usage:
+                driver.get_cookie('my_cookie')
+        """
+        cookies = self.get_cookies()
+        for cookie in cookies:
+            if cookie['name'] == name:
+                return cookie
+        return None
+
+    def get_cookies(self):
+        return self._send_message("getAllCookies", "value")
+
     @property
     def application_cache(self):
         return ApplicationCache(self)

testing/marionette/client/marionette/tests/unit/test_cookies.py

@@ -0,0 +1,92 @@
+import calendar
+import time
+import random
+from marionette_test import MarionetteTestCase
+
+
+class CookieTest(MarionetteTestCase):
+
+    def setUp(self):
+        MarionetteTestCase.setUp(self)
+        test_url = self.marionette.absolute_url('test.html')
+        self.marionette.navigate(test_url)
+        self.COOKIE_A = {"name": "foo",
+                         "value": "bar",
+                         "path": "/",
+                         "secure": False}
+
+    def tearDown(self):
+        self.marionette.delete_all_cookies()
+        MarionetteTestCase.tearDown(self)
+
+    def testAddCookie(self):
+        self.marionette.add_cookie(self.COOKIE_A)
+        import pdb
+        pdb.set_trace()
+        cookie_returned = str(self.marionette.execute_script("return document.cookie"))
+        self.assertTrue(self.COOKIE_A["name"] in cookie_returned)
+
+    def testAddingACookieThatExpiredInThePast(self):
+        cookie = self.COOKIE_A.copy()
+        cookie["expiry"] = calendar.timegm(time.gmtime()) - 1
+        self.marionette.add_cookie(cookie)
+        cookies = self.marionette.get_cookies()
+        self.assertEquals(0, len(cookies))
+
+    def testDeleteAllCookie(self):
+        self.marionette.add_cookie(self.COOKIE_A)
+        cookie_returned = str(self.marionette.execute_script("return document.cookie"))
+        print cookie_returned
+        self.assertTrue(self.COOKIE_A["name"] in cookie_returned)
+        self.marionette.delete_all_cookies()
+        self.assertFalse(self.marionette.get_cookies())
+
+    def testDeleteCookie(self):
+        self.marionette.add_cookie(self.COOKIE_A)
+        cookie_returned = str(self.marionette.execute_script("return document.cookie"))
+        self.assertTrue(self.COOKIE_A["name"] in cookie_returned)
+        self.marionette.delete_cookie("foo")
+        cookie_returned = str(self.marionette.execute_script("return document.cookie"))
+        self.assertFalse(self.COOKIE_A["name"] in cookie_returned)
+
+    def testShouldGetCookieByName(self): 
+        key = "key_%d" % int(random.random()*10000000)
+        self.marionette.execute_script("document.cookie = arguments[0] + '=set';", [key])
+
+        cookie = self.marionette.get_cookie(key)
+        self.assertEquals("set", cookie["value"])
+
+    def testGetAllCookies(self):
+        key1 = "key_%d" % int(random.random()*10000000)
+        key2 = "key_%d" % int(random.random()*10000000)
+
+        cookies = self.marionette.get_cookies()
+        count = len(cookies)
+
+        one = {"name" :key1,
+               "value": "value"}
+        two = {"name":key2,
+               "value": "value"}
+
+        self.marionette.add_cookie(one)
+        self.marionette.add_cookie(two)
+
+        test_url = self.marionette.absolute_url('test.html')
+        self.marionette.navigate(test_url)
+        cookies = self.marionette.get_cookies()
+        self.assertEquals(count + 2, len(cookies))
+
+    def testShouldNotDeleteCookiesWithASimilarName(self):
+        cookieOneName = "fish"
+        cookie1 = {"name" :cookieOneName,
+                    "value":"cod"}
+        cookie2 = {"name" :cookieOneName + "x",
+                    "value": "earth"}
+        self.marionette.add_cookie(cookie1)
+        self.marionette.add_cookie(cookie2)
+
+        self.marionette.delete_cookie(cookieOneName)
+        cookies = self.marionette.get_cookies()
+
+        self.assertFalse(cookie1["name"] == cookies[0]["name"], msg=str(cookies))
+        self.assertEquals(cookie2["name"] , cookies[0]["name"], msg=str(cookies))

testing/marionette/client/marionette/tests/unit/unit-tests.ini

@@ -52,3 +52,5 @@ b2g = false
 
 [test_appcache.py]
 [test_screenshot.py]
+[test_cookies.py]
+b2g = false

testing/marionette/marionette-actors.js

@@ -1564,6 +1564,40 @@ MarionetteDriverActor.prototype = {
                                           command_id: this.command_id});
   },
 
+  /**
+   * Add a cookie to the document.
+   */
+  addCookie: function MDA_addCookie(aRequest) {
+    this.command_id = this.getCommandId();
+    this.sendAsync("addCookie", {cookie:aRequest.cookie,
+                                 command_id: this.command_id});
+  },
+
+  /**
+   * Get all visible cookies for a document
+   */
+  getAllCookies: function MDA_getAllCookies() {
+    this.command_id = this.getCommandId();
+    this.sendAsync("getAllCookies", {command_id: this.command_id});
+  },
+
+  /**
+   * Delete all cookies that are visible to a document
+   */
+  deleteAllCookies: function MDA_deleteAllCookies() {
+    this.command_id = this.getCommandId();
+    this.sendAsync("deleteAllCookies", {command_id: this.command_id});
+  },
+
+  /**
+   * Delete a cookie by name
+   */
+  deleteCookie: function MDA_deleteCookie(aRequest) {
+    this.command_id = this.getCommandId();
+    this.sendAsync("deleteCookie", {name:aRequest.name,
+                                    command_id: this.command_id});
+  },
+
   /**
    * Closes the Browser Window.
    *
@@ -1916,7 +1950,11 @@ MarionetteDriverActor.prototype.requestTypes = {
   "getAppCacheStatus": MarionetteDriverActor.prototype.getAppCacheStatus,
   "closeWindow": MarionetteDriverActor.prototype.closeWindow,
   "setTestName": MarionetteDriverActor.prototype.setTestName,
-  "screenShot": MarionetteDriverActor.prototype.screenShot
+  "screenShot": MarionetteDriverActor.prototype.screenShot,
+  "addCookie": MarionetteDriverActor.prototype.addCookie,
+  "getAllCookies": MarionetteDriverActor.prototype.getAllCookies,
+  "deleteAllCookies": MarionetteDriverActor.prototype.deleteAllCookies,
+  "deleteCookie": MarionetteDriverActor.prototype.deleteCookie
 };
 
 /**

testing/marionette/marionette-listener.js

@@ -122,6 +122,10 @@ function startListeners() {
   addMessageListenerId("Marionette:setTestName", setTestName);
   addMessageListenerId("Marionette:setState", setState);
   addMessageListenerId("Marionette:screenShot", screenShot);
+  addMessageListenerId("Marionette:addCookie", addCookie);
+  addMessageListenerId("Marionette:getAllCookies", getAllCookies);
+  addMessageListenerId("Marionette:deleteAllCookies", deleteAllCookies);
+  addMessageListenerId("Marionette:deleteCookie", deleteCookie);
 }
 
 /**
@@ -203,6 +207,10 @@ function deleteSession(msg) {
   removeMessageListenerId("Marionette:setTestName", setTestName);
   removeMessageListenerId("Marionette:setState", setState);
   removeMessageListenerId("Marionette:screenShot", screenShot);
+  removeMessageListenerId("Marionette:addCookie", addCookie);
+  removeMessageListenerId("Marionette:getAllCookies", getAllCookies);
+  removeMessageListenerId("Marionette:deleteAllCookies", deleteAllCookies);
+  removeMessageListenerId("Marionette:deleteCookie", deleteCookie);
   this.elementManager.reset();
   // reset frame to the top-most frame
   curWindow = content;
@@ -921,6 +929,140 @@ function switchToFrame(msg) {
     checkTimer.initWithCallback(checkLoad, 100, Ci.nsITimer.TYPE_ONE_SHOT);
   }
 }
+ /**
+  * Add a cookie to the document
+  */
+function addCookie(msg) {
+  cookie = msg.json.cookie;
+
+  if (!cookie.expiry) {
+    var date = new Date();
+    var thePresent = new Date(Date.now());
+    date.setYear(thePresent.getFullYear() + 20);
+    cookie.expiry = date.getTime() / 1000;  // Stored in seconds.
+  }
+
+  if (!cookie.domain) {
+    var location = curWindow.document.location;
+    cookie.domain = location.hostname;
+  }
+  else {
+    var currLocation = curWindow.location;
+    var currDomain = currLocation.host;
+    if (currDomain.indexOf(cookie.domain) == -1) {
+      sendError("You may only set cookies for the current domain", 24, null, msg.json.command_id);
+    }
+  }
+
+  // The cookie's domain may include a port. Which is bad. Remove it
+  // We'll catch ip6 addresses by mistake. Since no-one uses those
+  // this will be okay for now. See Bug 814416
+  if (cookie.domain.match(/:\d+$/)) {
+    cookie.domain = cookie.domain.replace(/:\d+$/, '');
+  }
+
+  var document = curWindow.document;
+  if (!document || !document.contentType.match(/html/i)) {
+    sendError('You may only set cookies on html documents', 25, null, msg.json.command_id);
+  }
+  var cookieManager = Cc['@mozilla.org/cookiemanager;1'].
+                        getService(Ci.nsICookieManager2);
+  cookieManager.add(cookie.domain, cookie.path, cookie.name, cookie.value,
+                   cookie.secure, false, false, cookie.expiry);
+  sendOk(msg.json.command_id);
+}
+
+/**
+ * Get All the cookies for a location
+ */
+function getAllCookies(msg) {
+  var toReturn = [];
+  var cookies = getVisibleCookies(curWindow.location);
+  for (var i = 0; i < cookies.length; i++) {
+    var cookie = cookies[i];
+    var expires = cookie.expires;
+    if (expires == 0) {  // Session cookie, don't return an expiry.
+      expires = null;
+    } else if (expires == 1) { // Date before epoch time, cap to epoch.
+      expires = 0;
+    }
+    toReturn.push({
+      'name': cookie.name,
+      'value': cookie.value,
+      'path': cookie.path,
+      'domain': cookie.host,
+      'secure': cookie.isSecure,
+      'expiry': expires
+    });
+  }
+
+  sendResponse({value: toReturn}, msg.json.command_id);
+}
+
+/**
+ * Delete a cookie by name
+ */
+function deleteCookie(msg) {
+  var toDelete = msg.json.name;
+  var cookieManager = Cc['@mozilla.org/cookiemanager;1'].
+                        getService(Ci.nsICookieManager);
+
+  var cookies = getVisibleCookies(curWindow.location);
+  for (var i = 0; i < cookies.length; i++) {
+    var cookie = cookies[i];
+    if (cookie.name == toDelete) {
+      cookieManager.remove(cookie.host, cookie.name, cookie.path, false);
+    }
+  }
+
+  sendOk(msg.json.command_id);
+}
+
+/**
+ * Delete all the visibile cookies on a page
+ */
+function deleteAllCookies(msg) {
+  let cookieManager = Cc['@mozilla.org/cookiemanager;1'].
+                        getService(Ci.nsICookieManager);
+  let cookies = getVisibleCookies(curWindow.location);
+  for (let i = 0; i < cookies.length; i++) {
+    let cookie = cookies[i];
+    cookieManager.remove(cookie.host, cookie.name, cookie.path, false);
+  }
+  sendOk(msg.json.command_id);
+}
+
+/**
+ * Get all the visible cookies from a location
+ */
+function getVisibleCookies(location) {
+  let results = [];
+  let currentPath = location.pathname;
+  if (!currentPath) currentPath = '/';
+  let isForCurrentPath = function(aPath) {
+    return currentPath.indexOf(aPath) != -1;
+  }
+
+  let cookieManager = Cc['@mozilla.org/cookiemanager;1'].
+                        getService(Ci.nsICookieManager);
+  let enumerator = cookieManager.enumerator;
+  while (enumerator.hasMoreElements()) {
+    let cookie = enumerator.getNext().QueryInterface(Ci['nsICookie']);
+
+    // Take the hostname and progressively shorten
+    let hostname = location.hostname;
+    do {
+      if ((cookie.host == '.' + hostname || cookie.host == hostname)
+          && isForCurrentPath(cookie.path)) {
+          results.push(cookie);
+          break;
+      }
+      hostname = hostname.replace(/^.*?\./, '');
+    } while (hostname.indexOf('.') != -1);
+  }
+
+  return results;
+}
 
 function getAppCacheStatus(msg) {
   sendResponse({ value: curWindow.applicationCache.status },