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