Add --cookies option to pass in a cookie file

webkit2png

@@ -185,8 +185,28 @@ class WebkitLoad (Foundation.NSObject, WebKit.protocols.WebFrameLoadDelegate):
         scriptobject.setValue_forKey_(Webkit2PngScriptBridge.alloc().init(),
                                       'webkit2png')
 
-        webview.mainFrame().loadRequest_(
-            Foundation.NSURLRequest.requestWithURL_(nsurl))
+        req = Foundation.NSMutableURLRequest.requestWithURL_(nsurl)
+
+        if self.options.cookieFile or self.options.cookieData:
+            cookie = ''
+            if self.options.cookieFile:
+                try:
+                    cookie = cookieFromFile(self.options.cookieFile, str(nsurl))
+                except:
+                    print "cookie loading error: ", sys.exc_info()
+                    AppKit.NSApplication.sharedApplication().terminate_(None)
+            if self.options.cookieData:
+                if cookie:
+                    cookie += '; '
+                cookie += '; '.join(self.options.cookieData)
+
+            req.setValue_forHTTPHeaderField_(cookie, 'Cookie')
+
+        if self.options.no_cookies:
+            req.setValue_forHTTPHeaderField_('', 'Cookie')
+
+        webview.mainFrame().loadRequest_(req)
+
         if not webview.mainFrame().provisionalDataSource():
             print " ... not a proper url?"
             self.getURL(webview)
@@ -417,6 +437,18 @@ Examples:
         "--transparent", action="store_true",
         help="render output on a transparent background (requires a web "
         "page with a transparent background)", default=False)
+    group.add_option(
+        "--cookie-file", type="string", dest="cookieFile",
+        help="specify a Netscape cookie file",
+        metavar="FILENAME")
+    group.add_option(
+        "--cookie", action="append", type="string", dest="cookieData",
+        help="specify a cookie name-value pair (multiple --cookie is allowed)",
+        metavar="NAME=VALUE")
+    group.add_option(
+        "--no-cookies", action="store_true",
+        help="explicitly empty the Cookie: header (Safari's cookies are used by default)")
+
     cmdparser.add_option_group(group)
 
     (options, args) = cmdparser.parse_args()
@@ -503,5 +535,41 @@ Examples:
 
     app.run()
 
+# given a cookie filename and a url, return the value for an HTTP cookie
+# header
+def cookieFromFile(filename, url):
+    import cookielib
+    import urllib2
+    import tempfile
+    import os
+
+    # cookielib won't even look at the rest of a cookie file if the first line
+    # doesn't match a particular regex, so we always guarantee the magic_re
+    # will succeed
+    tmp = tempfile.NamedTemporaryFile(delete=False)
+    tmp.write("# Netscape HTTP Cookie File\n")
+    real = open(filename, "r")
+    while 1:
+        line = real.readline()
+        if line:
+            tmp.write(line)
+        else:
+            break
+    tmp.close()
+
+    # we only need this Request() long enough to have cookielib set the Cookie
+    # header for us, we never actually pass it urlopen() or the like
+    req = urllib2.Request(url)
+    jar = cookielib.MozillaCookieJar(tmp.name)
+    jar.load()
+    jar.add_cookie_header(req)
+
+    cookie = req.get_header("Cookie")
+    if not cookie:
+        cookie = ""
+
+    os.unlink(tmp.name)
+    return cookie
+
 if __name__ == '__main__':
     main()