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