2abe650718040207fc43f78b0d520dc7e0327b51
[wrw.git] / 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.cthread = None
80         self.freezetime = 3600
81
82     def clean(self):
83         now = int(time.time())
84         with self.lock:
85             dlist = []
86             for sess in self.live.itervalues():
87                 if sess.atime + self.freezetime < now:
88                     try:
89                         if sess.dirty():
90                             self.freeze(sess)
91                     except:
92                         if sess.atime + sess.expire < now:
93                             dlist.append(sess)
94                     else:
95                         dlist.append(sess)
96             for sess in dlist:
97                 del self.live[sess.id]
98
99     def cleanloop(self):
100         try:
101             while True:
102                 time.sleep(300)
103                 self.clean()
104                 if len(self.live) == 0:
105                     break
106         finally:
107             with self.lock:
108                 self.cthread = None
109
110     def fetch(self, req):
111         now = int(time.time())
112         sessid = cookie.get(req, self.cookiename)
113         new = False
114         with self.lock:
115             if self.cthread is None:
116                 self.cthread = threading.Thread(target = self.cleanloop)
117                 self.cthread.setDaemon(True)
118                 self.cthread.start()
119             try:
120                 if sessid is None:
121                     raise KeyError()
122                 elif sessid in self.live:
123                     sess = self.live[sessid]
124                 else:
125                     sess = self.thaw(sessid)
126                     self.live[sessid] = sess
127                 if sess.atime + sess.expire < now:
128                     raise KeyError()
129                 sess.atime = now
130             except KeyError:
131                 sess = session()
132                 self.live[sess.id] = sess
133                 new = True
134
135         def ckfreeze(req):
136             if sess.dirty():
137                 try:
138                     if new:
139                         cookie.add(req, self.cookiename, sess.id, self.path)
140                     self.freeze(sess)
141                 except:
142                     pass
143         req.oncommit(ckfreeze)
144         return sess
145
146     def thaw(self, sessid):
147         raise KeyError()
148
149     def freeze(self, sess):
150         raise TypeError()
151
152     def get(self, req):
153         return req.item(self.fetch)
154
155 class backeddb(db):
156     def __init__(self, backdb, *args, **kw):
157         super(backeddb, self).__init__(*args, **kw)
158         self.backdb = backdb
159
160     def thaw(self, sessid):
161         data = self.backdb[sessid]
162         try:
163             return pickle.loads(data)
164         except Exception, e:
165             raise KeyError()
166
167     def freeze(self, sess):
168         self.backdb[sess.id] = pickle.dumps(sess)
169         sess.frozen()
170
171 class dirback(object):
172     def __init__(self, path):
173         self.path = path
174
175     def __getitem__(self, key):
176         try:
177             with open(os.path.join(self.path, key)) as inf:
178                 return inf.read()
179         except IOError:
180             raise KeyError(key)
181
182     def __setitem__(self, key, value):
183         if not os.path.exists(self.path):
184             os.makedirs(self.path)
185         with open(os.path.join(self.path, key), "w") as out:
186             out.write(value)
187
188 default = backeddb(dirback(os.path.join("/tmp", "wrwsess-" + str(os.getuid()))))
189
190 def get(req):
191     return default.get(req)