wrw/session.py

   1 import threading, time, pickle, random, os
   2 import cookie
   3
   4 __all__ = ["db", "get"]
   5
   6 def hexencode(str):
   7     ret = ""
   8     for byte in str:
   9         ret += "%02X" % (ord(byte),)
  10     return ret
  11
  12 def gennonce(length):
  13     nonce = ""
  14     for i in xrange(length):
  15         nonce += chr(random.randint(0, 255))
  16     return nonce
  17
  18 class session(object):
  19     def __init__(self, expire = 86400 * 7):
  20         self.id = hexencode(gennonce(16))
  21         self.dict = {}
  22         self.lock = threading.Lock()
  23         self.ctime = self.atime = self.mtime = int(time.time())
  24         self.expire = expire
  25         self.dctl = set()
  26         self.dirtyp = False
  27
  28     def dirty(self):
  29         for d in self.dctl:
  30             if d.sessdirty():
  31                 return True
  32         return self.dirtyp
  33
  34     def frozen(self):
  35         for d in self.dctl:
  36             d.sessfrozen()
  37         self.dirtyp = False
  38
  39     def __getitem__(self, key):
  40         return self.dict[key]
  41
  42     def get(self, key, default = None):
  43         return self.dict.get(key, default)
  44
  45     def __setitem__(self, key, value):
  46         self.dict[key] = value
  47         if hasattr(value, "sessdirty"):
  48             self.dctl.add(value)
  49         else:
  50             self.dirtyp = True
  51
  52     def __delitem__(self, key):
  53         old = self.dict.pop(key)
  54         if old in self.dctl:
  55             self.dctl.remove(old)
  56         self.dirtyp = True
  57
  58     def __contains__(self, key):
  59         return key in self.dict
  60
  61     def __getstate__(self):
  62         ret = []
  63         for k, v in self.__dict__.items():
  64             if k == "lock": continue
  65             ret.append((k, v))
  66         return ret
  67
  68     def __setstate__(self, st):
  69         for k, v in st:
  70             self.__dict__[k] = v
  71         self.lock = threading.Lock()
  72
  73 class db(object):
  74     def __init__(self, cookiename = "wrwsess", path = "/"):
  75         self.live = {}
  76         self.cookiename = cookiename
  77         self.path = path
  78         self.lock = threading.Lock()
  79         self.lastuse = 0
  80         self.cthread = None
  81         self.freezetime = 3600
  82
  83     def clean(self):
  84         now = int(time.time())
  85         with self.lock:
  86             dlist = []
  87             for sess in self.live.itervalues():
  88                 if sess.atime + self.freezetime < now:
  89                     try:
  90                         if sess.dirty():
  91                             self.freeze(sess)
  92                     except:
  93                         if sess.atime + sess.expire < now:
  94                             dlist.append(sess)
  95                     else:
  96                         dlist.append(sess)
  97             for sess in dlist:
  98                 del self.live[sess.id]
  99
 100     def cleanloop(self):
 101         try:
 102             lastuse = self.lastuse
 103             while self.lastuse >= lastuse:
 104                 lastuse = self.lastuse
 105                 time.sleep(300)
 106                 self.clean()
 107         finally:
 108             with self.lock:
 109                 self.cthread = None
 110
 111     def fetch(self, req):
 112         now = int(time.time())
 113         self.lastuse = now
 114         sessid = cookie.get(req, self.cookiename)
 115         with self.lock:
 116             if self.cthread is None:
 117                 self.cthread = threading.Thread(target = self.cleanloop)
 118                 self.cthread.setDaemon(True)
 119                 self.cthread.start()
 120             try:
 121                 if sessid is None:
 122                     raise KeyError()
 123                 elif sessid in self.live:
 124                     sess = self.live[sessid]
 125                 else:
 126                     sess = self.thaw(sessid)
 127                     self.live[sessid] = sess
 128                 if sess.atime + sess.expire < now:
 129                     raise KeyError()
 130                 sess.atime = now
 131             except KeyError:
 132                 sess = session()
 133                 self.live[sess.id] = sess
 134                 req.oncommit(self.addcookie)
 135         req.oncommit(self.ckfreeze)
 136         return sess
 137
 138     def addcookie(self, req):
 139         sess = req.item(self.fetch)
 140         cookie.add(req, self.cookiename, sess.id, self.path)
 141
 142     def ckfreeze(self, req):
 143         sess = req.item(self.fetch)
 144         if sess.dirty():
 145             try:
 146                 self.freeze(sess)
 147             except:
 148                 pass
 149
 150     def thaw(self, sessid):
 151         raise KeyError()
 152
 153     def freeze(self, sess):
 154         raise TypeError()
 155
 156 class backeddb(db):
 157     def __init__(self, backdb, *args, **kw):
 158         super(backeddb, self).__init__(*args, **kw)
 159         self.backdb = backdb
 160
 161     def thaw(self, sessid):
 162         data = self.backdb[sessid]
 163         try:
 164             return pickle.loads(data)
 165         except Exception, e:
 166             raise KeyError()
 167
 168     def freeze(self, sess):
 169         self.backdb[sess.id] = pickle.dumps(sess)
 170         sess.frozen()
 171
 172 class dirback(object):
 173     def __init__(self, path):
 174         self.path = path
 175
 176     def __getitem__(self, key):
 177         try:
 178             with open(os.path.join(self.path, key)) as inf:
 179                 return inf.read()
 180         except IOError:
 181             raise KeyError(key)
 182
 183     def __setitem__(self, key, value):
 184         if not os.path.exists(self.path):
 185             os.makedirs(self.path)
 186         with open(os.path.join(self.path, key), "w") as out:
 187             out.write(value)
 188
 189 default = backeddb(dirback(os.path.join("/tmp", "wrwsess-" + str(os.getuid()))))
 190
 191 def get(req):
 192     return req.item(default.fetch)