Fix docstring
[uzbl-mobile] / examples / data / uzbl / scripts / cookie_daemon.py
index eb6f4af..3546d92 100755 (executable)
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
+'''
+The Python Cookie Daemon
+========================
 
-# Issues:
-#  - There is no easy way of stopping a running daemon.
-#  - Some users experience the broken pipe socket error seemingly randomly.
-#    Currently when a broken pipe error is received the daemon fatally fails.
+This application is a re-write of the original cookies.py script found in
+uzbl's master branch. This script provides more functionality than the original
+cookies.py by adding command line options to specify different cookie jar
+locations, socket locations, verbose output, etc.
 
+Keeping up to date
+==================
 
-# Todo list:
-#  - Use a pid file to make stopping a running daemon easy.
-#  - add {start|stop|restart} command line arguments to make the cookie_daemon
-#    functionally similar to the daemons found in /etc/init.d/ (in gentoo)
-#    or /etc/rc.d/ (in arch).
-#  - Recover from broken pipe socket errors.
-#  - Add option to create a throwaway cookie jar in /tmp and delete it upon
-#    closing the daemon.
+Check the cookie daemon uzbl-wiki page for more information on where to
+find the latest version of the cookie_daemon.py
 
+    http://www.uzbl.org/wiki/cookie_daemon.py
+
+Command line options
+====================
+
+Usage: cookie_daemon.py [options]
+
+Options:
+  -h, --help            show this help message and exit
+  -n, --no-daemon       don't daemonise the process.
+  -v, --verbose         print verbose output.
+  -t SECONDS, --daemon-timeout=SECONDS
+                        shutdown the daemon after x seconds inactivity.
+                        WARNING: Do not use this when launching the cookie
+                        daemon manually.
+  -s SOCKET, --cookie-socket=SOCKET
+                        manually specify the socket location.
+  -j FILE, --cookie-jar=FILE
+                        manually specify the cookie jar location.
+  -m, --memory          store cookies in memory only - do not write to disk
+
+Talking with uzbl
+=================
+
+In order to get uzbl to talk to a running cookie daemon you add the following
+to your uzbl config:
+
+  set cookie_handler = talk_to_socket $XDG_CACHE_HOME/uzbl/cookie_daemon_socket
+
+Or if you prefer using the $HOME variable:
+
+  set cookie_handler = talk_to_socket $HOME/.cache/uzbl/cookie_daemon_socket
+
+Issues
+======
+
+ - There is no easy way of stopping a running daemon.
+
+Todo list
+=========
+
+ - Use a pid file to make stopping a running daemon easy.
+ - add {start|stop|restart} command line arguments to make the cookie_daemon
+   functionally similar to the daemons found in /etc/init.d/ (in gentoo)
+   or /etc/rc.d/ (in arch).
+
+Reporting bugs / getting help
+=============================
+
+The best and fastest way to get hold of the maintainers of the cookie_daemon.py
+is to send them a message in the #uzbl irc channel found on the Freenode IRC
+network (irc.freenode.org).
+'''
 
 import cookielib
 import os
@@ -43,6 +95,7 @@ import select
 import socket
 import time
 import atexit
+from traceback import print_exc
 from signal import signal, SIGTERM
 from optparse import OptionParser
 
@@ -60,31 +113,30 @@ except ImportError:
 
 # Location of the uzbl cache directory.
 if 'XDG_CACHE_HOME' in os.environ.keys() and os.environ['XDG_CACHE_HOME']:
-    cache_dir = os.path.join(os.environ['XDG_CACHE_HOME'], 'uzbl/')
+    CACHE_DIR = os.path.join(os.environ['XDG_CACHE_HOME'], 'uzbl/')
 
 else:
-    cache_dir = os.path.join(os.environ['HOME'], '.cache/uzbl/')
+    CACHE_DIR = os.path.join(os.environ['HOME'], '.cache/uzbl/')
 
 # Location of the uzbl data directory.
 if 'XDG_DATA_HOME' in os.environ.keys() and os.environ['XDG_DATA_HOME']:
-    data_dir = os.path.join(os.environ['XDG_DATA_HOME'], 'uzbl/')
+    DATA_DIR = os.path.join(os.environ['XDG_DATA_HOME'], 'uzbl/')
 
 else:
-    data_dir = os.path.join(os.environ['HOME'], '.local/share/uzbl/')
+    DATA_DIR = os.path.join(os.environ['HOME'], '.local/share/uzbl/')
 
 # Default config
 config = {
 
   # Default cookie jar and daemon socket locations.
-  'cookie_socket': os.path.join(cache_dir, 'cookie_daemon_socket'),
-  'cookie_jar': os.path.join(data_dir, 'cookies.txt'),
+  'cookie_socket': os.path.join(CACHE_DIR, 'cookie_daemon_socket'),
+  'cookie_jar': os.path.join(DATA_DIR, 'cookies.txt'),
 
   # Time out after x seconds of inactivity (set to 0 for never time out).
   # Set to 0 by default until talk_to_socket is doing the spawning.
   'daemon_timeout': 0,
 
-  # Enable/disable daemonizing the process (useful when debugging).
-  # Set to False by default until talk_to_socket is doing the spawning.
+  # Tell process to daemonise
   'daemon_mode': True,
 
   # Set true to print helpful debugging messages to the terminal.
@@ -98,10 +150,12 @@ config = {
 # ============================================================================
 
 
-_scriptname = os.path.basename(sys.argv[0])
+_SCRIPTNAME = os.path.basename(sys.argv[0])
 def echo(msg):
+    '''Prints messages sent only if the verbose flag has been set.'''
+
     if config['verbose']:
-        print "%s: %s" % (_scriptname, msg)
+        print "%s: %s" % (_SCRIPTNAME, msg)
 
 
 def mkbasedir(filepath):
@@ -113,6 +167,73 @@ def mkbasedir(filepath):
         os.makedirs(dirname)
 
 
+def reclaim_socket():
+    '''Check if another process (hopefully a cookie_daemon.py) is listening
+    on the cookie daemon socket. If another process is found to be
+    listening on the socket exit the daemon immediately and leave the
+    socket alone. If the connect fails assume the socket has been abandoned
+    and delete it (to be re-created in the create socket function).'''
+
+    cookie_socket = config['cookie_socket']
+
+    try:
+        sock = socket.socket(socket.AF_UNIX, socket.SOCK_SEQPACKET)
+        sock.connect(cookie_socket)
+        sock.close()
+
+    except socket.error:
+        # Failed to connect to cookie_socket so assume it has been
+        # abandoned by another cookie daemon process.
+        echo("reclaiming abandoned cookie_socket %r." % cookie_socket)
+        if os.path.exists(cookie_socket):
+            os.remove(cookie_socket)
+
+        return
+
+    echo("detected another process listening on %r." % cookie_socket)
+    echo("exiting.")
+    # Use os._exit() to avoid tripping the atexit cleanup function.
+    sys.exit(1)
+
+
+def daemonize():
+    '''Daemonize the process using the Stevens' double-fork magic.'''
+
+    try:
+        if os.fork():
+            os._exit(0)
+
+    except OSError:
+        print_exc()
+        sys.stderr.write("fork #1 failed")
+        sys.exit(1)
+
+    os.chdir('/')
+    os.setsid()
+    os.umask(0)
+
+    try:
+        if os.fork():
+            os._exit(0)
+
+    except OSError:
+        print_exc()
+        sys.stderr.write("fork #2 failed")
+        sys.exit(1)
+
+    sys.stdout.flush()
+    sys.stderr.flush()
+
+    devnull = '/dev/null'
+    stdin = file(devnull, 'r')
+    stdout = file(devnull, 'a+')
+    stderr = file(devnull, 'a+', 0)
+
+    os.dup2(stdin.fileno(), sys.stdin.fileno())
+    os.dup2(stdout.fileno(), sys.stdout.fileno())
+    os.dup2(stderr.fileno(), sys.stderr.fileno())
+
+
 class CookieMonster:
     '''The uzbl cookie daemon class.'''
 
@@ -122,6 +243,7 @@ class CookieMonster:
         self.server_socket = None
         self.jar = None
         self.last_request = time.time()
+        self._running = False
 
 
     def run(self):
@@ -131,12 +253,12 @@ class CookieMonster:
         # exit if another daemon is detected listening on the cookie socket
         # and remove the abandoned socket if there isnt.
         if os.path.exists(config['cookie_socket']):
-            self.reclaim_socket()
+            reclaim_socket()
 
         # Daemonize process.
         if config['daemon_mode']:
             echo("entering daemon mode.")
-            self.daemonize()
+            daemonize()
 
         # Register a function to cleanup on exit.
         atexit.register(self.quit)
@@ -144,111 +266,60 @@ class CookieMonster:
         # Make SIGTERM act orderly.
         signal(SIGTERM, lambda signum, stack_frame: sys.exit(1))
 
-        # Create cookie daemon socket.
-        self.create_socket()
-
         # Create cookie jar object from file.
         self.open_cookie_jar()
 
-        try:
-            # Listen for incoming cookie puts/gets.
-            echo("listening on %r" % config['cookie_socket'])
-            self.listen()
+        # Creating a way to exit nested loops by setting a running flag.
+        self._running = True
 
-        except KeyboardInterrupt:
-            print
+        while self._running:
+            # Create cookie daemon socket.
+            self.create_socket()
 
-        except:
-            # Clean up
-            self.del_socket()
-
-            # Raise exception
-            raise
+            try:
+                # Enter main listen loop.
+                self.listen()
 
+            except KeyboardInterrupt:
+                self._running = False
+                print
 
-    def reclaim_socket(self):
-        '''Check if another process (hopefully a cookie_daemon.py) is listening
-        on the cookie daemon socket. If another process is found to be
-        listening on the socket exit the daemon immediately and leave the
-        socket alone. If the connect fails assume the socket has been abandoned
-        and delete it (to be re-created in the create socket function).'''
+            except socket.error:
+                print_exc()
 
-        cookie_socket = config['cookie_socket']
-
-        try:
-            sock = socket.socket(socket.AF_UNIX, socket.SOCK_SEQPACKET)
-            sock.connect(cookie_socket)
-            sock.close()
+            except:
+                # Clean up
+                self.del_socket()
 
-        except socket.error:
-            # Failed to connect to cookie_socket so assume it has been
-            # abandoned by another cookie daemon process.
-            echo("reclaiming abandoned cookie_socket %r." % cookie_socket)
-            if os.path.exists(cookie_socket):
-                os.remove(cookie_socket)
+                # Raise exception
+                raise
 
-            return
-
-        echo("detected another process listening on %r." % cookie_socket)
-        echo("exiting.")
-        # Use os._exit() to avoid tripping the atexit cleanup function.
-        os._exit(1)
-
-
-    def daemonize(function):
-        '''Daemonize the process using the Stevens' double-fork magic.'''
-
-        try:
-            if os.fork(): os._exit(0)
-
-        except OSError, e:
-            sys.stderr.write("fork #1 failed: %s\n" % e)
-            sys.exit(1)
-
-        os.chdir('/')
-        os.setsid()
-        os.umask(0)
-
-        try:
-            if os.fork(): os._exit(0)
-
-        except OSError, e:
-            sys.stderr.write("fork #2 failed: %s\n" % e)
-            sys.exit(1)
-
-        sys.stdout.flush()
-        sys.stderr.flush()
-
-        devnull = '/dev/null'
-        stdin = file(devnull, 'r')
-        stdout = file(devnull, 'a+')
-        stderr = file(devnull, 'a+', 0)
-
-        os.dup2(stdin.fileno(), sys.stdin.fileno())
-        os.dup2(stdout.fileno(), sys.stdout.fileno())
-        os.dup2(stderr.fileno(), sys.stderr.fileno())
+            # Always delete the socket before calling create again.
+            self.del_socket()
 
 
     def open_cookie_jar(self):
         '''Open the cookie jar.'''
 
         cookie_jar = config['cookie_jar']
-        mkbasedir(cookie_jar)
+        if cookie_jar:
+            mkbasedir(cookie_jar)
 
         # Create cookie jar object from file.
         self.jar = cookielib.MozillaCookieJar(cookie_jar)
 
-        try:
-            # Attempt to load cookies from the cookie jar.
-            self.jar.load(ignore_discard=True)
+        if cookie_jar:
+            try:
+                # Attempt to load cookies from the cookie jar.
+                self.jar.load(ignore_discard=True)
 
-            # Ensure restrictive permissions are set on the cookie jar
-            # to prevent other users on the system from hi-jacking your
-            # authenticated sessions simply by copying your cookie jar.
-            os.chmod(cookie_jar, 0600)
+                # Ensure restrictive permissions are set on the cookie jar
+                # to prevent other users on the system from hi-jacking your
+                # authenticated sessions simply by copying your cookie jar.
+                os.chmod(cookie_jar, 0600)
 
-        except:
-            pass
+            except:
+                pass
 
 
     def create_socket(self):
@@ -258,12 +329,12 @@ class CookieMonster:
         cookie_socket = config['cookie_socket']
         mkbasedir(cookie_socket)
 
-        self.server_socket = socket.socket(socket.AF_UNIX,\
+        self.server_socket = socket.socket(socket.AF_UNIX,
           socket.SOCK_SEQPACKET)
 
         if os.path.exists(cookie_socket):
             # Accounting for super-rare super-fast racetrack condition.
-            self.reclaim_socket()
+            reclaim_socket()
 
         self.server_socket.bind(cookie_socket)
 
@@ -275,13 +346,15 @@ class CookieMonster:
     def listen(self):
         '''Listen for incoming cookie PUT and GET requests.'''
 
-        while True:
+        echo("listening on %r" % config['cookie_socket'])
+
+        while self._running:
             # This line tells the socket how many pending incoming connections
             # to enqueue. I haven't had any broken pipe errors so far while
             # using the non-obvious value of 1 under heavy load conditions.
             self.server_socket.listen(1)
 
-            if bool(select.select([self.server_socket],[],[],1)[0]):
+            if bool(select.select([self.server_socket], [], [], 1)[0]):
                 client_socket, _ = self.server_socket.accept()
                 self.handle_request(client_socket)
                 self.last_request = time.time()
@@ -289,7 +362,8 @@ class CookieMonster:
 
             if config['daemon_timeout']:
                 idle = time.time() - self.last_request
-                if idle > config['daemon_timeout']: break
+                if idle > config['daemon_timeout']:
+                    self._running = False
 
 
     def handle_request(self, client_socket):
@@ -297,14 +371,21 @@ class CookieMonster:
 
         # Receive cookie request from client.
         data = client_socket.recv(8192)
-        if not data: return
+        if not data:
+            return
 
         # Cookie argument list in packet is null separated.
         argv = data.split("\0")
 
+        # Catch the EXIT command sent to kill the daemon.
+        if len(argv) == 1 and argv[0].strip() == "EXIT":
+            self._running = False
+            return None
+
         # Determine whether or not to print cookie data to terminal.
         print_cookie = (config['verbose'] and not config['daemon_mode'])
-        if print_cookie: print ' '.join(argv[:4])
+        if print_cookie:
+            print ' '.join(argv[:4])
 
         action = argv[0]
 
@@ -323,36 +404,40 @@ class CookieMonster:
             if req.has_header('Cookie'):
                 cookie = req.get_header('Cookie')
                 client_socket.send(cookie)
-                if print_cookie: print cookie
+                if print_cookie:
+                    print cookie
 
             else:
                 client_socket.send("\0")
 
         elif action == "PUT":
-            if len(argv) > 3:
-                set_cookie = argv[4]
-                if print_cookie: print set_cookie
+            cookie = argv[4] if len(argv) > 3 else None
+            if print_cookie:
+                print cookie
 
-            else:
-                set_cookie = None
+            self.put_cookie(req, cookie)
 
-            hdr = urllib2.httplib.HTTPMessage(\
-              StringIO.StringIO('Set-Cookie: %s' % set_cookie))
-            res = urllib2.addinfourl(StringIO.StringIO(), hdr,\
-              req.get_full_url())
-            self.jar.extract_cookies(res,req)
-            self.jar.save(ignore_discard=True)
+        if print_cookie:
+            print
 
-        if print_cookie: print
+
+    def put_cookie(self, req, cookie=None):
+        '''Put a cookie in the cookie jar.'''
+
+        hdr = urllib2.httplib.HTTPMessage(\
+          StringIO.StringIO('Set-Cookie: %s' % cookie))
+        res = urllib2.addinfourl(StringIO.StringIO(), hdr,
+          req.get_full_url())
+        self.jar.extract_cookies(res, req)
+        if config['cookie_jar']:
+            self.jar.save(ignore_discard=True)
 
 
-    def quit(self, *args):
+    def quit(self):
         '''Called on exit to make sure all loose ends are tied up.'''
 
-        # Only one loose end so far.
         self.del_socket()
-
-        os._exit(0)
+        sys.exit(0)
 
 
     def del_socket(self):
@@ -360,7 +445,13 @@ class CookieMonster:
         is the daemons pid file equivalent.'''
 
         if self.server_socket:
-            self.server_socket.close()
+            try:
+                self.server_socket.close()
+
+            except:
+                pass
+
+        self.server_socket = None
 
         cookie_socket = config['cookie_socket']
         if os.path.exists(cookie_socket):
@@ -368,47 +459,48 @@ class CookieMonster:
             os.remove(cookie_socket)
 
 
-if __name__ == "__main__":
-
+def main():
+    '''Main function.'''
 
+    # Define command line parameters.
     parser = OptionParser()
-    parser.add_option('-d', '--daemon-mode', dest='daemon_mode',\
-      action='store_true', help="daemonise the cookie handler.")
-
-    parser.add_option('-n', '--no-daemon', dest='no_daemon',\
+    parser.add_option('-n', '--no-daemon', dest='no_daemon',
       action='store_true', help="don't daemonise the process.")
 
-    parser.add_option('-v', '--verbose', dest="verbose",\
+    parser.add_option('-v', '--verbose', dest="verbose",
       action='store_true', help="print verbose output.")
 
-    parser.add_option('-t', '--daemon-timeout', dest='daemon_timeout',\
+    parser.add_option('-t', '--daemon-timeout', dest='daemon_timeout',
       action="store", metavar="SECONDS", help="shutdown the daemon after x "\
       "seconds inactivity. WARNING: Do not use this when launching the "\
       "cookie daemon manually.")
 
-    parser.add_option('-s', '--cookie-socket', dest="cookie_socket",\
+    parser.add_option('-s', '--cookie-socket', dest="cookie_socket",
       metavar="SOCKET", help="manually specify the socket location.")
 
-    parser.add_option('-j', '--cookie-jar', dest='cookie_jar',\
+    parser.add_option('-j', '--cookie-jar', dest='cookie_jar',
       metavar="FILE", help="manually specify the cookie jar location.")
 
+    parser.add_option('-m', '--memory', dest='memory', action='store_true',
+      help="store cookies in memory only - do not write to disk")
+
+    # Parse the command line arguments.
     (options, args) = parser.parse_args()
 
-    if options.daemon_mode and options.no_daemon:
+    if len(args):
         config['verbose'] = True
-        echo("fatal error: conflicting options --daemon-mode & --no-daemon")
+        for arg in args:
+            echo("unknown argument %r" % arg)
+
+        echo("exiting.")
         sys.exit(1)
 
     if options.verbose:
         config['verbose'] = True
         echo("verbose mode on.")
 
-    if options.daemon_mode:
-        echo("daemon mode on.")
-        config['daemon_mode'] = True
-
     if options.no_daemon:
-        echo("daemon mode off")
+        echo("daemon mode off.")
         config['daemon_mode'] = False
 
     if options.cookie_socket:
@@ -419,6 +511,10 @@ if __name__ == "__main__":
         echo("using cookie_jar %r" % options.cookie_jar)
         config['cookie_jar'] = options.cookie_jar
 
+    if options.memory:
+        echo("using memory %r" % options.memory)
+        config['cookie_jar'] = None
+
     if options.daemon_timeout:
         try:
             config['daemon_timeout'] = int(options.daemon_timeout)
@@ -429,7 +525,13 @@ if __name__ == "__main__":
             echo("fatal error: expected int argument for --daemon-timeout")
             sys.exit(1)
 
+    # Expand $VAR's in config keys that relate to paths.
     for key in ['cookie_socket', 'cookie_jar']:
-        config[key] = os.path.expandvars(config[key])
+        if config[key]:
+            config[key] = os.path.expandvars(config[key])
 
     CookieMonster().run()
+
+
+if __name__ == "__main__":
+    main()