vcs.maemo.org Git - uzbl-mobile/blob - examples/data/uzbl/scripts/cookie_daemon.py

   1 #!/usr/bin/env python
   2
   3 # Uzbl tabbing wrapper using a fifo socket interface
   4 # Copyright (c) 2009, Tom Adams <tom@holizz.com>
   5 # Copyright (c) 2009, Dieter Plaetinck <dieter AT plaetinck.be>
   6 # Copyright (c) 2009, Mason Larobina <mason.larobina@gmail.com>
   7 #
   8 # This program is free software: you can redistribute it and/or modify
   9 # it under the terms of the GNU General Public License as published by
  10 # the Free Software Foundation, either version 3 of the License, or
  11 # (at your option) any later version.
  12 #
  13 # This program is distributed in the hope that it will be useful,
  14 # but WITHOUT ANY WARRANTY; without even the implied warranty of
  15 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  16 # GNU General Public License for more details.
  17 #
  18 # You should have received a copy of the GNU General Public License
  19 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
  20 #
  21 # For configuration and cookie daemon usage examples check out the the
  22 # cookie daemon wiki page at http://www.uzbl.org/wiki/cookie_daemon.py
  23 #
  24 # Issues:
  25 #  - There is no easy way of stopping a running daemon.
  26 #
  27 # Todo list:
  28 #  - Use a pid file to make stopping a running daemon easy.
  29 #  - add {start|stop|restart} command line arguments to make the cookie_daemon
  30 #    functionally similar to the daemons found in /etc/init.d/ (in gentoo)
  31 #    or /etc/rc.d/ (in arch).
  32 #  - Add option to create a throwaway cookie jar in /tmp and delete it upon
  33 #    closing the daemon.
  34
  35
  36 import cookielib
  37 import os
  38 import sys
  39 import urllib2
  40 import select
  41 import socket
  42 import time
  43 import atexit
  44 from traceback import print_exc
  45 from signal import signal, SIGTERM
  46 from optparse import OptionParser
  47
  48 try:
  49     import cStringIO as StringIO
  50
  51 except ImportError:
  52     import StringIO
  53
  54
  55 # ============================================================================
  56 # ::: Default configuration section ::::::::::::::::::::::::::::::::::::::::::
  57 # ============================================================================
  58
  59
  60 # Location of the uzbl cache directory.
  61 if 'XDG_CACHE_HOME' in os.environ.keys() and os.environ['XDG_CACHE_HOME']:
  62     cache_dir = os.path.join(os.environ['XDG_CACHE_HOME'], 'uzbl/')
  63
  64 else:
  65     cache_dir = os.path.join(os.environ['HOME'], '.cache/uzbl/')
  66
  67 # Location of the uzbl data directory.
  68 if 'XDG_DATA_HOME' in os.environ.keys() and os.environ['XDG_DATA_HOME']:
  69     data_dir = os.path.join(os.environ['XDG_DATA_HOME'], 'uzbl/')
  70
  71 else:
  72     data_dir = os.path.join(os.environ['HOME'], '.local/share/uzbl/')
  73
  74 # Default config
  75 config = {
  76
  77   # Default cookie jar and daemon socket locations.
  78   'cookie_socket': os.path.join(cache_dir, 'cookie_daemon_socket'),
  79   'cookie_jar': os.path.join(data_dir, 'cookies.txt'),
  80
  81   # Time out after x seconds of inactivity (set to 0 for never time out).
  82   # Set to 0 by default until talk_to_socket is doing the spawning.
  83   'daemon_timeout': 0,
  84
  85   # Tell process to daemonise
  86   'daemon_mode': True,
  87
  88   # Set true to print helpful debugging messages to the terminal.
  89   'verbose': False,
  90
  91 } # End of config dictionary.
  92
  93
  94 # ============================================================================
  95 # ::: End of configuration section :::::::::::::::::::::::::::::::::::::::::::
  96 # ============================================================================
  97
  98
  99 _scriptname = os.path.basename(sys.argv[0])
 100 def echo(msg):
 101     if config['verbose']:
 102         print "%s: %s" % (_scriptname, msg)
 103
 104
 105 def mkbasedir(filepath):
 106     '''Create base directory of filepath if it doesn't exist.'''
 107
 108     dirname = os.path.dirname(filepath)
 109     if not os.path.exists(dirname):
 110         echo("creating dirs: %r" % dirname)
 111         os.makedirs(dirname)
 112
 113
 114 class CookieMonster:
 115     '''The uzbl cookie daemon class.'''
 116
 117     def __init__(self):
 118         '''Initialise class variables.'''
 119
 120         self.server_socket = None
 121         self.jar = None
 122         self.last_request = time.time()
 123         self._running = False
 124
 125
 126     def run(self):
 127         '''Start the daemon.'''
 128
 129         # Check if another daemon is running. The reclaim_socket function will
 130         # exit if another daemon is detected listening on the cookie socket
 131         # and remove the abandoned socket if there isnt.
 132         if os.path.exists(config['cookie_socket']):
 133             self.reclaim_socket()
 134
 135         # Daemonize process.
 136         if config['daemon_mode']:
 137             echo("entering daemon mode.")
 138             self.daemonize()
 139
 140         # Register a function to cleanup on exit.
 141         atexit.register(self.quit)
 142
 143         # Make SIGTERM act orderly.
 144         signal(SIGTERM, lambda signum, stack_frame: sys.exit(1))
 145
 146         # Create cookie jar object from file.
 147         self.open_cookie_jar()
 148
 149         # Creating a way to exit nested loops by setting a running flag.
 150         self._running = True
 151
 152         while self._running:
 153             # Create cookie daemon socket.
 154             self.create_socket()
 155
 156             try:
 157                 # Enter main listen loop.
 158                 self.listen()
 159
 160             except KeyboardInterrupt:
 161                 self._running = False
 162                 print
 163
 164             except socket.error:
 165                 print_exc()
 166
 167             except:
 168                 # Clean up
 169                 self.del_socket()
 170
 171                 # Raise exception
 172                 raise
 173
 174             # Always delete the socket before calling create again.
 175             self.del_socket()
 176
 177
 178     def reclaim_socket(self):
 179         '''Check if another process (hopefully a cookie_daemon.py) is listening
 180         on the cookie daemon socket. If another process is found to be
 181         listening on the socket exit the daemon immediately and leave the
 182         socket alone. If the connect fails assume the socket has been abandoned
 183         and delete it (to be re-created in the create socket function).'''
 184
 185         cookie_socket = config['cookie_socket']
 186
 187         try:
 188             sock = socket.socket(socket.AF_UNIX, socket.SOCK_SEQPACKET)
 189             sock.connect(cookie_socket)
 190             sock.close()
 191
 192         except socket.error:
 193             # Failed to connect to cookie_socket so assume it has been
 194             # abandoned by another cookie daemon process.
 195             echo("reclaiming abandoned cookie_socket %r." % cookie_socket)
 196             if os.path.exists(cookie_socket):
 197                 os.remove(cookie_socket)
 198
 199             return
 200
 201         echo("detected another process listening on %r." % cookie_socket)
 202         echo("exiting.")
 203         # Use os._exit() to avoid tripping the atexit cleanup function.
 204         os._exit(1)
 205
 206
 207     def daemonize(function):
 208         '''Daemonize the process using the Stevens' double-fork magic.'''
 209
 210         try:
 211             if os.fork(): os._exit(0)
 212
 213         except OSError, e:
 214             sys.stderr.write("fork #1 failed: %s\n" % e)
 215             sys.exit(1)
 216
 217         os.chdir('/')
 218         os.setsid()
 219         os.umask(0)
 220
 221         try:
 222             if os.fork(): os._exit(0)
 223
 224         except OSError, e:
 225             sys.stderr.write("fork #2 failed: %s\n" % e)
 226             sys.exit(1)
 227
 228         sys.stdout.flush()
 229         sys.stderr.flush()
 230
 231         devnull = '/dev/null'
 232         stdin = file(devnull, 'r')
 233         stdout = file(devnull, 'a+')
 234         stderr = file(devnull, 'a+', 0)
 235
 236         os.dup2(stdin.fileno(), sys.stdin.fileno())
 237         os.dup2(stdout.fileno(), sys.stdout.fileno())
 238         os.dup2(stderr.fileno(), sys.stderr.fileno())
 239
 240
 241     def open_cookie_jar(self):
 242         '''Open the cookie jar.'''
 243
 244         cookie_jar = config['cookie_jar']
 245         mkbasedir(cookie_jar)
 246
 247         # Create cookie jar object from file.
 248         self.jar = cookielib.MozillaCookieJar(cookie_jar)
 249
 250         try:
 251             # Attempt to load cookies from the cookie jar.
 252             self.jar.load(ignore_discard=True)
 253
 254             # Ensure restrictive permissions are set on the cookie jar
 255             # to prevent other users on the system from hi-jacking your
 256             # authenticated sessions simply by copying your cookie jar.
 257             os.chmod(cookie_jar, 0600)
 258
 259         except:
 260             pass
 261
 262
 263     def create_socket(self):
 264         '''Create AF_UNIX socket for interprocess uzbl instance <-> cookie
 265         daemon communication.'''
 266
 267         cookie_socket = config['cookie_socket']
 268         mkbasedir(cookie_socket)
 269
 270         self.server_socket = socket.socket(socket.AF_UNIX,\
 271           socket.SOCK_SEQPACKET)
 272
 273         if os.path.exists(cookie_socket):
 274             # Accounting for super-rare super-fast racetrack condition.
 275             self.reclaim_socket()
 276
 277         self.server_socket.bind(cookie_socket)
 278
 279         # Set restrictive permissions on the cookie socket to prevent other
 280         # users on the system from data-mining your cookies.
 281         os.chmod(cookie_socket, 0600)
 282
 283
 284     def listen(self):
 285         '''Listen for incoming cookie PUT and GET requests.'''
 286
 287         echo("listening on %r" % config['cookie_socket'])
 288
 289         while self._running:
 290             # This line tells the socket how many pending incoming connections
 291             # to enqueue. I haven't had any broken pipe errors so far while
 292             # using the non-obvious value of 1 under heavy load conditions.
 293             self.server_socket.listen(1)
 294
 295             if bool(select.select([self.server_socket],[],[],1)[0]):
 296                 client_socket, _ = self.server_socket.accept()
 297                 self.handle_request(client_socket)
 298                 self.last_request = time.time()
 299                 client_socket.close()
 300
 301             if config['daemon_timeout']:
 302                 idle = time.time() - self.last_request
 303                 if idle > config['daemon_timeout']:
 304                     self._running = False
 305
 306
 307     def handle_request(self, client_socket):
 308         '''Connection made, now to serve a cookie PUT or GET request.'''
 309
 310         # Receive cookie request from client.
 311         data = client_socket.recv(8192)
 312         if not data: return
 313
 314         # Cookie argument list in packet is null separated.
 315         argv = data.split("\0")
 316
 317         # Catch the EXIT command sent to kill the daemon.
 318         if len(argv) == 1 and argv[0].strip() == "EXIT":
 319             self._running = False
 320             return None
 321
 322         # Determine whether or not to print cookie data to terminal.
 323         print_cookie = (config['verbose'] and not config['daemon_mode'])
 324         if print_cookie: print ' '.join(argv[:4])
 325
 326         action = argv[0]
 327
 328         uri = urllib2.urlparse.ParseResult(
 329           scheme=argv[1],
 330           netloc=argv[2],
 331           path=argv[3],
 332           params='',
 333           query='',
 334           fragment='').geturl()
 335
 336         req = urllib2.Request(uri)
 337
 338         if action == "GET":
 339             self.jar.add_cookie_header(req)
 340             if req.has_header('Cookie'):
 341                 cookie = req.get_header('Cookie')
 342                 client_socket.send(cookie)
 343                 if print_cookie: print cookie
 344
 345             else:
 346                 client_socket.send("\0")
 347
 348         elif action == "PUT":
 349             if len(argv) > 3:
 350                 set_cookie = argv[4]
 351                 if print_cookie: print set_cookie
 352
 353             else:
 354                 set_cookie = None
 355
 356             hdr = urllib2.httplib.HTTPMessage(\
 357               StringIO.StringIO('Set-Cookie: %s' % set_cookie))
 358             res = urllib2.addinfourl(StringIO.StringIO(), hdr,\
 359               req.get_full_url())
 360             self.jar.extract_cookies(res,req)
 361             self.jar.save(ignore_discard=True)
 362
 363         if print_cookie: print
 364
 365
 366     def quit(self, *args):
 367         '''Called on exit to make sure all loose ends are tied up.'''
 368
 369         # Only one loose end so far.
 370         self.del_socket()
 371
 372         os._exit(0)
 373
 374
 375     def del_socket(self):
 376         '''Remove the cookie_socket file on exit. In a way the cookie_socket
 377         is the daemons pid file equivalent.'''
 378
 379         if self.server_socket:
 380             try:
 381                 self.server_socket.close()
 382
 383             except:
 384                 pass
 385
 386         self.server_socket = None
 387
 388         cookie_socket = config['cookie_socket']
 389         if os.path.exists(cookie_socket):
 390             echo("deleting socket %r" % cookie_socket)
 391             os.remove(cookie_socket)
 392
 393
 394 if __name__ == "__main__":
 395
 396
 397     parser = OptionParser()
 398     parser.add_option('-n', '--no-daemon', dest='no_daemon',\
 399       action='store_true', help="don't daemonise the process.")
 400
 401     parser.add_option('-v', '--verbose', dest="verbose",\
 402       action='store_true', help="print verbose output.")
 403
 404     parser.add_option('-t', '--daemon-timeout', dest='daemon_timeout',\
 405       action="store", metavar="SECONDS", help="shutdown the daemon after x "\
 406       "seconds inactivity. WARNING: Do not use this when launching the "\
 407       "cookie daemon manually.")
 408
 409     parser.add_option('-s', '--cookie-socket', dest="cookie_socket",\
 410       metavar="SOCKET", help="manually specify the socket location.")
 411
 412     parser.add_option('-j', '--cookie-jar', dest='cookie_jar',\
 413       metavar="FILE", help="manually specify the cookie jar location.")
 414
 415     (options, args) = parser.parse_args()
 416
 417     if options.verbose:
 418         config['verbose'] = True
 419         echo("verbose mode on.")
 420
 421     if options.no_daemon:
 422         echo("daemon mode off")
 423         config['daemon_mode'] = False
 424
 425     if options.cookie_socket:
 426         echo("using cookie_socket %r" % options.cookie_socket)
 427         config['cookie_socket'] = options.cookie_socket
 428
 429     if options.cookie_jar:
 430         echo("using cookie_jar %r" % options.cookie_jar)
 431         config['cookie_jar'] = options.cookie_jar
 432
 433     if options.daemon_timeout:
 434         try:
 435             config['daemon_timeout'] = int(options.daemon_timeout)
 436             echo("set timeout to %d seconds." % config['daemon_timeout'])
 437
 438         except ValueError:
 439             config['verbose'] = True
 440             echo("fatal error: expected int argument for --daemon-timeout")
 441             sys.exit(1)
 442
 443     for key in ['cookie_socket', 'cookie_jar']:
 444         config[key] = os.path.expandvars(config[key])
 445
 446     CookieMonster().run()