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
  33
  34 import cookielib
  35 import os
  36 import sys
  37 import urllib2
  38 import select
  39 import socket
  40 import time
  41 import atexit
  42 from traceback import print_exc
  43 from signal import signal, SIGTERM
  44 from optparse import OptionParser
  45
  46 try:
  47     import cStringIO as StringIO
  48
  49 except ImportError:
  50     import StringIO
  51
  52
  53 # ============================================================================
  54 # ::: Default configuration section ::::::::::::::::::::::::::::::::::::::::::
  55 # ============================================================================
  56
  57
  58 # Location of the uzbl cache directory.
  59 if 'XDG_CACHE_HOME' in os.environ.keys() and os.environ['XDG_CACHE_HOME']:
  60     cache_dir = os.path.join(os.environ['XDG_CACHE_HOME'], 'uzbl/')
  61
  62 else:
  63     cache_dir = os.path.join(os.environ['HOME'], '.cache/uzbl/')
  64
  65 # Location of the uzbl data directory.
  66 if 'XDG_DATA_HOME' in os.environ.keys() and os.environ['XDG_DATA_HOME']:
  67     data_dir = os.path.join(os.environ['XDG_DATA_HOME'], 'uzbl/')
  68
  69 else:
  70     data_dir = os.path.join(os.environ['HOME'], '.local/share/uzbl/')
  71
  72 # Default config
  73 config = {
  74
  75   # Default cookie jar and daemon socket locations.
  76   'cookie_socket': os.path.join(cache_dir, 'cookie_daemon_socket'),
  77   'cookie_jar': os.path.join(data_dir, 'cookies.txt'),
  78
  79   # Time out after x seconds of inactivity (set to 0 for never time out).
  80   # Set to 0 by default until talk_to_socket is doing the spawning.
  81   'daemon_timeout': 0,
  82
  83   # Tell process to daemonise
  84   'daemon_mode': True,
  85
  86   # Set true to print helpful debugging messages to the terminal.
  87   'verbose': False,
  88
  89 } # End of config dictionary.
  90
  91
  92 # ============================================================================
  93 # ::: End of configuration section :::::::::::::::::::::::::::::::::::::::::::
  94 # ============================================================================
  95
  96
  97 _scriptname = os.path.basename(sys.argv[0])
  98 def echo(msg):
  99     if config['verbose']:
 100         print "%s: %s" % (_scriptname, msg)
 101
 102
 103 def mkbasedir(filepath):
 104     '''Create base directory of filepath if it doesn't exist.'''
 105
 106     dirname = os.path.dirname(filepath)
 107     if not os.path.exists(dirname):
 108         echo("creating dirs: %r" % dirname)
 109         os.makedirs(dirname)
 110
 111
 112 class CookieMonster:
 113     '''The uzbl cookie daemon class.'''
 114
 115     def __init__(self):
 116         '''Initialise class variables.'''
 117
 118         self.server_socket = None
 119         self.jar = None
 120         self.last_request = time.time()
 121         self._running = False
 122
 123
 124     def run(self):
 125         '''Start the daemon.'''
 126
 127         # Check if another daemon is running. The reclaim_socket function will
 128         # exit if another daemon is detected listening on the cookie socket
 129         # and remove the abandoned socket if there isnt.
 130         if os.path.exists(config['cookie_socket']):
 131             self.reclaim_socket()
 132
 133         # Daemonize process.
 134         if config['daemon_mode']:
 135             echo("entering daemon mode.")
 136             self.daemonize()
 137
 138         # Register a function to cleanup on exit.
 139         atexit.register(self.quit)
 140
 141         # Make SIGTERM act orderly.
 142         signal(SIGTERM, lambda signum, stack_frame: sys.exit(1))
 143
 144         # Create cookie jar object from file.
 145         self.open_cookie_jar()
 146
 147         # Creating a way to exit nested loops by setting a running flag.
 148         self._running = True
 149
 150         while self._running:
 151             # Create cookie daemon socket.
 152             self.create_socket()
 153
 154             try:
 155                 # Enter main listen loop.
 156                 self.listen()
 157
 158             except KeyboardInterrupt:
 159                 self._running = False
 160                 print
 161
 162             except socket.error:
 163                 print_exc()
 164
 165             except:
 166                 # Clean up
 167                 self.del_socket()
 168
 169                 # Raise exception
 170                 raise
 171
 172             # Always delete the socket before calling create again.
 173             self.del_socket()
 174
 175
 176     def reclaim_socket(self):
 177         '''Check if another process (hopefully a cookie_daemon.py) is listening
 178         on the cookie daemon socket. If another process is found to be
 179         listening on the socket exit the daemon immediately and leave the
 180         socket alone. If the connect fails assume the socket has been abandoned
 181         and delete it (to be re-created in the create socket function).'''
 182
 183         cookie_socket = config['cookie_socket']
 184
 185         try:
 186             sock = socket.socket(socket.AF_UNIX, socket.SOCK_SEQPACKET)
 187             sock.connect(cookie_socket)
 188             sock.close()
 189
 190         except socket.error:
 191             # Failed to connect to cookie_socket so assume it has been
 192             # abandoned by another cookie daemon process.
 193             echo("reclaiming abandoned cookie_socket %r." % cookie_socket)
 194             if os.path.exists(cookie_socket):
 195                 os.remove(cookie_socket)
 196
 197             return
 198
 199         echo("detected another process listening on %r." % cookie_socket)
 200         echo("exiting.")
 201         # Use os._exit() to avoid tripping the atexit cleanup function.
 202         os._exit(1)
 203
 204
 205     def daemonize(function):
 206         '''Daemonize the process using the Stevens' double-fork magic.'''
 207
 208         try:
 209             if os.fork(): os._exit(0)
 210
 211         except OSError, e:
 212             sys.stderr.write("fork #1 failed: %s\n" % e)
 213             sys.exit(1)
 214
 215         os.chdir('/')
 216         os.setsid()
 217         os.umask(0)
 218
 219         try:
 220             if os.fork(): os._exit(0)
 221
 222         except OSError, e:
 223             sys.stderr.write("fork #2 failed: %s\n" % e)
 224             sys.exit(1)
 225
 226         sys.stdout.flush()
 227         sys.stderr.flush()
 228
 229         devnull = '/dev/null'
 230         stdin = file(devnull, 'r')
 231         stdout = file(devnull, 'a+')
 232         stderr = file(devnull, 'a+', 0)
 233
 234         os.dup2(stdin.fileno(), sys.stdin.fileno())
 235         os.dup2(stdout.fileno(), sys.stdout.fileno())
 236         os.dup2(stderr.fileno(), sys.stderr.fileno())
 237
 238
 239     def open_cookie_jar(self):
 240         '''Open the cookie jar.'''
 241
 242         cookie_jar = config['cookie_jar']
 243         if config['cookie_jar']:
 244             mkbasedir(cookie_jar)
 245
 246         # Create cookie jar object from file.
 247         self.jar = cookielib.MozillaCookieJar(cookie_jar)
 248
 249         if config['cookie_jar']:
 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             if config['cookie_jar']:
 362                 self.jar.save(ignore_discard=True)
 363
 364         if print_cookie: print
 365
 366
 367     def quit(self, *args):
 368         '''Called on exit to make sure all loose ends are tied up.'''
 369
 370         # Only one loose end so far.
 371         self.del_socket()
 372
 373         os._exit(0)
 374
 375
 376     def del_socket(self):
 377         '''Remove the cookie_socket file on exit. In a way the cookie_socket
 378         is the daemons pid file equivalent.'''
 379
 380         if self.server_socket:
 381             try:
 382                 self.server_socket.close()
 383
 384             except:
 385                 pass
 386
 387         self.server_socket = None
 388
 389         cookie_socket = config['cookie_socket']
 390         if os.path.exists(cookie_socket):
 391             echo("deleting socket %r" % cookie_socket)
 392             os.remove(cookie_socket)
 393
 394
 395 if __name__ == "__main__":
 396
 397
 398     parser = OptionParser()
 399     parser.add_option('-n', '--no-daemon', dest='no_daemon',\
 400       action='store_true', help="don't daemonise the process.")
 401
 402     parser.add_option('-v', '--verbose', dest="verbose",\
 403       action='store_true', help="print verbose output.")
 404
 405     parser.add_option('-t', '--daemon-timeout', dest='daemon_timeout',\
 406       action="store", metavar="SECONDS", help="shutdown the daemon after x "\
 407       "seconds inactivity. WARNING: Do not use this when launching the "\
 408       "cookie daemon manually.")
 409
 410     parser.add_option('-s', '--cookie-socket', dest="cookie_socket",\
 411       metavar="SOCKET", help="manually specify the socket location.")
 412
 413     parser.add_option('-j', '--cookie-jar', dest='cookie_jar',\
 414       metavar="FILE", help="manually specify the cookie jar location.")
 415
 416     parser.add_option('-m', '--memory', dest='memory', action='store_true',
 417             help="store cookies in memory only - do not write to disk")
 418
 419     (options, args) = parser.parse_args()
 420
 421     if options.verbose:
 422         config['verbose'] = True
 423         echo("verbose mode on.")
 424
 425     if options.no_daemon:
 426         echo("daemon mode off")
 427         config['daemon_mode'] = False
 428
 429     if options.cookie_socket:
 430         echo("using cookie_socket %r" % options.cookie_socket)
 431         config['cookie_socket'] = options.cookie_socket
 432
 433     if options.cookie_jar:
 434         echo("using cookie_jar %r" % options.cookie_jar)
 435         config['cookie_jar'] = options.cookie_jar
 436
 437     if options.memory:
 438         echo("using memory %r" % options.memory)
 439         config['cookie_jar'] = None
 440
 441     if options.daemon_timeout:
 442         try:
 443             config['daemon_timeout'] = int(options.daemon_timeout)
 444             echo("set timeout to %d seconds." % config['daemon_timeout'])
 445
 446         except ValueError:
 447             config['verbose'] = True
 448             echo("fatal error: expected int argument for --daemon-timeout")
 449             sys.exit(1)
 450
 451     for key in ['cookie_socket', 'cookie_jar']:
 452         if config[key]:
 453             config[key] = os.path.expandvars(config[key])
 454
 455     CookieMonster().run()