Commit | Line | Data |
---|---|---|
a2916c06 MW |
1 | ### -*-python-*- |
2 | ### | |
3 | ### Miscellaneous utilities | |
4 | ### | |
5 | ### (c) 2013 Mark Wooding | |
6 | ### | |
7 | ||
8 | ###----- Licensing notice --------------------------------------------------- | |
9 | ### | |
10 | ### This file is part of Chopwood: a password-changing service. | |
11 | ### | |
12 | ### Chopwood is free software; you can redistribute it and/or modify | |
13 | ### it under the terms of the GNU Affero General Public License as | |
14 | ### published by the Free Software Foundation; either version 3 of the | |
15 | ### License, or (at your option) any later version. | |
16 | ### | |
17 | ### Chopwood is distributed in the hope that it will be useful, | |
18 | ### but WITHOUT ANY WARRANTY; without even the implied warranty of | |
19 | ### MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the | |
20 | ### GNU Affero General Public License for more details. | |
21 | ### | |
22 | ### You should have received a copy of the GNU Affero General Public | |
23 | ### License along with Chopwood; if not, see | |
24 | ### <http://www.gnu.org/licenses/>. | |
25 | ||
26 | from __future__ import with_statement | |
27 | ||
28 | import base64 as BN | |
29 | import contextlib as CTX | |
30 | import fcntl as F | |
31 | import os as OS | |
32 | import re as RX | |
33 | import signal as SIG | |
34 | import sys as SYS | |
35 | import time as T | |
36 | ||
37 | try: import threading as TH | |
38 | except ImportError: import dummy_threading as TH | |
39 | ||
40 | ###-------------------------------------------------------------------------- | |
41 | ### Some basics. | |
42 | ||
43 | def identity(x): | |
44 | """The identity function: returns its argument.""" | |
45 | return x | |
46 | ||
47 | def constantly(x): | |
48 | """The function which always returns X.""" | |
49 | return lambda: x | |
50 | ||
51 | class struct (object): | |
52 | """A simple object for storing data in attributes.""" | |
53 | DEFAULTS = {} | |
54 | def __init__(me, *args, **kw): | |
55 | cls = me.__class__ | |
56 | for k, v in kw.iteritems(): setattr(me, k, v) | |
57 | try: | |
58 | slots = cls.__slots__ | |
59 | except AttributeError: | |
60 | if args: raise ValueError, 'no slots defined' | |
61 | else: | |
62 | if len(args) > len(slots): raise ValueError, 'too many arguments' | |
63 | for k, v in zip(slots, args): setattr(me, k, v) | |
64 | for k in slots: | |
65 | if hasattr(me, k): continue | |
66 | try: setattr(me, k, cls.DEFAULTS[k]) | |
67 | except KeyError: raise ValueError, "no value for `%s'" % k | |
68 | ||
69 | class Tag (object): | |
70 | """An object whose only purpose is to be distinct from other objects.""" | |
71 | def __init__(me, name): | |
72 | me._name = name | |
73 | def __repr__(me): | |
74 | return '#<%s %r>' % (type(me).__name__, me._name) | |
75 | ||
76 | class DictExpanderClass (type): | |
77 | """ | |
78 | Metaclass for classes with autogenerated members. | |
79 | ||
80 | If the class body defines a dictionary `__extra__' then the key/value pairs | |
81 | in this dictionary are promoted into attributes of the class. This is much | |
82 | easier -- and safer -- than fiddling about with `locals'. | |
83 | """ | |
84 | def __new__(cls, name, supers, dict): | |
85 | try: | |
86 | ex = dict['__extra__'] | |
87 | except KeyError: | |
88 | pass | |
89 | else: | |
90 | for k, v in ex.iteritems(): | |
91 | dict[k] = v | |
92 | del dict['__extra__'] | |
93 | return super(DictExpanderClass, cls).__new__(cls, name, supers, dict) | |
94 | ||
95 | class ExpectedError (Exception): | |
96 | """ | |
97 | A (concrete) base class for various errors we expect to encounter. | |
98 | ||
99 | The `msg' attribute carries a human-readable message explaining what the | |
100 | problem actually is. The really important bit, though, is the `code' | |
101 | attribute, which carries an HTTP status code to be reported to the user | |
102 | agent, if we're running through CGI. | |
103 | """ | |
104 | def __init__(me, code, msg): | |
105 | me.code = code | |
106 | me.msg = msg | |
107 | def __str__(me): | |
108 | return '%s (%d)' % (me.msg, me.code) | |
109 | ||
110 | def register(dict, name): | |
111 | """A decorator: add the decorated function to DICT, under the key NAME.""" | |
112 | def _(func): | |
113 | dict[name] = func | |
114 | return func | |
115 | return _ | |
116 | ||
117 | class StringSubst (object): | |
118 | """ | |
119 | A string substitution. Initialize with a dictionary mapping source strings | |
120 | to target strings. The object is callable, and maps strings in the obvious | |
121 | way. | |
122 | """ | |
123 | def __init__(me, map): | |
124 | me._map = map | |
125 | me._rx = RX.compile('|'.join(RX.escape(s) for s in map)) | |
126 | def __call__(me, s): | |
127 | return me._rx.sub(lambda m: me._map[m.group(0)], s) | |
128 | ||
129 | def readline(what, file = SYS.stdin): | |
130 | """Read a single line from FILE (default stdin) and return it.""" | |
131 | try: line = SYS.stdin.readline() | |
132 | except IOError, e: raise ExpectedError, (500, str(e)) | |
133 | if not line.endswith('\n'): | |
134 | raise ExpectedError, (500, "Failed to read %s" % what) | |
135 | return line[:-1] | |
136 | ||
137 | class EscapeHatch (BaseException): | |
138 | """Exception used by the `Escape' context manager""" | |
139 | def __init__(me): pass | |
140 | ||
141 | class Escape (object): | |
142 | """ | |
143 | A context manager. Executes its body until completion or the `Escape' | |
144 | object itself is invoked as a function. Other exceptions propagate | |
145 | normally. | |
146 | """ | |
147 | def __init__(me): | |
148 | me.exc = EscapeHatch() | |
149 | def __call__(me): | |
150 | raise me.exc | |
151 | def __enter__(me): | |
152 | return me | |
153 | def __exit__(me, exty, exval, extb): | |
154 | return exval is me.exc | |
155 | ||
156 | class Fluid (object): | |
157 | """ | |
158 | Stores `fluid' variables which can be temporarily bound to new values, and | |
159 | later restored. | |
160 | ||
161 | A caller may use the object's attributes for storing arbitrary values | |
162 | (though storing a `bind' value would be silly). The `bind' method provides | |
163 | a context manager which binds attributes to other values during its | |
164 | execution. This works even with multiple threads. | |
165 | """ | |
166 | ||
167 | ## We maintain two stores for variables. One is a global store, `_g'; the | |
168 | ## other is a thread-local store `_t'. We look for a variable first in the | |
169 | ## thread-local store, and then if necessary in the global store. Binding | |
170 | ## works by remembering the old state of the variable on entry, setting it | |
171 | ## in the thread-local store (always), and then restoring the old state on | |
172 | ## exit. | |
173 | ||
174 | ## A special marker for unbound variables. If a variable is bound to a | |
175 | ## value, rebound temporarily with `bind', and then deleted, we must | |
176 | ## pretend that it's not there, and then restore it again afterwards. We | |
177 | ## use this tag to mark variables which have been deleted while they're | |
178 | ## rebound. | |
179 | UNBOUND = Tag('unbound-variable') | |
180 | ||
181 | def __init__(me, **kw): | |
182 | """Create a new set of fluid variables, initialized from the keywords.""" | |
183 | me.__dict__.update(_g = struct(), | |
184 | _t = TH.local()) | |
185 | for k, v in kw.iteritems(): | |
186 | setattr(me._g, k, v) | |
187 | ||
188 | def __getattr__(me, k): | |
189 | """Return the current value stored with K, or raise AttributeError.""" | |
190 | try: v = getattr(me._t, k) | |
191 | except AttributeError: v = getattr(me._g, k) | |
192 | if v is Fluid.UNBOUND: raise AttributeError, k | |
193 | return v | |
194 | ||
195 | def __setattr__(me, k, v): | |
196 | """Associate the value V with the variable K.""" | |
197 | if hasattr(me._t, k): setattr(me._t, k, v) | |
198 | else: setattr(me._g, k, v) | |
199 | ||
200 | def __delattr__(me, k): | |
201 | """ | |
202 | Forget about the variable K, so that attempts to read it result in an | |
203 | AttributeError. | |
204 | """ | |
205 | if hasattr(me._t, k): setattr(me._t, k, Fluid.UNBOUND) | |
206 | else: delattr(me._g, k) | |
207 | ||
208 | def __dir__(me): | |
209 | """Return a list of the currently known variables.""" | |
210 | seen = set() | |
211 | keys = [] | |
212 | for s in [me._t, me._g]: | |
213 | for k in dir(s): | |
214 | if k in seen: continue | |
215 | seen.add(k) | |
216 | if getattr(s, k) is not Fluid.UNBOUND: keys.append(k) | |
217 | return keys | |
218 | ||
219 | @CTX.contextmanager | |
220 | def bind(me, **kw): | |
221 | """ | |
222 | A context manager: bind values to variables according to the keywords KW, | |
223 | and execute the body; when the body exits, restore the rebound variables | |
224 | to their previous values. | |
225 | """ | |
226 | ||
227 | ## A list of things to do when we finish. | |
228 | unwind = [] | |
229 | ||
230 | def _delattr(k): | |
231 | ## Remove K from the thread-local store. Only it might already have | |
232 | ## been deleted, so be careful. | |
233 | try: delattr(me._t, k) | |
234 | except AttributeError: pass | |
235 | ||
236 | def stash(k): | |
237 | ## Stash a function for restoring the old state of K. We do this here | |
238 | ## rather than inline only because Python's scoping rules are crazy and | |
239 | ## we need to ensure that all of the necessary variables are | |
240 | ## lambda-bound. | |
241 | try: ov = getattr(me._t, k) | |
242 | except AttributeError: unwind.append(lambda: _delattr(k)) | |
243 | else: unwind.append(lambda: setattr(me._t, k, ov)) | |
244 | ||
245 | ## Rebind the variables. | |
246 | for k, v in kw.iteritems(): | |
247 | stash(k) | |
248 | setattr(me._t, k, v) | |
249 | ||
250 | ## Run the body, and restore. | |
251 | try: yield me | |
252 | finally: | |
253 | for f in unwind: f() | |
254 | ||
255 | class Cleanup (object): | |
256 | """ | |
257 | A context manager for stacking other context managers. | |
258 | ||
259 | By itself, it does nothing. Attach other context managers with `enter' or | |
260 | loose cleanup functions with `add'. On exit, contexts are left and | |
261 | cleanups performed in reverse order. | |
262 | """ | |
263 | def __init__(me): | |
264 | me._cleanups = [] | |
265 | def __enter__(me): | |
266 | return me | |
267 | def __exit__(me, exty, exval, extb): | |
268 | trap = False | |
269 | for c in reversed(me._cleanups): | |
270 | if c(exty, exval, extb): trap = True | |
271 | return trap | |
272 | def enter(me, ctx): | |
273 | v = ctx.__enter__() | |
274 | me._cleanups.append(ctx.__exit__) | |
275 | return v | |
276 | def add(me, func): | |
277 | me._cleanups.append(lambda exty, exval, extb: func()) | |
278 | ||
279 | ###-------------------------------------------------------------------------- | |
280 | ### Encodings. | |
281 | ||
282 | class Encoding (object): | |
283 | """ | |
284 | A pairing of injective encoding on binary strings, with its appropriate | |
285 | partial inverse. | |
286 | ||
287 | The two functions are available in the `encode' and `decode' attributes. | |
288 | See also the `ENCODINGS' dictionary. | |
289 | """ | |
290 | def __init__(me, encode, decode): | |
291 | me.encode = encode | |
292 | me.decode = decode | |
293 | ||
294 | ENCODINGS = { | |
295 | 'base64': Encoding(lambda s: BN.b64encode(s), | |
296 | lambda s: BN.b64decode(s)), | |
297 | 'base32': Encoding(lambda s: BN.b32encode(s).lower(), | |
298 | lambda s: BN.b32decode(s, casefold = True)), | |
299 | 'hex': Encoding(lambda s: BN.b16encode(s).lower(), | |
300 | lambda s: BN.b16decode(s, casefold = True)), | |
301 | None: Encoding(identity, identity) | |
302 | } | |
303 | ||
304 | ###-------------------------------------------------------------------------- | |
305 | ### Time and timeouts. | |
306 | ||
307 | def update_time(): | |
308 | """ | |
309 | Reset our idea of the current time, as kept in the global variable `NOW'. | |
310 | """ | |
311 | global NOW | |
312 | NOW = int(T.time()) | |
313 | update_time() | |
314 | ||
315 | class Alarm (Exception): | |
316 | """ | |
317 | Exception used internally by the `timeout' context manager. | |
318 | ||
319 | If you're very unlucky, you might get one of these at top level. | |
320 | """ | |
321 | pass | |
322 | ||
323 | class Timeout (ExpectedError): | |
324 | """ | |
325 | Report a timeout, from the `timeout' context manager. | |
326 | """ | |
327 | def __init__(me, what): | |
328 | ExpectedError.__init__(me, 500, "Timeout %s" % what) | |
329 | ||
330 | ## Set `DEADLINE' to be the absolute time of the next alarm. We'll keep this | |
331 | ## up to date in `timeout'. | |
332 | delta, _ = SIG.getitimer(SIG.ITIMER_REAL) | |
333 | if delta == 0: DEADLINE = None | |
334 | else: DEADLINE = NOW + delta | |
335 | ||
336 | def _alarm(sig, tb): | |
337 | """If we receive `SIGALRM', raise the alarm.""" | |
338 | raise Alarm | |
339 | SIG.signal(SIG.SIGALRM, _alarm) | |
340 | ||
341 | @CTX.contextmanager | |
342 | def timeout(delta, what): | |
343 | """ | |
344 | A context manager which interrupts execution of its body after DELTA | |
345 | seconds, if it doesn't finish before then. | |
346 | ||
347 | If execution is interrupted, a `Timeout' exception is raised, carrying WHY | |
348 | (a gerund phrase) as part of its message. | |
349 | """ | |
350 | ||
351 | global DEADLINE | |
352 | when = NOW + delta | |
353 | if DEADLINE is not None and when >= DEADLINE: | |
354 | yield | |
355 | update_time() | |
356 | else: | |
357 | od = DEADLINE | |
358 | try: | |
359 | DEADLINE = when | |
360 | SIG.setitimer(SIG.ITIMER_REAL, delta) | |
361 | yield | |
362 | except Alarm: | |
363 | raise Timeout, what | |
364 | finally: | |
365 | update_time() | |
366 | DEADLINE = od | |
367 | if od is None: SIG.setitimer(SIG.ITIMER_REAL, 0) | |
368 | else: SIG.setitimer(SIG.ITIMER_REAL, DEADLINE - NOW) | |
369 | ||
370 | ###-------------------------------------------------------------------------- | |
371 | ### File locking. | |
372 | ||
373 | @CTX.contextmanager | |
374 | def lockfile(lock, t = None): | |
375 | """ | |
376 | Acquire an exclusive lock on a named file LOCK while executing the body. | |
377 | ||
378 | If T is zero, fail immediately if the lock can't be acquired; if T is none, | |
379 | then wait forever if necessary; otherwise give up after T seconds. | |
380 | """ | |
381 | fd = -1 | |
382 | try: | |
383 | fd = OS.open(lock, OS.O_WRONLY | OS.O_CREAT, 0600) | |
384 | if timeout is None: | |
385 | F.lockf(fd, F.LOCK_EX) | |
386 | elif timeout == 0: | |
387 | F.lockf(fd, F.LOCK_EX | F.LOCK_NB) | |
388 | else: | |
389 | with timeout(t, "waiting for lock file `%s'" % lock): | |
390 | F.lockf(fd, F.LOCK_EX) | |
391 | yield None | |
392 | finally: | |
393 | if fd != -1: OS.close(fd) | |
394 | ||
395 | ###-------------------------------------------------------------------------- | |
396 | ### Database utilities. | |
397 | ||
398 | ### Python's database API is dreadful: it exposes far too many | |
399 | ### implementation-specific details to the programmer, who may well want to | |
400 | ### write code which works against many different databases. | |
401 | ### | |
402 | ### One particularly frustrating problem is the variability of placeholder | |
403 | ### syntax in SQL statements: there's no universal convention, just a number | |
404 | ### of possible syntaxes, at least one of which will be implemented (and some | |
405 | ### of which are mutually incompatible). Because not doing this invites all | |
406 | ### sorts of misery such as SQL injection vulnerabilties, we introduce a | |
407 | ### simple abstraction. A database parameter-type object keeps track of one | |
408 | ### particular convention, providing the correct placeholders to be inserted | |
409 | ### into the SQL command string, and the corresponding arguments, in whatever | |
410 | ### way is necessary. | |
411 | ### | |
412 | ### The protocol is fairly simple. An object of the appropriate class is | |
413 | ### instantiated for each SQL statement, providing it with a dictionary | |
414 | ### mapping placeholder names to their values. The object's `sub' method is | |
415 | ### called for each placeholder found in the statement, with a match object | |
416 | ### as an argument; the match object picks out the name of the placeholder in | |
417 | ### question in group 1, and the method returns a piece of syntax appropriate | |
418 | ### to the database backend. Finally, the collected arguments are made | |
419 | ### available, in whatever format is required, in the object's `args' | |
420 | ### attribute. | |
421 | ||
422 | ## Turn simple Unix not-quite-glob patterns into SQL `LIKE' patterns. | |
423 | ## Match using: x LIKE y ESCAPE '\\' | |
424 | globtolike = StringSubst({ | |
425 | '\\*': '*', '%': '\\%', '*': '%', | |
426 | '\\?': '?', '_': '\\_', '?': '_' | |
427 | }) | |
428 | ||
429 | class LinearParam (object): | |
430 | """ | |
431 | Abstract parent class for `linear' parameter conventions. | |
432 | ||
433 | A linear convention is one where the arguments are supplied as a list, and | |
434 | placeholders are either all identical (with semantics `insert the next | |
435 | argument'), or identify their argument by its position within the list. | |
436 | """ | |
437 | def __init__(me, kw): | |
438 | me._i = 0 | |
439 | me.args = [] | |
440 | me._kw = kw | |
441 | def sub(me, match): | |
442 | name = match.group(1) | |
443 | me.args.append(me._kw[name]) | |
444 | marker = me._format() | |
445 | me._i += 1 | |
446 | return marker | |
447 | class QmarkParam (LinearParam): | |
448 | def _format(me): return '?' | |
449 | class NumericParam (LinearParam): | |
450 | def _format(me): return ':%d' % me._i | |
451 | class FormatParam (LinearParam): | |
452 | def _format(me): return '%s' | |
453 | ||
454 | class DictParam (object): | |
455 | """ | |
456 | Abstract parent class for `dictionary' parameter conventions. | |
457 | ||
458 | A dictionary convention is one where the arguments are provided as a | |
459 | dictionary, and placeholders contain a key name identifying the | |
460 | corresponding value in that dictionary. | |
461 | """ | |
462 | def __init__(me, kw): | |
463 | me.args = kw | |
464 | def sub(me, match): | |
465 | name = match.group(1) | |
466 | return me._format(name) | |
467 | def NamedParam (object): | |
468 | def _format(me, name): return ':%s' % name | |
469 | def PyFormatParam (object): | |
470 | def _format(me, name): return '%%(%s)s' % name | |
471 | ||
472 | ### Since we're doing a bunch of work to paper over idiosyncratic placeholder | |
473 | ### syntax, we might as well also sort out other problems. The `DB_FIXUPS' | |
474 | ### dictionary maps database module names to functions which might need to do | |
475 | ### clever stuff at connection setup time. | |
476 | ||
477 | DB_FIXUPS = {} | |
478 | ||
479 | @register(DB_FIXUPS, 'sqlite3') | |
480 | def fixup_sqlite3(db): | |
481 | """ | |
482 | Unfortunately, SQLite learnt about FOREIGN KEY constraints late, and so | |
483 | doesn't enforce them unless explicitly told to. | |
484 | """ | |
485 | c = db.cursor() | |
486 | c.execute("PRAGMA foreign_keys = ON") | |
487 | ||
488 | class SimpleDBConnection (object): | |
489 | """ | |
490 | Represents a database connection, while trying to hide the differences | |
491 | between various kinds of database backends. | |
492 | """ | |
493 | ||
494 | __metaclass__ = DictExpanderClass | |
495 | ||
496 | ## A map from placeholder convention names to classes implementing them. | |
497 | PLACECLS = { | |
498 | 'qmark': QmarkParam, | |
499 | 'numeric': NumericParam, | |
500 | 'named': NamedParam, | |
501 | 'format': FormatParam, | |
502 | 'pyformat': PyFormatParam | |
503 | } | |
504 | ||
505 | ## A pattern for our own placeholder syntax. | |
506 | R_PLACE = RX.compile(r'\$(\w+)') | |
507 | ||
508 | def __init__(me, modname, modargs): | |
509 | """ | |
510 | Make a new database connection, using the module MODNAME, and passing its | |
511 | `connect' function the MODARGS -- which may be either a list or a | |
512 | dictionary. | |
513 | """ | |
514 | ||
515 | ## Get the module, and create a connection. | |
516 | mod = __import__(modname) | |
517 | if isinstance(modargs, dict): me._db = mod.connect(**modargs) | |
518 | else: me._db = mod.connect(*modargs) | |
519 | ||
520 | ## Apply any necessary fixups. | |
521 | try: fixup = DB_FIXUPS[modname] | |
522 | except KeyError: pass | |
523 | else: fixup(me._db) | |
524 | ||
525 | ## Grab hold of other interesting things. | |
526 | me.Error = mod.Error | |
527 | me.Warning = mod.Warning | |
528 | me._placecls = me.PLACECLS[mod.paramstyle] | |
529 | ||
530 | def execute(me, command, **kw): | |
531 | """ | |
532 | Execute the SQL COMMAND. The keyword arguments are used to provide | |
533 | values corresponding to `$NAME' placeholders in the COMMAND. | |
534 | ||
535 | Return the receiver, so that iterator protocol is convenient. | |
536 | """ | |
537 | me._cur = me._db.cursor() | |
538 | plc = me._placecls(kw) | |
539 | subst = me.R_PLACE.sub(plc.sub, command) | |
540 | ##PRINT('*** %s : %r' % (subst, plc.args)) | |
541 | me._cur.execute(subst, plc.args) | |
542 | return me | |
543 | ||
544 | def __iter__(me): | |
545 | """Iterator protocol: simply return the receiver.""" | |
546 | return me | |
547 | def next(me): | |
548 | """Iterator protocol: return the next row from the current query.""" | |
549 | row = me.fetchone() | |
550 | if row is None: raise StopIteration | |
551 | return row | |
552 | ||
553 | def __enter__(me): | |
74b87214 | 554 | """Context protocol: begin a transaction.""" |
a2916c06 MW |
555 | ##PRINT('<<< BEGIN') |
556 | return | |
557 | def __exit__(me, exty, exval, tb): | |
558 | """Context protocol: commit or roll back a transaction.""" | |
559 | if exty: | |
560 | ##PRINT('>*> ROLLBACK') | |
561 | me.rollback() | |
562 | else: | |
563 | ##PRINT('>>> COMMIT') | |
564 | me.commit() | |
565 | ||
566 | ## Import a number of methods from the underlying connection. | |
567 | __extra__ = {} | |
568 | for _name in ['fetchone', 'fetchmany', 'fetchall']: | |
569 | def _(name, extra): | |
570 | extra[name] = lambda me, *args, **kw: \ | |
571 | getattr(me._cur, name)(*args, **kw) | |
572 | _(_name, __extra__) | |
573 | for _name in ['commit', 'rollback']: | |
574 | def _(name, extra): | |
575 | extra[name] = lambda me, *args, **kw: \ | |
576 | getattr(me._db, name)(*args, **kw) | |
577 | _(_name, __extra__) | |
578 | del _name, _ | |
579 | ||
580 | ###----- That's all, folks -------------------------------------------------- |