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 # Todo list:
  23 #  - Setup some option parsing so the daemon can take optional command line
  24 #    arguments.
  25
  26
  27 import cookielib
  28 import os
  29 import sys
  30 import urllib2
  31 import select
  32 import socket
  33 import time
  34 import atexit
  35 from signal import signal, SIGTERM
  36
  37 try:
  38     import cStringIO as StringIO
  39
  40 except ImportError:
  41     import StringIO
  42
  43
  44 # ============================================================================
  45 # ::: Default configuration section ::::::::::::::::::::::::::::::::::::::::::
  46 # ============================================================================
  47
  48 # Location of the uzbl cache directory.
  49 if 'XDG_CACHE_HOME' in os.environ.keys() and os.environ['XDG_CACHE_HOME']:
  50     cache_dir = os.path.join(os.environ['XDG_CACHE_HOME'], 'uzbl/')
  51
  52 else:
  53     cache_dir = os.path.join(os.environ['HOME'], '.cache/uzbl/')
  54
  55 # Location of the uzbl data directory.
  56 if 'XDG_DATA_HOME' in os.environ.keys() and os.environ['XDG_DATA_HOME']:
  57     data_dir = os.path.join(os.environ['XDG_DATA_HOME'], 'uzbl/')
  58
  59 else:
  60     data_dir = os.path.join(os.environ['HOME'], '.local/share/uzbl/')
  61
  62 # Create cache dir and data dir if they are missing.
  63 for path in [data_dir, cache_dir]:
  64     if not os.path.exists(path):
  65         os.makedirs(path)
  66
  67 # Default config
  68 cookie_socket = os.path.join(cache_dir, 'cookie_daemon_socket')
  69 cookie_jar = os.path.join(data_dir, 'cookies.txt')
  70
  71 # Time out after x seconds of inactivity (set to 0 for never time out).
  72 # Set to 0 by default until talk_to_socket is doing the spawning.
  73 daemon_timeout = 0
  74
  75 # Enable/disable daemonizing the process (useful when debugging).
  76 # Set to False by default until talk_to_socket is doing the spawning.
  77 daemon_mode = False
  78
  79 # Set true to print helpful debugging messages to the terminal.
  80 verbose = True
  81
  82 # ============================================================================
  83 # ::: End of configuration section :::::::::::::::::::::::::::::::::::::::::::
  84 # ============================================================================
  85
  86
  87 _scriptname = os.path.basename(sys.argv[0])
  88 def echo(msg):
  89     if verbose:
  90         print "%s: %s" % (_scriptname, msg)
  91
  92
  93 class CookieMonster:
  94     '''The uzbl cookie daemon class.'''
  95
  96     def __init__(self, cookie_socket, cookie_jar, daemon_timeout,\
  97       daemon_mode):
  98
  99         self.cookie_socket = os.path.expandvars(cookie_socket)
 100         self.server_socket = None
 101         self.cookie_jar = os.path.expandvars(cookie_jar)
 102         self.daemon_mode = daemon_mode
 103         self.jar = None
 104         self.daemon_timeout = daemon_timeout
 105         self.last_request = time.time()
 106
 107
 108     def run(self):
 109         '''Start the daemon.'''
 110
 111         # Check if another daemon is running. The reclaim_socket function will
 112         # exit if another daemon is detected listening on the cookie socket
 113         # and remove the abandoned socket if there isnt.
 114         if os.path.exists(self.cookie_socket):
 115             self.reclaim_socket()
 116
 117         # Daemonize process.
 118         if self.daemon_mode:
 119             echo("entering daemon mode.")
 120             self.daemonize()
 121
 122         # Register a function to cleanup on exit.
 123         atexit.register(self.quit)
 124
 125         # Make SIGTERM act orderly.
 126         signal(SIGTERM, lambda signum, stack_frame: sys.exit(1))
 127
 128         # Create cookie daemon socket.
 129         self.create_socket()
 130
 131         # Create cookie jar object from file.
 132         self.open_cookie_jar()
 133
 134         try:
 135             # Listen for incoming cookie puts/gets.
 136             self.listen()
 137
 138         except KeyboardInterrupt:
 139             print
 140
 141         except:
 142             # Clean up
 143             self.del_socket()
 144
 145             # Raise exception
 146             raise
 147
 148
 149     def reclaim_socket(self):
 150         '''Check if another process (hopefully a cookie_daemon.py) is listening
 151         on the cookie daemon socket. If another process is found to be
 152         listening on the socket exit the daemon immediately and leave the
 153         socket alone. If the connect fails assume the socket has been abandoned
 154         and delete it (to be re-created in the create socket function).'''
 155
 156         try:
 157             sock = socket.socket(socket.AF_UNIX, socket.SOCK_SEQPACKET)
 158             sock.connect(self.cookie_socket)
 159             sock.close()
 160
 161         except socket.error:
 162             # Failed to connect to cookie_socket so assume it has been
 163             # abandoned by another cookie daemon process.
 164             echo("reclaiming abandoned cookie_socket %r." % self.cookie_socket)
 165             if os.path.exists(self.cookie_socket):
 166                 os.remove(self.cookie_socket)
 167
 168             return
 169
 170         echo("detected another process listening on %r." % self.cookie_socket)
 171         echo("exiting.")
 172         # Use os._exit() to avoid tripping the atexit cleanup function.
 173         os._exit(1)
 174
 175
 176     def daemonize(function):
 177         '''Daemonize the process using the Stevens' double-fork magic.'''
 178
 179         try:
 180             if os.fork(): os._exit(0)
 181
 182         except OSError, e:
 183             sys.stderr.write("fork #1 failed: %s\n" % e)
 184             sys.exit(1)
 185
 186         os.chdir('/')
 187         os.setsid()
 188         os.umask(0)
 189
 190         try:
 191             if os.fork(): os._exit(0)
 192
 193         except OSError, e:
 194             sys.stderr.write("fork #2 failed: %s\n" % e)
 195             sys.exit(1)
 196
 197         sys.stdout.flush()
 198         sys.stderr.flush()
 199
 200         devnull = '/dev/null'
 201         stdin = file(devnull, 'r')
 202         stdout = file(devnull, 'a+')
 203         stderr = file(devnull, 'a+', 0)
 204
 205         os.dup2(stdin.fileno(), sys.stdin.fileno())
 206         os.dup2(stdout.fileno(), sys.stdout.fileno())
 207         os.dup2(stderr.fileno(), sys.stderr.fileno())
 208
 209
 210     def open_cookie_jar(self):
 211         '''Open the cookie jar.'''
 212
 213         # Create cookie jar object from file.
 214         self.jar = cookielib.MozillaCookieJar(self.cookie_jar)
 215
 216         try:
 217             # Attempt to load cookies from the cookie jar.
 218             self.jar.load(ignore_discard=True)
 219
 220             # Ensure restrictive permissions are set on the cookie jar
 221             # to prevent other users on the system from hi-jacking your
 222             # authenticated sessions simply by copying your cookie jar.
 223             os.chmod(self.cookie_jar, 0600)
 224
 225         except:
 226             pass
 227
 228
 229     def create_socket(self):
 230         '''Create AF_UNIX socket for interprocess uzbl instance <-> cookie
 231         daemon communication.'''
 232
 233         self.server_socket = socket.socket(socket.AF_UNIX,\
 234           socket.SOCK_SEQPACKET)
 235
 236         if os.path.exists(self.cookie_socket):
 237             # Accounting for super-rare super-fast racetrack condition.
 238             self.reclaim_socket()
 239
 240         self.server_socket.bind(self.cookie_socket)
 241
 242         # Set restrictive permissions on the cookie socket to prevent other
 243         # users on the system from data-mining your cookies.
 244         os.chmod(self.cookie_socket, 0600)
 245
 246
 247     def listen(self):
 248         '''Listen for incoming cookie PUT and GET requests.'''
 249
 250         while True:
 251             # This line tells the socket how many pending incoming connections
 252             # to enqueue. I haven't had any broken pipe errors so far while
 253             # using the non-obvious value of 1 under heavy load conditions.
 254             self.server_socket.listen(1)
 255
 256             if bool(select.select([self.server_socket],[],[],1)[0]):
 257                 client_socket, _ = self.server_socket.accept()
 258                 self.handle_request(client_socket)
 259                 self.last_request = time.time()
 260
 261             if self.daemon_timeout:
 262                 idle = time.time() - self.last_request
 263                 if idle > self.daemon_timeout: break
 264
 265
 266     def handle_request(self, client_socket):
 267         '''Connection made, now to serve a cookie PUT or GET request.'''
 268
 269         # Receive cookie request from client.
 270         data = client_socket.recv(8192)
 271         if not data: return
 272
 273         # Cookie argument list in packet is null separated.
 274         argv = data.split("\0")
 275
 276         # Determine whether or not to print cookie data to terminal.
 277         print_cookie = (verbose and not self.daemon_mode)
 278         if print_cookie: print ' '.join(argv[:4])
 279
 280         action = argv[0]
 281
 282         uri = urllib2.urlparse.ParseResult(
 283           scheme=argv[1],
 284           netloc=argv[2],
 285           path=argv[3],
 286           params='',
 287           query='',
 288           fragment='').geturl()
 289
 290         req = urllib2.Request(uri)
 291
 292         if action == "GET":
 293             self.jar.add_cookie_header(req)
 294             if req.has_header('Cookie'):
 295                 cookie = req.get_header('Cookie')
 296                 client_socket.send(cookie)
 297                 if print_cookie: print cookie
 298
 299             else:
 300                 client_socket.send("\0")
 301
 302         elif action == "PUT":
 303             if len(argv) > 3:
 304                 set_cookie = argv[4]
 305                 if print_cookie: print set_cookie
 306
 307             else:
 308                 set_cookie = None
 309
 310             hdr = urllib2.httplib.HTTPMessage(\
 311               StringIO.StringIO('Set-Cookie: %s' % set_cookie))
 312             res = urllib2.addinfourl(StringIO.StringIO(), hdr,\
 313               req.get_full_url())
 314             self.jar.extract_cookies(res,req)
 315             self.jar.save(ignore_discard=True)
 316
 317         if print_cookie: print
 318
 319         client_socket.close()
 320
 321
 322     def quit(self, *args):
 323         '''Called on exit to make sure all loose ends are tied up.'''
 324
 325         # Only one loose end so far.
 326         self.del_socket()
 327
 328         os._exit(0)
 329
 330
 331     def del_socket(self):
 332         '''Remove the cookie_socket file on exit. In a way the cookie_socket
 333         is the daemons pid file equivalent.'''
 334
 335         if self.server_socket:
 336             self.server_socket.close()
 337
 338         if os.path.exists(self.cookie_socket):
 339             os.remove(self.cookie_socket)
 340
 341
 342 if __name__ == "__main__":
 343
 344     CookieMonster(cookie_socket, cookie_jar, daemon_timeout,\
 345       daemon_mode).run()
 346