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