wrw/session.py

   1 import threading, time, pickle, random, os
   2 from . import cookie, env
   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 range(length):
  15         nonce += chr(random.randint(0, 255))
  16     return nonce
  17
  18 class session(object):
  19     def __init__(self, lock, expire=86400 * 7):
  20         self.id = hexencode(gennonce(16))
  21         self.dict = {}
  22         self.lock = 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         # The proper lock is set by the thawer
  72
  73     def __repr__(self):
  74         return "<session %s>" % self.id
  75
  76 class db(object):
  77     def __init__(self, backdb=None, cookiename="wrwsess", path="/"):
  78         self.live = {}
  79         self.cookiename = cookiename
  80         self.path = path
  81         self.lock = threading.Lock()
  82         self.cthread = None
  83         self.freezetime = 3600
  84         self.backdb = backdb
  85
  86     def clean(self):
  87         now = int(time.time())
  88         with self.lock:
  89             clist = list(self.live.keys())
  90         for sessid in clist:
  91             with self.lock:
  92                 try:
  93                     entry = self.live[sessid]
  94                 except KeyError:
  95                     continue
  96             with entry[0]:
  97                 rm = False
  98                 if entry[1] == "retired":
  99                     pass
 100                 elif entry[1] is None:
 101                     pass
 102                 else:
 103                     sess = entry[1]
 104                     if sess.atime + self.freezetime < now:
 105                         try:
 106                             if sess.dirty():
 107                                 self.freeze(sess)
 108                         except:
 109                             if sess.atime + sess.expire < now:
 110                                 rm = True
 111                         else:
 112                             rm = True
 113                 if rm:
 114                     entry[1] = "retired"
 115                     with self.lock:
 116                         del self.live[sessid]
 117
 118     def cleanloop(self):
 119         try:
 120             while True:
 121                 time.sleep(300)
 122                 self.clean()
 123                 if len(self.live) == 0:
 124                     break
 125         finally:
 126             with self.lock:
 127                 self.cthread = None
 128
 129     def _fetch(self, sessid):
 130         while True:
 131             now = int(time.time())
 132             with self.lock:
 133                 if sessid in self.live:
 134                     entry = self.live[sessid]
 135                 else:
 136                     entry = self.live[sessid] = [threading.RLock(), None]
 137             with entry[0]:
 138                 if isinstance(entry[1], session):
 139                     entry[1].atime = now
 140                     return entry[1]
 141                 elif entry[1] == "retired":
 142                     continue
 143                 elif entry[1] is None:
 144                     try:
 145                         thawed = self.thaw(sessid)
 146                         if thawed.atime + thawed.expire < now:
 147                             raise KeyError()
 148                         thawed.lock = entry[0]
 149                         thawed.atime = now
 150                         entry[1] = thawed
 151                         return thawed
 152                     finally:
 153                         if entry[1] is None:
 154                             entry[1] = "retired"
 155                             with self.lock:
 156                                 del self.live[sessid]
 157                 else:
 158                     raise Exception("Illegal session entry: " + repr(entry[1]))
 159
 160     def checkclean(self):
 161         with self.lock:
 162             if self.cthread is None:
 163                 self.cthread = threading.Thread(target = self.cleanloop)
 164                 self.cthread.setDaemon(True)
 165                 self.cthread.start()
 166
 167     def mksession(self, req):
 168         return session(threading.RLock())
 169
 170     def mkcookie(self, req, sess):
 171         cookie.add(req, self.cookiename, sess.id,
 172                    path=self.path,
 173                    expires=cookie.cdate(time.time() + sess.expire))
 174
 175     def fetch(self, req):
 176         now = int(time.time())
 177         sessid = cookie.get(req, self.cookiename)
 178         new = False
 179         try:
 180             if sessid is None:
 181                 raise KeyError()
 182             sess = self._fetch(sessid)
 183         except KeyError:
 184             sess = self.mksession(req)
 185             new = True
 186
 187         def ckfreeze(req):
 188             if sess.dirty():
 189                 if new:
 190                     self.mkcookie(req, sess)
 191                     with self.lock:
 192                         self.live[sess.id] = [sess.lock, sess]
 193                 try:
 194                     self.freeze(sess)
 195                 except:
 196                     pass
 197                 self.checkclean()
 198         req.oncommit(ckfreeze)
 199         return sess
 200
 201     def thaw(self, sessid):
 202         if self.backdb is None:
 203             raise KeyError()
 204         data = self.backdb[sessid]
 205         try:
 206             return pickle.loads(data)
 207         except:
 208             raise KeyError()
 209
 210     def freeze(self, sess):
 211         if self.backdb is None:
 212             raise TypeError()
 213         with sess.lock:
 214             data = pickle.dumps(sess, -1)
 215         self.backdb[sess.id] = data
 216         sess.frozen()
 217
 218     def get(self, req):
 219         return req.item(self.fetch)
 220
 221 class dirback(object):
 222     def __init__(self, path):
 223         self.path = path
 224
 225     def __getitem__(self, key):
 226         try:
 227             with open(os.path.join(self.path, key)) as inf:
 228                 return inf.read()
 229         except IOError:
 230             raise KeyError(key)
 231
 232     def __setitem__(self, key, value):
 233         if not os.path.exists(self.path):
 234             os.makedirs(self.path)
 235         with open(os.path.join(self.path, key), "w") as out:
 236             out.write(value)
 237
 238 default = env.var(db(backdb=dirback(os.path.join("/tmp", "wrwsess-" + str(os.getuid())))))
 239
 240 def get(req):
 241     return default.val.get(req)