diff --git a/webkit2png b/webkit2png index a870ba3..cda8cea 100755 --- a/webkit2png +++ b/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()