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                 client_socket.close()
 261
 262             if self.daemon_timeout:
 263                 idle = time.time() - self.last_request
 264                 if idle > self.daemon_timeout: break
 265
 266
 267     def handle_request(self, client_socket):
 268         '''Connection made, now to serve a cookie PUT or GET request.'''
 269
 270         # Receive cookie request from client.
 271         data = client_socket.recv(8192)
 272         if not data: return
 273
 274         # Cookie argument list in packet is null separated.
 275         argv = data.split("\0")
 276
 277         # Determine whether or not to print cookie data to terminal.
 278         print_cookie = (verbose and not self.daemon_mode)
 279         if print_cookie: print ' '.join(argv[:4])
 280
 281         action = argv[0]
 282
 283         uri = urllib2.urlparse.ParseResult(
 284           scheme=argv[1],
 285           netloc=argv[2],
 286           path=argv[3],
 287           params='',
 288           query='',
 289           fragment='').geturl()
 290
 291         req = urllib2.Request(uri)
 292
 293         if action == "GET":
 294             self.jar.add_cookie_header(req)
 295             if req.has_header('Cookie'):
 296                 cookie = req.get_header('Cookie')
 297                 client_socket.send(cookie)
 298                 if print_cookie: print cookie
 299
 300             else:
 301                 client_socket.send("\0")
 302
 303         elif action == "PUT":
 304             if len(argv) > 3:
 305                 set_cookie = argv[4]
 306                 if print_cookie: print set_cookie
 307
 308             else:
 309                 set_cookie = None
 310
 311             hdr = urllib2.httplib.HTTPMessage(\
 312               StringIO.StringIO('Set-Cookie: %s' % set_cookie))
 313             res = urllib2.addinfourl(StringIO.StringIO(), hdr,\
 314               req.get_full_url())
 315             self.jar.extract_cookies(res,req)
 316             self.jar.save(ignore_discard=True)
 317
 318         if print_cookie: print
 319
 320
 321     def quit(self, *args):
 322         '''Called on exit to make sure all loose ends are tied up.'''
 323
 324         # Only one loose end so far.
 325         self.del_socket()
 326
 327         os._exit(0)
 328
 329
 330     def del_socket(self):
 331         '''Remove the cookie_socket file on exit. In a way the cookie_socket
 332         is the daemons pid file equivalent.'''
 333
 334         if self.server_socket:
 335             self.server_socket.close()
 336
 337         if os.path.exists(self.cookie_socket):
 338             os.remove(self.cookie_socket)
 339
 340
 341 if __name__ == "__main__":
 342
 343     CookieMonster(cookie_socket, cookie_jar, daemon_timeout,\
 344       daemon_mode).run()
 345