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